agentu/internal/tools/patch.go

package tools

import (
	"fmt"
	"os"
	"path/filepath"
	"strings"

	"agentu/internal/safeio"
)

// ────────────────────────────────────────────────────────────────────────────
// V4 patch format data structures
// ────────────────────────────────────────────────────────────────────────────

type patchAction string

const (
	actionAdd    patchAction = "add"
	actionUpdate patchAction = "update"
	actionDelete patchAction = "delete"
)

type patchFile struct {
	action     patchAction
	path       string
	newContent string   // for "add" — full file content
	lines      []string // for "update" — lines to match/replace; for "delete" — unused
}

type parsedPatch struct {
	files []patchFile
}

// ────────────────────────────────────────────────────────────────────────────
// Parser
// ────────────────────────────────────────────────────────────────────────────

// parsePatch parses a V4-format patch string.
//
// Format:
//
//	*** Begin Patch
//	*** Add File: <path>
//	<file content>
//	*** Update File: <path>
//	@@
//	<context line>
//	@@
//	-<old line>
//	+<new line>
//	*** Delete File: <path>
//	*** End Patch
func parsePatch(patch string) (*parsedPatch, error) {
	patch = strings.TrimSpace(patch)
	if patch == "" {
		return nil, fmt.Errorf("patch is empty")
	}

	lines := strings.Split(patch, "\n")
	if len(lines) == 0 {
		return nil, fmt.Errorf("patch is empty")
	}

	if strings.TrimSpace(lines[0]) != "*** Begin Patch" {
		return nil, fmt.Errorf("patch must start with *** Begin Patch")
	}

	var files []patchFile
	seenPaths := make(map[string]struct{})
	seenEnd := false
	i := 1

	for i < len(lines) {
		line := strings.TrimSpace(lines[i])

		switch {
		case line == "*** End Patch":
			seenEnd = true
			i = len(lines)
		case strings.HasPrefix(line, "*** Add File:"):
			path := strings.TrimSpace(strings.TrimPrefix(line, "*** Add File:"))
			if err := validatePatchPath(path, seenPaths); err != nil {
				return nil, err
			}
			i++
			var contentLines []string
			for i < len(lines) {
				next := strings.TrimSpace(lines[i])
				if strings.HasPrefix(next, "***") {
					break
				}
				content := lines[i]
				if strings.HasPrefix(content, "+") {
					content = strings.TrimPrefix(content, "+")
				}
				contentLines = append(contentLines, content)
				i++
			}
			files = append(files, patchFile{
				action:     actionAdd,
				path:       path,
				newContent: strings.Join(contentLines, "\n") + "\n",
			})

		case strings.HasPrefix(line, "*** Delete File:"):
			path := strings.TrimSpace(strings.TrimPrefix(line, "*** Delete File:"))
			if err := validatePatchPath(path, seenPaths); err != nil {
				return nil, err
			}
			files = append(files, patchFile{
				action: actionDelete,
				path:   path,
			})
			i++

		case strings.HasPrefix(line, "*** Update File:"):
			path := strings.TrimSpace(strings.TrimPrefix(line, "*** Update File:"))
			if err := validatePatchPath(path, seenPaths); err != nil {
				return nil, err
			}
			i++
			var hunkLines []string
			for i < len(lines) {
				next := strings.TrimSpace(lines[i])
				if strings.HasPrefix(next, "***") {
					break
				}
				hunkLines = append(hunkLines, lines[i])
				i++
			}
			files = append(files, patchFile{
				action: actionUpdate,
				path:   path,
				lines:  hunkLines,
			})

		default:
			i++
		}
	}

	if len(files) == 0 {
		return nil, fmt.Errorf("patch contains no file operations")
	}
	if !seenEnd {
		return nil, fmt.Errorf("patch must end with *** End Patch")
	}

	return &parsedPatch{files: files}, nil
}

func validatePatchPath(path string, seen map[string]struct{}) error {
	if path == "" {
		return fmt.Errorf("patch file path is required")
	}
	if _, ok := seen[path]; ok {
		return fmt.Errorf("patch contains duplicate file path: %s", path)
	}
	seen[path] = struct{}{}
	return nil
}

// ────────────────────────────────────────────────────────────────────────────
// Applier
// ────────────────────────────────────────────────────────────────────────────

