// Package emit reconstructs Go source code from a lattice's bonded AST elements.
//
// The emitter reads bonded elements and projects them back into source code.
// The lattice's structure determines what code is emitted — it writes itself out.
//
// Body-level elements (assign, return, if, for, expr, etc.) carry their
// parent function name and rendered source text separated by \x00. The
// emitter groups these by parent to reconstruct function bodies.
//
// When body elements exist, the emitter produces actual executable code.
// When only declarations and literals exist (legacy mode), it falls back
// to self-logging skeleton code for behavioral equivalence.
package emit

import (
	"crypto/rand"
	"encoding/binary"
	"fmt"
	"go/parser"
	"go/token"
	"io"
	"os"
	"path/filepath"
	"sort"
	"strings"
	"sync"

	"git.mleku.dev/mleku/dendrite/pkg/causal"
	"git.mleku.dev/mleku/dendrite/pkg/enzyme"
	"git.mleku.dev/mleku/dendrite/pkg/lattice"
)

// Fragment is a bonded element extracted from the lattice for emission.
type Fragment struct {
	NodeID lattice.NodeID
	Type   string
	Value  string
	LockIn lattice.LockInDepth
}

// bodyTags are element types that represent statements inside function bodies.
var bodyTags = map[string]bool{
	"assign": true, "return": true, "if": true, "for": true,
	"switch": true, "select": true, "go": true, "send": true,
	"expr": true, "defer": true, "decl": true, "branch": true,
	"case": true, "comm": true,
}

// Harvest collects all bonded elements from a lattice, grouped by file
// (if file markers exist) or as a single unnamed group.
func Harvest(l *lattice.Lattice) map[string][]Fragment {
	files := make(map[string][]Fragment)
	currentFile := ""

	nodes := l.Nodes()
	sort.Slice(nodes, func(i, j int) bool {
		return nodes[i].ID() < nodes[j].ID()
	})

	for _, n := range nodes {
		if !n.Occupied() {
			continue
		}
		e := n.Occupant()

		// Skip reference material — these elements influence lattice
		// topology during growth but should not appear in emitted output.
		if enzyme.IsRef(e) {
			continue
		}

		val := ""
		if e.Value() != nil {
			val = fmt.Sprintf("%v", e.Value())
		}

		if e.Type() == "file" {
			currentFile = val
			continue
		}

		// Grammar rules are structural context, not emittable code.
		if e.Type() == "grammar-rule" {
			continue
		}

		f := Fragment{
			NodeID: n.ID(),
			Type:   e.Type(),
			Value:  val,
			LockIn: n.LockIn(),
		}
		files[currentFile] = append(files[currentFile], f)
	}

	return files
}

// bodyStmt holds a parsed body-level statement with its parent function.
type bodyStmt struct {
	Parent  string
	Source  string
	LineNum int // original source line number (0 = unknown)
	NodeID  lattice.NodeID
	LockIn  lattice.LockInDepth
	Tag     string
}

// parseBodyValue splits "parent\x00source" from a body element's value.
// Returns empty parent if no separator found (legacy element).
// Used for fields, directives, and other two-part values.
func parseBodyValue(val string) (parent, source string) {
	idx := strings.IndexByte(val, '\x00')
	if idx < 0 {
		return "", val
	}
	return val[:idx], val[idx+1:]
}

// parseBodyStmt splits body element values in two supported formats:
//   - Three-part: "parent\x00linenum\x00source" (new: carries source position)
//   - Two-part:   "parent\x00source" (legacy: no position info)
//
// Returns empty parent if no separator found.
func parseBodyStmt(val string) (parent string, lineNum int, source string) {
	idx := strings.IndexByte(val, '\x00')
	if idx < 0 {
		return "", 0, val
	}
	rest := val[idx+1:]
	parent = val[:idx]

	// Check for three-part format: is there another \x00?
	idx2 := strings.IndexByte(rest, '\x00')
	if idx2 >= 0 {
		// Three-part: parent\x00linenum\x00source
		lineStr := rest[:idx2]
		source = rest[idx2+1:]
		ln := 0
		for _, ch := range lineStr {
			if ch >= '0' && ch <= '9' {
				ln = ln*10 + int(ch-'0')
			}
		}
		return parent, ln, source
	}

	// Two-part: parent\x00source (legacy)
	return parent, 0, rest
}

// funcDecl holds information about a function or method declaration.
type funcDecl struct {
	Name      string // bare name for matching body statements
	Signature string // full value from the enzyme (e.g., "main()" or "Foo.Hello() string")
	IsMethod  bool
	RecvVar   string // receiver variable name (e.g., "c", "e") — empty means use "x"
	NodeID    lattice.NodeID
}

// parseFuncName extracts the bare function name from a func/method value.
// "main()" → "main", "Foo.Hello() string" → "Hello"
func parseFuncName(val string, isMethod bool) string {
	name := val
	if isMethod {
		// "Foo.Hello() string" → "Hello() string"
		if dot := strings.IndexByte(name, '.'); dot >= 0 {
			name = name[dot+1:]
		}
	}
	// "main()" → "main"
	if paren := strings.IndexByte(name, '('); paren >= 0 {
		name = name[:paren]
	}
	return name
}

// EmitGo writes reconstructed Go source from harvested fragments.
// Uses "main" as the package name. See EmitGoPackage for arbitrary packages.
func EmitGo(fragments []Fragment, w io.Writer) error {
	return EmitGoPackage(fragments, w, "main")
}

