// Package fitness evaluates how closely emitted code reproduces the original.
//
// Three dimensions:
//   - Source: structural similarity of Go AST (types, functions, methods)
//   - Binary: size ratio and shared content of compiled binaries
//   - Behavior: same inputs → same outputs
//
// The fitness score feeds back into the spore, driving evolution:
// generations with higher fitness produce spores that nucleate
// better-adapted lattices.
package fitness

import (
	"bytes"
	"go/ast"
	"go/parser"
	"go/token"
	"os"
	"os/exec"
	"path/filepath"
	"strings"

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

// Score holds the three-dimensional fitness evaluation.
// All fields are exact rationals — no floating-point nondeterminism.
type Score struct {
	Source       ratio.Ratio `json:"source"`        // 0..1 structural AST similarity
	Binary       ratio.Ratio `json:"binary"`        // 0..1 compiled binary similarity
	Behav        ratio.Ratio `json:"behav"`         // 0..1 behavioral equivalence
	Overall      ratio.Ratio `json:"overall"`       // weighted combination
	CompileError string      `json:"compile_error,omitempty"`
}

// Compute sets the Overall score as a weighted combination.
// Behavioral equivalence dominates — "does the code do the same thing?"
// Weights: 3/20 source + 1/20 binary + 16/20 behavioral = 1.
func (s *Score) Compute() {
	s.Overall = ratio.New(3, 20).Mul(s.Source).
		Add(ratio.New(1, 20).Mul(s.Binary)).
		Add(ratio.New(16, 20).Mul(s.Behav))
}

// SourceSimilarity compares two Go source strings structurally.
// It parses both into ASTs and measures the overlap of declarations:
// package name, type names, function names, method names.
func SourceSimilarity(original, emitted string) ratio.Ratio {
	origDecls := extractDecls(original)
	emitDecls := extractDecls(emitted)

	if len(origDecls) == 0 {
		return ratio.Zero
	}

	// Count how many original declarations appear in the emitted code.
	matches := 0
	for decl := range origDecls {
		if emitDecls[decl] {
			matches++
		}
	}

	return ratio.New(int64(matches), int64(len(origDecls)))
}

// extractDecls parses Go source and returns a set of declaration signatures.
func extractDecls(src string) map[string]bool {
	decls := make(map[string]bool)

	fset := token.NewFileSet()
	f, err := parser.ParseFile(fset, "", src, parser.SkipObjectResolution)
	if err != nil {
		// If it doesn't parse, try to extract what we can from text.
		return extractDeclsFromText(src)
	}

	if f.Name != nil {
		decls["package:"+f.Name.Name] = true
	}

	for _, d := range f.Decls {
		switch decl := d.(type) {
		case *ast.FuncDecl:
			if decl.Recv != nil && len(decl.Recv.List) > 0 {
				// Method.
				decls["method:"+decl.Name.Name] = true
			} else {
				decls["func:"+decl.Name.Name] = true
			}
		case *ast.GenDecl:
			for _, spec := range decl.Specs {
				switch s := spec.(type) {
				case *ast.TypeSpec:
					decls["type:"+s.Name.Name] = true
				case *ast.ImportSpec:
					if s.Path != nil {
						decls["import:"+s.Path.Value] = true
					}
				}
			}
		}
	}

	return decls
}

// extractDeclsFromText does a rough text-based extraction for non-parseable source.
func extractDeclsFromText(src string) map[string]bool {
	decls := make(map[string]bool)
	for _, line := range strings.Split(src, "\n") {
		line = strings.TrimSpace(line)
		if strings.HasPrefix(line, "package ") {
			decls["package:"+strings.Fields(line)[1]] = true
		}
		if strings.HasPrefix(line, "func ") {
			// Extract function name.
			rest := strings.TrimPrefix(line, "func ")
			if idx := strings.IndexByte(rest, '('); idx > 0 {
				name := strings.TrimSpace(rest[:idx])
				if strings.Contains(name, ")") {
					// Method: "func (x T) Name("
					parts := strings.SplitAfter(name, ")")
					if len(parts) > 1 {
						decls["method:"+strings.TrimSpace(parts[1])] = true
					}
				} else {
					decls["func:"+name] = true
				}
			}
		}
		if strings.HasPrefix(line, "type ") {
			fields := strings.Fields(line)
			if len(fields) >= 2 {
				decls["type:"+fields[1]] = true
			}
		}
	}
	return decls
}

// BinarySimilarity compares two compiled binaries.
// Measures: size ratio and shared byte sequences.
func BinarySimilarity(originalPath, emittedPath string) ratio.Ratio {
	origData, err := os.ReadFile(originalPath)
	if err != nil {
		return ratio.Zero
	}
	emitData, err := os.ReadFile(emittedPath)
	if err != nil {
		return ratio.Zero
	}

	if len(origData) == 0 || len(emitData) == 0 {
		return ratio.Zero
	}

	// Size ratio — min/max so it's always <= 1.
	small, large := int64(len(emitData)), int64(len(origData))
	if small > large {
		small, large = large, small
	}
	sizeRatio := ratio.New(small, large)

	// Shared 4-byte sequences (rough structural similarity).
	// Sample to keep it fast.
	const chunkSize = 4     // 2^2 byte chunks
	const maxSamples = 25600 // 10^2 × 2^8 — epoch-aligned with binary chunk boundary

	origChunks := make(map[string]bool)
	step := intMax(1, (len(origData)-chunkSize)/maxSamples)
	for i := 0; i+chunkSize <= len(origData); i += step {
		origChunks[string(origData[i:i+chunkSize])] = true
	}

	shared := 0
	total := 0
	step = intMax(1, (len(emitData)-chunkSize)/maxSamples)
	for i := 0; i+chunkSize <= len(emitData); i += step {
		total++
		if origChunks[string(emitData[i:i+chunkSize])] {
			shared++
		}
	}

	chunkSim := ratio.Zero
	if total > 0 {
		chunkSim = ratio.New(int64(shared), int64(total))
	}

	return ratio.Half.Mul(sizeRatio).Add(ratio.Half.Mul(chunkSim))
}

// BehavioralSimilarity runs both binaries with the same input and
// compares their stdout output. This is the real fitness test.
func BehavioralSimilarity(originalBin, emittedBin string, args []string, timeout string) ratio.Ratio {
	origOut := runBinary(originalBin, args, timeout)
	emitOut := runBinary(emittedBin, args, timeout)

	if len(origOut) == 0 && len(emitOut) == 0 {
		// Both produce no output — trivially equivalent.
		return ratio.One
	}
	if len(origOut) == 0 || len(emitOut) == 0 {
		return ratio.Zero
	}

	// Exact match.
	if bytes.Equal(origOut, emitOut) {
		return ratio.One
	}

	// Line-level similarity.
	origLines := strings.Split(string(origOut), "\n")
	emitLines := strings.Split(string(emitOut), "\n")

	return lineSimilarity(origLines, emitLines)
}

// runBinary executes a binary in a temp directory and captures stdout.
// Using a temp dir prevents the binary from clobbering files in the project.
func runBinary(binPath string, args []string, timeout string) []byte {
	if timeout == "" {
		timeout = "10s"
	}

	tmpDir, err := os.MkdirTemp("", "fitness-run-*")
	if err != nil {
		return nil
	}
	defer os.RemoveAll(tmpDir)

	cmd := exec.Command("timeout", append([]string{timeout, binPath}, args...)...)
	cmd.Dir = tmpDir
	cmd.Env = cleanGoEnv("/home/mleku/sdk/go1.24.6")
	out, _ := cmd.CombinedOutput()
	return out
}

// cleanGoEnv returns an environment with GOROOT, PATH, and GOTOOLCHAIN
// set to use the specified Go root, preventing contamination from the
// system Go installation.
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
}

