package main

import (
	"bytes"
	"encoding/json"
	"fmt"
	"io"
	"os"
	"os/exec"
	"path/filepath"
	"regexp"
	"slices"
	"strings"
)

func main() {
	if len(os.Args) != 3 {
		fmt.Fprintf(os.Stderr, "Usage: %s <path> <path>\n\nSee the vdf(1) manual page for more information.\n", os.Args[0])
		os.Exit(1)
	}
	err := innerMain(os.Stdout, os.Args[1], os.Args[2])
	if err != nil {
		fmt.Fprintf(os.Stderr, "%s", err)
		os.Exit(1)
	}
}

func innerMain(w io.Writer, oldRoot string, newRoot string) error {
	// We resolve symlinks so that we can filter out the roots from the dependencies.
	oldRoot, err := filepath.EvalSymlinks(oldRoot)
	if err != nil {
		return fmt.Errorf("failed to eval symlinks in %q: %v", oldRoot, err)
	}
	newRoot, err = filepath.EvalSymlinks(newRoot)
	if err != nil {
		return fmt.Errorf("failed to eval symlinks in %q: %v", newRoot, err)
	}

	oldPathInfos, err := queryPaths(oldRoot)
	if err != nil {
		return err
	}
	newPathInfos, err := queryPaths(newRoot)
	if err != nil {
		return err
	}

	var nameRe *regexp.Regexp
	if strings.HasSuffix(oldRoot, ".drv") && strings.HasSuffix(newRoot, ".drv") {
		nameRe = derivationNameRe
	} else if strings.HasSuffix(oldRoot, ".drv") || strings.HasSuffix(newRoot, ".drv") {
		return fmt.Errorf("error: one derivation path given (either both paths should be output paths or both should be derivation paths)")
	} else {
		nameRe = outputNameRe
	}

	var oldDependencies []Dependency
	for path := range oldPathInfos {
		if path == oldRoot {
			continue
		}
		output, err := parseStorePath(nameRe, path)
		if err != nil {
			return err
		}
		oldDependencies = append(oldDependencies, *output)
	}

	var newDependencies []Dependency
	for path := range newPathInfos {
		if path == newRoot {
			continue
		}
		output, err := parseStorePath(nameRe, path)
		if err != nil {
			return err
		}
		newDependencies = append(newDependencies, *output)
	}

	// TODO: exclude package names defined in TOML config

	added, removed, changed := comparePackages(oldDependencies, newDependencies)

	systemPathPath := findReference(newPathInfos[newRoot], "system-path")
	etcPath := findReference(newPathInfos[newRoot], "etc")

	directDependencies := make(map[string]bool)
	if systemPathPath != nil && etcPath != nil {
		// We assume this is a NixOS system.

		// Detect environment.systemPackages as direct dependencies.
		for _, path := range newPathInfos[*systemPathPath].References {
			directDependencies[path] = true
		}

		// Detect users.users.<name>.packages as direct dependencies.
		for _, path := range newPathInfos[*etcPath].References {
			name, _ := stripDigest(filepath.Base(path))
			if name == "user-environment" {
				for _, path := range newPathInfos[path].References {
					directDependencies[path] = true
				}
			}
		}
	} else {
		for _, path := range newPathInfos[newRoot].References {
			directDependencies[path] = true
		}

	}

	printChanges(w, added, removed, changed, directDependencies)

	return nil
}

type PathInfo struct {
	References []string `json:"references"`
}

func queryPaths(path string) (map[string]PathInfo, error) {
	// We're using the experimental CLI because the stable CLI does not support JSON output.
	cmd := exec.Command("nix", "path-info", "--json", "--recursive", "--extra-experimental-features", "nix-command", "--", path)

	var stderr bytes.Buffer
	cmd.Stderr = &stderr

	output, err := cmd.Output()
	if err != nil {
		return nil, fmt.Errorf(
			"error: nix path-info --json failed: %v\n    %s",
			err, strings.TrimSpace(stderr.String()))
	}

	var paths map[string]PathInfo
	if err := json.Unmarshal(output, &paths); err != nil {
		return nil, fmt.Errorf("error parsing JSON: %v", err)
	}
	return paths, nil
}