// EmitGoPackage writes reconstructed Go source from harvested fragments
// using the specified package name.
//
// If body elements (assign, return, if, expr, etc.) are present, the emitter
// reconstructs actual function bodies from the rendered source text each
// element carries. Otherwise it falls back to self-logging skeleton code.
func EmitGoPackage(fragments []Fragment, w io.Writer, pkgName string) error {
	// Separate fragments by role.
	var (
		funcDecls   []funcDecl
		structDecls []Fragment // type+struct pairs
		ifaceDecls  []Fragment // type+interface pairs
		typeDecls   []Fragment // standalone types
		fields      []Fragment
		imports     []Fragment
		bodyStmts   []bodyStmt
		literals    []Fragment
		directives  []Fragment
		varDecls    []Fragment // top-level var/const declarations
	)

	// Track struct and interface names from their dedicated elements.
	structNames := make(map[string]bool)
	ifaceNames := make(map[string]bool)

	for _, f := range fragments {
		switch f.Type {
		case "func":
			name := parseFuncName(f.Value, false)
			funcDecls = append(funcDecls, funcDecl{
				Name: name, Signature: f.Value, IsMethod: false, NodeID: f.NodeID,
			})
		case "method":
			methodVal := f.Value
			recvVar := ""
			// Method value format: "recvVar\x00RecvType.MethodName(...) returns"
			if idx := strings.IndexByte(methodVal, '\x00'); idx >= 0 {
				recvVar = methodVal[:idx]
				methodVal = methodVal[idx+1:]
			}
			name := parseFuncName(methodVal, true)
			funcDecls = append(funcDecls, funcDecl{
				Name: name, Signature: methodVal, IsMethod: true,
				RecvVar: recvVar, NodeID: f.NodeID,
			})
		case "struct":
			structDecls = append(structDecls, f)
			structNames[f.Value] = true
		case "interface":
			ifaceDecls = append(ifaceDecls, f)
			ifaceNames[f.Value] = true
		case "type":
			typeDecls = append(typeDecls, f)
		case "field":
			fields = append(fields, f)
		case "import":
			imports = append(imports, f)
		case "directive":
			directives = append(directives, f)
		case "var":
			varDecls = append(varDecls, f)
		default:
			if strings.HasPrefix(f.Type, "literal:") {
				literals = append(literals, f)
			} else if strings.HasPrefix(f.Type, "ident:") {
				// Ident subtypes contribute to naming but aren't directly emitted.
			} else if bodyTags[f.Type] {
				parent, lineNum, source := parseBodyStmt(f.Value)
				if source != "" {
					bodyStmts = append(bodyStmts, bodyStmt{
						Parent:  parent,
						Source:  source,
						LineNum: lineNum,
						NodeID:  f.NodeID,
						LockIn:  f.LockIn,
						Tag:     f.Type,
					})
				}
			}
		}
	}

	// Group body statements by parent function, maintaining node ID order.
	sort.Slice(bodyStmts, func(i, j int) bool {
		return bodyStmts[i].NodeID < bodyStmts[j].NodeID
	})
	funcBodies := make(map[string][]bodyStmt)
	for _, bs := range bodyStmts {
		funcBodies[bs.Parent] = append(funcBodies[bs.Parent], bs)
	}

	// Determine if we have real body content or need legacy self-logging.
	hasBodyContent := len(bodyStmts) > 0

	// === Build the source ===
	var b strings.Builder
	// Always use the caller's package name. When ingesting multi-package
	// source (e.g. -self), fragments contain package elements from every
	// file parsed ("axiom", "enzyme", "ed25519", etc). Using those would
	// produce a non-main package that can't compile as a binary.
	fmt.Fprintf(&b, "package %s\n\n", pkgName)

	// Determine imports (with alias resolution for name collisions).
	// Build a set of known package names for reference validation.
	knownPkgs := make(map[string]bool)
	if hasBodyContent {
		neededImports := inferImports(bodyStmts, imports, funcDecls, varDecls, fields)
		neededImports = aliasCollisions(neededImports)
		for _, imp := range neededImports {
			knownPkgs[extractPkgName(imp)] = true
		}
		if len(neededImports) > 0 {
			b.WriteString("import (\n")
			for _, imp := range neededImports {
				fmt.Fprintf(&b, "\t%s\n", imp)
			}
			b.WriteString(")\n\n")
		}
	} else {
		needsFmt := hasOutputLiterals(literals)
		if needsFmt {
			b.WriteString("import \"fmt\"\n\n")
			knownPkgs["fmt"] = true
		}
	}

	// Top-level var/const declarations (with associated directives).
	// Validate each declaration: must parse, and must not reference
	// unknown packages (e.g., "hash.Len" when "hash" isn't imported).
	emittedVars := make(map[string]bool)
	emittedNames := make(map[string]bool) // track declared names for duplicate detection
	for _, v := range varDecls {
		if v.Value == "" || emittedVars[v.Value] {
			continue
		}
		if !isValidDecl(v.Value) {
			continue
		}
		name := extractVarName(v.Value)
		if emittedNames[name] {
			continue // duplicate declaration name
		}
		if hasUndefinedRefs(v.Value, knownPkgs) {
			continue
		}
		emittedVars[v.Value] = true
		emittedNames[name] = true
		emitDirectives(&b, directives, name)
		fmt.Fprintf(&b, "%s\n\n", v.Value)
	}

	// Type declarations.
	emittedTypes := make(map[string]bool)
	ifaceAliases := make(map[string]bool) // types emitted as "= interface{}"
	for _, t := range typeDecls {
		name := t.Value
		if name == "" || emittedTypes[name] || emittedNames[name] {
			continue
		}
		emittedTypes[name] = true

		// Emit any directives associated with this type.
		emitDirectives(&b, directives, name)

		if structNames[name] {
			fmt.Fprintf(&b, "type %s struct {\n", name)
			for _, f := range fields {
				parent, fieldDef := parseBodyValue(f.Value)
				if parent != name {
					continue // skip fields from other structs
				}
				if fieldDef == "" {
					continue
				}
				fmt.Fprintf(&b, "\t%s\n", fieldDef)
			}
			b.WriteString("}\n\n")
		} else if ifaceNames[name] {
			fmt.Fprintf(&b, "type %s interface {\n", name)
			// Only include methods whose receiver type matches this
			// interface name (or has no body — interface method stubs).
			seen := make(map[string]bool)
			for _, fd := range funcDecls {
				if !fd.IsMethod {
					continue
				}
				// Extract receiver type from "RecvType.MethodName(...) returns"
				sig := fd.Signature
				recv := ""
				if dot := strings.IndexByte(sig, '.'); dot >= 0 {
					recv = sig[:dot]
					sig = sig[dot+1:]
				}
				// Only include if receiver matches this interface, or
				// there's no body (could be an interface method stub).
				if recv != name && recv != "*"+name {
					continue
				}
				if _, hasBod := funcBodies[fd.Name]; hasBod {
					continue
				}
				if seen[fd.Name] {
					continue
				}
				seen[fd.Name] = true
				fmt.Fprintf(&b, "\t%s\n", sig)
			}
			b.WriteString("}\n\n")
		} else {
			// Type with unknown underlying definition — alias to interface{}.
			// Track these so we can skip methods with this receiver type
			// (interface aliases can't be method receivers).
			ifaceAliases[name] = true
			fmt.Fprintf(&b, "type %s = interface{}\n\n", name)
		}
	}

	// Functions and methods — with actual bodies when available.
	// Validate signatures before emitting to avoid garbled output.
	emittedFuncs := make(map[string]bool)
	for _, fd := range funcDecls {
		if emittedFuncs[fd.Name] {
			continue
		}
		// Skip non-method functions whose name clashes with a type or var/const.
		// Go doesn't allow a function and type/var with the same name in one package.
		if !fd.IsMethod && (emittedTypes[fd.Name] || emittedNames[fd.Name]) {
			continue
		}

		var funcLine string
		var returnSig string // portion after ')' for zero-value return generation

		if fd.IsMethod {
			sig := fd.Signature
			recv := ""
			rest := sig
			if dot := strings.IndexByte(sig, '.'); dot >= 0 {
				recv = sig[:dot]
				rest = sig[dot+1:]
			}
			if !strings.Contains(rest, "(") {
				continue
			}
			if recv != "" && !isValidReceiver(recv) {
				continue
			}
			// Skip methods on interface-aliased types (can't have receivers).
			bareRecv := strings.TrimPrefix(recv, "*")
			if ifaceAliases[bareRecv] {
				continue
			}
			// Skip if signature references undefined packages.
			if hasUndefinedRefs(rest, knownPkgs) {
				continue
			}
			if recv != "" {
				rv := fd.RecvVar
				if rv == "" {
					rv = "x"
				}
				funcLine = fmt.Sprintf("func (%s %s) %s {\n", rv, recv, rest)
			} else {
				funcLine = fmt.Sprintf("func %s {\n", rest)
			}
			returnSig = extractReturnSig(rest)
		} else {
			if !strings.Contains(fd.Signature, "(") {
				continue
			}
			if hasUndefinedRefs(fd.Signature, knownPkgs) {
				continue
			}
			funcLine = fmt.Sprintf("func %s {\n", fd.Signature)
			returnSig = extractReturnSig(fd.Signature)
		}

		emittedFuncs[fd.Name] = true
		emitDirectives(&b, directives, fd.Name)
		b.WriteString(funcLine)

		hasBody := false
		if stmts, ok := funcBodies[fd.Name]; ok && hasBodyContent {
			emitBody(&b, stmts)
			hasBody = len(stmts) > 0
		} else if !hasBodyContent && fd.Name == "main" {
			emitLegacyMain(&b, literals, typeDecls, funcDecls)
			hasBody = true
		}

		// Add zero-value return for functions with return types and no body.
		if !hasBody && returnSig != "" {
			zr := zeroReturn(returnSig)
			if zr != "" {
				fmt.Fprintf(&b, "\t%s\n", zr)
			}
		}
		b.WriteString("}\n\n")
	}

	// Ensure main exists (only for package main).
	if pkgName == "main" && !emittedFuncs["main"] {
		b.WriteString("func main() {\n")
		if stmts, ok := funcBodies["main"]; ok {
			emitBody(&b, stmts)
		} else if !hasBodyContent {
			emitLegacyMain(&b, literals, typeDecls, funcDecls)
		}
		b.WriteString("}\n\n")
	}

	// Post-process: try to parse with go/parser and format with go/format.
	// If parsing succeeds, emit the formatted version (fixes import ordering,
	// whitespace). If it fails, emit the raw version — the compiler will
	// report specific errors that feed back into fitness.
	source := b.String()
	if formatted, err := sanitizeGoSource(source); err == nil {
		source = formatted
	}

	_, err := io.WriteString(w, source)
	return err
}

