// Package train implements the training loop: pick a function with tests,
// hide the original, regenerate from description, run tests to verify
// functional equivalence.
//
// The organism proves it understands a function by reproducing it from its
// own English description. The tests are the oracle of truth — if they pass,
// the generated code is functionally equivalent.
//
// The Run function accepts a CodeGenerator interface so the generation
// strategy is pluggable (lattice-based, deterministic, etc.).
package train

import (
	"context"
	"fmt"
	"os"
	"os/exec"
	"path/filepath"
	"sort"
	"strings"

	"git.mleku.dev/mleku/dendrite/pkg/cartography"
)

// Target is a function selected for training.
type Target struct {
	Entry       *cartography.Entry
	OrigSource  string // original function source
	Description string // English description from atlas
	TestPkg     string // package to test (e.g., "./describe/...")
	TypeDefs    string // type definitions from the same package that this function uses
	TestSource  string // actual test code for the function (so oracle can see expectations)
}

// Result captures the outcome of a training attempt.
type Result struct {
	Target      Target
	Attempts    int    // how many oracle calls were needed
	MaxAttempts int    // limit
	Passed      bool   // all tests passed
	Generated   string // the winning (or last) generated source
	LastError   string // last test failure output
}

// Config controls the training loop.
type Config struct {
	MaxAttempts int    // max oracle calls per target (default 5)
	GoRoot      string // path to Go installation
	ProjectRoot string // path to project root
	Verbose     bool   // print detailed output
}

// FindCandidates returns atlas entries that are good training targets:
// exported functions with tests, in packages that can be tested independently.
func FindCandidates(atlas *cartography.Atlas) []*cartography.Entry {
	ids := sortedEntryIDs(atlas)
	var candidates []*cartography.Entry
	for _, id := range ids {
		e := atlas.Entries[id]
		if !e.HasTest {
			continue
		}
		if !e.Exported {
			continue
		}
		if e.Kind != "func" {
			continue // start with standalone functions, not methods
		}
		if e.Description == "" {
			continue
		}
		if e.Signature == "" {
			continue
		}
		candidates = append(candidates, e)
	}
	return candidates
}

// CodeGenerator produces Go source from a query string.
// This replaces the former oracle dependency with a pluggable interface.
type CodeGenerator interface {
	// GenerateCode takes a prompt and returns Go source code.
	GenerateCode(ctx context.Context, query string) (string, error)
	// SetGeneration advances the generation counter (for rate limiting).
	SetGeneration(gen int)
}

