package emit

import (
	"fmt"
	"go/parser"
	"go/token"
	"sort"
	"strings"
	"testing"

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

type tagConstraint struct{ tag string }

func (c tagConstraint) Tag() string                 { return c.tag }
func (c tagConstraint) Admits(e axiom.Element) bool  { return e.Type() == c.tag }

// bondAll ingests Go source into a lattice and bonds elements to matching sites.
func bondAll(t *testing.T, l *lattice.Lattice, src string) int {
	t.Helper()
	ge := enzyme.GoSource{}
	elements := ge.Digest(strings.NewReader(src))
	bonded := 0
	for e := range elements {
		for _, n := range l.Nodes() {
			if !n.Occupied() && n.Admits(e) {
				n.Bond(e)
				bonded++
				break
			}
		}
	}
	return bonded
}

// makeLattice creates a lattice with sites for all Go AST element types.
func makeLattice(t *testing.T) *lattice.Lattice {
	t.Helper()
	l := lattice.New()
	tags := []string{
		"package", "import", "type", "struct", "interface", "field",
		"func", "method", "comment",
		"assign", "return", "if", "for", "switch", "select",
		"go", "send", "expr", "defer", "decl", "branch", "case", "comm",
		// Declaration-level ident subtypes.
		"ident:func-name", "ident:method-name", "ident:type-name",
		"ident:field-name", "ident:param", "ident:result",
		"ident:receiver", "ident:var-name",
	}
	for _, tag := range tags {
		for range 20 {
			n := l.AddNode([]axiom.Constraint{tagConstraint{tag}})
			n.SetEnergy(true)
		}
	}
	// Wire up neighbors.
	nodes := l.Nodes()
	for i := range nodes {
		if i > 0 {
			l.Connect(nodes[i-1], nodes[i])
		}
	}
	return l
}

func TestHarvestAndEmit(t *testing.T) {
	l := lattice.New()

	tags := []string{"package", "import", "type", "struct", "field",
		"func", "ident:var-name", "ident:func-name", "method", "return"}
	nodes := make(map[string]*lattice.Node)
	for _, tag := range tags {
		n := l.AddNode([]axiom.Constraint{tagConstraint{tag}})
		nodes[tag] = n
	}
	for range 5 {
		l.AddNode([]axiom.Constraint{tagConstraint{"ident:var-name"}})
	}

	nodes["package"].Bond(enzyme.Elem("package", "main"))
	nodes["import"].Bond(enzyme.Elem("import", `"fmt"`))
	nodes["type"].Bond(enzyme.Elem("type", "MyType"))
	nodes["struct"].Bond(enzyme.Elem("struct", "MyType"))
	nodes["field"].Bond(enzyme.Elem("field", "MyType\x00Name string"))
	nodes["func"].Bond(enzyme.Elem("func", "main()"))
	nodes["ident:var-name"].Bond(enzyme.Elem("ident:var-name", "foo"))
	nodes["ident:func-name"].Bond(enzyme.Elem("ident:func-name", "main"))
	nodes["method"].Bond(enzyme.Elem("method", "MyType.Run()"))
	nodes["return"].Bond(enzyme.Elem("return", ""))

	files := Harvest(l)
	frags := files[""]
	t.Logf("harvested %d fragments", len(frags))

	var buf strings.Builder
	err := EmitGo(frags, &buf)
	if err != nil {
		t.Fatal(err)
	}

	src := buf.String()
	t.Log("--- emitted source ---")
	t.Log(src)

	if !strings.Contains(src, "package main") {
		t.Error("missing package declaration")
	}
	if !strings.Contains(src, "type MyType struct") {
		t.Error("missing type declaration")
	}
	if !strings.Contains(src, "func main()") {
		t.Error("missing func declaration")
	}
}

func TestEmitFromSelfIngest(t *testing.T) {
	l := makeLattice(t)

	src := `package main

import "fmt"

type Thing struct {
	Name string
}

func main() {
	f := Thing{Name: "world"}
	fmt.Println(f.Name)
}
`
	bonded := bondAll(t, l, src)
	t.Logf("bonded %d elements", bonded)

	files := Harvest(l)
	frags := files[""]

	var buf strings.Builder
	err := EmitGo(frags, &buf)
	if err != nil {
		t.Fatal(err)
	}

	emitted := buf.String()
	t.Log("--- emitted from self-ingest ---")
	t.Log(emitted)

	if !strings.Contains(emitted, "package main") {
		t.Error("missing package declaration")
	}
	if !strings.Contains(emitted, "func main()") {
		t.Error("expected main function")
	}
	// Should have the assignment in main's body.
	if !strings.Contains(emitted, "Thing{") {
		t.Error("expected Thing struct literal in main body")
	}
	// Should have the fmt.Println call.
	if !strings.Contains(emitted, "fmt.Println") {
		t.Error("expected fmt.Println call in body")
	}
	// Should have fmt import.
	if !strings.Contains(emitted, `"fmt"`) {
		t.Error("expected fmt import")
	}
}

func TestEmitReconstructsBody(t *testing.T) {
	l := makeLattice(t)

	src := `package main

import "fmt"

func greet(name string) string {
	if name == "" {
		return "hello, stranger"
	}
	return fmt.Sprintf("hello, %s", name)
}

func main() {
	msg := greet("world")
	fmt.Println(msg)
}
`
	bonded := bondAll(t, l, src)
	t.Logf("bonded %d elements", bonded)

	files := Harvest(l)
	frags := files[""]

	var buf strings.Builder
	err := EmitGo(frags, &buf)
	if err != nil {
		t.Fatal(err)
	}

	emitted := buf.String()
	t.Log("--- emitted with body reconstruction ---")
	t.Log(emitted)

	// Check function declarations include signatures.
	if !strings.Contains(emitted, "func greet(") {
		t.Error("expected greet function with parameters")
	}
	// Check greet body has the if statement.
	if !strings.Contains(emitted, `name == ""`) {
		t.Error("expected if condition in greet body")
	}
	// Check return statement.
	if !strings.Contains(emitted, `return fmt.Sprintf`) || !strings.Contains(emitted, `return "hello, stranger"`) {
		t.Error("expected return statements in greet body")
	}
	// Check main body has assignment.
	if !strings.Contains(emitted, `greet("world")`) {
		t.Error("expected greet call in main body")
	}
}

func TestEmitMethodReconstruction(t *testing.T) {
	l := makeLattice(t)

	src := `package main

import "fmt"

type Server struct {
	Port int
}

func (s Server) Start() error {
	fmt.Printf("listening on %d\n", s.Port)
	return nil
}

func main() {
	s := Server{Port: 8080}
	s.Start()
}
`
	bonded := bondAll(t, l, src)
	t.Logf("bonded %d elements", bonded)

	files := Harvest(l)
	frags := files[""]

	var buf strings.Builder
	err := EmitGo(frags, &buf)
	if err != nil {
		t.Fatal(err)
	}

	emitted := buf.String()
	t.Log("--- emitted with methods ---")
	t.Log(emitted)

	// Method should have receiver with original variable name.
	if !strings.Contains(emitted, "func (s Server) Start()") {
		t.Error("expected method with receiver")
	}
	// Method body should have the Printf call.
	if !strings.Contains(emitted, "fmt.Printf") {
		t.Error("expected fmt.Printf in method body")
	}
	// Struct declaration.
	if !strings.Contains(emitted, "type Server struct") {
		t.Error("expected Server struct")
	}
}

func TestEmitCompilesClean(t *testing.T) {
	l := makeLattice(t)

	src := `package main

import "fmt"

type Score struct {
	Value int
}

func compute() int {
	return 42
}

func main() {
	x := compute()
	fmt.Printf("score: %d\n", x)
}
`
	bondAll(t, l, src)

	files := Harvest(l)
	frags := files[""]

	var buf strings.Builder
	err := EmitGo(frags, &buf)
	if err != nil {
		t.Fatal(err)
	}

	emitted := buf.String()
	t.Log("--- emitted source ---")
	t.Log(emitted)

	// Verify it parses as valid Go.
	fset := token.NewFileSet()
	_, parseErr := parser.ParseFile(fset, "emitted.go", emitted, parser.SkipObjectResolution)
	if parseErr != nil {
		t.Errorf("emitted source does not parse as valid Go: %v\nsource:\n%s", parseErr, emitted)
	}
}

func TestEmitLegacyMode(t *testing.T) {
	// Test that the legacy self-logging mode still works when only
	// declarations and literals are present (no body elements).
	l := lattice.New()

	nodes := make(map[string]*lattice.Node)
	for _, tag := range []string{"package", "type", "struct", "field", "func", "literal:string", "method"} {
		n := l.AddNode([]axiom.Constraint{tagConstraint{tag}})
		nodes[tag] = n
	}

	nodes["package"].Bond(enzyme.Elem("package", "main"))
	nodes["type"].Bond(enzyme.Elem("type", "Score"))
	nodes["struct"].Bond(enzyme.Elem("struct", "Score"))
	nodes["field"].Bond(enzyme.Elem("field", "Score\x00Value int"))
	nodes["func"].Bond(enzyme.Elem("func", "main()"))
	nodes["literal:string"].Bond(enzyme.Elem("literal:string", `"score: %d nodes\n"`))
	nodes["method"].Bond(enzyme.Elem("method", "Score.Compute()"))

	files := Harvest(l)
	frags := files[""]

	var buf strings.Builder
	err := EmitGo(frags, &buf)
	if err != nil {
		t.Fatal(err)
	}

	emitted := buf.String()
	t.Log("--- legacy emitted source ---")
	t.Log(emitted)

	// Should have self-knowledge in legacy mode.
	if !strings.Contains(emitted, "nTypes") {
		t.Error("expected self-knowledge declaration in legacy mode")
	}
	if !strings.Contains(emitted, "fmt.Printf") {
		t.Error("expected fmt.Printf for format strings in legacy mode")
	}
}

func TestParseBodyValue(t *testing.T) {
	parent, source := parseBodyValue("main\x00x := 42")
	if parent != "main" {
		t.Errorf("expected parent 'main', got %q", parent)
	}
	if source != "x := 42" {
		t.Errorf("expected source 'x := 42', got %q", source)
	}

	// Legacy element with no separator.
	parent, source = parseBodyValue("")
	if parent != "" {
		t.Errorf("expected empty parent, got %q", parent)
	}
}

func TestParseFuncName(t *testing.T) {
	tests := []struct {
		val      string
		isMethod bool
		want     string
	}{
		{"main()", false, "main"},
		{"greet(name string) string", false, "greet"},
		{"Foo.Hello() string", true, "Hello"},
		{"*Server.Start() error", true, "Start"},
	}
	for _, tt := range tests {
		got := parseFuncName(tt.val, tt.isMethod)
		if got != tt.want {
			t.Errorf("parseFuncName(%q, %v) = %q, want %q", tt.val, tt.isMethod, got, tt.want)
		}
	}
}

func TestEmitMultipleStructs(t *testing.T) {
	l := makeLattice(t)

	src := `package main

import "fmt"

type Server struct {
	Port int
	Host string
}

type Client struct {
	Name    string
	Timeout int
}

func main() {
	s := Server{Port: 8080, Host: "localhost"}
	c := Client{Name: "test", Timeout: 30}
	fmt.Println(s, c)
}
`
	bonded := bondAll(t, l, src)
	t.Logf("bonded %d elements", bonded)

	files := Harvest(l)
	frags := files[""]

	var buf strings.Builder
	err := EmitGo(frags, &buf)
	if err != nil {
		t.Fatal(err)
	}

	emitted := buf.String()
	t.Log("--- emitted with multiple structs ---")
	t.Log(emitted)

	// Server struct should have Port and Host but NOT Name or Timeout.
	if !strings.Contains(emitted, "type Server struct") {
		t.Error("expected Server struct")
	}
	if !strings.Contains(emitted, "type Client struct") {
		t.Error("expected Client struct")
	}

	// Find Server struct block and verify its fields.
	serverIdx := strings.Index(emitted, "type Server struct")
	clientIdx := strings.Index(emitted, "type Client struct")
	if serverIdx < 0 || clientIdx < 0 {
		t.Fatal("missing struct declarations")
	}

	// Extract the Server struct block.
	serverEnd := strings.Index(emitted[serverIdx:], "}\n")
	serverBlock := emitted[serverIdx : serverIdx+serverEnd+2]
	t.Logf("Server block: %q", serverBlock)

	if !strings.Contains(serverBlock, "Port") {
		t.Error("Server should contain Port field")
	}
	if !strings.Contains(serverBlock, "Host") {
		t.Error("Server should contain Host field")
	}
	if strings.Contains(serverBlock, "Name") {
		t.Error("Server should NOT contain Name field (belongs to Client)")
	}
	if strings.Contains(serverBlock, "Timeout") {
		t.Error("Server should NOT contain Timeout field (belongs to Client)")
	}

	// Extract the Client struct block.
	clientEnd := strings.Index(emitted[clientIdx:], "}\n")
	clientBlock := emitted[clientIdx : clientIdx+clientEnd+2]
	t.Logf("Client block: %q", clientBlock)

	if !strings.Contains(clientBlock, "Name") {
		t.Error("Client should contain Name field")
	}
	if !strings.Contains(clientBlock, "Timeout") {
		t.Error("Client should contain Timeout field")
	}
	if strings.Contains(clientBlock, "Port") {
		t.Error("Client should NOT contain Port field (belongs to Server)")
	}

	// Verify it parses as valid Go.
	fset := token.NewFileSet()
	_, parseErr := parser.ParseFile(fset, "emitted.go", emitted, parser.SkipObjectResolution)
	if parseErr != nil {
		t.Errorf("emitted source does not parse: %v\nsource:\n%s", parseErr, emitted)
	}
}

func TestEmitDirective(t *testing.T) {
	l := makeLattice(t)
	// Add sites for directive and var element types.
	for range 4 {
		n := l.AddNode([]axiom.Constraint{tagConstraint{"directive"}})
		n.SetEnergy(true)
	}
	for range 4 {
		n := l.AddNode([]axiom.Constraint{tagConstraint{"var"}})
		n.SetEnergy(true)
	}

	src := `package main

import "embed"

//go:embed hello.txt
var content string

func main() {
	println(content)
}
`
	bonded := bondAll(t, l, src)
	t.Logf("bonded %d elements", bonded)

	files := Harvest(l)
	frags := files[""]

	var buf strings.Builder
	err := EmitGo(frags, &buf)
	if err != nil {
		t.Fatal(err)
	}

	emitted := buf.String()
	t.Log("--- emitted with directive ---")
	t.Log(emitted)

	// The //go:embed directive should appear in the output.
	if !strings.Contains(emitted, "//go:embed") {
		t.Error("expected //go:embed directive in emitted source")
	}

	// The var declaration should appear.
	if !strings.Contains(emitted, "var content") {
		t.Error("expected var content declaration in emitted source")
	}

	// The directive should appear before the var declaration.
	embedIdx := strings.Index(emitted, "//go:embed")
	varIdx := strings.Index(emitted, "var content")
	if embedIdx >= 0 && varIdx >= 0 && embedIdx > varIdx {
		t.Error("//go:embed directive should appear before var declaration")
	}
}

func TestImportPruning(t *testing.T) {
	// Verify unused imports are pruned.
	stmts := []bodyStmt{
		{Source: `fmt.Println("hello")`},
	}
	imports := []Fragment{
		{Type: "import", Value: `"fmt"`},
		{Type: "import", Value: `"os"`},       // unused — should be pruned
		{Type: "import", Value: `"strings"`},   // unused — should be pruned
	}
	result := inferImports(stmts, imports, nil, nil, nil)
	if len(result) != 1 || result[0] != `"fmt"` {
		t.Errorf("expected only fmt import, got %v", result)
	}
}

func TestRejectUntrustedImport(t *testing.T) {
	// External imports must be rejected by trust policy.
	stmts := []bodyStmt{
		{Source: `evil.DoSomething()`},
		{Source: `axiom.NewElement()`},
	}
	imports := []Fragment{
		{Type: "import", Value: `"github.com/evil/pkg"`},                       // untrusted
		{Type: "import", Value: `"git.mleku.dev/mleku/dendrite/pkg/axiom"`},        // trusted self-import
		{Type: "import", Value: `"fmt"`},                                        // trusted stdlib (unused)
	}
	result := inferImports(stmts, imports, nil, nil, nil)

	// Only the self-import should survive (used + trusted).
	// fmt is trusted but unused, evil is used but untrusted.
	found := make(map[string]bool)
	for _, imp := range result {
		found[imp] = true
	}

	if found[`"github.com/evil/pkg"`] {
		t.Error("untrusted external import should have been rejected")
	}
	if !found[`"git.mleku.dev/mleku/dendrite/pkg/axiom"`] {
		t.Error("trusted self-import should have been kept")
	}
	if found[`"fmt"`] {
		t.Error("unused import should have been pruned")
	}
}

func TestInferImportsFromSignatures(t *testing.T) {
	// Imports should be inferred from function signatures and var declarations,
	// not just body statements.
	stmts := []bodyStmt{
		{Source: `fmt.Println("hello")`},
	}
	funcs := []funcDecl{
		{Name: "foo", Signature: "foo(s *spore.Spore) ratio.Ratio", IsMethod: false},
	}
	vars := []Fragment{
		{Type: "var", Value: `var Name = reflect.TypeOf(Block{}).Name()`},
	}
	result := inferImports(stmts, nil, funcs, vars, nil)
	found := make(map[string]bool)
	for _, imp := range result {
		found[imp] = true
	}
	if !found[`"fmt"`] {
		t.Error("expected fmt import from body, got", result)
	}
	if !found[`"reflect"`] {
		t.Error("expected reflect import from var declaration, got", result)
	}
	if !found[`"git.mleku.dev/mleku/dendrite/pkg/spore"`] {
		t.Error("expected spore import from func signature, got", result)
	}
	if !found[`"git.mleku.dev/mleku/dendrite/pkg/ratio"`] {
		t.Error("expected ratio import from func signature, got", result)
	}
}

func TestDependencyOrdering(t *testing.T) {
	// Strict dependency chain: x → y → Println.
	// x must come before y, y before Println.
	stmts := []bodyStmt{
		{Source: `fmt.Println(y)`, Tag: "expr"},
		{Source: `y := x + 1`, Tag: "assign"},
		{Source: `x := 1`, Tag: "assign"},
	}

	for range 10 {
		ordered := orderBody(stmts)
		xIdx, yIdx, printIdx := -1, -1, -1
		for i, s := range ordered {
			switch {
			case strings.Contains(s.Source, "x := 1"):
				xIdx = i
			case strings.Contains(s.Source, "y := x"):
				yIdx = i
			case strings.Contains(s.Source, "Println"):
				printIdx = i
			}
		}
		if xIdx > yIdx {
			t.Errorf("x := 1 (idx %d) should come before y := x + 1 (idx %d)", xIdx, yIdx)
		}
		if yIdx > printIdx {
			t.Errorf("y := x + 1 (idx %d) should come before Println(y) (idx %d)", yIdx, printIdx)
		}
	}
}

func TestNonDependentShuffling(t *testing.T) {
	// Three independent assignments + one dependent Println.
	// a, b, c have no dependencies on each other — should be shuffled.
	// Println depends on all three — must come last.
	stmts := []bodyStmt{
		{Source: `a := 1`, Tag: "assign"},
		{Source: `b := 2`, Tag: "assign"},
		{Source: `c := 3`, Tag: "assign"},
		{Source: `fmt.Println(a, b, c)`, Tag: "expr"},
	}

	orderings := make(map[string]bool)
	for range 100 {
		ordered := orderBody(stmts)

		// Println must always be last.
		last := ordered[len(ordered)-1]
		if !strings.Contains(last.Source, "Println") {
			t.Fatal("Println must always be last (depends on a, b, c)")
		}

		// Collect the order of a, b, c.
		var order string
		for _, s := range ordered[:3] {
			switch {
			case strings.Contains(s.Source, "a :="):
				order += "a"
			case strings.Contains(s.Source, "b :="):
				order += "b"
			case strings.Contains(s.Source, "c :="):
				order += "c"
			}
		}
		orderings[order] = true
	}

	if len(orderings) < 2 {
		t.Errorf("expected at least 2 distinct orderings of independent statements, got %d: %v",
			len(orderings), orderings)
	}
	t.Logf("observed %d distinct orderings: %v", len(orderings), orderings)
}

func TestEmitProject(t *testing.T) {
	// Create fragments for two files: main package and a library package.
	mainFrags := []Fragment{
		{Type: "package", Value: "main"},
		{Type: "import", Value: `"fmt"`},
		{Type: "import", Value: `"git.mleku.dev/mleku/dendrite/pkg/axiom"`},
		{Type: "func", Value: "main()"},
	}
	// Simulate body content so imports get resolved.
	mainFrags = append(mainFrags, Fragment{Type: "expr", Value: "main\x00fmt.Println(axiom.New())"})

	libFrags := []Fragment{
		{Type: "package", Value: "axiom"},
		{Type: "type", Value: "Element"},
		{Type: "interface", Value: "Element"},
		{Type: "func", Value: "New() Element"},
	}

	files := map[string][]Fragment{
		"main.go":       mainFrags,
		"axiom/axiom.go": libFrags,
	}

	project := EmitProject(files, "git.mleku.dev/mleku/dendrite")

	// Should have 3 files: main.go, axiom/axiom.go, go.mod.
	if len(project) != 3 {
		t.Errorf("expected 3 files, got %d: %v", len(project), keys(project))
	}

	// Check go.mod exists and has replace directive.
	gomod := project["go.mod"]
	if gomod == "" {
		t.Fatal("missing go.mod")
	}
	t.Log("--- go.mod ---")
	t.Log(gomod)

	if !strings.Contains(gomod, "module git.mleku.dev/mleku/dendrite") {
		t.Error("go.mod missing module declaration")
	}
	if !strings.Contains(gomod, "git.mleku.dev/mleku/dendrite/axiom => ./axiom") {
		t.Error("go.mod missing replace directive for axiom")
	}

	// Check main.go has package main.
	mainSrc := project["main.go"]
	if !strings.Contains(mainSrc, "package main") {
		t.Error("main.go missing package main")
	}
	t.Log("--- main.go ---")
	t.Log(mainSrc)

	// Check axiom/axiom.go has package axiom.
	axiomSrc := project["axiom/axiom.go"]
	if !strings.Contains(axiomSrc, "package axiom") {
		t.Error("axiom/axiom.go missing package axiom")
	}
	t.Log("--- axiom/axiom.go ---")
	t.Log(axiomSrc)
}

func TestGoModReplace(t *testing.T) {
	internalPkgs := map[string]bool{
		"axiom":   true,
		"emit":    true,
		"enzyme":  true,
		"lattice": true,
	}
	gomod := emitGoMod("git.mleku.dev/mleku/dendrite", internalPkgs)
	t.Log(gomod)

	if !strings.Contains(gomod, "module git.mleku.dev/mleku/dendrite") {
		t.Error("missing module declaration")
	}
	if !strings.Contains(gomod, "go 1.24") {
		t.Error("missing go version")
	}

	// All internal packages should have replace directives.
	for pkg := range internalPkgs {
		expected := fmt.Sprintf("git.mleku.dev/mleku/dendrite/%s => ./%s", pkg, pkg)
		if !strings.Contains(gomod, expected) {
			t.Errorf("missing replace directive for %s", pkg)
		}
	}
}

func keys(m map[string]string) []string {
	var ks []string
	for k := range m {
		ks = append(ks, k)
	}
	sort.Strings(ks)
	return ks
}

func TestEmitGoPackage(t *testing.T) {
	// Verify non-main package emission doesn't force a main function.
	frags := []Fragment{
		{Type: "package", Value: "axiom"},
		{Type: "type", Value: "Element"},
		{Type: "interface", Value: "Element"},
	}

	var buf strings.Builder
	err := EmitGoPackage(frags, &buf, "axiom")
	if err != nil {
		t.Fatal(err)
	}

	emitted := buf.String()
	t.Log("--- non-main package ---")
	t.Log(emitted)

	if !strings.Contains(emitted, "package axiom") {
		t.Error("expected package axiom")
	}
	if strings.Contains(emitted, "func main()") {
		t.Error("non-main package should NOT have func main()")
	}
}