// EmitProject reconstructs a multi-file Go module from harvested fragments.
// It emits source files and a go.mod with replace directives for all
// internal packages. Returns filepath → source code.
func EmitProject(files map[string][]Fragment, modPath string) map[string]string {
	result := make(map[string]string)
	internalPkgs := make(map[string]bool)

	for filename, frags := range files {
		// Determine package name from fragments.
		pkgName := "main"
		for _, f := range frags {
			if f.Type == "package" {
				pkgName = f.Value
				break
			}
		}

		var buf strings.Builder
		EmitGoPackage(frags, &buf, pkgName)

		// Use filename if provided, otherwise derive from package name.
		outName := filename
		if outName == "" {
			outName = pkgName + ".go"
		}
		result[outName] = buf.String()

		// Track subpackage directories for go.mod replace directives.
		if dir := filepath.Dir(outName); dir != "." && dir != "" {
			internalPkgs[dir] = true
		}
	}

	// Emit go.mod with replace directives.
	result["go.mod"] = emitGoMod(modPath, internalPkgs)
	return result
}

// emitGoMod generates a go.mod file with replace directives for all
// internal subpackages. All imports resolve locally — no network access.
func emitGoMod(modPath string, internalPkgs map[string]bool) string {
	var b strings.Builder
	fmt.Fprintf(&b, "module %s\n\ngo 1.24\n", modPath)

	if len(internalPkgs) > 0 {
		var pkgs []string
		for pkg := range internalPkgs {
			pkgs = append(pkgs, pkg)
		}
		sort.Strings(pkgs)

		b.WriteString("\nreplace (\n")
		for _, pkg := range pkgs {
			fmt.Fprintf(&b, "\t%s/%s => ./%s\n", modPath, pkg, pkg)
		}
		b.WriteString(")\n")
	}
	return b.String()
}

// emitBody writes the body statements for a function.
// It deduplicates, filters compound sub-statements, orders by
// dependency DAG, and shuffles non-dependent statements for
// anti-fingerprinting.
func emitBody(b *strings.Builder, stmts []bodyStmt) {
	// Deduplicate.
	var unique []bodyStmt
	seen := make(map[string]bool)
	for _, s := range stmts {
		if seen[s.Source] {
			continue
		}
		seen[s.Source] = true
		unique = append(unique, s)
	}

	// Remove inner statements that are contained within compound statements.
	// Use line numbers when available: a statement is "inner" if its line
	// number falls within the line range of a compound (multi-line) statement.
	// Falls back to substring containment when line numbers are absent.
	// Also filters case/comm clauses that belong inside a switch/select.
	compoundTags := map[string]bool{"if": true, "for": true, "switch": true, "select": true}
	innerTags := map[string]bool{"case": true, "comm": true}
	var filtered []bodyStmt
	for i, s := range unique {
		contained := false

		// Case/comm clauses are always inner to switch/select — filter them
		// if any enclosing compound statement exists.
		if innerTags[s.Tag] {
			for j, other := range unique {
				if i == j {
					continue
				}
				if (s.Tag == "case" && (other.Tag == "switch")) ||
					(s.Tag == "comm" && other.Tag == "select") {
					// Check containment by line range or substring.
					if s.LineNum > 0 && other.LineNum > 0 {
						otherLines := strings.Count(other.Source, "\n") + 1
						if s.LineNum > other.LineNum && s.LineNum <= other.LineNum+otherLines-1 {
							contained = true
							break
						}
					}
					if strings.Contains(other.Source, s.Source) {
						contained = true
						break
					}
				}
			}
		}

		// Line-number-based filtering for other inner statements.
		if !contained && s.LineNum > 0 {
			for j, other := range unique {
				if i == j || !compoundTags[other.Tag] || other.LineNum == 0 {
					continue
				}
				otherLines := strings.Count(other.Source, "\n") + 1
				otherEnd := other.LineNum + otherLines - 1
				if s.LineNum > other.LineNum && s.LineNum <= otherEnd {
					contained = true
					break
				}
			}
		}

		// Fallback: substring containment (original heuristic).
		if !contained {
			for j, other := range unique {
				if i != j && len(other.Source) > len(s.Source) && strings.Contains(other.Source, s.Source) {
					contained = true
					break
				}
			}
		}

		if !contained {
			filtered = append(filtered, s)
		}
	}

	// Strip orphan break/continue/fallthrough — these are branch statements
	// from loop/switch bodies that got placed at the top level of a function.
	// They're only valid inside for/switch/select.
	hasLoop := false
	for _, s := range filtered {
		if s.Tag == "for" || s.Tag == "switch" || s.Tag == "select" {
			hasLoop = true
			break
		}
		if strings.HasPrefix(strings.TrimSpace(s.Source), "for ") ||
			strings.HasPrefix(strings.TrimSpace(s.Source), "switch ") ||
			strings.HasPrefix(strings.TrimSpace(s.Source), "select {") {
			hasLoop = true
			break
		}
	}
	if !hasLoop {
		var clean []bodyStmt
		for _, s := range filtered {
			trimmed := strings.TrimSpace(s.Source)
			if trimmed == "break" || trimmed == "continue" || trimmed == "fallthrough" ||
				strings.HasPrefix(trimmed, "break ") || strings.HasPrefix(trimmed, "continue ") ||
				s.Tag == "branch" {
				continue
			}
			clean = append(clean, s)
		}
		filtered = clean
	}

	// Remove duplicate declarations within the function scope.
	// Go allows re-declaration with := when at least one variable is new,
	// but duplicate definitions are a common source of compile errors in
	// emitted code. Keep the first definition, drop subsequent ones.
	{
		declaredIds := make(map[string]bool)
		var deduped []bodyStmt
		for _, s := range filtered {
			defs := causal.ExtractDefines(s.Source)
			if len(defs) == 0 {
				deduped = append(deduped, s)
				continue
			}
			conflict := false
			for id := range defs {
				if declaredIds[id] {
					conflict = true
					break
				}
			}
			if conflict {
				continue // drop duplicate declaration
			}
			for id := range defs {
				declaredIds[id] = true
			}
			deduped = append(deduped, s)
		}
		filtered = deduped
	}

	// Order by dependency DAG with non-dependent shuffling.
	ordered := orderBody(filtered)

	for _, s := range ordered {
		sublines := strings.Split(s.Source, "\n")
		for _, sl := range sublines {
			fmt.Fprintf(b, "\t%s\n", sl)
		}
	}
}