func findReference(info PathInfo, name string) *string {
	pathIdx := slices.IndexFunc(info.References, func(path string) bool {
		pathName, _ := stripDigest(filepath.Base(path))
		return pathName == name
	})
	if pathIdx == -1 {
		return nil
	}
	return &info.References[pathIdx]
}

var digestRe = regexp.MustCompile(`^[a-z0-9]{32}-`)

func stripDigest(name string) (string, error) {
	indices := digestRe.FindStringIndex(name)
	if indices == nil {
		return name, fmt.Errorf("path %q did not start with digest", name)
	}
	return name[indices[1]:], nil
}

type Dependency struct {
	Name    string
	Version string
	Output  string
	Path    string
}

// We assume that versions start with a digit (nixpkgs also requires this.)
// We only recognize outputs after a version and assume it starts with a non-digit and doesn't contain a dash.
var outputNameRe = regexp.MustCompile(`(.+?)(?:-([0-9].*?)(?:-([^0-9].+))?)?$`)

// For derivations there is no output name and we strip .drv if it exists.
var derivationNameRe = regexp.MustCompile(`(.+?)(?:-([0-9].*?)())?(?:\.drv)?$`)

func parseStorePath(nameRe *regexp.Regexp, path string) (*Dependency, error) {
	name, err := stripDigest(filepath.Base(path))
	if err != nil {
		return nil, fmt.Errorf("path %s did not start with digest", path)
	}
	// fixed output derivations have a name starting with two digests
	name, _ = stripDigest(name)
	matches := nameRe.FindStringSubmatch(name)
	if matches == nil {
		return nil, fmt.Errorf("path %s did not match regex", path)
	}
	return &Dependency{
		Name:    matches[1],
		Version: matches[2],
		Output:  matches[3],
		Path:    path,
	}, nil
}

type VersionChange struct {
	Name        string
	Output      string
	OldVersions []Version
	NewVersions []Version
}

type Version struct {
	Version string
	Paths   []string
}

func (vc *Dependency) nameOutput() NameOutput {
	return NameOutput{vc.Name, vc.Output}
}

type NameOutput struct {
	Name   string
	Output string
}

func comparePackages(oldDependencies, newDependencies []Dependency) ([]Dependency, []Dependency, []VersionChange) {
	// For one package name there can be different versions in the same dependency tree.
	oldMap := make(map[NameOutput][]Dependency)
	newMap := make(map[NameOutput][]Dependency)

	for _, dep := range oldDependencies {
		versions, exists := oldMap[dep.nameOutput()]
		if exists {
			oldMap[dep.nameOutput()] = append(versions, dep)
		} else {
			oldMap[dep.nameOutput()] = []Dependency{dep}
		}
	}

	for _, dep := range newDependencies {
		versions, exists := newMap[dep.nameOutput()]
		if exists {
			newMap[dep.nameOutput()] = append(versions, dep)
		} else {
			newMap[dep.nameOutput()] = []Dependency{dep}
		}
	}

	var added []Dependency
	var removed []Dependency
	var changed []VersionChange

	for nameOut, newDeps := range newMap {
		if oldDeps, exists := oldMap[nameOut]; exists {
			// There can be different paths in the same dependency tree with the same package name and version.
			oldVersions := groupByVersion(oldDeps)
			newVersions := groupByVersion(newDeps)

			if !slices.EqualFunc(newVersions, oldVersions, sameVersion) {
				changed = append(changed, VersionChange{
					Name:        nameOut.Name,
					Output:      nameOut.Output,
					OldVersions: oldVersions,
					NewVersions: newVersions,
				})
			}
		} else {
			added = append(added, newDeps...)
		}
	}

	for nameOut, oldDeps := range oldMap {
		if _, exists := newMap[nameOut]; !exists {
			removed = append(removed, oldDeps...)
		}
	}

	return added, removed, changed
}