// Run executes the training loop for a single target.
//
// Steps:
//  1. Copy project to temp dir
//  2. Remove the target function from the temp copy
//  3. Generate replacement from description via CodeGenerator
//  4. Write generated code to temp copy
//  5. Run tests — if pass, done
//  6. If fail, compose feedback query with test errors and retry
func Run(ctx context.Context, gen CodeGenerator, target Target, cfg Config) *Result {
	if cfg.MaxAttempts <= 0 {
		cfg.MaxAttempts = 5
	}

	result := &Result{
		Target:      target,
		MaxAttempts: cfg.MaxAttempts,
	}

	// Create temp copy of project.
	tmpDir, err := os.MkdirTemp("", "train-*")
	if err != nil {
		result.LastError = fmt.Sprintf("create temp dir: %v", err)
		return result
	}
	defer os.RemoveAll(tmpDir)

	if err := copyProject(cfg.ProjectRoot, tmpDir); err != nil {
		result.LastError = fmt.Sprintf("copy project: %v", err)
		return result
	}

	// Remove the target function from the temp copy and fix imports.
	if err := removeFunction(tmpDir, target.Entry); err != nil {
		result.LastError = fmt.Sprintf("remove function: %v", err)
		return result
	}

	// Run goimports to clean up unused imports in the modified file.
	goimportsCleanup(tmpDir, target.Entry.FilePath, cfg.GoRoot)

	if cfg.Verbose {
		fmt.Printf("  removed %s from temp copy\n", target.Entry.ID)
	}

	// Iterative generation loop.
	var lastTestOutput string
	for attempt := 1; attempt <= cfg.MaxAttempts; attempt++ {
		result.Attempts = attempt

		// Compose the generation query.
		query := composeTrainQuery(target, lastTestOutput, attempt)

		// Call generator.
		gen.SetGeneration(attempt)
		answer, err := gen.GenerateCode(ctx, query)
		if err != nil {
			result.LastError = fmt.Sprintf("generate call %d: %v", attempt, err)
			if cfg.Verbose {
				fmt.Printf("  attempt %d: generate error: %v\n", attempt, err)
			}
			continue
		}

		// Extract Go source from response.
		// We use extractTrainSource instead of describe.ExtractGoSource
		// because the target function may itself contain backtick markers
		// that confuse the standard extraction regex.
		goSource := extractTrainSource(answer)
		if goSource == "" {
			result.LastError = "no Go source in oracle response"
			if cfg.Verbose {
				fmt.Printf("  attempt %d: no source extracted\n", attempt)
			}
			continue
		}

		// Fix package declaration if oracle used wrong package.
		goSource = fixPackage(goSource, target.Entry.Package)
		result.Generated = goSource

		if cfg.Verbose {
			lines := strings.Split(goSource, "\n")
			n := 10
			if len(lines) < n {
				n = len(lines)
			}
			fmt.Printf("  attempt %d: generated %d bytes (%d lines)\n", attempt, len(goSource), len(lines))
			for _, line := range lines[:n] {
				fmt.Printf("    > %s\n", line)
			}
			if len(lines) > n {
				fmt.Printf("    > ... (%d more lines)\n", len(lines)-n)
			}
		}

		// Write the generated code to the temp copy.
		genPath := filepath.Join(tmpDir, target.Entry.FilePath)
		// The generated code needs to go into the same package.
		// Write as a new file alongside the original (with the function removed).
		genFile := filepath.Join(filepath.Dir(genPath), "train_generated.go")
		if err := os.WriteFile(genFile, []byte(goSource), 0o644); err != nil {
			result.LastError = fmt.Sprintf("write generated: %v", err)
			continue
		}

		// Run tests.
		testPkg := target.TestPkg
		if testPkg == "" {
			testPkg = "./" + filepath.Dir(target.Entry.FilePath) + "/..."
		}
		testOutput, testErr := runTests(tmpDir, testPkg, cfg.GoRoot, target.Entry.TestFuncs)

		if testErr == nil {
			// Tests passed — functional equivalence achieved.
			result.Passed = true
			if cfg.Verbose {
				fmt.Printf("  attempt %d: PASSED\n", attempt)
			}
			return result
		}

		lastTestOutput = testOutput
		result.LastError = testOutput
		if cfg.Verbose {
			// Show first 5 lines of test output.
			lines := strings.Split(testOutput, "\n")
			n := 5
			if len(lines) < n {
				n = len(lines)
			}
			fmt.Printf("  attempt %d: FAILED\n", attempt)
			for _, line := range lines[:n] {
				fmt.Printf("    %s\n", line)
			}
			if len(lines) > n {
				fmt.Printf("    ... (%d more lines)\n", len(lines)-n)
			}
		}

		// Clean up generated file for next attempt.
		os.Remove(genFile)
	}

	return result
}

