package emit

import (
	"context"
	"fmt"
	"os"
	"os/exec"
	"path/filepath"
	"regexp"
	"sort"
	"strings"
	"sync"
	"time"
	"unicode"
)

// compileMu serializes CompileAndRepair calls. Each go build invocation
// uses ~200MB+ per compiler process; running multiple concurrently from
// colony instances blows past 16GB RAM. One at a time.
var compileMu sync.Mutex

// CompileAndRepair attempts to compile source, and if compilation fails,
// iteratively fixes errors by adding type/var stubs and removing unused
// imports. Returns the repaired source or error if it cannot be fixed
// within maxPasses attempts.
func CompileAndRepair(source, goRoot string, maxPasses int) (string, error) {
	compileMu.Lock()
	defer compileMu.Unlock()

	// Hard deadline for the entire repair process.
	deadline := time.Now().Add(2 * time.Minute)
	originalSize := len(source)

	if maxPasses < 8 {
		maxPasses = 8
	}

	// Create a persistent workspace for all compile attempts in this repair
	// session. Reusing the same directory avoids re-running go mod tidy and
	// lets the Go build cache warm up across iterations.
	workDir, err := os.MkdirTemp("", "repair-*")
	if err != nil {
		return source, err
	}
	defer os.RemoveAll(workDir)

	// Write go.mod once — it doesn't change between iterations.
	modContent := buildRepairGoMod(source)
	os.WriteFile(filepath.Join(workDir, "go.mod"), []byte(modContent), 0o644)
	if strings.Contains(source, selfModule) {
		modRoot := findRepairModuleRoot()
		if modRoot != "" {
			if data, err := os.ReadFile(filepath.Join(modRoot, "go.sum")); err == nil {
				os.WriteFile(filepath.Join(workDir, "go.sum"), data, 0o644)
			}
		}
	}

	goBin := filepath.Join(goRoot, "bin", "go")
	env := repairCleanEnv(goRoot)

	// Write initial source and run go mod tidy once. Go 1.24+ requires
	// the go.sum to be current before building. We write the source first
	// so tidy can see which imports are actually used.
	os.WriteFile(filepath.Join(workDir, "main.go"), []byte(source), 0o644)
	tidyCtx, tidyCancel := context.WithTimeout(context.Background(), 30*time.Second)
	tidy := exec.CommandContext(tidyCtx, goBin, "mod", "tidy")
	tidy.Dir = workDir
	tidy.Env = env
	tidy.CombinedOutput()
	tidyCancel()

	// tryCompileLocal is the fast path: only rewrites main.go and builds.
	// Uses -p 1 to limit build parallelism (each compiler process uses
	// ~200MB; with 16 cores that's 3.2GB just for the compiler).
	// Each compile gets a 60-second timeout to prevent hanging.
	tryCompileLocal := func(src string) (string, error) {
		os.WriteFile(filepath.Join(workDir, "main.go"), []byte(src), 0o644)
		ctx, cancel := context.WithTimeout(context.Background(), 60*time.Second)
		defer cancel()
		cmd := exec.CommandContext(ctx, goBin, "build", "-p", "1", "-gcflags=-e", "-o", filepath.Join(workDir, "offspring"), ".")
		cmd.Dir = workDir
		cmd.Env = env
		out, err := cmd.CombinedOutput()
		return string(out), err
	}

	// Pre-scan: find names used as method receivers so we stub them as
	// struct types instead of vars.
	receiverTypes := findReceiverTypes(source)

	// Phase A: Iterative stub + import/redecl fixes.
	// Run until no more stubbable errors or maxPasses reached.
	for pass := range maxPasses * 2 {
		if time.Now().After(deadline) || len(source) > originalSize*3 {
			break
		}
		compileOutput, compileErr := tryCompileLocal(source)
		if compileErr == nil {
			return source, nil
		}

		fixes := parseCompileErrors(compileOutput)
		promoteReceiverFixes(fixes, receiverTypes)

		if len(fixes) == 0 {
			break // no more stubbable errors — move to removal
		}

		prev := source
		source = applyFixes(source, fixes)
		if source == prev {
			break // no progress from stubs
		}
		_ = pass
	}

	// Phase B: Incremental repair — neutralize before removing.
	// Strategy: (1) try to neutralize broken lines by replacing RHS with
	// zero values, preserving variable declarations; (2) if neutralization
	// makes no progress, remove error lines; (3) as last resort, replace
	// entire broken functions with signature-preserving stubs.
	for pass := range maxPasses * 4 {
		if time.Now().After(deadline) || len(source) > originalSize*3 {
			break
		}
		_ = pass
		compileOutput, compileErr := tryCompileLocal(source)
		if compileErr == nil {
			return source, nil
		}

		prev := source

		// Step 1: Try to neutralize error lines — replace broken RHS
		// with zero values while keeping the variable declaration.
		source = neutralizeErrorLines(source, compileOutput)

		// Step 2: If neutralization didn't help, remove error lines.
		if source == prev {
			source = removeErrorLines(source, compileOutput)
		}

		// Step 3: If line removal didn't help, replace broken functions
		// with signature-preserving stubs.
		if source == prev {
			source = removeErrorFunctions(source, compileOutput)
		}

		if source == prev {
			break // no progress
		}

		// After each change, run stub/import fixes until stable.
		// Cap at 3 rounds and abort if source grows beyond 3x original.
		for innerPass := range 3 {
			if len(source) > originalSize*3 {
				break
			}
			compileOutput, compileErr = tryCompileLocal(source)
			if compileErr == nil {
				return source, nil
			}
			fixes := parseCompileErrors(compileOutput)
			promoteReceiverFixes(fixes, receiverTypes)
			if len(fixes) == 0 {
				break
			}
			prevInner := source
			source = applyFixes(source, fixes)
			if source == prevInner {
				break
			}
			_ = innerPass
		}
	}

	// Ensure main() exists — it may have been emptied or removed.
	source = ensureMain(source)

	// Remove broken top-level declarations (unclosed braces, etc.)
	source = removeIncompleteDecls(source)

	// Final cleanup: remove any remaining orphaned blocks.
	source = cleanOrphanedBlocks(source)

	// One last compile attempt.
	_, finalErr := tryCompileLocal(source)
	if finalErr == nil {
		return source, nil
	}
	return source, fmt.Errorf("still has errors after %d passes", maxPasses)
}


