package iskra

import (
	"bytes"
	"slices"
)

type Symbol struct {
	Name  string
	Depth int32
	Kind  uint8
}

const (
	SymDecl  uint8 = 1
	SymRef   uint8 = 2
	SymParam uint8 = 3
)

type SymbolTable struct {
	Symbols     []Symbol
	Counts      map[string]int32
	ParamTypes  []string
	ResultTypes []string
	RecvType    string
	RefSeqHash  uint64
}

// Binding represents a single variable's geometric footprint:
// where it was declared, and the depth of each reference to it.
type Binding struct {
	DeclDepth int32
	RefDepths []int32
}

// BindingGraph captures the full binding structure of a function.
// Each declaration creates a binding; each reference resolves to
// the nearest enclosing declaration of that name.
type BindingGraph struct {
	Bindings []Binding
	// Externals are symbols referenced but never declared in this scope
	// (package-level constants, functions from other packages, etc).
	Externals []ExternalRef
}

type ExternalRef struct {
	Name     string
	RefCount int32
	Depths   []int32
}

// BindingSignature is the geometric fingerprint with type tags.
// Two functions are isomorphic if their signatures match.
type BindingSignature struct {
	Locals      []LocalSig
	Externs     []ExternSig
	ParamTypes  []string
	ResultTypes []string
	RecvType    string
	ExternNames []string
	RefSeqHash  uint64
}

type LocalSig struct {
	DeclDepth int32
	RefDepths []int32
}

type ExternSig struct {
	RefCount int32
	Depths   []int32
}

func ExtractSymbols(astDump string) *SymbolTable {
	st := &SymbolTable{
		Counts: map[string]int32{},
	}
	lines := bytes.Split([]byte(astDump), []byte("\n"))
	inParams := false
	inResults := false
	sectionDepth := int32(-1)
	paramOrd := map[string]int32{}
	resultOrd := map[string]int32{}
	localOrd := map[string]int32{}
	var refSeq []byte

	for _, line := range lines {
		if len(line) == 0 {
			continue
		}
		d := countIndent(line)
		trimmed := bytes.TrimSpace(line)

		if (inParams || inResults) && d <= sectionDepth {
			inParams = false
			inResults = false
		}

		if bytes.HasPrefix(trimmed, []byte("FuncDecl ")) {
			name := string(trimmed[9:])
			if idx := bytes.IndexByte(trimmed[9:], ' '); idx >= 0 {
				name = string(trimmed[9 : 9+idx])
			}
			addSym(st, name, d, SymDecl)
			if ri := bytes.Index(trimmed, []byte("recv=")); ri >= 0 {
				rt := trimmed[ri+5:]
				if si := bytes.IndexByte(rt, ' '); si >= 0 {
					rt = rt[:si]
				}
				st.RecvType = string(rt)
			}
		} else if bytes.HasPrefix(trimmed, []byte("Params")) {
			inParams = true
			inResults = false
			sectionDepth = d
		} else if bytes.HasPrefix(trimmed, []byte("Results")) {
			inResults = true
			inParams = false
			sectionDepth = d
		} else if inParams && d > sectionDepth {
			typeName := extractTypeFromParamLine(trimmed)
			if typeName != "" {
				st.ParamTypes = append(st.ParamTypes, typeName)
			}
			names := extractParamNames(trimmed)
			for _, n := range names {
				paramOrd[n] = int32(len(paramOrd))
			}
			parseParamLine(st, trimmed, d)
		} else if inResults && d > sectionDepth {
			typeName := extractTypeFromParamLine(trimmed)
			if typeName != "" {
				st.ResultTypes = append(st.ResultTypes, typeName)
			}
			if isParamLine(trimmed) {
				names := extractParamNames(trimmed)
				for _, n := range names {
					resultOrd[n] = int32(len(resultOrd))
				}
				parseParamLine(st, trimmed, d)
			}
		} else if bytes.HasPrefix(trimmed, []byte("Assign ")) {
			rest := trimmed[7:]
			spIdx := bytes.IndexByte(rest, ' ')
			if spIdx > 0 {
				lhs := rest[:spIdx]
				parts := bytes.Split(lhs, []byte(","))
				for _, p := range parts {
					p = bytes.TrimSpace(p)
					nm := string(p)
					if len(nm) > 0 {
						if _, ok := localOrd[nm]; !ok {
							if _, ok2 := paramOrd[nm]; !ok2 {
								if _, ok3 := resultOrd[nm]; !ok3 {
									localOrd[nm] = int32(len(localOrd))
								}
							}
						}
					}
				}
			}
			refSeq = appendRefFingerprint(refSeq, trimmed, paramOrd, resultOrd, localOrd)
			parseAssign(st, trimmed, d)
		} else if bytes.HasPrefix(trimmed, []byte("Return")) {
			refSeq = appendRefFingerprint(refSeq, trimmed, paramOrd, resultOrd, localOrd)
			parseRefs(st, trimmed, d)
		} else if bytes.HasPrefix(trimmed, []byte("If")) {
			refSeq = appendRefFingerprint(refSeq, trimmed, paramOrd, resultOrd, localOrd)
			parseRefs(st, trimmed, d)
		} else if bytes.HasPrefix(trimmed, []byte("For")) {
			refSeq = appendRefFingerprint(refSeq, trimmed, paramOrd, resultOrd, localOrd)
			parseRefs(st, trimmed, d)
		} else if bytes.HasPrefix(trimmed, []byte("Range ")) {
			refSeq = appendRefFingerprint(refSeq, trimmed, paramOrd, resultOrd, localOrd)
			parseRefs(st, trimmed, d)
		} else if bytes.HasPrefix(trimmed, []byte("Expr ")) {
			refSeq = appendRefFingerprint(refSeq, trimmed, paramOrd, resultOrd, localOrd)
			parseRefs(st, trimmed, d)
		} else if bytes.HasPrefix(trimmed, []byte("Case")) {
			refSeq = appendRefFingerprint(refSeq, trimmed, paramOrd, resultOrd, localOrd)
			parseRefs(st, trimmed, d)
		} else if bytes.HasPrefix(trimmed, []byte("Switch")) {
			refSeq = appendRefFingerprint(refSeq, trimmed, paramOrd, resultOrd, localOrd)
			parseRefs(st, trimmed, d)
		} else if bytes.HasPrefix(trimmed, []byte("Value ")) {
			parseParamLine(st, trimmed[6:], d)
		} else if isParamLine(trimmed) {
			parseParamLine(st, trimmed, d)
		}
	}
	if len(refSeq) > 0 && !allSeparators(refSeq) {
		st.RefSeqHash = fnvHash64(refSeq)
	}
	return st
}