// composeTrainQuery builds the oracle prompt for generating a function.
func composeTrainQuery(target Target, lastTestError string, attempt int) string {
	var b strings.Builder

	b.WriteString("Generate a Go function that is functionally equivalent to the following specification.\n\n")

	// Include dependent type definitions so the oracle knows the data model.
	if target.TypeDefs != "" {
		b.WriteString("The following types are used by this function. They are already defined — do NOT redefine them.\n")
		b.WriteString("Types from other packages are imported; use their package qualifier (e.g., axiom.Element).\n")
		b.WriteString("IMPORTANT: Use the FULL import path shown in the comments (e.g., \"git.mleku.dev/mleku/dendrite/pkg/axiom\"), NOT a short path like \"axiom\".\n")
		b.WriteString("Types from the same package are used directly (e.g., Description).\n\n")
		b.WriteString(target.TypeDefs)
		b.WriteString("\n")
	}

	// Description from atlas.
	fmt.Fprintf(&b, "Function: %s\n", target.Entry.Name)
	fmt.Fprintf(&b, "Package: %s\n", target.Entry.Package)
	fmt.Fprintf(&b, "Signature: %s\n", target.Entry.Signature)
	fmt.Fprintf(&b, "\nDescription: %s\n", target.Description)

	if target.Entry.Contract != "" {
		fmt.Fprintf(&b, "\nContract: %s\n", target.Entry.Contract)
	}
	if target.Entry.EdgeCases != "" {
		fmt.Fprintf(&b, "\nEdge cases: %s\n", target.Entry.EdgeCases)
	}
	if target.Entry.DocComment != "" {
		fmt.Fprintf(&b, "\nDoc comment: %s\n", target.Entry.DocComment)
	}

	// Test specifications — what the function must satisfy.
	if len(target.Entry.TestFuncs) > 0 {
		fmt.Fprintf(&b, "\nTests that must pass: %s\n", strings.Join(target.Entry.TestFuncs, ", "))
	}
	if len(target.Entry.TestPatterns) > 0 {
		fmt.Fprintf(&b, "Test cases: %s\n", strings.Join(target.Entry.TestPatterns, ", "))
	}
	if target.Entry.Validation != "" {
		fmt.Fprintf(&b, "Validation: %s\n", target.Entry.Validation)
	}

	// Include the actual test source if available, so the oracle can see
	// the exact behavioral expectations.
	if target.TestSource != "" {
		b.WriteString("\nActual test code (for reference — your function must pass these):\n```go\n")
		b.WriteString(target.TestSource)
		b.WriteString("\n```\n")
	}

	// Parameter details.
	if len(target.Entry.Params) > 0 {
		b.WriteString("\nParameters:\n")
		for _, p := range target.Entry.Params {
			fmt.Fprintf(&b, "  - %s %s", p.Name, p.Type)
			if p.Semantic != "" {
				fmt.Fprintf(&b, " (%s)", p.Semantic)
			}
			b.WriteString("\n")
		}
	}
	if len(target.Entry.Returns) > 0 {
		b.WriteString("Returns:\n")
		for _, r := range target.Entry.Returns {
			fmt.Fprintf(&b, "  - %s", r.Type)
			if r.IsError {
				b.WriteString(" (error)")
			}
			if r.Semantic != "" {
				fmt.Fprintf(&b, " — %s", r.Semantic)
			}
			b.WriteString("\n")
		}
	}

	// Side effects.
	if len(target.Entry.SideEffects) > 0 {
		fmt.Fprintf(&b, "\nSide effects: %s\n", strings.Join(target.Entry.SideEffects, ", "))
	}

	// Dependencies — what other functions/types this calls.
	if len(target.Entry.DependsOn) > 0 {
		fmt.Fprintf(&b, "Dependencies: %s\n", strings.Join(target.Entry.DependsOn, ", "))
	}

	// Feedback from previous failed attempt.
	if attempt > 1 && lastTestError != "" {
		b.WriteString("\n--- PREVIOUS ATTEMPT FAILED ---\n")
		b.WriteString("The following test failures occurred:\n\n")
		// Only include FAIL lines and error messages — filter out PASS lines.
		filtered := filterTestFailures(lastTestError)
		b.WriteString(filtered)
		b.WriteString("\n\nFix the issues and generate a corrected version.\n")
	}

	b.WriteString("\nRequirements:\n")
	fmt.Fprintf(&b, "- Use package %s\n", target.Entry.Package)
	b.WriteString("- Include all necessary imports\n")
	b.WriteString("- The function must have the EXACT same signature\n")
	b.WriteString("- Return the complete Go source inside a single ```go code block\n")
	b.WriteString("- Include ONLY the function — no test code, no main function\n")
	b.WriteString("- IMPORTANT: If the function body contains backtick characters, use raw string literals (backtick) carefully — the code block must still be valid\n")

	return b.String()
}

// removeFunction removes a specific function from its source file in the
// temp project copy. It reads the file, finds the function by line number
// and brace matching, and writes the file back without it.
func removeFunction(tmpDir string, entry *cartography.Entry) error {
	path := filepath.Join(tmpDir, entry.FilePath)
	data, err := os.ReadFile(path)
	if err != nil {
		return err
	}

	lines := strings.Split(string(data), "\n")
	start := entry.Line - 1 // 0-indexed
	if start < 0 || start >= len(lines) {
		return fmt.Errorf("line %d out of range (file has %d lines)", entry.Line, len(lines))
	}

	// Include the doc comment above the function.
	for start > 0 && strings.HasPrefix(strings.TrimSpace(lines[start-1]), "//") {
		start--
	}

	// Find the end by brace matching.
	end := start
	depth := 0
	foundOpen := false
	for i := start; i < len(lines); i++ {
		for _, ch := range lines[i] {
			if ch == '{' {
				depth++
				foundOpen = true
			}
			if ch == '}' {
				depth--
				if depth == 0 && foundOpen {
					end = i + 1
					goto found
				}
			}
		}
	}
	// No braces found — single-line declaration or type.
	end = start + 1

found:
	// Remove lines [start:end].
	var result []string
	result = append(result, lines[:start]...)
	result = append(result, lines[end:]...)

	return os.WriteFile(path, []byte(strings.Join(result, "\n")), 0o644)
}

// runTests runs `go test` for a specific package in the temp directory.
// If testFuncs is non-empty, uses -run to filter to those specific tests.
// Returns the test output and any error.
func runTests(tmpDir, testPkg, goRoot string, testFuncs []string) (string, error) {
	goBin := filepath.Join(goRoot, "bin", "go")
	args := []string{"test", "-v", "-count=1"}
	if len(testFuncs) > 0 {
		args = append(args, "-run", "^("+strings.Join(testFuncs, "|")+")$")
	}
	args = append(args, testPkg)
	cmd := exec.Command(goBin, args...)
	cmd.Dir = tmpDir
	cmd.Env = cleanGoEnv(goRoot)
	out, err := cmd.CombinedOutput()
	return string(out), err
}