const selfModule = "git.mleku.dev/mleku/dendrite"

// buildRepairGoMod generates a minimal go.mod for repair compilation.
func buildRepairGoMod(source string) string {
	var b strings.Builder
	b.WriteString("module offspring\n\ngo 1.24\n")

	if !strings.Contains(source, selfModule) {
		return b.String()
	}

	modRoot := findRepairModuleRoot()
	if modRoot == "" {
		return b.String()
	}

	// Read parent go.mod for require blocks.
	parentMod, err := os.ReadFile(filepath.Join(modRoot, "go.mod"))
	if err != nil {
		return b.String()
	}

	// Extract require blocks.
	lines := strings.Split(string(parentMod), "\n")
	inBlock := false
	for _, line := range lines {
		trimmed := strings.TrimSpace(line)
		if trimmed == "require (" {
			inBlock = true
			b.WriteString(line + "\n")
			continue
		}
		if inBlock {
			b.WriteString(line + "\n")
			if trimmed == ")" {
				inBlock = false
			}
			continue
		}
		if strings.HasPrefix(trimmed, "require ") && !strings.HasPrefix(trimmed, "require (") {
			b.WriteString(line + "\n")
		}
	}

	b.WriteString("\nrequire " + selfModule + " v0.0.0\n")
	b.WriteString("\nreplace " + selfModule + " => " + modRoot + "\n")

	return b.String()
}

// findRepairModuleRoot walks up from cwd looking for go.mod declaring selfModule.
func findRepairModuleRoot() string {
	dir, err := os.Getwd()
	if err != nil {
		return ""
	}
	for {
		data, err := os.ReadFile(filepath.Join(dir, "go.mod"))
		if err == nil {
			for _, line := range strings.Split(string(data), "\n") {
				line = strings.TrimSpace(line)
				if strings.HasPrefix(line, "module ") {
					mod := strings.TrimSpace(strings.TrimPrefix(line, "module"))
					if mod == selfModule {
						return dir
					}
				}
			}
		}
		parent := filepath.Dir(dir)
		if parent == dir {
			break
		}
		dir = parent
	}
	return ""
}

// repairCleanEnv returns a clean Go environment.
func repairCleanEnv(root string) []string {
	env := os.Environ()
	clean := make([]string, 0, len(env)+3)
	for _, e := range env {
		if strings.HasPrefix(e, "GOROOT=") ||
			strings.HasPrefix(e, "GOTOOLCHAIN=") ||
			strings.HasPrefix(e, "PATH=") {
			continue
		}
		clean = append(clean, e)
	}
	clean = append(clean,
		"GOROOT="+root,
		"GOTOOLCHAIN=local",
		"GOMAXPROCS=1",
		"PATH="+filepath.Join(root, "bin")+":"+os.Getenv("PATH"),
	)
	return clean
}

// reReceiverType matches method declarations to extract receiver type names.
// Matches: func (x TypeName), func (x *TypeName)
var reReceiverType = regexp.MustCompile(`func\s+\(\s*\w+\s+\*?(\w+)\)`)

// findReceiverTypes scans source for method declarations and returns a set
// of type names used as receivers.
func findReceiverTypes(source string) map[string]bool {
	types := make(map[string]bool)
	for _, line := range strings.Split(source, "\n") {
		if m := reReceiverType.FindStringSubmatch(line); m != nil {
			types[m[1]] = true
		}
	}
	return types
}

// promoteReceiverFixes upgrades "undefined_var" fixes to "undefined_type"
// if the name is used as a method receiver.
func promoteReceiverFixes(fixes []fix, receiverTypes map[string]bool) {
	for i := range fixes {
		if fixes[i].kind == "undefined_var" && receiverTypes[fixes[i].name] {
			fixes[i].kind = "undefined_type"
		}
	}
}

// ensureMain adds an empty func main() if one is not present.
func ensureMain(source string) string {
	for _, line := range strings.Split(source, "\n") {
		trimmed := strings.TrimSpace(line)
		if strings.HasPrefix(trimmed, "func main()") {
			return source
		}
	}
	return source + "\nfunc main() {}\n"
}

// fix represents a single repair action.
type fix struct {
	kind string // "undefined_type", "undefined_var", "unused_import", "redeclared", "missing_return"
	name string // the identifier, import path, or line number (for missing_return)
}

// Patterns for parsing Go compiler errors.
var (
	reUndefined      = regexp.MustCompile(`undefined:\s+([\w.]+)`)
	reUnusedImport   = regexp.MustCompile(`"([^"]+)" imported and not used`)
	reRedeclared     = regexp.MustCompile(`(\w+) redeclared in this block`)
	reMissingReturn  = regexp.MustCompile(`^(.+):(\d+):\d+: missing return`)
	reNotImplement   = regexp.MustCompile(`does not implement (\w+)`)
	reErrorAtLine    = regexp.MustCompile(`\./main\.go:(\d+):\d+:`)
	reAssignMismatch = regexp.MustCompile(`assignment mismatch: (\d+) variables but .* returns (\d+) values?`)
	reMultiValue     = regexp.MustCompile(`multiple-value (\w+)\(`)
	reDeclNotUsed    = regexp.MustCompile(`declared and not used: (\w+)`)
)