// orderBody builds a dependency DAG from define-use analysis,
// topologically sorts into levels, and crypto/rand shuffles
// within each level for anti-fingerprinting.
func orderBody(stmts []bodyStmt) []bodyStmt {
	if len(stmts) <= 1 {
		return stmts
	}

	n := len(stmts)

	// Extract defines and uses for each statement.
	defines := make([]map[string]bool, n)
	uses := make([]map[string]bool, n)
	for i, s := range stmts {
		defines[i] = causal.ExtractDefines(s.Source)
		uses[i] = causal.ExtractUses(s.Source)
	}

	// Build dependency graph: deps[j] contains indices that j depends on.
	deps := make([][]int, n)
	for j := range n {
		for i := range n {
			if i == j {
				continue
			}
			// j depends on i if j uses something i defines.
			for id := range uses[j] {
				if defines[i][id] {
					deps[j] = append(deps[j], i)
					break
				}
			}
		}
	}

	// Topological sort into levels.
	levels := topoLevels(n, deps)

	// Sort within each level by source line number when available,
	// falling back to crypto/rand shuffle when line numbers are absent.
	for _, level := range levels {
		hasLineNums := false
		for _, idx := range level {
			if stmts[idx].LineNum > 0 {
				hasLineNums = true
				break
			}
		}
		if hasLineNums {
			sort.Slice(level, func(a, b int) bool {
				la, lb := stmts[level[a]].LineNum, stmts[level[b]].LineNum
				if la != lb {
					return la < lb
				}
				return level[a] < level[b]
			})
		} else {
			cryptoShuffle(level)
		}
	}

	// Flatten levels into ordered result.
	result := make([]bodyStmt, 0, n)
	for _, level := range levels {
		for _, idx := range level {
			result = append(result, stmts[idx])
		}
	}
	return result
}

// topoLevels performs topological sort and groups nodes into levels.
// Level 0 has no dependencies, level 1 depends only on level 0, etc.
func topoLevels(n int, deps [][]int) [][]int {
	// Compute in-degree.
	inDeg := make([]int, n)
	for j := range n {
		inDeg[j] = len(deps[j])
	}

	// Collect level 0 (no dependencies).
	var levels [][]int
	remaining := make([]bool, n)
	for i := range n {
		remaining[i] = true
	}

	for {
		var level []int
		for i := range n {
			if !remaining[i] {
				continue
			}
			// Check if all dependencies are already placed.
			allResolved := true
			for _, dep := range deps[i] {
				if remaining[dep] {
					allResolved = false
					break
				}
			}
			if allResolved {
				level = append(level, i)
			}
		}
		if len(level) == 0 {
			// Remaining nodes have circular dependencies.
			// Add them in original order.
			for i := range n {
				if remaining[i] {
					level = append(level, i)
				}
			}
			levels = append(levels, level)
			break
		}
		levels = append(levels, level)
		for _, idx := range level {
			remaining[idx] = false
		}
	}
	return levels
}

// cryptoShuffle performs Fisher-Yates shuffle using crypto/rand.
func cryptoShuffle(indices []int) {
	for i := len(indices) - 1; i > 0; i-- {
		var buf [8]byte
		rand.Read(buf[:])
		j := int(binary.LittleEndian.Uint64(buf[:]) % uint64(i+1))
		indices[i], indices[j] = indices[j], indices[i]
	}
}

// modulePath is the self-import prefix. Imports matching this are trusted.
const modulePath = "git.mleku.dev/mleku/dendrite"

// moduleRootOnce lazily finds the module root directory on disk.
var (
	moduleRootOnce sync.Once
	moduleRootDir  string          // absolute path to the dendrite module root, or ""
	subPkgCache    map[string]bool // cache: directory name → exists as sub-package
)

// findModuleRootDir locates the dendrite module root by walking up from cwd
// looking for a go.mod that declares the module.
func findModuleRootDir() string {
	dir, err := os.Getwd()
	if err != nil {
		return ""
	}
	for {
		modFile := filepath.Join(dir, "go.mod")
		data, err := os.ReadFile(modFile)
		if err == nil && strings.Contains(string(data), "module "+modulePath) {
			return dir
		}
		parent := filepath.Dir(dir)
		if parent == dir {
			return "" // reached filesystem root
		}
		dir = parent
	}
}

// isValidSubPkg checks if w is an actual sub-package directory in the dendrite
// module that contains Go source files (not just sub-directories).
// Caches results after the first filesystem scan.
func isValidSubPkg(w string) bool {
	moduleRootOnce.Do(func() {
		moduleRootDir = findModuleRootDir()
		subPkgCache = make(map[string]bool)
		if moduleRootDir != "" {
			entries, err := os.ReadDir(filepath.Join(moduleRootDir, "pkg"))
			if err == nil {
				for _, e := range entries {
					if !e.IsDir() || strings.HasPrefix(e.Name(), ".") || strings.HasPrefix(e.Name(), "_") {
						continue
					}
					// Only count as a package if it contains .go files.
					subEntries, err := os.ReadDir(filepath.Join(moduleRootDir, "pkg", e.Name()))
					if err != nil {
						continue
					}
					for _, se := range subEntries {
						if !se.IsDir() && strings.HasSuffix(se.Name(), ".go") {
							subPkgCache[e.Name()] = true
							break
						}
					}
				}
			}
		}
	})
	return subPkgCache[w]
}