// copyProject copies a Go project to a temp directory, skipping non-essential
// files for speed.
func copyProject(src, dst string) error {
	return filepath.Walk(src, func(path string, info os.FileInfo, err error) error {
		if err != nil {
			return err
		}

		rel, err := filepath.Rel(src, path)
		if err != nil {
			return err
		}

		if info.IsDir() {
			base := filepath.Base(path)
			switch base {
			case ".git", "_output", "node_modules", "vendor",
				".svelte-kit", ".next", "dist", "build", "coverage":
				return filepath.SkipDir
			}
			return os.MkdirAll(filepath.Join(dst, rel), info.Mode())
		}

		// Only copy Go source and module files.
		ext := filepath.Ext(path)
		if ext != ".go" && ext != ".mod" && ext != ".sum" {
			return nil
		}

		data, err := os.ReadFile(path)
		if err != nil {
			return err
		}
		return os.WriteFile(filepath.Join(dst, rel), data, info.Mode())
	})
}

// cleanGoEnv returns an environment with GOROOT, PATH, and GOTOOLCHAIN set.
func cleanGoEnv(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",
		"PATH="+filepath.Join(root, "bin")+":"+os.Getenv("PATH"),
	)
	return clean
}

// goimportsCleanup removes unused imports from a Go file after a function
// has been deleted. It does a simple text scan: finds all import paths,
// checks if each is referenced in the remaining code, and removes unreferenced ones.
func goimportsCleanup(tmpDir, relPath, goRoot string) {
	absPath := filepath.Join(tmpDir, relPath)
	data, err := os.ReadFile(absPath)
	if err != nil {
		return
	}

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

	// Find import block boundaries and individual imports.
	type importLine struct {
		lineIdx int
		path    string // e.g., "go/ast"
		alias   string // e.g., "ast" or custom alias
	}

	var imports []importLine
	inImportBlock := false
	importBlockStart := -1
	importBlockEnd := -1

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

		if trimmed == "import (" {
			inImportBlock = true
			importBlockStart = i
			continue
		}
		if inImportBlock && trimmed == ")" {
			inImportBlock = false
			importBlockEnd = i
			continue
		}
		if inImportBlock && trimmed != "" && !strings.HasPrefix(trimmed, "//") {
			// Parse import line: could be `"path"` or `alias "path"`
			path := ""
			alias := ""
			parts := strings.Fields(trimmed)
			for _, p := range parts {
				clean := strings.Trim(p, `"`)
				if clean != p { // it was quoted — this is the path
					path = clean
				} else {
					alias = p
				}
			}
			if path != "" {
				if alias == "" {
					alias = importAlias(path)
				}
				imports = append(imports, importLine{lineIdx: i, path: path, alias: alias})
			}
		}
	}

	if len(imports) == 0 || importBlockStart < 0 {
		return
	}

	// Build the code text WITHOUT import block for searching.
	var codeLines []string
	for i, line := range lines {
		if i >= importBlockStart && i <= importBlockEnd {
			continue
		}
		codeLines = append(codeLines, line)
	}
	codeText := strings.Join(codeLines, "\n")

	// Check which imports are used by looking for the alias as a package
	// qualifier. We need word-boundary awareness: "ast." must not match
	// "fast." or "last.". A package qualifier in Go is always preceded by
	// a non-alphanumeric character (space, tab, paren, star, ampersand, etc.)
	// or starts a line.
	var unusedLineIdxs []int
	for _, imp := range imports {
		if imp.alias == "_" {
			continue
		}
		used := isImportUsed(codeText, imp.alias)
		if !used {
			unusedLineIdxs = append(unusedLineIdxs, imp.lineIdx)
		}
	}

	if len(unusedLineIdxs) == 0 {
		return
	}

	// Remove unused import lines.
	skip := make(map[int]bool)
	for _, idx := range unusedLineIdxs {
		skip[idx] = true
	}

	var result []string
	for i, line := range lines {
		if skip[i] {
			continue
		}
		result = append(result, line)
	}

	os.WriteFile(absPath, []byte(strings.Join(result, "\n")), 0o644)
}