// lineSimilarity computes how well the emitted output reproduces the original.
// Uses a combination of exact line matches and token overlap to measure
// partial matches (e.g., format strings with verbs stripped).
func lineSimilarity(orig, emitted []string) ratio.Ratio {
	if len(orig) == 0 {
		return ratio.Zero
	}

	// First pass: exact line matches (highest confidence).
	emitSet := make(map[string]bool, len(emitted))
	for _, line := range emitted {
		emitSet[strings.TrimSpace(line)] = true
	}

	exactMatches := 0
	for _, line := range orig {
		if emitSet[strings.TrimSpace(line)] {
			exactMatches++
		}
	}

	// Second pass: token overlap — do the emitted lines share words
	// with the original? This catches format-string matches where
	// "abiogenesis: nodes" partially matches "abiogenesis: 72 nodes".
	emitTokens := make(map[string]bool)
	for _, line := range emitted {
		for _, tok := range strings.Fields(strings.TrimSpace(line)) {
			if len(tok) >= 3 { // skip short tokens
				emitTokens[strings.ToLower(tok)] = true
			}
		}
	}

	tokenScore := ratio.Zero
	for _, line := range orig {
		trimmed := strings.TrimSpace(line)
		if trimmed == "" {
			continue
		}
		toks := strings.Fields(trimmed)
		if len(toks) == 0 {
			continue
		}
		hits := 0
		for _, tok := range toks {
			if len(tok) >= 3 && emitTokens[strings.ToLower(tok)] {
				hits++
			}
		}
		total := 0
		for _, tok := range toks {
			if len(tok) >= 3 {
				total++
			}
		}
		if total > 0 {
			tokenScore = tokenScore.Add(ratio.New(int64(hits), int64(total)))
		}
	}

	nonEmpty := 0
	for _, line := range orig {
		if strings.TrimSpace(line) != "" {
			nonEmpty++
		}
	}
	if nonEmpty == 0 {
		nonEmpty = 1
	}
	avgTokenSim := tokenScore.Div(ratio.FromInt(int64(nonEmpty)))

	// Combine: exact matches are worth more, token overlap adds partial credit.
	exactRat := ratio.New(int64(exactMatches), int64(len(orig)))
	return ratio.Half.Mul(exactRat).Add(ratio.Half.Mul(avgTokenSim))
}