// inferImports determines which imports are needed based on body statements,
// function signatures, var declarations, and field types. Three passes:
//  1. Collect candidates (explicit from lattice + inferred from all text)
//  2. Trust filter: reject anything that isn't stdlib or a self-import
//  3. Usage prune: drop imports whose package name doesn't appear in text
func inferImports(stmts []bodyStmt, importFrags []Fragment, funcs []funcDecl, vars []Fragment, fields []Fragment) []string {
	// Pass 1: Collect candidates.
	candidates := make(map[string]bool)

	// Explicit imports from the lattice.
	for _, imp := range importFrags {
		candidates[imp.Value] = true
	}

	// Collect ALL text that may reference packages: body statements,
	// function signatures, var declarations, and struct field types.
	allText := &strings.Builder{}
	for _, s := range stmts {
		allText.WriteString(s.Source)
		allText.WriteByte('\n')
	}
	for _, fd := range funcs {
		allText.WriteString(fd.Signature)
		allText.WriteByte('\n')
	}
	for _, v := range vars {
		allText.WriteString(v.Value)
		allText.WriteByte('\n')
	}
	for _, f := range fields {
		allText.WriteString(f.Value)
		allText.WriteByte('\n')
	}
	text := allText.String()

	// Infer stdlib imports from body text.
	// Only single-segment packages belong here. Multi-segment paths
	// like "path/filepath" and "os/exec" are in multiPkgs below.
	stdPkgs := map[string]string{
		"fmt": "fmt.", "os": "os.", "io": "io.",
		"strings": "strings.", "strconv": "strconv.",
		"log": "log.", "time": "time.", "sync": "sync.",
		"context": "context.", "math": "math.", "sort": "sort.",
		"bytes": "bytes.", "errors": "errors.", "path": "path.",
		"net": "net.", "regexp": "regexp.", "reflect": "reflect.",
		"testing": "testing.", "embed": "embed.",
		"bufio": "bufio.", "unicode": "unicode.",
		"crypto": "crypto.", "hash": "hash.", "encoding": "encoding.",
		"flag": "flag.", "slices": "slices.", "maps": "maps.",
		"unsafe": "unsafe.",
	}
	// Multi-segment stdlib paths — package name differs from import path.
	multiPkgs := map[string]string{
		"net/http":        "http.",
		"encoding/json":   "json.",
		"encoding/hex":    "hex.",
		"encoding/binary": "binary.",
		"os/exec":         "exec.",
		"os/signal":       "signal.",
		"path/filepath":   "filepath.",
		"crypto/rand":     "rand.",
		"crypto/sha256":   "sha256.",
		"crypto/ed25519":  "ed25519.",
		"go/ast":          "ast.",
		"go/parser":       "parser.",
		"go/token":        "token.",
		"go/format":       "format.",
		"go/printer":      "printer.",
		"math/big":        "big.",
		"math/rand":       "rand.",
		"io/fs":           "fs.",
	}

	for pkg, marker := range stdPkgs {
		if strings.Contains(text, marker) {
			candidates[`"`+pkg+`"`] = true
		}
	}
	for path, marker := range multiPkgs {
		if strings.Contains(text, marker) {
			candidates[`"`+path+`"`] = true
		}
	}

	// Infer self-module sub-package imports from body text.
	// If the body uses "ratio.New(...)" and "ratio" is an actual sub-package
	// directory, infer "git.mleku.dev/mleku/dendrite/pkg/ratio" as a candidate.
	// Validates against the filesystem to avoid false positives from struct
	// field access (e.g., "score.Behav.Float64()" where Behav is a field).
	knownStdPkg := make(map[string]bool)
	for pkg := range stdPkgs {
		knownStdPkg[pkg] = true
	}
	for _, marker := range multiPkgs {
		knownStdPkg[strings.TrimSuffix(marker, ".")] = true
	}
	words := identifiersInText(text)
	for _, w := range words {
		if knownStdPkg[w] {
			continue
		}
		if !strings.Contains(text, w+".") {
			continue
		}
		// Must be an actual sub-package directory on disk.
		if !isValidSubPkg(w) {
			continue
		}
		candidate := `"` + modulePath + "/pkg/" + w + `"`
		if !candidates[candidate] {
			candidates[candidate] = true
		}
	}

	// Pass 2+3: Trust filter + usage prune.
	var result []string
	for imp := range candidates {
		bare := strings.Trim(imp, `"`)

		// Trust check: only stdlib and self-imports allowed.
		if !isStdlib(bare) && !strings.HasPrefix(bare, modulePath) {
			continue // untrusted external — drop
		}

		// Usage check: package name must appear as "pkg." in body text.
		pkgName := extractPkgName(imp)
		if strings.Contains(text, pkgName+".") {
			result = append(result, imp)
		}
	}
	sort.Strings(result)
	return result
}

// isStdlib returns true if the import path is a Go standard library package.
// stdlib paths have no dots in the first path segment.
func isStdlib(importPath string) bool {
	first := importPath
	if idx := strings.IndexByte(importPath, '/'); idx >= 0 {
		first = importPath[:idx]
	}
	return !strings.Contains(first, ".")
}

// extractPkgName returns the package name from an import path.
// "fmt" → "fmt", "net/http" → "http", "git.mleku.dev/mleku/dendrite/pkg/axiom" → "axiom"
func extractPkgName(importPath string) string {
	bare := strings.Trim(importPath, `"`)
	if idx := strings.LastIndex(bare, "/"); idx >= 0 {
		return bare[idx+1:]
	}
	return bare
}

// hasOutputLiterals checks if any literal fragments look like program output.
func hasOutputLiterals(literals []Fragment) bool {
	for _, lit := range literals {
		cat, _ := classifyLiteral(lit.Value)
		if cat != "" {
			return true
		}
	}
	return false
}

// emitLegacyMain writes the self-logging main body (legacy mode).
// Used when no body elements are available.
func emitLegacyMain(b *strings.Builder, literals []Fragment, types []Fragment, funcs []funcDecl) {
	typeNames := dedupFragNames(types)
	funcNames := make([]string, 0)
	methodNames := make([]string, 0)
	for _, fd := range funcs {
		if fd.IsMethod {
			methodNames = append(methodNames, fd.Name)
		} else {
			funcNames = append(funcNames, fd.Name)
		}
	}

	type scoredLiteral struct {
		bare   string
		cat    string
		lockIn lattice.LockInDepth
	}
	var scored []scoredLiteral
	seen := make(map[string]bool)
	for _, f := range literals {
		cat, bare := classifyLiteral(f.Value)
		if cat == "" || seen[bare] {
			continue
		}
		seen[bare] = true
		scored = append(scored, scoredLiteral{bare, cat, f.LockIn})
	}
	sort.Slice(scored, func(i, j int) bool {
		return scored[j].lockIn.Less(scored[i].lockIn)
	})

	if len(scored) == 0 {
		return
	}

	b.WriteString("\t// Self-knowledge: the offspring counts its own structure.\n")
	fmt.Fprintf(b, "\tnTypes := %d\n", len(typeNames))
	fmt.Fprintf(b, "\tnFuncs := %d\n", len(funcNames))
	fmt.Fprintf(b, "\tnMethods := %d\n", len(methodNames))
	b.WriteString("\tnTotal := nTypes + nFuncs + nMethods\n")
	b.WriteString("\t_ = nTotal\n")

	for _, lit := range scored {
		switch lit.cat {
		case "format":
			verbs := parseFormatVerbs(lit.bare)
			if len(verbs) == 0 {
				stripped := stripEscapes(lit.bare)
				if stripped != "" {
					fmt.Fprintf(b, "\tfmt.Println(%q)\n", stripped)
				}
				continue
			}
			args := make([]string, len(verbs))
			intArgIdx := 0
			intArgs := []string{"nTypes", "nFuncs", "nMethods", "nTotal"}
			for i, v := range verbs {
				switch v.verb {
				case 'd':
					args[i] = intArgs[intArgIdx%len(intArgs)]
					intArgIdx++
				case 's':
					if i < len(typeNames) {
						args[i] = fmt.Sprintf("%q", typeNames[i])
					} else if i < len(funcNames) {
						args[i] = fmt.Sprintf("%q", funcNames[i%len(funcNames)])
					} else {
						args[i] = `""`
					}
				case 'f':
					args[i] = "float64(nTypes) / float64(nTotal+1) * 100"
				case 'v':
					args[i] = "nTotal"
				case 'p':
					args[i] = "0"
				default:
					args[i] = "nTotal"
				}
			}
			fmtStr := cleanFormatString(lit.bare)
			fmt.Fprintf(b, "\tfmt.Printf(\"%s\", %s)\n", fmtStr, strings.Join(args, ", "))
		case "output":
			stripped := stripEscapes(lit.bare)
			if stripped != "" {
				fmt.Fprintf(b, "\tfmt.Println(%q)\n", stripped)
			}
		}
	}
}