// filterTestFailures extracts only the failure-relevant lines from test output.
// Removes PASS lines and keeps FAIL lines, error messages, and compiler errors.
func filterTestFailures(output string) string {
	lines := strings.Split(output, "\n")
	var result []string
	inFailBlock := false

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

		// Always include compiler errors.
		if strings.HasPrefix(trimmed, "#") || strings.Contains(line, ": undefined:") ||
			strings.Contains(line, "imported and not used") ||
			strings.Contains(line, "redeclared") {
			result = append(result, line)
			continue
		}

		// Skip PASS lines.
		if strings.HasPrefix(trimmed, "--- PASS:") || strings.HasPrefix(trimmed, "=== PAUSE") {
			inFailBlock = false
			continue
		}

		// Include FAIL lines and their context.
		if strings.HasPrefix(trimmed, "--- FAIL:") || strings.HasPrefix(trimmed, "FAIL") {
			result = append(result, line)
			inFailBlock = true
			continue
		}

		// Include test error output (indented lines after a RUN or FAIL).
		if strings.HasPrefix(trimmed, "=== RUN") {
			// Check if this test fails — peek ahead is hard, so include the RUN.
			inFailBlock = true
			result = append(result, line)
			continue
		}

		// Include assertion failures and error messages.
		if inFailBlock && trimmed != "" {
			result = append(result, line)
		}

		// Cap output.
		if len(result) >= 40 {
			result = append(result, "... (truncated)")
			break
		}
	}

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

// isImportUsed checks if a package alias is used as a qualifier in Go code.
// Handles word boundaries: "ast." must not match "fast." or "last.".
func isImportUsed(code, alias string) bool {
	needle := alias + "."
	idx := 0
	for {
		pos := strings.Index(code[idx:], needle)
		if pos < 0 {
			return false
		}
		absPos := idx + pos
		// Check that the character before the match is not alphanumeric.
		if absPos == 0 {
			return true // starts at beginning of code
		}
		prev := code[absPos-1]
		if !isAlphaNum(prev) {
			return true // word boundary before alias
		}
		idx = absPos + len(needle)
		if idx >= len(code) {
			return false
		}
	}
}

func isAlphaNum(b byte) bool {
	return (b >= 'a' && b <= 'z') || (b >= 'A' && b <= 'Z') || (b >= '0' && b <= '9') || b == '_'
}

// importAlias derives the Go package name from an import path.
// Handles versioned module paths: "github.com/foo/bar/v2" → "bar",
// "github.com/btcsuite/btcd/btcec/v2" → "btcec".
func importAlias(path string) string {
	parts := strings.Split(path, "/")
	last := parts[len(parts)-1]
	// If the last component is a Go module version suffix (v2, v3, etc.),
	// use the second-to-last component.
	if len(parts) >= 2 && len(last) >= 2 && last[0] == 'v' && last[1] >= '0' && last[1] <= '9' {
		return parts[len(parts)-2]
	}
	return last
}

// extractTrainSource extracts Go source from an oracle response, handling
// the case where the generated function body itself contains backtick markers
// (e.g., when regenerating ExtractGoSource). Unlike describe.ExtractGoSource
// which uses non-greedy regex, this finds code blocks by locating the first
// ```go marker and the LAST ``` marker, which handles nested backticks.
func extractTrainSource(answer string) string {
	// Find the opening marker.
	openIdx := strings.Index(answer, "```go")
	if openIdx < 0 {
		openIdx = strings.Index(answer, "```")
		if openIdx < 0 {
			// No code fences — if it starts with "package", use the whole thing.
			trimmed := strings.TrimSpace(answer)
			if strings.HasPrefix(trimmed, "package ") {
				return trimmed
			}
			return ""
		}
	}

	// Move past the opening marker line.
	startOfCode := strings.Index(answer[openIdx:], "\n")
	if startOfCode < 0 {
		return ""
	}
	startOfCode += openIdx + 1

	// Find the LAST ``` in the response (closing marker).
	// This handles cases where the generated code contains backticks.
	lastClose := strings.LastIndex(answer, "```")
	if lastClose <= openIdx {
		// No closing marker — use everything after opening.
		return strings.TrimSpace(answer[startOfCode:])
	}

	source := answer[startOfCode:lastClose]
	source = strings.TrimSpace(source)

	// Safety: strip any residual opening backtick markers at the start.
	for strings.HasPrefix(source, "```") {
		nl := strings.Index(source, "\n")
		if nl < 0 {
			break
		}
		source = strings.TrimSpace(source[nl+1:])
	}

	return source
}

// fixPackage ensures the generated source uses the correct package name.
// If the oracle produced "package main" but we need "package describe",
// this fixes it.
func fixPackage(source, wantPkg string) string {
	lines := strings.Split(source, "\n")
	for i, line := range lines {
		trimmed := strings.TrimSpace(line)
		if strings.HasPrefix(trimmed, "package ") {
			parts := strings.Fields(trimmed)
			if len(parts) >= 2 && parts[1] != wantPkg {
				lines[i] = "package " + wantPkg
			}
			break
		}
	}
	return strings.Join(lines, "\n")
}