// parseCompileErrors extracts fixable errors from compiler output.
func parseCompileErrors(output string) []fix {
	seen := make(map[string]bool)
	var fixes []fix

	for _, line := range strings.Split(output, "\n") {
		if m := reUndefined.FindStringSubmatch(line); m != nil {
			name := m[1]
			key := "undef:" + name
			if !seen[key] {
				seen[key] = true
				if strings.Contains(name, ".") {
					// Qualified reference (e.g., hash.Len) — can't stub,
					// will be handled by removeErrorFunctions.
				} else if isTypeName(name) {
					fixes = append(fixes, fix{kind: "undefined_type", name: name})
				} else {
					fixes = append(fixes, fix{kind: "undefined_var", name: name})
				}
			}
		}
		if m := reUnusedImport.FindStringSubmatch(line); m != nil {
			path := m[1]
			key := "unused:" + path
			if !seen[key] {
				seen[key] = true
				fixes = append(fixes, fix{kind: "unused_import", name: path})
			}
		}
		if m := reRedeclared.FindStringSubmatch(line); m != nil {
			name := m[1]
			key := "redecl:" + name
			if !seen[key] {
				seen[key] = true
				fixes = append(fixes, fix{kind: "redeclared", name: name})
			}
		}
		// "missing return" at a line — remove the enclosing function.
		if reMissingReturn.MatchString(line) {
			// Extract line number and mark for function removal.
			if m := reMissingReturn.FindStringSubmatch(line); m != nil {
				key := "misret:" + m[2]
				if !seen[key] {
					seen[key] = true
					fixes = append(fixes, fix{kind: "missing_return", name: m[2]})
				}
			}
		}
		// "does not implement X" — the struct is used as an interface
		// it doesn't satisfy. Remove the offending statement.
		if reNotImplement.MatchString(line) {
			// Not easily fixable — skip for now, will be handled by
			// removal of the enclosing function if it also has other errors.
		}

		// "X declared and not used" — suppress by adding "_ = X".
		if m := reDeclNotUsed.FindStringSubmatch(line); m != nil {
			name := m[1]
			key := "declnotused:" + name
			if !seen[key] {
				seen[key] = true
				fixes = append(fixes, fix{kind: "decl_not_used", name: name})
			}
		}

		// "assignment mismatch: N variables but F returns M values" or
		// "multiple-value X() used in single-value context" — the func
		// stub has wrong return arity. Fix by regenerating the stub.
		if m := reMultiValue.FindStringSubmatch(line); m != nil {
			name := m[1]
			key := "multiret:" + name
			if !seen[key] {
				seen[key] = true
				fixes = append(fixes, fix{kind: "fix_return_arity", name: name})
			}
		}
		if m := reAssignMismatch.FindStringSubmatch(line); m != nil {
			// The error is on this line — we need to find the function name
			// being called. Extract from the error line reference.
			if lineM := reErrorAtLine.FindStringSubmatch(line); lineM != nil {
				key := "assignmis:" + lineM[1]
				if !seen[key] {
					seen[key] = true
					fixes = append(fixes, fix{kind: "fix_assign_mismatch", name: lineM[1]})
				}
			}
		}
	}

	return fixes
}

// isTypeName heuristically determines if an identifier is likely a type name
// (starts with uppercase letter).
func isTypeName(name string) bool {
	if len(name) == 0 {
		return false
	}
	return unicode.IsUpper(rune(name[0]))
}

// applyFixes modifies the source to fix the given errors.
func applyFixes(source string, fixes []fix) string {
	for _, f := range fixes {
		switch f.kind {
		case "undefined_type":
			source = addTypeStub(source, f.name)
		case "undefined_var":
			source = addVarStub(source, f.name)
		case "unused_import":
			source = removeImport(source, f.name)
		case "redeclared":
			source = removeRedeclaration(source, f.name)
		case "missing_return":
			source = removeFuncAtLine(source, f.name)
		case "fix_return_arity":
			source = fixReturnArity(source, f.name)
		case "fix_assign_mismatch":
			source = fixAssignMismatchAtLine(source, f.name)
		case "decl_not_used":
			source = suppressUnusedVar(source, f.name)
		}
	}
	return source
}

// fixReturnArity fixes a func stub whose return arity doesn't match call sites.
// Removes the existing stub and re-adds it with correct arity detection.
func fixReturnArity(source, name string) string {
	// Remove the existing func stub declaration.
	lines := strings.Split(source, "\n")
	var cleaned []string
	for _, line := range lines {
		trimmed := strings.TrimSpace(line)
		if strings.HasPrefix(trimmed, "func "+name+"(args ...interface{})") {
			continue // remove old stub
		}
		cleaned = append(cleaned, line)
	}
	source = strings.Join(cleaned, "\n")

	// Re-add with correct arity.
	return addVarStub(source, name)
}

// fixAssignMismatchAtLine handles "assignment mismatch" errors by finding
// the function call on the error line and adjusting its stub.
func fixAssignMismatchAtLine(source, lineNumStr string) string {
	lineNum := 0
	for _, ch := range lineNumStr {
		if ch >= '0' && ch <= '9' {
			lineNum = lineNum*10 + int(ch-'0')
		}
	}
	if lineNum == 0 {
		return source
	}

	lines := strings.Split(source, "\n")
	if lineNum > len(lines) {
		return source
	}

	// Extract the line and find the function name being called.
	line := lines[lineNum-1]
	trimmed := strings.TrimSpace(line)

	// Look for patterns: "a, b := funcName(" or "a, b = funcName("
	assignOps := []string{":=", "="}
	for _, op := range assignOps {
		idx := strings.Index(trimmed, op)
		if idx < 0 {
			continue
		}
		rhs := strings.TrimSpace(trimmed[idx+len(op):])
		// rhs should start with funcName(
		parenIdx := strings.IndexByte(rhs, '(')
		if parenIdx <= 0 {
			continue
		}
		funcName := strings.TrimSpace(rhs[:parenIdx])
		// Check if this function has a stub we can fix.
		if strings.Contains(source, "func "+funcName+"(args ...interface{})") {
			return fixReturnArity(source, funcName)
		}
	}

	return source
}

