package jsbackend

import (
	"encoding/json"
	"fmt"
	"go/token"
	"go/types"
	"os"
	"path/filepath"
	"sort"
	"strings"
	"time"

	"golang.org/x/tools/go/ssa"
)

// Config holds JS backend configuration.
type Config struct {
	OutputDir  string // directory to write .mjs files
	RuntimeDir string // path to jsruntime/ directory (embedded in output)
	DumpSSA    bool   // print SSA for debugging
}

// PackageInfo holds the minimal info we need about a package from the loader.
// This avoids a dependency on the loader package (which pulls in LLVM via cgo).
type PackageInfo struct {
	Path    string         // import path
	TypePkg *types.Package // typechecker package
}

// ProgramCompiler compiles an entire Go program to JavaScript ES modules.
type ProgramCompiler struct {
	config     *Config
	program    *ssa.Program
	typeMapper *TypeMapper
	compiled   map[string]bool
	asyncFuncs map[*ssa.Function]bool // functions that need to be async
	errors     []error
	manifest   []string // relative paths of all output files
}

// CompileProgram is the main entry point.
// Takes an SSA program, a list of packages in dependency order, and the main package path.
func CompileProgram(config *Config, ssaProg *ssa.Program, packages []PackageInfo, mainPkgPath string) error {
	pc := &ProgramCompiler{
		config:     config,
		program:    ssaProg,
		typeMapper: NewTypeMapper(),
		compiled:   make(map[string]bool),
		asyncFuncs: make(map[*ssa.Function]bool),
	}

	if err := os.MkdirAll(config.OutputDir, 0o755); err != nil {
		return fmt.Errorf("create output dir: %w", err)
	}

	// Build all SSA packages.
	ssaProg.Build()

	// Compute which functions need to be async (transitive closure).
	pc.computeAsyncFuncs(packages)

	// Compile packages in dependency order.
	for _, pkgInfo := range packages {
		ssaPkg := ssaProg.Package(pkgInfo.TypePkg)
		if ssaPkg == nil {
			continue
		}
		if err := pc.compilePackage(ssaPkg); err != nil {
			pc.errors = append(pc.errors, err)
		}
	}

	// Write the main entry point.
	if mainPkgPath != "" {
		if err := pc.writeEntryPoint(mainPkgPath); err != nil {
			pc.errors = append(pc.errors, err)
		}
	}

	// Copy runtime files to output.
	if err := pc.copyRuntime(); err != nil {
		pc.errors = append(pc.errors, err)
	}

	// Write manifest of all output files.
	if err := pc.writeManifest(); err != nil {
		pc.errors = append(pc.errors, err)
	}

	if len(pc.errors) > 0 {
		return pc.errors[0]
	}
	return nil
}