// ExtractTypeDefs finds type definitions that the target function's parameters
// or returns reference. Handles both same-package types (e.g., Description)
// and cross-package types (e.g., axiom.Element, lattice.Lattice).
// Also extracts method signatures for cross-package types so the oracle knows
// the API surface (e.g., Lattice.Nodes(), Node.Bond()).
// Reads actual source from disk to get full struct/interface definitions.
func ExtractTypeDefs(entry *cartography.Entry, atlas *cartography.Atlas, projectRoot string) string {
	var b strings.Builder

	// Read module path from go.mod for constructing full import paths.
	modulePath := readModulePath(projectRoot)

	// Collect type references from params and returns.
	// samePackage: {"Description": true}
	// crossPackage: {"axiom.Element": true, "lattice.Lattice": true}
	samePackage := make(map[string]bool)
	crossPackage := make(map[string]bool)

	for _, p := range entry.Params {
		classifyType(p.Type, entry.Package, samePackage, crossPackage)
	}
	for _, r := range entry.Returns {
		classifyType(r.Type, entry.Package, samePackage, crossPackage)
	}

	// Iteratively expand types: for each type found, also collect types
	// referenced by its struct fields and method return types.
	// Run up to 3 rounds to follow chains like Lattice → Node → axiom.Constraint.
	for range 3 {
		before := len(samePackage) + len(crossPackage)
		expandFieldTypes(samePackage, crossPackage, entry.Package, atlas)
		after := len(samePackage) + len(crossPackage)
		if after == before {
			break // no new types discovered
		}
	}

	if len(samePackage) == 0 && len(crossPackage) == 0 {
		return ""
	}

	// Collect already-emitted type IDs to avoid duplicates.
	emitted := make(map[string]bool)

	// Same-package types: emit ALL exported type definitions from the package.
	// The oracle needs to see every type it might construct or reference.
	// Same-package types are always accessible, so this is safe.
	ids := sortedEntryIDs(atlas)
	for _, id := range ids {
		e := atlas.Entries[id]
		if e.Package != entry.Package {
			continue
		}
		if e.Kind != "type" && e.Kind != "interface" {
			continue
		}
		if !e.Exported {
			continue
		}
		if emitted[e.ID] {
			continue
		}
		src := readTypeSource(e, projectRoot)
		if src != "" {
			fmt.Fprintf(&b, "// (same package %s)\n%s\n\n", e.Package, src)
			emitted[e.ID] = true
		}
	}

	// Cross-package types: emit struct/interface definitions + method signatures.
	crossNames := make([]string, 0, len(crossPackage))
	for qn := range crossPackage {
		crossNames = append(crossNames, qn)
	}
	sort.Strings(crossNames)
	for _, qualifiedName := range crossNames {
		parts := strings.SplitN(qualifiedName, ".", 2)
		if len(parts) != 2 {
			continue
		}
		pkg, name := parts[0], parts[1]

		// Find the atlas entry for this type.
		for _, id := range ids {
			e := atlas.Entries[id]
			if e.Package != pkg || e.Name != name {
				continue
			}
			if e.Kind != "type" && e.Kind != "interface" {
				continue
			}
			if emitted[e.ID] {
				continue
			}
			src := readTypeSource(e, projectRoot)
			fullImport := modulePath + "/" + filepath.Dir(e.FilePath)
			if src != "" {
				fmt.Fprintf(&b, "// (from package %s — import path: %q, do NOT redefine)\n%s\n", pkg, fullImport, src)
				emitted[e.ID] = true
			} else if e.Description != "" {
				fmt.Fprintf(&b, "// (from package %s — import path: %q) %s\n", pkg, fullImport, e.Description)
				emitted[e.ID] = true
			}

			// Append exported method signatures for this type.
			methods := collectMethods(pkg, name, atlas)
			if len(methods) > 0 {
				fmt.Fprintf(&b, "// Exported methods on %s.%s:\n", pkg, name)
				for _, sig := range methods {
					fmt.Fprintf(&b, "//   %s\n", sig)
				}
			}
			b.WriteString("\n")
			break
		}
	}

	return b.String()
}