func allSeparators(b []byte) bool {
	for _, c := range b {
		if c != '|' {
			return false
		}
	}
	return true
}

func extractParamNames(line []byte) []string {
	idx := bytes.IndexByte(line, ' ')
	if idx <= 0 {
		return nil
	}
	namesPart := line[:idx]
	parts := bytes.Split(namesPart, []byte(","))
	var names []string
	for _, p := range parts {
		p = bytes.TrimSpace(p)
		if len(p) > 0 && p[0] >= 'a' && p[0] <= 'z' {
			names = append(names, string(p))
		}
	}
	return names
}

func appendRefFingerprint(seq, line []byte, paramOrd, resultOrd, localOrd map[string]int32) []byte {
	start := bytes.IndexByte(line, '[')
	if start < 0 {
		seq = append(seq, '|')
		return seq
	}
	end := bytes.IndexByte(line[start:], ']')
	if end < 0 {
		seq = append(seq, '|')
		return seq
	}
	refs := line[start+1 : start+end]
	parts := bytes.Split(refs, []byte(","))
	for i, p := range parts {
		p = bytes.TrimSpace(p)
		if i > 0 {
			seq = append(seq, ',')
		}
		name := string(p)
		if ord, ok := paramOrd[name]; ok {
			seq = append(seq, 'P')
			seq = append(seq, []byte(intToStr(int(ord)))...)
		} else if ord, ok := resultOrd[name]; ok {
			seq = append(seq, 'R')
			seq = append(seq, []byte(intToStr(int(ord)))...)
		} else if ord, ok := localOrd[name]; ok {
			seq = append(seq, 'L')
			seq = append(seq, []byte(intToStr(int(ord)))...)
		} else {
			seq = append(seq, p...)
		}
	}
	seq = append(seq, '|')
	return seq
}

func fnvHash64(data []byte) uint64 {
	h := uint64(0xcbf29ce484222325)
	for _, b := range data {
		h ^= uint64(b)
		h *= 0x100000001b3
	}
	return h
}

func extractTypeFromParamLine(line []byte) string {
	// Param: "name type" or "name,name type" or result: "name type" or just "type"
	// e.g. "r rune", "s,prefix []byte", "f float64", "bool", "error"
	idx := bytes.IndexByte(line, ' ')
	if idx < 0 {
		return string(line) // no space = bare type (result)
	}
	// Check if what's before the space could be a type name (starts with [, *, map, etc.)
	before := line[:idx]
	if before[0] == '[' || before[0] == '*' || bytes.HasPrefix(before, []byte("map")) || bytes.HasPrefix(before, []byte("chan")) || bytes.HasPrefix(before, []byte("func")) {
		return string(line) // entire thing is a type
	}
	return string(line[idx+1:])
}