// addTypeStub adds a type stub after the import block.
// Uses struct{} so methods can be defined on the type.
func addTypeStub(source, name string) string {
	stub := fmt.Sprintf("type %s struct{}\n", name)

	// Check if already declared (any declaration form).
	if nameIsDeclared(source, name) {
		return source
	}

	return insertAfterImports(source, stub)
}

// addVarStub adds a var or func stub after the import block.
// If the name is used as a function call in the source, it emits a
// func stub. When another function in the source calls this name and
// assigns the result, the stub uses types inferred from the assignment
// context (e.g., "s, err := name()" → (string, error) if the vars
// are used with string/error operations). Falls back to interface{}.
func addVarStub(source, name string) string {
	// Check if already declared (any declaration form).
	if nameIsDeclared(source, name) {
		return source
	}

	// If this name is called as a function, emit a func stub.
	// Use word-boundary matching: the char before name( must not be
	// an identifier char (to avoid matching "runGeneration(" for "gen(").
	if isCalledAsFunc(source, name) {
		retSig := inferReturnSignature(source, name)
		if retSig != "" {
			zr := zeroReturn(retSig)
			if zr == "" {
				zr = "return"
			}
			stub := fmt.Sprintf("func %s(args ...interface{}) %s { %s }\n", name, retSig, zr)
			return insertAfterImports(source, stub)
		}
		// Fallback: infer arity from assignment LHS count.
		n := inferReturnArity(source, name)
		if n <= 1 {
			stub := fmt.Sprintf("func %s(args ...interface{}) interface{} { return nil }\n", name)
			return insertAfterImports(source, stub)
		}
		retTypes := make([]string, n)
		retVals := make([]string, n)
		for i := range n {
			retTypes[i] = "interface{}"
			retVals[i] = "nil"
		}
		stub := fmt.Sprintf("func %s(args ...interface{}) (%s) { return %s }\n",
			name, strings.Join(retTypes, ", "), strings.Join(retVals, ", "))
		return insertAfterImports(source, stub)
	}

	stub := fmt.Sprintf("var %s interface{}\n", name)
	return insertAfterImports(source, stub)
}

// findWordBoundaryCall finds "name(" in text where name starts at a word
// boundary (not preceded by an identifier character). Returns the index of
// name in text, or -1 if not found.
func findWordBoundaryCall(text, name string) int {
	callPat := name + "("
	idx := 0
	for {
		pos := strings.Index(text[idx:], callPat)
		if pos < 0 {
			return -1
		}
		absPos := idx + pos
		if absPos > 0 {
			prev := text[absPos-1]
			if (prev >= 'a' && prev <= 'z') || (prev >= 'A' && prev <= 'Z') ||
				(prev >= '0' && prev <= '9') || prev == '_' {
				idx = absPos + len(callPat)
				continue
			}
		}
		return absPos
	}
}

// isCalledAsFunc checks if name appears as a function call "name(" in the
// source, using word-boundary matching to avoid false positives like matching
// "runGeneration(" when looking for "gen(".
func isCalledAsFunc(source, name string) bool {
	callPat := name + "("
	idx := 0
	for {
		pos := strings.Index(source[idx:], callPat)
		if pos < 0 {
			return false
		}
		absPos := idx + pos
		// Check character before the match — must be a non-identifier char.
		if absPos > 0 {
			prev := source[absPos-1]
			if (prev >= 'a' && prev <= 'z') || (prev >= 'A' && prev <= 'Z') ||
				(prev >= '0' && prev <= '9') || prev == '_' {
				idx = absPos + len(callPat)
				continue // false positive: part of a longer identifier
			}
		}
		return true
	}
}

// inferReturnSignature tries to determine the return type signature for a
// function call by examining assignment context. Looks for patterns like:
//
//	s, err := name(...)   → checks if "err" is used with error comparisons
//	bin, err := name(...) → (string, error) if bin is used in string context
//	ok := name(...)       → bool if "ok" is used as a bool
//
// Returns a Go return type string like "string", "error", "(string, error)"
// or "" if inference fails.
func inferReturnSignature(source, name string) string {
	// Scan for assignment call sites with word-boundary matching.
	for _, line := range strings.Split(source, "\n") {
		trimmed := strings.TrimSpace(line)
		callIdx := findWordBoundaryCall(trimmed, name)
		if callIdx < 0 {
			continue
		}
		prefix := strings.TrimSpace(trimmed[:callIdx])
		isAssign := false
		if strings.HasSuffix(prefix, ":=") {
			prefix = strings.TrimSuffix(prefix, ":=")
			isAssign = true
		} else if strings.HasSuffix(prefix, "=") && !strings.HasSuffix(prefix, "!=") && !strings.HasSuffix(prefix, "==") {
			prefix = strings.TrimSuffix(prefix, "=")
			isAssign = true
		}
		if !isAssign {
			continue
		}
		prefix = strings.TrimSpace(prefix)
		if prefix == "" {
			continue
		}
		parts := strings.Split(prefix, ",")
		var varNames []string
		for _, p := range parts {
			p = strings.TrimSpace(p)
			if p != "" {
				varNames = append(varNames, p)
			}
		}
		if len(varNames) == 0 {
			continue
		}
		// Infer types from variable names and usage context.
		types := make([]string, len(varNames))
		for i, v := range varNames {
			types[i] = inferVarType(source, v)
		}
		if len(types) == 1 {
			return types[0]
		}
		return "(" + strings.Join(types, ", ") + ")"
	}
	return ""
}