// compilePackage compiles a single Go package to a JS module file.
func (pc *ProgramCompiler) compilePackage(ssaPkg *ssa.Package) error {
	pkgPath := ssaPkg.Pkg.Path()
	if pc.compiled[pkgPath] {
		return nil
	}
	pc.compiled[pkgPath] = true

	if shouldSkipPackage(pkgPath) {
		return nil
	}

	e := NewEmitter()

	e.Comment("Package %s", pkgPath)
	e.Comment("Generated by MoxieJS — Moxie JavaScript backend")
	e.Newline()

	e.ImportAll("$rt", pc.runtimeImportPath(pkgPath))
	e.Newline()

	pc.emitImports(e, ssaPkg)

	// Collect and sort members.
	var memberNames []string
	for name := range ssaPkg.Members {
		memberNames = append(memberNames, name)
	}
	sort.Slice(memberNames, func(i, j int) bool {
		iPos := ssaPkg.Members[memberNames[i]].Pos()
		jPos := ssaPkg.Members[memberNames[j]].Pos()
		if iPos == jPos {
			return memberNames[i] < memberNames[j]
		}
		return iPos < jPos
	})

	// Package-level variables.
	e.Comment("Package-level variables")
	for _, name := range memberNames {
		if g, ok := ssaPkg.Members[name].(*ssa.Global); ok {
			pc.emitGlobal(e, g)
		}
	}
	e.Newline()

	// Type registrations.
	for _, name := range memberNames {
		if t, ok := ssaPkg.Members[name].(*ssa.Type); ok {
			pc.typeMapper.EmitTypeRegistration(e, t.Type())
		}
	}

	// Functions.
	for _, name := range memberNames {
		if fn, ok := ssaPkg.Members[name].(*ssa.Function); ok {
			if fn.TypeParams() != nil {
				continue
			}
			pc.compileFunction(e, fn)
		}
	}

	// Methods on named types.
	for _, name := range memberNames {
		if t, ok := ssaPkg.Members[name].(*ssa.Type); ok {
			if types.IsInterface(t.Type()) {
				continue
			}
			pc.compileTypeMethods(e, ssaPkg, t)
		}
	}

	filename := PackageModuleName(pkgPath)
	outputPath := filepath.Join(pc.config.OutputDir, filename)
	pc.manifest = append(pc.manifest, filename)
	return os.WriteFile(outputPath, []byte(e.String()), 0o644)
}

func (pc *ProgramCompiler) compileFunction(e *Emitter, fn *ssa.Function) {
	if fn.Blocks == nil {
		pc.emitExternalStub(e, fn)
		return
	}

	// Bridge functions: if the package path contains "/jsbridge/",
	// emit a pass-through to the JS runtime instead of compiling Go.
	if pc.emitBridgeFunction(e, fn) {
		return
	}

	if pc.config.DumpSSA {
		fmt.Fprintf(os.Stderr, "=== %s ===\n", fn.String())
		fn.WriteTo(os.Stderr)
		fmt.Fprintln(os.Stderr)
	}

	fc := &FunctionCompiler{
		pc:     pc,
		fn:     fn,
		e:      e,
		locals: make(map[ssa.Value]string),
		varGen: 0,
	}
	fc.compile()
}

// emitBridgeFunction checks if fn belongs to a jsbridge package and emits
// a pass-through to the JS runtime. Returns true if handled.
func (pc *ProgramCompiler) emitBridgeFunction(e *Emitter, fn *ssa.Function) bool {
	pkg := fn.Package()
	if pkg == nil {
		return false
	}
	pkgPath := pkg.Pkg.Path()

	// Find the jsbridge subpackage name.
	idx := strings.Index(pkgPath, "/jsbridge/")
	if idx < 0 {
		return false
	}
	bridgePkg := pkgPath[idx+len("/jsbridge/"):]
	// bridgePkg is e.g. "dom", "ws", "localstorage"

	goName := fn.Name()
	// Don't bridge the init function — let it compile normally.
	if goName == "init" {
		return false
	}
	jsName := functionJsName(fn)
	params := functionParams(fn)

	// Check if function returns a value.
	results := fn.Signature.Results()
	hasReturn := results != nil && results.Len() > 0

	e.Comment("jsbridge: %s.%s", bridgePkg, goName)
	e.Block("export function %s(%s)", jsName, params)
	if hasReturn {
		e.Line("return $rt.%s.%s(%s);", bridgePkg, goName, params)
	} else {
		e.Line("$rt.%s.%s(%s);", bridgePkg, goName, params)
	}
	e.EndBlock()
	e.Newline()
	return true
}