// ExtractBindings builds a BindingGraph from a SymbolTable.
// It resolves each reference to the nearest prior declaration of
// the same name at equal or lesser depth (lexical scoping).
func ExtractBindings(st *SymbolTable) *BindingGraph {
	bg := &BindingGraph{}

	// Index: name → binding index (for locals).
	localIdx := map[string]int32{}
	// Track external refs by name.
	extIdx := map[string]int32{}

	// First pass: register all declarations.
	for _, sym := range st.Symbols {
		if sym.Kind == SymDecl || sym.Kind == SymParam {
			idx := int32(len(bg.Bindings))
			bg.Bindings = append(bg.Bindings, Binding{
				DeclDepth: sym.Depth,
			})
			localIdx[sym.Name] = idx
		}
	}

	// Second pass: resolve references.
	for _, sym := range st.Symbols {
		if sym.Kind != SymRef {
			continue
		}
		if idx, ok := localIdx[sym.Name]; ok {
			bg.Bindings[idx].RefDepths = append(bg.Bindings[idx].RefDepths, sym.Depth)
		} else {
			if ei, ok := extIdx[sym.Name]; ok {
				bg.Externals[ei].RefCount++
				bg.Externals[ei].Depths = append(bg.Externals[ei].Depths, sym.Depth)
			} else {
				extIdx[sym.Name] = int32(len(bg.Externals))
				bg.Externals = append(bg.Externals, ExternalRef{
					Name:     sym.Name,
					RefCount: 1,
					Depths:   []int32{sym.Depth},
				})
			}
		}
	}

	return bg
}

// Signature erases names and produces a sorted geometric fingerprint.
func (bg *BindingGraph) Signature() *BindingSignature {
	sig := &BindingSignature{}

	for _, b := range bg.Bindings {
		depths := []int32{:0:len(b.RefDepths)}
		depths = append(depths, b.RefDepths...)
		slices.SortFunc(depths, cmpInt32)
		sig.Locals = append(sig.Locals, LocalSig{
			DeclDepth: b.DeclDepth,
			RefDepths: depths,
		})
	}

	for _, e := range bg.Externals {
		depths := []int32{:0:len(e.Depths)}
		depths = append(depths, e.Depths...)
		slices.SortFunc(depths, cmpInt32)
		sig.Externs = append(sig.Externs, ExternSig{
			RefCount: e.RefCount,
			Depths:   depths,
		})
	}

	slices.SortFunc(sig.Locals, cmpLocalSig)
	slices.SortFunc(sig.Externs, cmpExternSig)
	return sig
}

func SignatureWithTypes(st *SymbolTable) *BindingSignature {
	bg := ExtractBindings(st)
	sig := bg.Signature()
	sig.ParamTypes = append(sig.ParamTypes, st.ParamTypes...)
	sig.ResultTypes = append(sig.ResultTypes, st.ResultTypes...)
	sig.RecvType = st.RecvType
	for _, e := range bg.Externals {
		sig.ExternNames = append(sig.ExternNames, e.Name)
	}
	slices.Sort(sig.ExternNames)
	sig.RefSeqHash = st.RefSeqHash
	return sig
}

func (a *BindingSignature) IsIsomorphic(b *BindingSignature) bool {
	if len(a.Locals) != len(b.Locals) {
		return false
	}
	if len(a.Externs) != len(b.Externs) {
		return false
	}
	for i := range a.Locals {
		if a.Locals[i].DeclDepth != b.Locals[i].DeclDepth {
			return false
		}
		if !int32sEqual(a.Locals[i].RefDepths, b.Locals[i].RefDepths) {
			return false
		}
	}
	for i := range a.Externs {
		if a.Externs[i].RefCount != b.Externs[i].RefCount {
			return false
		}
		if !int32sEqual(a.Externs[i].Depths, b.Externs[i].Depths) {
			return false
		}
	}
	if !stringsEqual(a.ParamTypes, b.ParamTypes) {
		return false
	}
	if !stringsEqual(a.ResultTypes, b.ResultTypes) {
		return false
	}
	if a.RecvType != b.RecvType {
		return false
	}
	if !stringsEqual(a.ExternNames, b.ExternNames) {
		return false
	}
	if a.RefSeqHash != 0 && b.RefSeqHash != 0 && a.RefSeqHash != b.RefSeqHash {
		return false
	}
	return true
}