// inferVarType guesses the type of a variable from its name and how it's
// used in the source. Common patterns:
//
//	err       → error
//	ok        → bool
//	_         → interface{}
//	n, count  → int
//	s, str    → string
//	b, buf    → []byte
func inferVarType(source, name string) string {
	if name == "_" {
		return "interface{}"
	}
	if name == "err" {
		return "error"
	}
	if name == "ok" {
		return "bool"
	}
	// Check usage patterns in source.
	if strings.Contains(source, name+" != nil") || strings.Contains(source, name+" == nil") {
		// Nil-comparable — could be error, pointer, slice, map, interface.
		if strings.Contains(source, "return "+name) {
			// If returned alone or last, likely error.
			return "error"
		}
		return "interface{}"
	}
	if strings.Contains(source, "string("+name+")") || strings.Contains(source, name+` + "`) || strings.Contains(source, `" + `+name) {
		return "string"
	}
	if strings.Contains(source, "len("+name+")") {
		return "string" // could be []byte too, string is safer
	}
	// Default to interface{} — most permissive.
	return "interface{}"
}

// inferReturnArity scans source for assignment patterns that call name()
// and returns the number of LHS variables. Handles:
//   - a, b := name(...)         → 2
//   - a, b, c = name(...)       → 3
//   - a := name(...)            → 1
//   - name(...)  (no assignment) → 1
func inferReturnArity(source, name string) int {
	maxArity := 1
	for _, line := range strings.Split(source, "\n") {
		trimmed := strings.TrimSpace(line)

		// Find call sites with word-boundary matching.
		callIdx := findWordBoundaryCall(trimmed, name)
		if callIdx < 0 {
			continue
		}

		// Check if this is part of an assignment (not just a standalone call).
		// Look for := or = before the call.
		prefix := trimmed[:callIdx]
		prefix = strings.TrimSpace(prefix)

		// Strip trailing := or =
		isAssign := false
		if strings.HasSuffix(prefix, ":=") {
			prefix = strings.TrimSuffix(prefix, ":=")
			isAssign = true
		} else if strings.HasSuffix(prefix, "=") && !strings.HasSuffix(prefix, "!=") && !strings.HasSuffix(prefix, "==") {
			prefix = strings.TrimSuffix(prefix, "=")
			isAssign = true
		}

		if !isAssign {
			continue
		}

		// Count comma-separated LHS variables.
		prefix = strings.TrimSpace(prefix)
		if prefix == "" {
			continue
		}
		parts := strings.Split(prefix, ",")
		arity := 0
		for _, p := range parts {
			p = strings.TrimSpace(p)
			if p != "" && p != "_" || p == "_" {
				arity++
			}
		}
		if arity > maxArity {
			maxArity = arity
		}
	}
	return maxArity
}

// nameIsDeclared checks if a name is already declared in the source
// as a type, var, const, or function.
func nameIsDeclared(source, name string) bool {
	for _, line := range strings.Split(source, "\n") {
		trimmed := strings.TrimSpace(line)
		if strings.HasPrefix(trimmed, "type "+name+" ") ||
			strings.HasPrefix(trimmed, "type "+name+"=") ||
			strings.HasPrefix(trimmed, "var "+name+" ") ||
			strings.HasPrefix(trimmed, "const "+name+" ") ||
			strings.HasPrefix(trimmed, "func "+name+"(") ||
			strings.HasPrefix(trimmed, name+" :=") {
			return true
		}
	}
	return false
}

// insertAfterImports inserts text after the import block (or after package line).
func insertAfterImports(source, text string) string {
	lines := strings.Split(source, "\n")
	insertIdx := -1

	// Find end of import block.
	inImport := false
	for i, line := range lines {
		trimmed := strings.TrimSpace(line)
		if trimmed == "import (" {
			inImport = true
		}
		if inImport && trimmed == ")" {
			insertIdx = i + 1
			break
		}
		// Single-line import.
		if strings.HasPrefix(trimmed, "import ") && !strings.Contains(trimmed, "(") {
			insertIdx = i + 1
		}
	}

	// Fallback: after package line.
	if insertIdx < 0 {
		for i, line := range lines {
			if strings.HasPrefix(strings.TrimSpace(line), "package ") {
				insertIdx = i + 1
				break
			}
		}
	}
	if insertIdx < 0 {
		insertIdx = 0
	}

	// Insert.
	result := make([]string, 0, len(lines)+2)
	result = append(result, lines[:insertIdx]...)
	result = append(result, text)
	result = append(result, lines[insertIdx:]...)
	return strings.Join(result, "\n")
}

// removeImport removes an import path from the source.
func removeImport(source, importPath string) string {
	lines := strings.Split(source, "\n")
	result := make([]string, 0, len(lines))

	for _, line := range lines {
		trimmed := strings.TrimSpace(line)
		// Match: "path" or alias "path"
		if strings.Contains(trimmed, `"`+importPath+`"`) {
			// Skip single-line import or import block entry.
			if strings.HasPrefix(trimmed, "import ") || !strings.HasPrefix(trimmed, "import") {
				continue
			}
		}
		result = append(result, line)
	}

	return strings.Join(result, "\n")
}

// removeRedeclaration removes the SECOND occurrence of a declaration.
func removeRedeclaration(source, name string) string {
	lines := strings.Split(source, "\n")
	seen := false
	result := make([]string, 0, len(lines))

	patterns := []string{
		"type " + name + " ",
		"type " + name + "=",
		"var " + name + " ",
		"func " + name + "(",
	}

	for i := 0; i < len(lines); i++ {
		trimmed := strings.TrimSpace(lines[i])
		isDecl := false
		for _, pat := range patterns {
			if strings.HasPrefix(trimmed, pat) {
				isDecl = true
				break
			}
		}

		if isDecl {
			if seen {
				// Skip this declaration (and its body if it has braces).
				depth := strings.Count(lines[i], "{") - strings.Count(lines[i], "}")
				for depth > 0 && i+1 < len(lines) {
					i++
					depth += strings.Count(lines[i], "{") - strings.Count(lines[i], "}")
				}
				continue
			}
			seen = true
		}

		result = append(result, lines[i])
	}

	return strings.Join(result, "\n")
}

