package compiler

import (
	"fmt"
	"strings"

	"github.com/vic/godnet/pkg/lambda"
)

// CodeGenerator translates lambda AST to Go code.
type CodeGenerator struct {
	SourceFile string
	SourceText string
	buf        strings.Builder
	nodeCount  int
	vars       map[string]*varInfo
}

type varInfo struct {
	nodeName string
	port     int // Next available aux port (for replicators)
	level    int
	uses     int // Number of times variable has been used
}

// Generate produces Go source code from a lambda term.
func (g *CodeGenerator) Generate(term lambda.Term) string {
	g.vars = make(map[string]*varInfo)

	g.writeHeader()
	g.writeBuildNetFunction(term)
	g.writeMainFunction()

	return g.buf.String()
}

func (g *CodeGenerator) writeHeader() {
	g.writeLine("package main")
	g.writeLine("")
	g.writeLine("// Generated by godnet compiler")
	g.writeLine("// Source: %s", g.SourceFile)
	g.writeLine("")
	g.writeLine("import (")
	g.writeLine("\t\"fmt\"")
	g.writeLine("\t\"os\"")
	g.writeLine("\t\"time\"")
	g.writeLine("\t\"github.com/vic/godnet/pkg/deltanet\"")
	g.writeLine("\t\"github.com/vic/godnet/pkg/lambda\"")
	g.writeLine(")")
	g.writeLine("")
}

func (g *CodeGenerator) writeBuildNetFunction(term lambda.Term) {
	g.writeLine("func buildNet(net *deltanet.Network) (deltanet.Node, int, map[uint64]string) {")
	g.writeLine("\tvarNames := make(map[uint64]string)")
	g.writeLine("")

	// Translate the term
	rootNode, rootPort := g.translateTerm(term, 0, 0)

	g.writeLine("")
	g.writeLine("\treturn %s, %d, varNames", rootNode, rootPort)
	g.writeLine("}")
	g.writeLine("")
}

func (g *CodeGenerator) writeMainFunction() {
	g.writeLine("func main() {")
	g.writeLine("\tnet := deltanet.NewNetwork()")
	g.writeLine("\troot, port, varNames := buildNet(net)")
	g.writeLine("")
	g.writeLine("\toutput := net.NewVar()")
	g.writeLine("\tnet.Link(root, port, output, 0)")
	g.writeLine("")
	g.writeLine("\tstart := time.Now()")
	g.writeLine("\tnet.ReduceAll()")
	g.writeLine("\telapsed := time.Since(start)")
	g.writeLine("")
	g.writeLine("\tresultNode, resultPort := net.GetLink(output, 0)")
	g.writeLine("\tresult := lambda.FromDeltaNet(net, resultNode, resultPort, varNames)")
	g.writeLine("\tfmt.Println(result)")
	g.writeLine("")
	g.writeLine("\tstats := net.GetStats()")
	g.writeLine("\tseconds := elapsed.Seconds()")
	g.writeLine("\tfmt.Fprintf(os.Stderr, \"\\nStats:\\n\")")
	g.writeLine("\tfmt.Fprintf(os.Stderr, \"Time: %%v\\n\", elapsed)")
	g.writeLine("\tfmt.Fprintf(os.Stderr, \"Total Reductions: %%d\", stats.TotalReductions)")
	g.writeLine("\tif seconds > 0 {")
	g.writeLine("\t\tfmt.Fprintf(os.Stderr, \" (%%.2f ops/sec)\", float64(stats.TotalReductions)/seconds)")
	g.writeLine("\t}")
	g.writeLine("\tfmt.Fprintf(os.Stderr, \"\\n\")")
	g.writeLine("}")
}

func (g *CodeGenerator) translateTerm(term lambda.Term, level int, depth uint64) (string, int) {
	switch t := term.(type) {
	case lambda.Var:
		return g.genVar(t, level, depth)
	case lambda.Abs:
		return g.genAbs(t, level, depth)
	case lambda.App:
		return g.genApp(t, level, depth)
	case lambda.Let:
		// Desugar: let x = v in b  →  (λx. b) v
		desugared := lambda.App{
			Fun: lambda.Abs{Arg: t.Name, Body: t.Body},
			Arg: t.Val,
		}
		return g.translateTerm(desugared, level, depth)
	default:
		panic(fmt.Sprintf("unknown term type: %T", term))
	}
}