// isOutputLiteral returns true if a string literal looks like program output.
func isOutputLiteral(bare string) bool {
	if bare == "" {
		return false
	}
	if strings.Contains(bare, "/") {
		return false
	}
	if len(bare) <= 2 {
		return false
	}
	if !strings.ContainsAny(bare, " :=,.(){}[]!?%") && len(bare) < 20 {
		return false
	}
	if strings.HasPrefix(bare, "json:") || strings.HasPrefix(bare, "yaml:") {
		return false
	}
	return true
}

// classifyLiteral categorizes a string literal for emission.
func classifyLiteral(raw string) (category string, bare string) {
	if !strings.HasPrefix(raw, `"`) && !strings.HasPrefix(raw, "`") {
		return "", ""
	}
	bare = strings.Trim(raw, `"`+"`")
	if !isOutputLiteral(bare) {
		return "", ""
	}
	if strings.Contains(bare, "%") {
		return "format", bare
	}
	return "output", bare
}

// verbInfo describes a format verb found in a format string.
type verbInfo struct {
	verb byte
}

// parseFormatVerbs extracts format verbs from a printf-style format string.
func parseFormatVerbs(s string) []verbInfo {
	var verbs []verbInfo
	i := 0
	for i < len(s) {
		if s[i] == '%' && i+1 < len(s) {
			i++
			if s[i] == '%' {
				i++
				continue
			}
			for i < len(s) && strings.ContainsRune("-+# 0", rune(s[i])) {
				i++
			}
			for i < len(s) && s[i] >= '0' && s[i] <= '9' {
				i++
			}
			if i < len(s) && s[i] == '.' {
				i++
				for i < len(s) && s[i] >= '0' && s[i] <= '9' {
					i++
				}
			}
			if i < len(s) {
				verbs = append(verbs, verbInfo{verb: s[i]})
				i++
			}
		} else {
			i++
		}
	}
	return verbs
}

// cleanFormatString converts a raw format string into a safe Go string
// literal body for embedding inside double quotes.
func cleanFormatString(bare string) string {
	var b strings.Builder
	b.Grow(len(bare))
	for i := 0; i < len(bare); i++ {
		switch bare[i] {
		case '\\':
			b.WriteString(`\\`)
		case '\n':
			b.WriteString(`\n`)
		case '\r':
			b.WriteString(`\r`)
		case '\t':
			b.WriteString(`\t`)
		case '"':
			b.WriteString(`\"`)
		default:
			b.WriteByte(bare[i])
		}
	}
	return b.String()
}

// stripEscapes removes Go source escape sequences from a string.
func stripEscapes(s string) string {
	s = strings.ReplaceAll(s, `\n`, "")
	s = strings.ReplaceAll(s, `\t`, "")
	s = strings.ReplaceAll(s, `\r`, "")
	for strings.Contains(s, "  ") {
		s = strings.ReplaceAll(s, "  ", " ")
	}
	return strings.TrimSpace(s)
}

// dedupFragNames returns unique non-empty names from fragments.
func dedupFragNames(frags []Fragment) []string {
	seen := make(map[string]bool)
	var names []string
	for _, f := range frags {
		if f.Value != "" && !seen[f.Value] {
			seen[f.Value] = true
			names = append(names, f.Value)
		}
	}
	return names
}

// dedupNames returns unique non-empty names from a slice of fragments.
// Kept for backward compatibility.
func dedupNames(frags []Fragment) []string {
	return dedupFragNames(frags)
}

// extractVarName pulls the variable name from a rendered var declaration.
// "var ownSources embed.FS" → "ownSources"
// Handles leading comments: "var // comment\nGitRef string" → "GitRef"
func extractVarName(decl string) string {
	// Process each line looking for the var/const name.
	for _, line := range strings.Split(decl, "\n") {
		line = strings.TrimSpace(line)
		// Skip comment lines.
		if strings.HasPrefix(line, "//") || line == "" {
			continue
		}
		// Strip "var " or "const " prefix.
		for _, prefix := range []string{"var ", "const "} {
			if strings.HasPrefix(line, prefix) {
				line = line[len(prefix):]
				break
			}
		}
		// First identifier token is the name (strip trailing comma for multi-var).
		if idx := strings.IndexAny(line, " =,"); idx > 0 {
			return line[:idx]
		}
		if line != "" {
			return line
		}
	}
	return ""
}

// emitDirectives writes any //go: directives associated with the named
// declaration. Each directive's value is "assocName\x00//go:embed ..."
func emitDirectives(b *strings.Builder, directives []Fragment, name string) {
	for _, d := range directives {
		parent, text := parseBodyValue(d.Value)
		if parent == name && text != "" {
			fmt.Fprintf(b, "%s\n", text)
		}
	}
}

// isValidDecl checks whether a var/const declaration string is syntactically
// valid Go. It wraps the declaration in a minimal package and tries to parse it.
func isValidDecl(decl string) bool {
	// Quick reject: const/var without a value (iota members outside block).
	trimmed := strings.TrimSpace(decl)
	if strings.HasPrefix(trimmed, "const ") {
		// "const Foo" with no = and no type — iota member, invalid standalone.
		rest := strings.TrimPrefix(trimmed, "const ")
		// Strip trailing comment.
		if ci := strings.Index(rest, "//"); ci >= 0 {
			rest = strings.TrimSpace(rest[:ci])
		}
		// Must have '=' for a standalone const, or be a typed const like "const X Type = val".
		if !strings.Contains(rest, "=") {
			return false
		}
	}
	// Try to parse the declaration.
	src := "package p\n" + decl + "\n"
	fset := token.NewFileSet()
	_, err := parser.ParseFile(fset, "", src, parser.SkipObjectResolution)
	return err == nil
}

// extractReturnSig extracts the return type portion from a function signature.
// "main()" → "", "Foo() string" → "string", "Bar() (int, error)" → "(int, error)"
func extractReturnSig(sig string) string {
	// Find the FIRST depth-0 closing ')' — this closes the parameter list.
	// Everything after it is the return signature.
	depth := 0
	paramClose := -1
	for i, c := range sig {
		if c == '(' {
			depth++
		} else if c == ')' {
			depth--
			if depth == 0 {
				paramClose = i
				break
			}
		}
	}
	if paramClose < 0 || paramClose >= len(sig)-1 {
		return ""
	}
	ret := strings.TrimSpace(sig[paramClose+1:])
	if ret == "" || ret == "{" {
		return ""
	}
	return ret
}

// zeroReturn generates a return statement with zero values for the given
// return type signature. Handles single types, named returns, and tuples.
func zeroReturn(retSig string) string {
	retSig = strings.TrimSpace(retSig)
	if retSig == "" {
		return ""
	}

	// Named returns: "(x int, err error)" — just "return" suffices.
	if strings.HasPrefix(retSig, "(") {
		inner := strings.Trim(retSig, "()")
		parts := strings.Split(inner, ",")
		// Check if these are named (have both name and type).
		for _, p := range parts {
			fields := strings.Fields(strings.TrimSpace(p))
			if len(fields) >= 2 {
				return "return" // named returns — zero-initialized
			}
		}
		// Unnamed tuple: generate zero values for each.
		var zeros []string
		for _, p := range parts {
			zeros = append(zeros, zeroValue(strings.TrimSpace(p)))
		}
		return "return " + strings.Join(zeros, ", ")
	}

	// Single return type.
	return "return " + zeroValue(retSig)
}