// expandFieldTypes looks at each already-collected type's struct fields AND
// method return types in the atlas and adds any new type references.
// One level only — prevents explosion. This ensures that if Lattice has
// method Nodes() []*Node, the Node type is also collected.
func expandFieldTypes(samePackage, crossPackage map[string]bool, currentPkg string, atlas *cartography.Atlas) {
	// Snapshot current sets so we don't iterate while modifying.
	sameNames := make([]string, 0, len(samePackage))
	for k := range samePackage {
		sameNames = append(sameNames, k)
	}
	sort.Strings(sameNames)

	crossNames := make([]string, 0, len(crossPackage))
	for k := range crossPackage {
		crossNames = append(crossNames, k)
	}
	sort.Strings(crossNames)

	ids := sortedEntryIDs(atlas)

	// For each same-package type, collect types from struct fields.
	for _, name := range sameNames {
		for _, id := range ids {
			e := atlas.Entries[id]
			if e.Package != currentPkg || e.Name != name {
				continue
			}
			if e.Kind != "type" {
				continue
			}
			for _, f := range e.Params {
				classifyType(f.Type, currentPkg, samePackage, crossPackage)
			}
			break
		}
	}

	// For each cross-package type, collect types from struct fields
	// AND from method return types.
	for _, qualifiedName := range crossNames {
		parts := strings.SplitN(qualifiedName, ".", 2)
		if len(parts) != 2 {
			continue
		}
		pkg, name := parts[0], parts[1]

		// Struct fields.
		for _, id := range ids {
			e := atlas.Entries[id]
			if e.Package != pkg || e.Name != name || e.Kind != "type" {
				continue
			}
			for _, f := range e.Params {
				classifyType(f.Type, currentPkg, samePackage, crossPackage)
			}
			break
		}

		// Method return types: if Lattice.Nodes() returns []*Node,
		// add lattice.Node to crossPackage so its methods are also extracted.
		// Note: return types are recorded unqualified within their own package,
		// so "[]*Node" in a lattice method means "lattice.Node" for our purposes.
		for _, id := range ids {
			e := atlas.Entries[id]
			if e.Package != pkg || e.Kind != "method" || !e.Exported {
				continue
			}
			recv := strings.TrimPrefix(e.Receiver, "*")
			if recv != name {
				continue
			}
			for _, r := range e.Returns {
				// Classify relative to the method's own package so that
				// unqualified types like "Node" resolve to same-package
				// of the method (i.e., "lattice"), then we promote them
				// to cross-package from the target's perspective.
				methodSame := make(map[string]bool)
				methodCross := make(map[string]bool)
				classifyType(r.Type, pkg, methodSame, methodCross)
				// Same-package of the method = cross-package of the target.
				for typeName := range methodSame {
					if pkg != currentPkg {
						crossPackage[pkg+"."+typeName] = true
					} else {
						samePackage[typeName] = true
					}
				}
				for k := range methodCross {
					crossPackage[k] = true
				}
			}
		}
	}
}

// collectMethods finds all exported method signatures for a given type
// in the atlas. Returns signature strings like "Nodes() []*Node".
func collectMethods(pkg, typeName string, atlas *cartography.Atlas) []string {
	ids := sortedEntryIDs(atlas)
	var sigs []string
	// Match methods whose receiver is the type (pointer or value).
	for _, id := range ids {
		e := atlas.Entries[id]
		if e.Package != pkg || e.Kind != "method" {
			continue
		}
		if !e.Exported {
			continue
		}
		// Receiver is stored as "*Lattice" or "Lattice".
		recv := strings.TrimPrefix(e.Receiver, "*")
		if recv != typeName {
			continue
		}
		// Use the signature, stripping the "func " prefix and receiver.
		// The atlas Signature looks like: "func (l *Lattice) Nodes() []*Node"
		// We want just: "Nodes() []*Node"
		sig := e.Signature
		if i := strings.Index(sig, ") "); i >= 0 {
			sig = strings.TrimSpace(sig[i+2:])
		}
		sigs = append(sigs, sig)
	}
	return sigs
}

// sortedEntryIDs returns atlas entry IDs in sorted order for deterministic iteration.
func sortedEntryIDs(atlas *cartography.Atlas) []string {
	ids := make([]string, 0, len(atlas.Entries))
	for id := range atlas.Entries {
		ids = append(ids, id)
	}
	sort.Strings(ids)
	return ids
}