func (pc *ProgramCompiler) compileTypeMethods(e *Emitter, ssaPkg *ssa.Package, t *ssa.Type) {
	methods := getAllMethods(ssaPkg.Prog, t.Type())
	methods = append(methods, getAllMethods(ssaPkg.Prog, types.NewPointer(t.Type()))...)

	for _, method := range methods {
		fn := ssaPkg.Prog.MethodValue(method)
		if fn == nil || fn.Blocks == nil {
			continue
		}
		if fn.Synthetic != "" && fn.Synthetic != "package initializer" {
			continue
		}
		pc.compileFunction(e, fn)

		typeID := pc.typeMapper.TypeID(t.Type())
		jsName := methodJsName(fn)
		e.Line("$rt.types.getType(%s)?.methods?.set(%s, %s);",
			JsString(typeID), JsString(method.Obj().Name()), jsName)
	}
}

func (pc *ProgramCompiler) emitGlobal(e *Emitter, g *ssa.Global) {
	name := JsIdentifier(g.Name())
	elemType := g.Type().(*types.Pointer).Elem()
	zero := pc.typeMapper.ZeroExpr(elemType)
	e.Line("export let %s = { $value: %s, $get() { return this.$value; }, $set(v) { this.$value = v; } };", name, zero)
}

func (pc *ProgramCompiler) emitImports(e *Emitter, pkg *ssa.Package) {
	imports := make(map[string]bool)

	for _, member := range pkg.Members {
		if fn, ok := member.(*ssa.Function); ok {
			collectImports(fn, imports)
		}
	}

	// Also scan methods on named types.
	for _, member := range pkg.Members {
		if t, ok := member.(*ssa.Type); ok {
			for _, mset := range []*types.MethodSet{
				pkg.Prog.MethodSets.MethodSet(t.Type()),
				pkg.Prog.MethodSets.MethodSet(types.NewPointer(t.Type())),
			} {
				for i := 0; i < mset.Len(); i++ {
					fn := pkg.Prog.MethodValue(mset.At(i))
					if fn != nil {
						collectImports(fn, imports)
					}
				}
			}
		}
	}

	delete(imports, pkg.Pkg.Path())
	for path := range imports {
		if shouldSkipPackage(path) {
			delete(imports, path)
		}
	}

	paths := make([]string, 0, len(imports))
	for p := range imports {
		paths = append(paths, p)
	}
	sort.Strings(paths)

	for _, p := range paths {
		alias := JsIdentifier(p)
		modFile := PackageModuleName(p)
		e.ImportAll(alias, "./"+modFile)
	}
	if len(paths) > 0 {
		e.Newline()
	}
}

func (pc *ProgramCompiler) emitExternalStub(e *Emitter, fn *ssa.Function) {
	name := functionJsName(fn)
	params := functionParams(fn)

	fullName := fn.String()
	switch {
	case fullName == "fmt.Println" || fullName == "(*fmt.pp).doPrintln":
		e.Block("export function %s(%s)", name, params)
		e.Line("$rt.runtime.println(...arguments);")
		e.EndBlock()
	case fullName == "fmt.Printf" || fullName == "(*fmt.pp).doPrintf":
		e.Block("export function %s(%s)", name, params)
		e.Line("$rt.runtime.print(...arguments);")
		e.EndBlock()
	case strings.HasPrefix(fullName, "runtime."):
		e.Comment("runtime stub: %s", fullName)
		e.Block("export function %s(%s)", name, params)
		e.Comment("intrinsic: %s", fullName)
		e.EndBlock()
	default:
		e.Comment("external: %s", fullName)
		e.Block("export function %s(%s)", name, params)
		e.Line("throw new Error('external function not implemented: %s');", fullName)
		e.EndBlock()
	}
	e.Newline()
}