// CompileTo compiles a Go source file into a binary in the given directory.
// Returns the path to the binary, or error.
func CompileTo(sourceFile, outputBin, goRoot string) error {
	dir := filepath.Dir(outputBin)
	os.MkdirAll(dir, 0o755)

	// Read the source to check for self-imports.
	src, err := os.ReadFile(sourceFile)
	if err != nil {
		return err
	}

	// Write a go.mod. If the source imports dendrite packages, add
	// require + replace directives so the build can resolve them.
	modPath := filepath.Join(dir, "go.mod")
	if _, err := os.Stat(modPath); os.IsNotExist(err) {
		modContent := buildGoMod(string(src), dir)
		os.WriteFile(modPath, []byte(modContent), 0o644)
	}

	// Copy source to dir/main.go.
	mainPath := filepath.Join(dir, "main.go")
	if err := os.WriteFile(mainPath, src, 0o644); err != nil {
		return err
	}

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

	// Run go mod tidy to reconcile dependencies before building.
	tidy := exec.Command(goBin, "mod", "tidy")
	tidy.Dir = dir
	tidy.Env = env
	tidy.CombinedOutput() // best-effort; build will report real errors

	cmd := exec.Command(goBin, "build", "-o", outputBin, ".")
	cmd.Dir = dir
	cmd.Env = env
	out, err := cmd.CombinedOutput()
	if err != nil {
		return &CompileError{Output: string(out), Err: err}
	}

	return nil
}

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