func groupByVersion(deps []Dependency) []Version {
	versionMap := make(map[string][]string)
	for _, dep := range deps {
		if dep.Version == "" {
			continue
		}
		paths, exists := versionMap[dep.Version]
		if exists {
			versionMap[dep.Version] = append(paths, dep.Path)
		} else {
			versionMap[dep.Version] = []string{dep.Path}
		}
	}
	versions := make([]Version, 0, len(versionMap))
	for version, paths := range versionMap {
		versions = append(versions, Version{version, paths})
	}
	slices.SortFunc(versions, func(a Version, b Version) int {
		return strings.Compare(a.Version, b.Version)
	})
	return versions
}

func sameVersion(v1, v2 Version) bool {
	return v1.Version == v2.Version
}

func printChanges(w io.Writer, added, removed []Dependency, changed []VersionChange, newDirectDeps map[string]bool) {
	if len(changed) > 0 {
		fmt.Fprintln(w, "Version changed:")
		// TODO: grouping changes into categories defined in TOML config
		changedSelected := make([]VersionChange, 0)
		changedUnselected := make([]VersionChange, 0)
		for _, change := range changed {
			// If any path of any new version is a direct dependency in the new dependency tree
			// we show the whole version change in the direct dependency section because figuring
			// out which old version exactly changed to which new version is non-trivial
			// since there can be several versions.
			if slices.IndexFunc(change.NewVersions, func(v Version) bool {
				for _, path := range v.Paths {
					_, exists := newDirectDeps[path]
					if exists {
						return true
					}
				}
				return false
			}) != -1 {
				changedSelected = append(changedSelected, change)
			} else {
				changedUnselected = append(changedUnselected, change)
			}
		}
		printVersionChanges(w, filterVersionChanges(changedSelected))
		if len(changedSelected) > 0 && len(changedUnselected) > 0 {
			fmt.Fprintln(w)
		}
		printVersionChanges(w, filterVersionChanges(changedUnselected))
	}
	if len(added) > 0 {
		if len(changed) > 0 {
			fmt.Fprintln(w)
		}
		fmt.Fprintln(w, "Added packages:")
		printDependencyList(w, added)
	}
	if len(removed) > 0 {
		if len(added) > 0 || len(changed) > 0 {
			fmt.Fprintln(w)
		}
		fmt.Fprintln(w, "Removed packages:")
		printDependencyList(w, removed)
	}
}

func filterVersionChanges(versionChanges []VersionChange) []VersionChange {
	slices.SortFunc(versionChanges, func(a, b VersionChange) int {
		if a.Name == b.Name {
			return strings.Compare(a.Output, b.Output)
		}
		return strings.Compare(a.Name, b.Name)
	})
	var outChange VersionChange
	filteredVersionChanges := []VersionChange{}
	for _, change := range versionChanges {
		if change.Name == outChange.Name &&
			slices.EqualFunc(change.OldVersions, outChange.OldVersions, sameVersion) &&
			slices.EqualFunc(change.NewVersions, outChange.NewVersions, sameVersion) {
			continue
		}

		filteredVersionChanges = append(filteredVersionChanges, change)

		if change.Output == "" {
			outChange = change
		}
	}
	return filteredVersionChanges
}

func printVersionChanges(w io.Writer, changed []VersionChange) {
	for _, change := range changed {
		// TODO: link changelogs defined in TOML config
		fmt.Fprintf(w, "  %s%s: %s -> %s\n", change.Name, formatOutput(change.Output),
			strings.Join(mapSlice(change.OldVersions, version), ", "),
			strings.Join(mapSlice(change.NewVersions, version), ", "),
		)
	}
}

func version(v Version) string {
	return v.Version
}

func printDependencyList(w io.Writer, deps []Dependency) {
	slices.SortFunc(deps, func(a, b Dependency) int {
		if a.Name == b.Name {
			return strings.Compare(a.Output, b.Output)
		}
		return strings.Compare(a.Name, b.Name)
	})
	for _, dep := range deps {
		fmt.Fprintf(w, "  %s%s %s\n", dep.Name, formatOutput(dep.Output), dep.Version)
	}
}

func formatOutput(output string) string {
	if output == "" {
		return ""
	} else {
		return fmt.Sprintf("(%s)", output)
	}
}

func mapSlice[T, U any](slice []T, fn func(T) U) []U {
	result := make([]U, len(slice))
	for i, v := range slice {
		result[i] = fn(v)
	}
	return result
}