// classifyType parses a Go type expression and classifies referenced custom
// types as same-package or cross-package.
// "*RevenueSummary" → samePackage["RevenueSummary"]
// "[]axiom.Element" → crossPackage["axiom.Element"]
// "*lattice.Lattice" → crossPackage["lattice.Lattice"]
func classifyType(typ, currentPkg string, samePackage, crossPackage map[string]bool) {
	// Strip pointer, slice, variadic, map prefixes.
	t := typ
	t = strings.TrimPrefix(t, "*")
	t = strings.TrimPrefix(t, "[]")
	t = strings.TrimPrefix(t, "...")
	t = strings.TrimPrefix(t, "*") // **T
	// Handle map types: map[K]V — extract V.
	if strings.HasPrefix(t, "map[") {
		if i := strings.LastIndex(t, "]"); i >= 0 {
			t = t[i+1:]
			t = strings.TrimPrefix(t, "*")
		}
	}

	if t == "" || isBuiltinType(t) {
		return
	}

	// Check for package qualifier: "axiom.Element"
	if i := strings.Index(t, "."); i > 0 {
		crossPackage[t] = true
	} else if t[0] >= 'A' && t[0] <= 'Z' {
		// Unqualified uppercase → same package type.
		samePackage[t] = true
	}
}

// isBuiltinType returns true for Go builtin types.
func isBuiltinType(name string) bool {
	switch name {
	case "string", "int", "int8", "int16", "int32", "int64",
		"uint", "uint8", "uint16", "uint32", "uint64",
		"float32", "float64", "bool", "byte", "rune",
		"error", "any", "interface{}", "uintptr",
		"context.Context": // treat as builtin — everyone knows it
		return true
	}
	return false
}

// readTypeSource reads the actual type definition from source.
func readTypeSource(entry *cartography.Entry, projectRoot string) string {
	path := filepath.Join(projectRoot, entry.FilePath)
	data, err := os.ReadFile(path)
	if err != nil {
		return ""
	}

	lines := strings.Split(string(data), "\n")
	start := entry.Line - 1
	if start < 0 || start >= len(lines) {
		return ""
	}

	// Find the end by brace matching.
	depth := 0
	for i := start; i < len(lines) && i < start+50; i++ {
		for _, ch := range lines[i] {
			if ch == '{' {
				depth++
			}
			if ch == '}' {
				depth--
				if depth == 0 {
					return strings.Join(lines[start:i+1], "\n")
				}
			}
		}
	}

	// No braces — single-line type.
	return lines[start]
}

// readModulePath reads the Go module path from go.mod.
func readModulePath(projectRoot string) string {
	data, err := os.ReadFile(filepath.Join(projectRoot, "go.mod"))
	if err != nil {
		return ""
	}
	for _, line := range strings.Split(string(data), "\n") {
		line = strings.TrimSpace(line)
		if strings.HasPrefix(line, "module ") {
			return strings.TrimSpace(strings.TrimPrefix(line, "module"))
		}
	}
	return ""
}

// firstSentence returns the first sentence of a string.
func firstSentence(s string) string {
	s = strings.TrimSpace(s)
	if i := strings.IndexByte(s, '.'); i >= 0 && i < 120 {
		return s[:i+1]
	}
	if len(s) > 120 {
		return s[:120] + "..."
	}
	return s
}

// ExtractTestSource reads the test functions for an entry from the test file.
func ExtractTestSource(entry *cartography.Entry, projectRoot string) string {
	if entry.TestFile == "" || len(entry.TestFuncs) == 0 {
		return ""
	}

	path := filepath.Join(projectRoot, entry.TestFile)
	data, err := os.ReadFile(path)
	if err != nil {
		return ""
	}

	source := string(data)
	lines := strings.Split(source, "\n")
	var result []string

	for _, testFunc := range entry.TestFuncs {
		// Find the test function in the file.
		prefix := "func " + testFunc + "("
		for i, line := range lines {
			if !strings.Contains(line, prefix) {
				continue
			}
			// Found start — extract until closing brace.
			depth := 0
			for j := i; j < len(lines) && j < i+100; j++ {
				result = append(result, lines[j])
				for _, ch := range lines[j] {
					if ch == '{' {
						depth++
					}
					if ch == '}' {
						depth--
						if depth == 0 {
							goto nextFunc
						}
					}
				}
			}
		nextFunc:
			result = append(result, "")
			break
		}
	}

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

// BuildDescription composes a rich English description of a function
// from its atlas entry, suitable for training. This is what the oracle
// receives instead of the source code.
func BuildDescription(entry *cartography.Entry) string {
	var b strings.Builder

	// Start with the atlas description.
	if entry.Description != "" {
		b.WriteString(entry.Description)
	}

	// Add contract details.
	if entry.Contract != "" {
		fmt.Fprintf(&b, "\n\nContract: %s", entry.Contract)
	}

	// Add edge cases.
	if entry.EdgeCases != "" {
		fmt.Fprintf(&b, "\n\nEdge cases: %s", entry.EdgeCases)
	}

	// Add doc comment (original author's intent).
	if entry.DocComment != "" {
		fmt.Fprintf(&b, "\n\nDoc: %s", entry.DocComment)
	}

	return b.String()
}