package diff

import (
	"fmt"
	"slices"
	"strings"
)

// patchHunk is one hunk of a patch to apply: the lines it expects and the
// lines it puts in their place.
type patchHunk struct {
	header   string
	old, new []string
	ops      []Line
	// oldEnd and newEnd mark a "\ No newline at end of file" after a line
	// of the old or new text.
	oldEnd, newEnd bool
}

// parse reads the hunks of a unified diff. File headers (---, +++, diff,
// index) are skipped. Models often get the line numbers in @@ headers
// wrong, so hunks are placed by their content.
func parse(patch string) ([]patchHunk, error) {
	var hunks []patchHunk
	var cur *patchHunk
	flush := func() {
		if cur == nil {
			return
		}
		for _, op := range cur.ops {
			if op.Kind != Add {
				cur.old = append(cur.old, op.Text)
			}
			if op.Kind != Del {
				cur.new = append(cur.new, op.Text)
			}
		}
		if len(cur.ops) > 0 {
			hunks = append(hunks, *cur)
		}
		cur = nil
	}
	// Lines keep any "\r": a file's lines may end in one, and the patch
	// must say exactly what to write. Matching may ignore it (see Apply).
	lines := strings.Split(patch, "\n")
	if len(lines) > 0 && lines[len(lines)-1] == "" {
		lines = lines[:len(lines)-1]
	}
	for i, line := range lines {
		switch {
		case strings.HasPrefix(line, "@@"):
			flush()
			cur = &patchHunk{header: line}
		case cur == nil && (strings.HasPrefix(line, "---") || strings.HasPrefix(line, "+++") ||
			strings.HasPrefix(line, "diff ") || strings.HasPrefix(line, "index ")):
		case strings.HasPrefix(line, `\`): // "\ No newline at end of file"
			if cur != nil && len(cur.ops) > 0 {
				switch cur.ops[len(cur.ops)-1].Kind {
				case Same:
					cur.oldEnd, cur.newEnd = true, true
				case Del:
					cur.oldEnd = true
				case Add:
					cur.newEnd = true
				}
			}
		case cur == nil && strings.TrimSpace(line) == "":
		case cur == nil:
			// A patch without @@ headers is one hunk.
			cur = &patchHunk{header: "@@"}
			fallthrough
		default:
			if line == "" || line == "\r" {
				line = " " + line // an empty context line whose leading space was lost
			}
			k := Kind(line[0])
			if k != Same && k != Del && k != Add {
				return nil, fmt.Errorf("line %d of the diff starts with %q; every line of a hunk must start with a space, '-' or '+'", i+1, line[:1])
			}
			cur.ops = append(cur.ops, Line{k, line[1:]})
		}
	}
	flush()
	if len(hunks) == 0 {
		return nil, fmt.Errorf("the diff has no hunks")
	}
	return hunks, nil
}

// find returns every place want occurs in lines at or after from.
func find(lines, want []string, from int, eq func(a, b string) bool) []int {
	var at []int
	for i := from; i+len(want) <= len(lines); i++ {
		if slices.EqualFunc(lines[i:i+len(want)], want, eq) {
			at = append(at, i)
		}
	}
	return at
}

// Apply applies a unified diff to text. Each hunk is found by its context
// and removed lines, in order; trailing whitespace (and a CRLF line's
// "\r") is ignored if an exact match fails. A hunk must match one place
// only, and lines added with no context only to an empty file: the @@
// line numbers never place a hunk, because a wrong one would silently
// change the wrong lines.
func Apply(text, patch string) (string, error) {
	hunks, err := parse(patch)
	if err != nil {
		return "", err
	}
	trailingNewline := text == "" || strings.HasSuffix(text, "\n")
	lines := SplitLines(text)
	from := 0
	shift := 0 // how far earlier hunks moved later lines
	exact := func(a, b string) bool { return a == b }
	loose := func(a, b string) bool { return strings.TrimRight(a, " \t\r") == strings.TrimRight(b, " \t\r") }
	for n, h := range hunks {
		var at int
		if len(h.old) == 0 {
			// Lines added with no context have nowhere to go but an empty
			// file: the header's line number would place them unchecked.
			if len(lines) > 0 {
				return "", fmt.Errorf("hunk %d (%s) adds lines without any context lines to place them; include a few unchanged lines around the change", n+1, h.header)
			}
		} else {
			found := find(lines, h.old, from, exact)
			if len(found) == 0 {
				found = find(lines, h.old, from, loose)
			}
			switch len(found) {
			case 0:
				return "", fmt.Errorf("hunk %d (%s) does not match the file. Read the file again and make the context and '-' lines match it exactly. These lines were not found in order:\n%s", n+1, h.header, quote(h.old))
			case 1:
				at = found[0]
			default:
				return "", fmt.Errorf("hunk %d (%s) matches the file in %d places (%s); add context lines from around the one to change, so that it matches only there", n+1, h.header, len(found), lineList(found, shift))
			}
		}
		out := append([]string{}, lines[:at]...)
		// Context lines keep the file's own text, which a loose match may
		// differ from.
		k := at
		for _, op := range h.ops {
			switch op.Kind {
			case Same:
				out = append(out, lines[k])
				k++
			case Del:
				k++
			case Add:
				out = append(out, op.Text)
			}
		}
		out = append(out, lines[at+len(h.old):]...)
		lines = out
		from = at + len(h.new)
		shift += len(h.new) - len(h.old)
		// A hunk that marks the end of the old or new text changes
		// whether the result ends in a newline.
		if from == len(lines) {
			switch {
			case h.newEnd:
				trailingNewline = false
			case h.oldEnd:
				trailingNewline = true
			}
		}
	}
	result := strings.Join(lines, "\n")
	if len(lines) > 0 && trailingNewline {
		result += "\n"
	}
	return result, nil
}

func quote(lines []string) string {
	var b strings.Builder
	for i, l := range lines {
		if i >= 12 {
			fmt.Fprintf(&b, "    ... (%d more)\n", len(lines)-i)
			break
		}
		b.WriteString("    " + l + "\n")
	}
	return strings.TrimSuffix(b.String(), "\n")
}

// lineList names the lines where matches start, as numbered in the file
// before earlier hunks were applied, as the model read it. Matches after
// an earlier hunk are shifted by what it added or removed.
func lineList(at []int, shift int) string {
	var b strings.Builder
	for i, a := range at {
		if i == 5 {
			fmt.Fprintf(&b, " and %d more", len(at)-i)
			break
		}
		if i > 0 {
			b.WriteString(", ")
		}
		fmt.Fprintf(&b, "line %d", a+1-shift)
	}
	return b.String()
}