func (pc *ProgramCompiler) writeEntryPoint(mainPkgPath string) error {
	e := NewEmitter()
	e.Comment("MoxieJS Entry Point")
	e.Comment("Generated by Moxie JavaScript backend")
	// Build stamp ensures the entry script is byte-different on every build,
	// so the browser detects the SW as updated even when only imports changed.
	e.Comment(fmt.Sprintf("Build: %d", time.Now().UnixNano()))
	e.Newline()

	e.ImportAll("$rt", "./$runtime/index.mjs")
	mainMod := PackageModuleName(mainPkgPath)
	e.ImportAll("$main", "./"+mainMod)
	e.Newline()

	e.Block("async function $start()")
	e.Line("if ($rt.crypto && $rt.crypto.init) await $rt.crypto.init();")
	e.Line("if ($main.init) await $main.init();")
	e.Line("if ($main.main) await $main.main();")
	e.EndBlock()
	e.Newline()

	e.Line("$rt.goroutine.runMain($start);")

	outputPath := filepath.Join(pc.config.OutputDir, "$entry.mjs")
	pc.manifest = append(pc.manifest, "$entry.mjs")
	return os.WriteFile(outputPath, []byte(e.String()), 0o644)
}

func (pc *ProgramCompiler) copyRuntime() error {
	rtDir := filepath.Join(pc.config.OutputDir, "$runtime")
	if err := os.MkdirAll(rtDir, 0o755); err != nil {
		return err
	}

	if pc.config.RuntimeDir != "" {
		entries, err := os.ReadDir(pc.config.RuntimeDir)
		if err != nil {
			return fmt.Errorf("read runtime dir: %w", err)
		}
		for _, entry := range entries {
			if entry.IsDir() || !strings.HasSuffix(entry.Name(), ".mjs") {
				continue
			}
			data, err := os.ReadFile(filepath.Join(pc.config.RuntimeDir, entry.Name()))
			if err != nil {
				return err
			}
			if err := os.WriteFile(filepath.Join(rtDir, entry.Name()), data, 0o644); err != nil {
				return err
			}
			pc.manifest = append(pc.manifest, "$runtime/"+entry.Name())
		}
	}
	return nil
}

func (pc *ProgramCompiler) writeManifest() error {
	sort.Strings(pc.manifest)
	// JS module version.
	var b strings.Builder
	b.WriteString("// MoxieJS Manifest — all output files.\n")
	b.WriteString("// Generated by moxiejs. Do not edit.\n\n")
	b.WriteString("export const files = [\n")
	for _, f := range pc.manifest {
		b.WriteString("  ")
		b.WriteString(JsString(f))
		b.WriteString(",\n")
	}
	b.WriteString("];\n")
	if err := os.WriteFile(filepath.Join(pc.config.OutputDir, "$manifest.mjs"), []byte(b.String()), 0o644); err != nil {
		return err
	}
	// JSON version for runtime consumption.
	data, err := json.Marshal(pc.manifest)
	if err != nil {
		return err
	}
	return os.WriteFile(filepath.Join(pc.config.OutputDir, "$manifest.json"), data, 0o644)
}

func (pc *ProgramCompiler) runtimeImportPath(pkgPath string) string {
	return "./$runtime/index.mjs"
}

// computeAsyncFuncs determines which functions need to be async.
// A function is async if it directly contains channel/select/go operations,
// or if it calls any function that is async (transitive closure).
func (pc *ProgramCompiler) computeAsyncFuncs(packages []PackageInfo) {
	// Phase 1: find direct async roots and build call graph.
	callers := make(map[*ssa.Function][]*ssa.Function) // callee -> list of callers
	var allFuncs []*ssa.Function

	for _, pkgInfo := range packages {
		ssaPkg := pc.program.Package(pkgInfo.TypePkg)
		if ssaPkg == nil {
			continue
		}
		for _, mem := range ssaPkg.Members {
			fn, ok := mem.(*ssa.Function)
			if ok {
				allFuncs = append(allFuncs, fn)
				// Also collect anonymous functions.
				allFuncs = append(allFuncs, fn.AnonFuncs...)
			}
			if t, ok := mem.(*ssa.Type); ok {
				methods := getAllMethods(ssaPkg.Prog, t.Type())
				methods = append(methods, getAllMethods(ssaPkg.Prog, types.NewPointer(t.Type()))...)
				for _, m := range methods {
					mfn := ssaPkg.Prog.MethodValue(m)
					if mfn != nil && mfn.Blocks != nil {
						allFuncs = append(allFuncs, mfn)
						allFuncs = append(allFuncs, mfn.AnonFuncs...)
					}
				}
			}
		}
	}

	// Mark direct async roots.
	for _, fn := range allFuncs {
		if isDirectAsync(fn) {
			pc.asyncFuncs[fn] = true
		}
	}

	// Build caller graph.
	for _, fn := range allFuncs {
		if fn.Blocks == nil {
			continue
		}
		for _, block := range fn.Blocks {
			for _, instr := range block.Instrs {
				if call, ok := instr.(*ssa.Call); ok {
					if callee := call.Common().StaticCallee(); callee != nil {
						callers[callee] = append(callers[callee], fn)
					}
				}
			}
		}
	}

	// Phase 2: propagate — if a callee is async, all callers become async.
	changed := true
	for changed {
		changed = false
		for fn := range pc.asyncFuncs {
			for _, caller := range callers[fn] {
				if !pc.asyncFuncs[caller] {
					pc.asyncFuncs[caller] = true
					changed = true
				}
			}
		}
	}
}