func stringsEqual(a, b []string) bool {
	if len(a) != len(b) {
		return false
	}
	for i := range a {
		if a[i] != b[i] {
			return false
		}
	}
	return true
}

// GeometricDistance measures how far two signatures diverge.
// Returns 0 for perfect isomorphism, higher for greater divergence.
// The score counts: missing/extra bindings + depth mismatches + ref count diffs.
func GeometricDistance(a, b *BindingSignature) int32 {
	dist := int32(0)

	// Compare locals.
	la := len(a.Locals)
	lb := len(b.Locals)
	minL := la
	if lb < minL {
		minL = lb
	}
	for i := 0; i < minL; i++ {
		if a.Locals[i].DeclDepth != b.Locals[i].DeclDepth {
			dist++
		}
		dist += int32SliceDistance(a.Locals[i].RefDepths, b.Locals[i].RefDepths)
	}
	// Penalty for unmatched bindings.
	if la > lb {
		dist += int32(la - lb)
	} else {
		dist += int32(lb - la)
	}

	// Compare externals.
	ea := len(a.Externs)
	eb := len(b.Externs)
	minE := ea
	if eb < minE {
		minE = eb
	}
	for i := 0; i < minE; i++ {
		d := a.Externs[i].RefCount - b.Externs[i].RefCount
		if d < 0 {
			d = -d
		}
		dist += d
		dist += int32SliceDistance(a.Externs[i].Depths, b.Externs[i].Depths)
	}
	if ea > eb {
		dist += int32(ea - eb)
	} else {
		dist += int32(eb - ea)
	}

	return dist
}

func int32SliceDistance(a, b []int32) int32 {
	dist := int32(0)
	la := len(a)
	lb := len(b)
	minN := la
	if lb < minN {
		minN = lb
	}
	for i := 0; i < minN; i++ {
		d := a[i] - b[i]
		if d < 0 {
			d = -d
		}
		dist += d
	}
	if la > lb {
		dist += int32(la - lb)
	} else {
		dist += int32(lb - la)
	}
	return dist
}

func int32sEqual(a, b []int32) bool {
	if len(a) != len(b) {
		return false
	}
	for i := range a {
		if a[i] != b[i] {
			return false
		}
	}
	return true
}

func cmpInt32(a, b int32) int {
	if a < b {
		return -1
	}
	if a > b {
		return 1
	}
	return 0
}

func cmpLocalSig(a, b LocalSig) int {
	if a.DeclDepth != b.DeclDepth {
		if a.DeclDepth < b.DeclDepth {
			return -1
		}
		return 1
	}
	la := int32(len(a.RefDepths))
	lb := int32(len(b.RefDepths))
	if la != lb {
		if la < lb {
			return -1
		}
		return 1
	}
	for i := int32(0); i < la; i++ {
		if a.RefDepths[i] != b.RefDepths[i] {
			if a.RefDepths[i] < b.RefDepths[i] {
				return -1
			}
			return 1
		}
	}
	return 0
}

func cmpExternSig(a, b ExternSig) int {
	if a.RefCount != b.RefCount {
		if a.RefCount < b.RefCount {
			return -1
		}
		return 1
	}
	la := int32(len(a.Depths))
	lb := int32(len(b.Depths))
	if la != lb {
		if la < lb {
			return -1
		}
		return 1
	}
	for i := int32(0); i < la; i++ {
		if a.Depths[i] != b.Depths[i] {
			if a.Depths[i] < b.Depths[i] {
				return -1
			}
			return 1
		}
	}
	return 0
}

// SignatureHash24 produces a 24-bit hash of a binding signature for fast
// rejection during signature lookup. Collisions are expected but rare enough
// (1 in 16M) to make the full IsIsomorphic check infrequent.
func SignatureHash24(sig *BindingSignature) uint32 {
	h := uint32(0x811c9dc5) // FNV-1a offset basis
	mix := func(b byte) {
		h ^= uint32(b)
		h *= 0x01000193
	}
	mix(byte(len(sig.Locals)))
	for _, l := range sig.Locals {
		mix(byte(l.DeclDepth))
		mix(byte(l.DeclDepth >> 8))
		mix(byte(len(l.RefDepths)))
		for _, d := range l.RefDepths {
			mix(byte(d))
			mix(byte(d >> 8))
		}
	}
	mix(byte(len(sig.Externs)))
	for _, e := range sig.Externs {
		mix(byte(e.RefCount))
		mix(byte(e.RefCount >> 8))
		for _, d := range e.Depths {
			mix(byte(d))
			mix(byte(d >> 8))
		}
	}
	for _, c := range []byte(sig.RecvType) {
		mix(c)
	}
	for _, name := range sig.ExternNames {
		for _, c := range []byte(name) {
			mix(c)
		}
		mix(0xFF)
	}
	rsb := uint64ToBytes(sig.RefSeqHash)
	for _, c := range rsb {
		mix(c)
	}
	return h & 0xFFFFFF
}