// zeroValue returns the zero-value literal for a Go type.
func zeroValue(typ string) string {
	typ = strings.TrimSpace(typ)
	switch {
	case typ == "string":
		return `""`
	case typ == "bool":
		return "false"
	case typ == "error":
		return "nil"
	case typ == "int" || typ == "int8" || typ == "int16" || typ == "int32" || typ == "int64" ||
		typ == "uint" || typ == "uint8" || typ == "uint16" || typ == "uint32" || typ == "uint64" ||
		typ == "float32" || typ == "float64" || typ == "byte" || typ == "rune":
		return "0"
	case strings.HasPrefix(typ, "*") || strings.HasPrefix(typ, "[]") ||
		strings.HasPrefix(typ, "map[") || strings.HasPrefix(typ, "chan ") ||
		strings.HasPrefix(typ, "func(") || strings.HasPrefix(typ, "<-chan"):
		return "nil"
	case strings.Contains(typ, "."):
		// Package-qualified type — assume it's a struct or interface.
		return typ + "{}"
	default:
		// Unknown type — use zero value by name.
		return typ + "{}"
	}
}

// identifiersInText extracts unique identifiers from Go source text
// that appear before a dot (potential package references like "ratio.New").
func identifiersInText(text string) []string {
	seen := make(map[string]bool)
	var result []string
	for i := 0; i < len(text)-1; i++ {
		if text[i] == '.' && i > 0 {
			// Walk back to find the identifier.
			j := i - 1
			for j >= 0 && isIdentChar(rune(text[j])) {
				j--
			}
			word := text[j+1 : i]
			if len(word) >= 2 && !seen[word] {
				// Skip single-char receiver variables and common Go keywords.
				seen[word] = true
				result = append(result, word)
			}
		}
	}
	return result
}

// hasUndefinedRefs checks whether a var/const declaration references
// package-qualified identifiers (like "hash.Len" or "_l.Get(...)") where
// the package isn't in the known imports. Returns true if there are
// references that can't be resolved.
func hasUndefinedRefs(decl string, knownPkgs map[string]bool) bool {
	// Strip the "var name" / "const name" prefix to get the initializer.
	// Look for patterns like "identifier." that suggest package references.
	for i := 0; i < len(decl)-1; i++ {
		if decl[i] == '.' && i > 0 {
			// Walk back to find the identifier before the dot.
			j := i - 1
			for j >= 0 && isIdentChar(rune(decl[j])) {
				j--
			}
			pkg := decl[j+1 : i]
			if pkg == "" || pkg == "x" {
				continue // receiver variable, not a package
			}
			// Skip if it looks like a method call on a known variable
			// (single lowercase letter or "err", "ctx", etc.)
			if len(pkg) == 1 && pkg[0] >= 'a' && pkg[0] <= 'z' {
				continue
			}
			// Check if this package is in our import set.
			if pkg != "" && !knownPkgs[pkg] {
				// Check if it's a commonly available identifier
				// (builtin types like "reflect", etc.)
				if !isBuiltinIdent(pkg) {
					return true
				}
			}
		}
	}
	return false
}

// isBuiltinIdent returns true for Go built-in identifiers and common
// receiver variable names that aren't package references.
func isBuiltinIdent(name string) bool {
	switch name {
	case "true", "false", "nil", "iota",
		"int", "int8", "int16", "int32", "int64",
		"uint", "uint8", "uint16", "uint32", "uint64",
		"float32", "float64", "complex64", "complex128",
		"string", "bool", "byte", "rune", "error", "any",
		"make", "new", "len", "cap", "append", "copy", "delete",
		"close", "panic", "recover", "print", "println",
		"err", "ctx", "ok", "self":
		return true
	}
	return false
}

// isValidReceiver checks that a method receiver type is a simple identifier
// or pointer to identifier (e.g. "Foo" or "*Foo"), not a garbled signature.
func isValidReceiver(recv string) bool {
	r := strings.TrimPrefix(recv, "*")
	if r == "" {
		return false
	}
	for _, c := range r {
		if !isIdentChar(c) {
			return false
		}
	}
	return true
}

func isIdentChar(c rune) bool {
	return (c >= 'a' && c <= 'z') || (c >= 'A' && c <= 'Z') || (c >= '0' && c <= '9') || c == '_'
}

// aliasCollisions detects import paths that resolve to the same package name
// and adds aliases to resolve collisions. For example, if both "crypto" and
// "git.mleku.dev/mleku/dendrite/pkg/crypto" are imported, the latter becomes
// "dcrypto" aliased.
func aliasCollisions(imports []string) []string {
	type entry struct {
		path  string
		pkg   string
		index int
	}
	entries := make([]entry, len(imports))
	byPkg := make(map[string][]int)

	for i, imp := range imports {
		bare := strings.Trim(imp, `"`)
		pkg := extractPkgName(imp)
		entries[i] = entry{path: bare, pkg: pkg, index: i}
		byPkg[pkg] = append(byPkg[pkg], i)
	}

	result := make([]string, len(imports))
	copy(result, imports)

	for pkg, indices := range byPkg {
		if len(indices) <= 1 {
			continue
		}
		// Keep the shortest path (likely stdlib) without alias.
		// Alias the others with a prefix.
		sort.Slice(indices, func(a, b int) bool {
			return len(entries[indices[a]].path) < len(entries[indices[b]].path)
		})
		for k := 1; k < len(indices); k++ {
			idx := indices[k]
			alias := fmt.Sprintf("%s%d", pkg, k)
			result[idx] = fmt.Sprintf("%s %q", alias, entries[idx].path)
		}
	}
	return result
}

// sanitizeGoSource post-processes emitted Go source:
// 1. Always strip orphan break/continue/goto (valid syntax, invalid semantics)
// 2. If parsing still fails, remove invalid top-level declarations
func sanitizeGoSource(source string) (string, error) {
	// Always strip orphan branches — go/parser accepts them as valid syntax,
	// but the Go compiler rejects break/continue outside loops.
	source = stripOrphanBranches(source)

	fset := token.NewFileSet()
	_, err := parser.ParseFile(fset, "output.go", source, parser.SkipObjectResolution|parser.AllErrors|parser.ParseComments)
	if err == nil {
		return source, nil
	}

	// Parse failed. Salvage by removing invalid top-level declarations.
	return repairSource(source)
}