// isDirectAsync checks if a function directly contains channel/select/go ops.
func isDirectAsync(fn *ssa.Function) bool {
	if fn.Blocks == nil {
		return false
	}
	for _, block := range fn.Blocks {
		for _, instr := range block.Instrs {
			switch instr.(type) {
			case *ssa.Send, *ssa.Select, *ssa.Go:
				return true
			}
			if unop, ok := instr.(*ssa.UnOp); ok && unop.Op == token.ARROW {
				return true
			}
		}
	}
	return false
}

// --- Helpers ---

func shouldSkipPackage(path string) bool {
	skip := []string{
		"runtime",
		"runtime/internal",
		"internal/task",
		"internal/abi",
		"internal/reflectlite",
		"unsafe",
		"internal/goarch",
		"internal/goos",
	}
	for _, s := range skip {
		if path == s || strings.HasPrefix(path, s+"/") {
			return true
		}
	}
	return false
}

func functionJsName(fn *ssa.Function) string {
	name := fn.Name()
	if fn.Signature.Recv() != nil {
		recv := fn.Signature.Recv().Type()
		if ptr, ok := recv.(*types.Pointer); ok {
			recv = ptr.Elem()
		}
		typeName := "unknown"
		if named, ok := recv.(*types.Named); ok {
			typeName = named.Obj().Name()
		}
		name = typeName + "$" + name
	}
	return JsIdentifier(name)
}

func methodJsName(fn *ssa.Function) string {
	return functionJsName(fn)
}

func functionParams(fn *ssa.Function) string {
	var params []string
	// Free variables come first (bound via .bind() in MakeClosure).
	for _, fv := range fn.FreeVars {
		params = append(params, JsIdentifier(fv.Name()))
	}
	for _, p := range fn.Params {
		params = append(params, JsIdentifier(p.Name()))
	}
	return strings.Join(params, ", ")
}

func collectImports(fn *ssa.Function, imports map[string]bool) {
	if fn.Blocks == nil {
		return
	}
	for _, block := range fn.Blocks {
		for _, instr := range block.Instrs {
			if call, ok := instr.(*ssa.Call); ok {
				if callee := call.Common().StaticCallee(); callee != nil {
					if pkg := callee.Package(); pkg != nil {
						imports[pkg.Pkg.Path()] = true
					}
				}
			}
		}
	}
	// Recurse into anonymous functions (closures).
	for _, anon := range fn.AnonFuncs {
		collectImports(anon, imports)
	}
}

func getAllMethods(prog *ssa.Program, t types.Type) []*types.Selection {
	mset := prog.MethodSets.MethodSet(t)
	var methods []*types.Selection
	for i := 0; i < mset.Len(); i++ {
		methods = append(methods, mset.At(i))
	}
	return methods
}