// removeFuncAtLine removes the function whose closing brace is at the
// given line number (as a string). Walks backward to find "func " and
// forward to find the matching "}".
func removeFuncAtLine(source, lineNumStr string) string {
	lineNum := 0
	for _, ch := range lineNumStr {
		if ch >= '0' && ch <= '9' {
			lineNum = lineNum*10 + int(ch-'0')
		}
	}
	if lineNum == 0 {
		return source
	}

	lines := strings.Split(source, "\n")
	if lineNum > len(lines) {
		return source
	}

	// Walk backward from the error line to find the function declaration.
	funcStart := -1
	for i := lineNum - 1; i >= 0; i-- {
		trimmed := strings.TrimSpace(lines[i])
		if strings.HasPrefix(trimmed, "func ") {
			funcStart = i
			break
		}
	}
	if funcStart < 0 {
		return source
	}

	// Walk forward from funcStart to find the matching closing brace.
	depth := 0
	funcEnd := -1
	for i := funcStart; i < len(lines); i++ {
		depth += strings.Count(lines[i], "{") - strings.Count(lines[i], "}")
		if depth <= 0 && strings.Count(lines[i], "{") > 0 || (depth == 0 && i > funcStart) {
			funcEnd = i
			break
		}
	}
	if funcEnd < 0 {
		funcEnd = len(lines) - 1
	}

	// Remove the function.
	result := make([]string, 0, len(lines)-(funcEnd-funcStart+1))
	result = append(result, lines[:funcStart]...)
	result = append(result, lines[funcEnd+1:]...)
	return strings.Join(result, "\n")
}

// removeErrorFunctions parses compiler output to find which line numbers
// have errors, then removes the enclosing functions for those lines.
func removeErrorFunctions(source, compileOutput string) string {
	// Extract error line numbers from compiler output.
	errorLines := make(map[int]bool)
	for _, line := range strings.Split(compileOutput, "\n") {
		if m := reErrorAtLine.FindStringSubmatch(line); m != nil {
			lineNum := 0
			for _, ch := range m[1] {
				lineNum = lineNum*10 + int(ch-'0')
			}
			errorLines[lineNum] = true
		}
	}

	if len(errorLines) == 0 {
		return source
	}

	// Find which functions contain errors and remove them.
	lines := strings.Split(source, "\n")
	type funcRange struct {
		start, end int
	}

	// Build a list of all function ranges.
	var funcs []funcRange
	for i := 0; i < len(lines); i++ {
		trimmed := strings.TrimSpace(lines[i])
		if !strings.HasPrefix(trimmed, "func ") {
			continue
		}
		start := i
		depth := 0
		foundBody := false
		for j := i; j < len(lines); j++ {
			depth += strings.Count(lines[j], "{") - strings.Count(lines[j], "}")
			if strings.Contains(lines[j], "{") {
				foundBody = true
			}
			if foundBody && depth <= 0 {
				funcs = append(funcs, funcRange{start, j})
				i = j // skip past this function
				break
			}
		}
	}

	// Find functions that contain error lines.
	removeFuncs := make(map[int]bool)
	for _, fr := range funcs {
		for lineNum := range errorLines {
			// Line numbers are 1-based; our slice is 0-based.
			if lineNum-1 >= fr.start && lineNum-1 <= fr.end {
				removeFuncs[fr.start] = true
				break
			}
		}
	}

	if len(removeFuncs) == 0 {
		return source
	}

	// Rebuild source, replacing error functions with signature-preserving stubs.
	// This keeps the function callable (correct name, params, return types)
	// while removing the broken body.
	result := make([]string, 0, len(lines))
	for i := 0; i < len(lines); i++ {
		// Check if this is a function we're replacing.
		skip := false
		for _, fr := range funcs {
			if fr.start == i && removeFuncs[fr.start] {
				trimmed := strings.TrimSpace(lines[i])
				if trimmed == "func main() {" || strings.HasPrefix(trimmed, "func main()") {
					// Replace main with empty body.
					result = append(result, "func main() {}")
					i = fr.end
					skip = true
					break
				}
				// Extract the full func signature line and generate a stub
				// with the correct return types.
				stub := stubFromSignature(trimmed)
				if stub != "" {
					result = append(result, stub)
				}
				i = fr.end // skip to end of function
				skip = true
				break
			}
		}
		if !skip {
			result = append(result, lines[i])
		}
	}

	return strings.Join(result, "\n")
}

// neutralizeErrorLines attempts to preserve variable declarations from error
// lines by replacing broken RHS expressions with zero values. For example:
//
//	x, err := brokenCall()  →  var x interface{}; var err error
//	s := undefined + "foo"  →  var s string
//
// This prevents cascading "undefined" errors from variables that were declared
// on the removed line. Returns unchanged source if no neutralization is possible.
func neutralizeErrorLines(source, compileOutput string) string {
	errorLines := make(map[int]bool)
	for _, line := range strings.Split(compileOutput, "\n") {
		if m := reErrorAtLine.FindStringSubmatch(line); m != nil {
			lineNum := 0
			for _, ch := range m[1] {
				lineNum = lineNum*10 + int(ch-'0')
			}
			if lineNum > 0 {
				errorLines[lineNum] = true
			}
		}
	}
	if len(errorLines) == 0 {
		return source
	}

	lines := strings.Split(source, "\n")
	changed := false
	for lineNum := range errorLines {
		idx := lineNum - 1
		if idx < 0 || idx >= len(lines) {
			continue
		}
		line := lines[idx]
		trimmed := strings.TrimSpace(line)

		// Only neutralize short-assign statements: "a, b := expr"
		assignIdx := strings.Index(trimmed, ":=")
		if assignIdx <= 0 {
			continue
		}

		// Don't neutralize structural lines.
		if strings.HasPrefix(trimmed, "func ") || strings.HasPrefix(trimmed, "if ") ||
			strings.HasPrefix(trimmed, "for ") || strings.HasPrefix(trimmed, "switch ") {
			continue
		}

		lhs := strings.TrimSpace(trimmed[:assignIdx])
		parts := strings.Split(lhs, ",")
		var decls []string
		for _, p := range parts {
			p = strings.TrimSpace(p)
			if p == "" || p == "_" {
				continue
			}
			typ := inferVarType(source, p)
			decls = append(decls, fmt.Sprintf("var %s %s", p, typ))
		}
		if len(decls) > 0 {
			// Preserve indentation.
			indent := line[:len(line)-len(strings.TrimLeft(line, " \t"))]
			lines[idx] = indent + strings.Join(decls, "; ")
			changed = true
		}
	}
	if !changed {
		return source
	}
	return strings.Join(lines, "\n")
}

