package main

import (
	"bufio"
	"os/exec"
	"path/filepath"
	"strconv"
	"strings"
)

// diffIndex answers "is this file:line part of the unstaged working-tree diff?"
// It resolves each file to its enclosing git repo and, once per repo, parses
// `git diff --unified=0` (added/modified lines) plus untracked files (which count
// as fully changed). Everything is cached; files with no enclosing repo, or no
// matching change, are treated as unchanged.
type diffIndex struct {
	repoOf map[string]string    // dir -> git repo root ("" = none)
	repos  map[string]*repoDiff // repo root -> parsed diff
}

type repoDiff struct {
	changed   map[string]map[int]bool // repo-relative slash path -> changed line set
	untracked map[string]bool         // repo-relative slash path (fully changed)
}

func newDiffIndex() *diffIndex {
	return &diffIndex{repoOf: map[string]string{}, repos: map[string]*repoDiff{}}
}

// isChanged reports whether absPath:line lies in the unstaged working-tree diff.
func (d *diffIndex) isChanged(absPath string, line int) bool {
	repo := d.repoRoot(filepath.Dir(absPath))
	if repo == "" {
		return false
	}
	rd := d.repoDiff(repo)
	rel, err := filepath.Rel(repo, absPath)
	if err != nil {
		return false
	}
	rel = filepath.ToSlash(rel)
	if rd.untracked[rel] {
		return true
	}
	if lines, ok := rd.changed[rel]; ok {
		return lines[line]
	}
	return false
}

// repoRoot walks up from dir to find the enclosing git repo root, cached.
func (d *diffIndex) repoRoot(dir string) string {
	if root, ok := d.repoOf[dir]; ok {
		return root
	}
	out, err := exec.Command("git", "-C", dir, "rev-parse", "--show-toplevel").Output()
	root := ""
	if err == nil {
		root = strings.TrimSpace(string(out))
	}
	d.repoOf[dir] = root
	return root
}

// repoDiff parses (and caches) the unstaged diff + untracked files for a repo.
func (d *diffIndex) repoDiff(repo string) *repoDiff {
	if rd, ok := d.repos[repo]; ok {
		return rd
	}
	rd := &repoDiff{changed: map[string]map[int]bool{}, untracked: map[string]bool{}}

	if out, err := exec.Command("git", "-C", repo, "diff", "--unified=0", "--no-color").Output(); err == nil {
		parseUnifiedDiff(string(out), rd)
	}
	if out, err := exec.Command("git", "-C", repo, "ls-files", "--others", "--exclude-standard").Output(); err == nil {
		for _, f := range strings.Split(strings.TrimSpace(string(out)), "\n") {
			if f != "" {
				rd.untracked[filepath.ToSlash(f)] = true
			}
		}
	}

	d.repos[repo] = rd
	return rd
}

// parseUnifiedDiff records added/modified line numbers per file from a
// `git diff --unified=0` payload. Only the new-file side (`+c,d`) matters.
func parseUnifiedDiff(diff string, rd *repoDiff) {
	var cur string
	sc := bufio.NewScanner(strings.NewReader(diff))
	sc.Buffer(make([]byte, 0, 64*1024), 4*1024*1024)
	for sc.Scan() {
		line := sc.Text()
		switch {
		case strings.HasPrefix(line, "+++ b/"):
			cur = strings.TrimPrefix(line, "+++ b/")
			if _, ok := rd.changed[cur]; !ok {
				rd.changed[cur] = map[int]bool{}
			}
		case strings.HasPrefix(line, "@@"):
			start, count := parseHunkNewRange(line)
			if start <= 0 || cur == "" {
				continue
			}
			if rd.changed[cur] == nil {
				rd.changed[cur] = map[int]bool{}
			}
			for i := 0; i < count; i++ {
				rd.changed[cur][start+i] = true
			}
		}
	}
}

// parseHunkNewRange extracts (start, count) from the `+c,d` part of a hunk
// header like `@@ -a,b +c,d @@`. A missing `,d` means count 1.
func parseHunkNewRange(header string) (int, int) {
	plus := strings.IndexByte(header, '+')
	if plus < 0 {
		return 0, 0
	}
	rest := header[plus+1:]
	if sp := strings.IndexByte(rest, ' '); sp >= 0 {
		rest = rest[:sp]
	}
	start, count := 0, 1
	if comma := strings.IndexByte(rest, ','); comma >= 0 {
		start, _ = strconv.Atoi(rest[:comma])
		count, _ = strconv.Atoi(rest[comma+1:])
	} else {
		start, _ = strconv.Atoi(rest)
	}
	return start, count
}