func (g *CodeGenerator) genVar(v lambda.Var, level int, depth uint64) (string, int) {
	if info, ok := g.vars[v.Name]; ok {
		// Bound variable - expand replicator
		g.writeComment("Variable: %s (bound)", v.Name)

		if strings.HasPrefix(info.nodeName, "rep_") {
			// Already a replicator, expand it
			oldRepName := info.nodeName
			newRepName := g.nextNode("rep")
			delta := level - (info.level + 1)
			nextPort := info.uses + 1 // Next aux port index (1-based, since port 0 is principal)

			g.writeLine("\t// Expand replicator for variable '%s' (use #%d)", v.Name, info.uses+1)
			g.writeLine("\t%s := net.NewReplicator(%d, append(%s.Deltas(), %d))",
				newRepName, info.level, oldRepName, delta)

			// Move principal connection
			g.writeLine("\tsourceNode, sourcePort := net.GetLink(%s, 0)", oldRepName)
			g.writeLine("\tnet.LinkAt(%s, 0, sourceNode, sourcePort, %d)", newRepName, depth)

			// Move existing aux ports
			g.writeLine("\tfor i := 0; i < len(%s.Deltas()); i++ {", oldRepName)
			g.writeLine("\t\tdestNode, destPort := net.GetLink(%s, i+1)", oldRepName)
			g.writeLine("\t\tif destNode != nil {")
			g.writeLine("\t\t\tnet.LinkAt(%s, i+1, destNode, destPort, %d)", newRepName, depth)
			g.writeLine("\t\t}")
			g.writeLine("\t}")

			info.nodeName = newRepName
			info.uses++
			return newRepName, nextPort
		} else {
			// First use after eraser - create replicator
			g.writeLine("\t// First use of variable '%s'", v.Name)
			repName := g.nextNode("rep")
			delta := level - (info.level + 1)
			repLevel := info.level + 1

			g.writeLine("\t%s := net.NewReplicator(%d, []int{%d})", repName, repLevel, delta)
			g.writeLine("\tnet.LinkAt(%s, 0, %s, %d, %d)", repName, info.nodeName, info.port, depth)

			info.nodeName = repName
			info.uses = 1
			return repName, 1
		}
	} else {
		// Free variable
		g.writeComment("Variable: %s (free)", v.Name)
		varName := g.nextNode("var")
		repName := g.nextNode("rep")

		g.writeLine("\t%s := net.NewVar()", varName)
		g.writeLine("\tvarNames[%s.ID()] = \"%s\"", varName, v.Name)
		g.writeLine("\t%s := net.NewReplicator(0, []int{%d})", repName, level-1)
		g.writeLine("\tnet.LinkAt(%s, 0, %s, 0, %d)", repName, varName, depth)

		g.vars[v.Name] = &varInfo{
			nodeName: repName,
			port:     0,
			level:    0,
		}

		return repName, 1
	}
}

func (g *CodeGenerator) genAbs(abs lambda.Abs, level int, depth uint64) (string, int) {
	g.writeComment("Abstraction: λ%s. ...", abs.Arg)

	fanName := g.nextNode("fan")
	eraName := g.nextNode("era")

	g.writeLine("\t%s := net.NewFan()", fanName)
	g.writeLine("\t%s := net.NewEraser()", eraName)
	g.writeLine("\tnet.LinkAt(%s, 0, %s, 2, %d)", eraName, fanName, depth)

	// Save old binding if shadowing
	oldVar := g.vars[abs.Arg]
	g.vars[abs.Arg] = &varInfo{
		nodeName: fanName,
		port:     2,
		level:    level,
	}

	// Generate body
	bodyNode, bodyPort := g.translateTerm(abs.Body, level, depth)
	g.writeLine("\tnet.LinkAt(%s, 1, %s, %d, %d)", fanName, bodyNode, bodyPort, depth)

	// Restore old binding
	if oldVar != nil {
		g.vars[abs.Arg] = oldVar
	} else {
		delete(g.vars, abs.Arg)
	}

	return fanName, 0
}

func (g *CodeGenerator) genApp(app lambda.App, level int, depth uint64) (string, int) {
	g.writeComment("Application")

	fanName := g.nextNode("fan")
	g.writeLine("\t%s := net.NewFan()", fanName)

	// Generate function
	funNode, funPort := g.translateTerm(app.Fun, level, depth)
	g.writeLine("\tnet.LinkAt(%s, 0, %s, %d, %d)", fanName, funNode, funPort, depth)

	// Generate argument (level + 1)
	argNode, argPort := g.translateTerm(app.Arg, level+1, depth+1)
	g.writeLine("\tnet.LinkAt(%s, 2, %s, %d, %d)", fanName, argNode, argPort, depth+1)

	return fanName, 1
}

func (g *CodeGenerator) nextNode(prefix string) string {
	g.nodeCount++
	return fmt.Sprintf("%s_%d", prefix, g.nodeCount)
}

func (g *CodeGenerator) writeLine(format string, args ...interface{}) {
	g.buf.WriteString(fmt.Sprintf(format, args...))
	g.buf.WriteString("\n")
}

func (g *CodeGenerator) writeComment(format string, args ...interface{}) {
	comment := "\t// " + fmt.Sprintf(format, args...)
	g.writeLine("%s", comment)
}