// removeErrorLines removes specific lines that the compiler reports errors on,
// then cleans up orphaned syntax (empty var/const/type blocks, etc.).
func removeErrorLines(source, compileOutput string) string {
	// Extract error line numbers.
	errorLines := make(map[int]bool)
	for _, line := range strings.Split(compileOutput, "\n") {
		if m := reErrorAtLine.FindStringSubmatch(line); m != nil {
			lineNum := 0
			for _, ch := range m[1] {
				lineNum = lineNum*10 + int(ch-'0')
			}
			if lineNum > 0 {
				errorLines[lineNum] = true
			}
		}
	}

	if len(errorLines) == 0 {
		return source
	}

	lines := strings.Split(source, "\n")

	// First, find which statements the error lines belong to.
	// A "statement" spans from the line with a top-level expression to its
	// closing paren/brace. If an error is inside a multi-line call like
	// runColony(arg1,\n arg2,\n arg3), remove the entire call.
	removeStmts := findStatementsToRemove(lines, errorLines)

	result := make([]string, 0, len(lines))
	for i, line := range lines {
		lineNum := i + 1 // 1-based
		if removeStmts[lineNum] || errorLines[lineNum] {
			// Don't remove structural lines (func, closing braces, package, import).
			trimmed := strings.TrimSpace(line)
			if trimmed == "}" || trimmed == "{" ||
				strings.HasPrefix(trimmed, "func ") ||
				strings.HasPrefix(trimmed, "package ") ||
				strings.HasPrefix(trimmed, "import") {
				result = append(result, line)
				continue
			}
			// Skip this error line or its containing statement.
			continue
		}
		result = append(result, line)
	}

	source = strings.Join(result, "\n")

	// Clean up orphaned group blocks: var\n), const\n), type\n)
	// Also clean up stray "var" or "const" or "type" keywords on their own line.
	source = cleanOrphanedBlocks(source)

	return source
}

// cleanOrphanedBlocks removes broken group declarations like:
//
//	var
//	}
//
// or empty groups like:
//
//	var (
//	)
func cleanOrphanedBlocks(source string) string {
	lines := strings.Split(source, "\n")
	result := make([]string, 0, len(lines))

	for i := 0; i < len(lines); i++ {
		trimmed := strings.TrimSpace(lines[i])

		// Check for a lone keyword followed by } or ) on the next non-empty line.
		if trimmed == "var" || trimmed == "const" || trimmed == "type" {
			// Look ahead for closing brace or paren.
			j := i + 1
			for j < len(lines) && strings.TrimSpace(lines[j]) == "" {
				j++
			}
			if j < len(lines) {
				next := strings.TrimSpace(lines[j])
				if next == "}" || next == ")" {
					i = j // skip both
					continue
				}
			}
		}

		// Check for empty group: "var (" followed (after blanks) by ")".
		if (trimmed == "var (" || trimmed == "const (" || trimmed == "type (") {
			j := i + 1
			for j < len(lines) && strings.TrimSpace(lines[j]) == "" {
				j++
			}
			if j < len(lines) && strings.TrimSpace(lines[j]) == ")" {
				i = j // skip the empty group
				continue
			}
		}

		// Check for a "var // comment" pattern (broken var decl).
		if strings.HasPrefix(trimmed, "var //") || strings.HasPrefix(trimmed, "const //") || strings.HasPrefix(trimmed, "type //") {
			continue // remove broken declaration
		}

		result = append(result, lines[i])
	}

	return strings.Join(result, "\n")
}

// removeIncompleteDecls scans top-level declarations and removes those
// with unbalanced braces (e.g., "var X = func() int {" with stripped body).
func removeIncompleteDecls(source string) string {
	lines := strings.Split(source, "\n")
	result := make([]string, 0, len(lines))

	for i := 0; i < len(lines); i++ {
		trimmed := strings.TrimSpace(lines[i])

		// Detect top-level declaration lines with opening brace.
		isDecl := strings.HasPrefix(trimmed, "var ") ||
			strings.HasPrefix(trimmed, "const ") ||
			strings.HasPrefix(trimmed, "type ") ||
			strings.HasPrefix(trimmed, "func ")

		if isDecl && strings.Contains(trimmed, "{") {
			// Count braces for this declaration.
			depth := 0
			j := i
			for j < len(lines) {
				depth += strings.Count(lines[j], "{") - strings.Count(lines[j], "}")
				if depth <= 0 {
					break
				}
				j++
			}
			if j >= len(lines) && depth > 0 {
				// Unclosed declaration — remove from i to end of file.
				// But don't discard everything; stop at the next top-level decl.
				for k := i + 1; k < len(lines); k++ {
					next := strings.TrimSpace(lines[k])
					if next == "" {
						continue
					}
					nextIsDecl := strings.HasPrefix(next, "var ") ||
						strings.HasPrefix(next, "const ") ||
						strings.HasPrefix(next, "type ") ||
						strings.HasPrefix(next, "func ")
					if nextIsDecl {
						// Skip lines i through k-1 (the broken decl).
						i = k - 1
						break
					}
				}
				continue // skip this broken declaration line
			}
		}

		result = append(result, lines[i])
	}

	return strings.Join(result, "\n")
}