// buildGoMod generates a go.mod for offspring compilation. If the source
// imports dendrite sub-packages, the go.mod includes require/replace
// directives pointing to the local source tree and copies go.sum for
// transitive dependency resolution. buildDir is the directory where
// go.mod and go.sum will be written.
func buildGoMod(source string, buildDir string) string {
	var b strings.Builder
	b.WriteString("module offspring\n\ngo 1.24\n")

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

	// Find the dendrite module root by walking up from cwd.
	modRoot := findModuleRoot(selfModule)
	if modRoot == "" {
		return b.String()
	}

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

	// Extract require blocks from parent go.mod.
	requires := extractRequireBlocks(string(parentMod))
	if requires != "" {
		b.WriteString("\n")
		b.WriteString(requires)
	}

	// Add self-require + replace.
	b.WriteString("\nrequire " + selfModule + " v0.0.0\n")
	b.WriteString("\nreplace " + selfModule + " => " + modRoot + "\n")

	// Copy go.sum for transitive dependency resolution.
	sumSrc := filepath.Join(modRoot, "go.sum")
	sumDst := filepath.Join(buildDir, "go.sum")
	if sumData, err := os.ReadFile(sumSrc); err == nil {
		os.WriteFile(sumDst, sumData, 0o644)
	}

	return b.String()
}

// findModuleRoot walks up from the current working directory looking for
// a go.mod that declares the given module path.
func findModuleRoot(modulePath string) string {
	dir, err := os.Getwd()
	if err != nil {
		return ""
	}
	for {
		modFile := filepath.Join(dir, "go.mod")
		data, err := os.ReadFile(modFile)
		if err == nil {
			// Check if this go.mod declares our module.
			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 == modulePath {
						return dir
					}
				}
			}
		}
		parent := filepath.Dir(dir)
		if parent == dir {
			break
		}
		dir = parent
	}
	return ""
}

// extractRequireBlocks extracts all require(...) blocks and single require
// lines from a go.mod string, excluding the module declaration.
func extractRequireBlocks(gomod string) string {
	var b strings.Builder
	lines := strings.Split(gomod, "\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
		}
		// Single-line require (but not "module" line).
		if strings.HasPrefix(trimmed, "require ") && !strings.HasPrefix(trimmed, "require (") {
			b.WriteString(line + "\n")
		}
	}
	return b.String()
}

// CompileError wraps a compilation failure with the compiler output.
type CompileError struct {
	Output string
	Err    error
}

func (e *CompileError) Error() string {
	return e.Err.Error() + ": " + e.Output
}

// Evaluate runs the full fitness evaluation pipeline.
func Evaluate(originalSource, emittedSource, originalBin, goRoot string, testArgs []string) Score {
	var s Score

	// 1. Source similarity.
	origSrc, err := os.ReadFile(originalSource)
	if err == nil {
		emitSrc, err2 := os.ReadFile(emittedSource)
		if err2 == nil {
			s.Source = SourceSimilarity(string(origSrc), string(emitSrc))
		}
	}

	// 2. Compile the emitted source and compare binaries.
	absEmitted, _ := filepath.Abs(emittedSource)
	emitBinDir := filepath.Join(filepath.Dir(absEmitted), "_fitness_build")
	emitBin := filepath.Join(emitBinDir, "offspring")
	absOrigBin, _ := filepath.Abs(originalBin)
	defer os.RemoveAll(emitBinDir)

	compileErr := CompileTo(absEmitted, emitBin, goRoot)
	if compileErr != nil {
		s.CompileError = compileErr.Error()
	} else {
		s.Binary = BinarySimilarity(absOrigBin, emitBin)
		s.Behav = BehavioralSimilarity(absOrigBin, emitBin, testArgs, "15s")
	}

	s.Compute()
	return s
}

// intMax returns the larger of two ints.
func intMax(a, b int) int {
	if a > b {
		return a
	}
	return b
}