// applyPatchFile applies a single patch file operation.
func applyPatchFile(baseDir string, pf patchFile) (string, error) {
	switch pf.action {
	case actionAdd:
		return applyAdd(baseDir, pf)
	case actionUpdate:
		return applyUpdate(baseDir, pf)
	case actionDelete:
		return applyDelete(baseDir, pf)
	default:
		return "", fmt.Errorf("unknown patch action: %s", pf.action)
	}
}

func applyAdd(baseDir string, pf patchFile) (string, error) {
	path, err := resolvePathForWrite(baseDir, pf.path)
	if err != nil {
		return "", err
	}
	if _, err := os.Stat(path); err == nil {
		return "", fmt.Errorf("cannot add existing file: %s", pf.path)
	} else if !os.IsNotExist(err) {
		return "", fmt.Errorf("stat new file %s: %w", pf.path, err)
	}
	if err := safeio.AtomicWriteFile(path, []byte(pf.newContent), 0o644); err != nil {
		return "", fmt.Errorf("write new file %s: %w", pf.path, err)
	}
	return fmt.Sprintf("+ %s", pf.path), nil
}

func applyDelete(baseDir string, pf patchFile) (string, error) {
	path, err := resolvePath(baseDir, pf.path)
	if err != nil {
		return "", err
	}
	if info, err := os.Stat(path); err != nil {
		if os.IsNotExist(err) {
			return "", fmt.Errorf("cannot delete non-existent file: %s", pf.path)
		}
		return "", err
	} else if info.IsDir() {
		return "", fmt.Errorf("cannot delete directory: %s", pf.path)
	}
	if err := os.Remove(path); err != nil {
		return "", fmt.Errorf("delete file %s: %w", pf.path, err)
	}
	// Try to remove empty parent directories.
	dir := filepath.Dir(path)
	for dir != baseDir && dir != "." {
		if entries, err := os.ReadDir(dir); err == nil && len(entries) == 0 {
			os.Remove(dir)
			dir = filepath.Dir(dir)
		} else {
			break
		}
	}
	return fmt.Sprintf("- %s", pf.path), nil
}

func applyUpdate(baseDir string, pf patchFile) (string, error) {
	path, err := resolvePath(baseDir, pf.path)
	if err != nil {
		return "", fmt.Errorf("cannot update non-existent file: %s: %w", pf.path, err)
	}
	existing, err := os.ReadFile(path)
	if err != nil {
		return "", fmt.Errorf("cannot update non-existent file: %s", pf.path)
	}
	info, err := os.Stat(path)
	if err != nil {
		return "", fmt.Errorf("stat file %s: %w", pf.path, err)
	}

	fileLines := strings.Split(string(existing), "\n")

	// Parse the hunk lines into: context blocks + old lines + new lines.
	result, err := applyHunkLines(fileLines, pf.lines, pf.path)
	if err != nil {
		return "", err
	}

	newContent := strings.Join(result, "\n")
	if err := safeio.AtomicWriteFile(path, []byte(newContent), info.Mode().Perm()); err != nil {
		return "", fmt.Errorf("write updated file %s: %w", pf.path, err)
	}
	return fmt.Sprintf("~ %s", pf.path), nil
}