func uint64ToBytes(v uint64) [8]byte {
	var b [8]byte
	b[0] = byte(v)
	b[1] = byte(v >> 8)
	b[2] = byte(v >> 16)
	b[3] = byte(v >> 24)
	b[4] = byte(v >> 32)
	b[5] = byte(v >> 40)
	b[6] = byte(v >> 48)
	b[7] = byte(v >> 56)
	return b
}

// --- Flat name-based comparison (kept for backward compat) ---

type SymbolMatch struct {
	TotalA     int32
	TotalB     int32
	Matched    int32
	OnlyInA    []string
	OnlyInB    []string
	CountDiffs []string
}

func CompareSymbols(a, b *SymbolTable) *SymbolMatch {
	m := &SymbolMatch{}
	namesA := map[string]bool{}
	namesB := map[string]bool{}
	for name := range a.Counts {
		namesA[name] = true
		m.TotalA++
	}
	for name := range b.Counts {
		namesB[name] = true
		m.TotalB++
	}
	for name := range namesA {
		if !namesB[name] {
			m.OnlyInA = append(m.OnlyInA, name)
		} else {
			m.Matched++
			if a.Counts[name] != b.Counts[name] {
				m.CountDiffs = append(m.CountDiffs, name)
			}
		}
	}
	for name := range namesB {
		if !namesA[name] {
			m.OnlyInB = append(m.OnlyInB, name)
		}
	}
	return m
}

func (m *SymbolMatch) IsEquivalent() bool {
	return len(m.OnlyInA) == 0 && len(m.OnlyInB) == 0
}

func (m *SymbolMatch) Score() int32 {
	if m.TotalA == 0 && m.TotalB == 0 {
		return 100
	}
	total := m.TotalA
	if m.TotalB > total {
		total = m.TotalB
	}
	if total == 0 {
		return 0
	}
	return (m.Matched * 100) / total
}

// --- Helpers ---

func countIndent(line []byte) int32 {
	n := int32(0)
	i := int32(0)
	for i+1 < int32(len(line)) && line[i] == ' ' && line[i+1] == ' ' {
		n++
		i += 2
	}
	return n
}

func isParamLine(trimmed []byte) bool {
	if len(trimmed) == 0 {
		return false
	}
	if trimmed[0] >= 'a' && trimmed[0] <= 'z' {
		return bytes.IndexByte(trimmed, ' ') > 0
	}
	return false
}

func parseParamLine(st *SymbolTable, trimmed []byte, depth int32) {
	idx := bytes.IndexByte(trimmed, ' ')
	if idx <= 0 {
		return
	}
	names := trimmed[:idx]
	parts := bytes.Split(names, []byte(","))
	for _, p := range parts {
		p = bytes.TrimSpace(p)
		if len(p) > 0 {
			addSym(st, string(p), depth, SymParam)
		}
	}
}

func parseAssign(st *SymbolTable, trimmed []byte, depth int32) {
	rest := trimmed[7:]
	spIdx := bytes.IndexByte(rest, ' ')
	if spIdx <= 0 {
		return
	}
	lhs := rest[:spIdx]
	parts := bytes.Split(lhs, []byte(","))
	for _, p := range parts {
		p = bytes.TrimSpace(p)
		if len(p) > 0 {
			addSym(st, string(p), depth, SymDecl)
		}
	}
	parseRefs(st, trimmed, depth)
}

func parseRefs(st *SymbolTable, trimmed []byte, depth int32) {
	start := bytes.IndexByte(trimmed, '[')
	if start < 0 {
		return
	}
	end := bytes.IndexByte(trimmed[start:], ']')
	if end < 0 {
		return
	}
	refs := trimmed[start+1 : start+end]
	parts := bytes.Split(refs, []byte(","))
	for _, p := range parts {
		p = bytes.TrimSpace(p)
		if len(p) > 0 {
			addSym(st, string(p), depth, SymRef)
		}
	}
}

func addSym(st *SymbolTable, name string, depth int32, kind uint8) {
	st.Symbols = append(st.Symbols, Symbol{Name: name, Depth: depth, Kind: kind})
	st.Counts[name] = st.Counts[name] + 1
}