// findStatementsToRemove identifies multi-line statements that contain error
// lines. When an error occurs inside a multi-line function call like
// runColony(arg1,\n arg2,\n arg3), all lines of that call should be removed
// together rather than one at a time (which leaves broken syntax).
func findStatementsToRemove(lines []string, errorLines map[int]bool) map[int]bool {
	removeLines := make(map[int]bool)

	// For each error line, check if it's inside a multi-line statement
	// (unclosed parentheses from above). If so, find the statement boundaries
	// and mark all lines for removal.
	for errLine := range errorLines {
		idx := errLine - 1 // 0-based
		if idx < 0 || idx >= len(lines) {
			continue
		}

		// Walk backward to find where the unclosed paren starts.
		parenDepth := 0
		for i := idx; i >= 0; i-- {
			line := lines[i]
			for j := len(line) - 1; j >= 0; j-- {
				if line[j] == ')' {
					parenDepth++
				} else if line[j] == '(' {
					parenDepth--
				}
			}
			if parenDepth < 0 {
				// Found the start of the multi-line expression.
				// This line has an unmatched '(' — find the matching ')'.
				stmtStart := i + 1 // line AFTER the opening paren
				stmtEnd := idx

				// Walk forward from the error line to find the closing ')'.
				depth := 0
				for j := i; j < len(lines); j++ {
					for _, c := range lines[j] {
						if c == '(' {
							depth++
						} else if c == ')' {
							depth--
						}
					}
					if depth <= 0 {
						stmtEnd = j
						break
					}
				}

				// Mark ALL lines of this statement for removal.
				// Include the statement start line too (the call itself).
				for k := stmtStart; k <= stmtEnd; k++ {
					removeLines[k+1] = true // 1-based
				}
				// Also mark the line with the function call.
				removeLines[i+1] = true
				break
			}
		}
	}

	return removeLines
}

// suppressUnusedVar adds "_ = name" after the declaration line of an unused
// variable to suppress the "declared and not used" error without removing it.
func suppressUnusedVar(source, name string) string {
	lines := strings.Split(source, "\n")
	for i, line := range lines {
		trimmed := strings.TrimSpace(line)
		// Match ":=" assignments: "name := ..." or "name, other := ..."
		if strings.Contains(trimmed, name) && strings.Contains(trimmed, ":=") {
			// Check if this line declares the variable.
			lhs := trimmed
			if eqIdx := strings.Index(lhs, ":="); eqIdx > 0 {
				lhs = strings.TrimSpace(lhs[:eqIdx])
			}
			// Check all LHS parts.
			parts := strings.Split(lhs, ",")
			for _, p := range parts {
				if strings.TrimSpace(p) == name {
					// Insert "_ = name" after this line.
					suppression := "\t_ = " + name
					result := make([]string, 0, len(lines)+1)
					result = append(result, lines[:i+1]...)
					result = append(result, suppression)
					result = append(result, lines[i+1:]...)
					return strings.Join(result, "\n")
				}
			}
		}
		// Also check "var name ..." declarations.
		if strings.HasPrefix(trimmed, "var "+name+" ") || strings.HasPrefix(trimmed, "var "+name+"=") {
			suppression := "\t_ = " + name
			result := make([]string, 0, len(lines)+1)
			result = append(result, lines[:i+1]...)
			result = append(result, suppression)
			result = append(result, lines[i+1:]...)
			return strings.Join(result, "\n")
		}
	}
	return source
}

// stubFromSignature takes a function declaration line (e.g.,
// "func buildSelf() (string, error) {") and returns a stub with an
// empty body and zero-value return. This preserves the function's
// call signature so callers don't break.
func stubFromSignature(funcLine string) string {
	// Strip trailing "{" and whitespace.
	sig := strings.TrimSpace(strings.TrimSuffix(strings.TrimSpace(funcLine), "{"))
	if !strings.HasPrefix(sig, "func ") {
		return ""
	}

	// Extract the return signature.
	retSig := extractReturnSig(sig)
	zr := zeroReturn(retSig)

	if zr != "" {
		return sig + " { " + zr + " }"
	}
	return sig + " {}"
}

// truncate shortens a string to maxLen characters.
func truncate(s string, maxLen int) string {
	if len(s) <= maxLen {
		return s
	}
	return s[:maxLen] + "..."
}

// StubNames returns the names of all type, var, and func stubs added by CompileAndRepair.
// Detects: "type X struct{}", "var x interface{}", and "func x(args ...interface{}) ..."
func StubNames(source string) []string {
	var names []string
	for _, line := range strings.Split(source, "\n") {
		trimmed := strings.TrimSpace(line)
		if strings.HasPrefix(trimmed, "type ") && strings.HasSuffix(trimmed, "struct{}") {
			fields := strings.Fields(trimmed)
			if len(fields) >= 2 {
				names = append(names, fields[1])
			}
		}
		if strings.HasPrefix(trimmed, "var ") && strings.HasSuffix(trimmed, "interface{}") {
			fields := strings.Fields(trimmed)
			if len(fields) >= 2 {
				names = append(names, fields[1])
			}
		}
		// Detect func stubs: "func name(args ...interface{})"
		if strings.HasPrefix(trimmed, "func ") && strings.Contains(trimmed, "(args ...interface{})") {
			// Extract name from "func name(args ..."
			rest := strings.TrimPrefix(trimmed, "func ")
			if paren := strings.IndexByte(rest, '('); paren > 0 {
				name := rest[:paren]
				if name != "" && !strings.Contains(name, " ") {
					names = append(names, name)
				}
			}
		}
	}
	sort.Strings(names)
	return names
}