// applyHunkLines applies a sequence of hunk lines to the file content.
// The hunk lines are a mix of:
//
//	@@           — hunk separator (a context line block starts)
//	<plain line> — context line (must appear in both old and new)
//	-<line>      — line to remove
//	+<line>      — line to add
//
// Hunks are matched greedily: for each group of contiguous context + - + + lines,
// find the first position in the file where context + - lines match,
// replace with context + + lines.
func applyHunkLines(fileLines []string, hunkLines []string, path string) ([]string, error) {
	if len(hunkLines) == 0 {
		return fileLines, nil
	}

	// Split the hunk lines into separate hunks at @@ markers.
	type hunk struct {
		lines []string
	}
	var hunks []hunk
	var current hunk

	for _, hl := range hunkLines {
		trimmed := strings.TrimSpace(hl)
		if strings.HasPrefix(trimmed, "@@") {
			if len(current.lines) > 0 {
				hunks = append(hunks, current)
				current = hunk{}
			}
			continue
		}
		current.lines = append(current.lines, hl)
	}
	if len(current.lines) > 0 {
		hunks = append(hunks, current)
	}

	if len(hunks) == 0 {
		return fileLines, nil
	}

	result := make([]string, len(fileLines))
	copy(result, fileLines)
	offset := 0 // tracks net line count change from previous hunks

	for hi, h := range hunks {
		// Separate the hunk lines into a search pattern and replacement.
		searchLines, replaceLines := buildHunkSearchReplace(h.lines)
		if len(searchLines) == 0 && len(replaceLines) == 0 {
			continue
		}

		matchStart := findHunkMatch(result, searchLines, offset)
		if matchStart < 0 {
			return nil, fmt.Errorf("hunk context not found in %s", path)
		}

		// Replace: remove searchLines, insert replaceLines.
		matchEnd := matchStart + len(searchLines)
		newResult := make([]string, 0, len(result)-len(searchLines)+len(replaceLines))
		newResult = append(newResult, result[:matchStart]...)
		newResult = append(newResult, replaceLines...)
		newResult = append(newResult, result[matchEnd:]...)
		result = newResult

		offset = matchStart + len(replaceLines)

		// Multiple matches check for this hunk only.
		secondMatch := findHunkMatch(result, searchLines, offset)
		if secondMatch >= 0 {
			return nil, fmt.Errorf("hunk context matched multiple times in %s", path)
		}

		_ = hi
	}

	return result, nil
}

// buildHunkSearchReplace splits a hunk's lines into the search pattern
// (context + old lines) and replacement (context + new lines).
func buildHunkSearchReplace(lines []string) (search, replace []string) {
	for _, l := range lines {
		if strings.HasPrefix(l, "-") {
			search = append(search, l[1:])
		} else if strings.HasPrefix(l, "+") {
			replace = append(replace, l[1:])
		} else {
			contextLine := strings.TrimPrefix(l, " ")
			search = append(search, contextLine)
			replace = append(replace, contextLine)
		}
	}
	return
}

// findHunkMatch finds the first position >= offset in fileLines where
// searchLines matches exactly (line-by-line comparison).
func findHunkMatch(fileLines []string, searchLines []string, offset int) int {
	for i := offset; i <= len(fileLines)-len(searchLines); i++ {
		match := true
		for j, sl := range searchLines {
			if fileLines[i+j] != sl {
				match = false
				break
			}
		}
		if match {
			return i
		}
	}
	return -1
}

// ────────────────────────────────────────────────────────────────────────────
// Preview
// ────────────────────────────────────────────────────────────────────────────

// summarizePatch returns a human-readable summary of the patch.
func summarizePatch(patch string) (string, error) {
	pp, err := parsePatch(patch)
	if err != nil {
		return "", err
	}

	var b strings.Builder
	added, updated, deleted := 0, 0, 0
	for _, pf := range pp.files {
		switch pf.action {
		case actionAdd:
			added++
			fmt.Fprintf(&b, "+ add %s\n", pf.path)
		case actionUpdate:
			updated++
			fmt.Fprintf(&b, "~ update %s\n", pf.path)
		case actionDelete:
			deleted++
			fmt.Fprintf(&b, "- delete %s\n", pf.path)
		}
	}

	summary := fmt.Sprintf("patch: %d file(s) changed", len(pp.files))
	if added > 0 {
		summary += fmt.Sprintf(" (%d added)", added)
	}
	if updated > 0 {
		summary += fmt.Sprintf(" (%d updated)", updated)
	}
	if deleted > 0 {
		summary += fmt.Sprintf(" (%d deleted)", deleted)
	}

	return summary + "\n" + b.String(), nil
}

// extractTargetPaths returns all file paths referenced in the patch.
func extractTargetPaths(baseDir, patch string) ([]string, error) {
	pp, err := parsePatch(patch)
	if err != nil {
		return nil, err
	}

	var paths []string
	for _, pf := range pp.files {
		// Use resolvePathForWrite for all files — TargetPaths needs paths
		// even for files that don't exist yet (rollback snapshots).
		resolved, err := resolvePathForWrite(baseDir, pf.path)
		if err != nil {
			return nil, fmt.Errorf("file %s: %w", pf.path, err)
		}
		paths = append(paths, resolved)
	}
	return paths, nil
}