// stripOrphanBranches removes break/continue/goto statements that appear
// outside of for/switch/select blocks. Uses character-level brace counting
// that ignores braces inside string literals and comments to avoid being
// confused by struct literals like `Type{}`.
func stripOrphanBranches(source string) string {
	lines := strings.Split(source, "\n")
	var result []string

	// Track nesting of for loops (for continue) and all loop-like constructs
	// (for/switch/select, for break). Continue is ONLY valid in for loops.
	forDepth := 0   // for loops only
	breakDepth := 0 // for + switch + select
	braceDepth := 0
	var forBraceStack []int   // brace depths where for loops start
	var breakBraceStack []int // brace depths where any break-accepting block starts
	inBlockComment := false

	for _, line := range lines {
		trimmed := strings.TrimSpace(line)

		// Classify block-opening statements.
		isFor := strings.HasPrefix(trimmed, "for ") || trimmed == "for {" ||
			strings.HasPrefix(trimmed, "for range ")
		isBreakable := isFor ||
			strings.HasPrefix(trimmed, "switch ") || trimmed == "switch {" ||
			strings.HasPrefix(trimmed, "select ") || trimmed == "select {"

		// Count braces at character level, skipping strings and comments.
		lineOpens, lineCloses := countBraces(line, &inBlockComment)

		if isFor && lineOpens > 0 {
			forBraceStack = append(forBraceStack, braceDepth+1)
			forDepth++
		}
		if isBreakable && lineOpens > 0 {
			breakBraceStack = append(breakBraceStack, braceDepth+1)
			breakDepth++
		}

		braceDepth += lineOpens - lineCloses

		// Pop stacks when we close past their level.
		for len(forBraceStack) > 0 && braceDepth < forBraceStack[len(forBraceStack)-1] {
			forBraceStack = forBraceStack[:len(forBraceStack)-1]
			forDepth--
		}
		for len(breakBraceStack) > 0 && braceDepth < breakBraceStack[len(breakBraceStack)-1] {
			breakBraceStack = breakBraceStack[:len(breakBraceStack)-1]
			breakDepth--
		}

		// Strip orphan branches and goto.
		if strings.HasPrefix(trimmed, "goto ") {
			continue // always strip goto — labels rarely survive emission
		}
		// continue is only valid in for loops
		if forDepth <= 0 && (trimmed == "continue" || strings.HasPrefix(trimmed, "continue ")) {
			continue
		}
		// break is valid in for/switch/select
		if breakDepth <= 0 && (trimmed == "break" || strings.HasPrefix(trimmed, "break ")) {
			continue
		}
		// fallthrough and case clauses only valid in switch
		if breakDepth <= 0 {
			if trimmed == "fallthrough" || strings.HasPrefix(trimmed, "case ") ||
				trimmed == "default:" {
				continue
			}
		}

		result = append(result, line)
	}

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

// countBraces counts `{` and `}` in a line, skipping those inside string
// literals (both "" and ``) and comments. Tracks block comment state
// across lines via the inBlockComment pointer.
func countBraces(line string, inBlockComment *bool) (opens, closes int) {
	inString := false
	inRawString := false
	inLineComment := false
	escaped := false

	for i := 0; i < len(line); i++ {
		c := line[i]

		if escaped {
			escaped = false
			continue
		}

		if *inBlockComment {
			if c == '*' && i+1 < len(line) && line[i+1] == '/' {
				*inBlockComment = false
				i++ // skip '/'
			}
			continue
		}

		if inLineComment {
			continue
		}

		if inString {
			if c == '\\' {
				escaped = true
			} else if c == '"' {
				inString = false
			}
			continue
		}

		if inRawString {
			if c == '`' {
				inRawString = false
			}
			continue
		}

		// Not inside any string or comment.
		switch c {
		case '"':
			inString = true
		case '`':
			inRawString = true
		case '/':
			if i+1 < len(line) {
				if line[i+1] == '/' {
					inLineComment = true
					i++
				} else if line[i+1] == '*' {
					*inBlockComment = true
					i++
				}
			}
		case '{':
			opens++
		case '}':
			closes++
		}
	}
	return
}

// repairSource splits Go source into top-level declaration blocks and
// re-assembles only the ones that parse successfully.
func repairSource(source string) (string, error) {
	lines := strings.Split(source, "\n")

	var header strings.Builder // package + import
	var decls []string         // individual top-level declarations
	var current strings.Builder

	inImport := false
	headerDone := false
	braceDepth := 0

	for _, line := range lines {
		trimmed := strings.TrimSpace(line)

		// Package and import go into the header.
		if !headerDone {
			if strings.HasPrefix(trimmed, "package ") {
				header.WriteString(line + "\n")
				continue
			}
			if strings.HasPrefix(trimmed, "import") {
				inImport = true
				header.WriteString(line + "\n")
				if strings.Contains(trimmed, "(") && !strings.Contains(trimmed, ")") {
					continue
				}
				if strings.Contains(trimmed, ")") || !strings.Contains(trimmed, "(") {
					inImport = false
					headerDone = true
				}
				continue
			}
			if inImport {
				header.WriteString(line + "\n")
				if trimmed == ")" {
					inImport = false
					headerDone = true
				}
				continue
			}
			headerDone = true
		}

		// Track brace depth to find declaration boundaries.
		for _, c := range line {
			if c == '{' {
				braceDepth++
			} else if c == '}' {
				braceDepth--
			}
		}

		current.WriteString(line + "\n")

		// At brace depth 0 and a non-empty line, we've completed a declaration.
		if braceDepth <= 0 && trimmed != "" {
			decl := current.String()
			if strings.TrimSpace(decl) != "" {
				decls = append(decls, decl)
			}
			current.Reset()
			braceDepth = 0
		}
	}
	// Flush remaining.
	if s := current.String(); strings.TrimSpace(s) != "" {
		decls = append(decls, s)
	}

	// Validate each declaration by attempting to parse it.
	// Track declared names to prevent redeclarations across types, vars, and funcs.
	var validDecls strings.Builder
	hdr := header.String()
	declaredNames := make(map[string]bool)
	for _, decl := range decls {
		test := hdr + "\n" + decl
		fset := token.NewFileSet()
		_, err := parser.ParseFile(fset, "", test, parser.SkipObjectResolution|parser.AllErrors|parser.ParseComments)
		if err != nil {
			continue
		}
		// Extract the declaration name and check for redeclarations.
		name := extractDeclName(strings.TrimSpace(decl))
		if name != "" && declaredNames[name] {
			continue // skip redeclaration
		}
		if name != "" {
			declaredNames[name] = true
		}
		validDecls.WriteString(decl)
		validDecls.WriteString("\n")
	}

	result := hdr + "\n" + validDecls.String()

	// Don't use go/format — it strips unused imports which breaks the
	// dependency chain. Return the filtered result directly.
	if validDecls.Len() > 0 {
		return result, nil
	}

	return "", fmt.Errorf("no valid declarations found")
}

// extractDeclName extracts the declared name from a top-level declaration.
// "type Foo struct {" → "Foo", "func Bar() {" → "Bar", "var x int" → "x",
// "func (r *T) Method() {" → "" (methods don't conflict with types/vars).
func extractDeclName(decl string) string {
	decl = strings.TrimSpace(decl)
	switch {
	case strings.HasPrefix(decl, "type "):
		rest := decl[5:]
		if idx := strings.IndexAny(rest, " ={"); idx > 0 {
			return rest[:idx]
		}
		return rest
	case strings.HasPrefix(decl, "var "):
		rest := decl[4:]
		if idx := strings.IndexAny(rest, " ="); idx > 0 {
			return rest[:idx]
		}
		return rest
	case strings.HasPrefix(decl, "const "):
		rest := decl[6:]
		if idx := strings.IndexAny(rest, " ="); idx > 0 {
			return rest[:idx]
		}
		return rest
	case strings.HasPrefix(decl, "func "):
		rest := decl[5:]
		// Method: "func (x *T) Name(..." → skip (methods don't conflict)
		if strings.HasPrefix(rest, "(") {
			return ""
		}
		// Plain function: "func Name(..." → extract Name
		if idx := strings.IndexByte(rest, '('); idx > 0 {
			return rest[:idx]
		}
		return ""
	}
	return ""
}