package grammar

import (
	"log"
	"math/rand/v2"
	"sort"

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

// BuildGrammarLattice creates a lattice with grammar-shaped topology.
//
// Each tag gets counts[tag] nodes. Within each tag group, nodes are
// connected in a ring (preserving locality). Between groups, connections
// are made according to the grammar's adjacency rules: a node with tag A
// is connected to nodes with tags that Grammar.CanNeighbor(A, B) permits.
//
// The instanceSeed provides per-instance variation: different seeds produce
// different selections of which grammar-permitted connections are made.
// Same grammar rules, different topological realization. This is what
// differentiates colony instances — the grammar defines the rigid backbone,
// the seed selects the specific innervation.
func BuildGrammarLattice(
	g *Grammar,
	counts map[string]int,
	instanceSeed [32]byte,
	constraintFactory func(string) axiom.Constraint,
) *lattice.Lattice {
	l := lattice.New()

	// Sort tags for deterministic node creation order.
	tags := make([]string, 0, len(counts))
	for tag := range counts {
		if counts[tag] > 0 {
			tags = append(tags, tag)
		}
	}
	sort.Strings(tags)

	// Create nodes grouped by tag.
	type tagGroup struct {
		tag   string
		nodes []*lattice.Node
	}
	groups := make([]tagGroup, 0, len(tags))
	groupIndex := make(map[string]int) // tag -> index in groups

	for _, tag := range tags {
		n := counts[tag]
		tg := tagGroup{tag: tag, nodes: make([]*lattice.Node, n)}
		for i := range n {
			node := l.AddNode([]axiom.Constraint{constraintFactory(tag)})
			node.SetEnergy(true)
			tg.nodes[i] = node
		}
		groupIndex[tag] = len(groups)
		groups = append(groups, tg)
	}

	// Intra-group connectivity: ring within each tag group.
	for _, tg := range groups {
		if len(tg.nodes) < 2 {
			continue
		}
		for i := range tg.nodes {
			l.Connect(tg.nodes[i], tg.nodes[(i+1)%len(tg.nodes)])
		}
	}

	// Inter-group connectivity: grammar-shaped cross-connections.
	// Seed a deterministic PRNG from the instance seed.
	var seed [32]byte
	copy(seed[:], instanceSeed[:])
	rng := rand.New(rand.NewChaCha8(seed))

	for i, tgA := range groups {
		for j := i + 1; j < len(groups); j++ {
			tgB := groups[j]

			// Check if grammar permits this pair (either direction).
			canAB := g.CanNeighbor(tgA.tag, tgB.tag)
			canBA := g.CanNeighbor(tgB.tag, tgA.tag)
			if !canAB && !canBA {
				continue
			}

			// Number of cross-connections: proportional to the smaller
			// group, with a minimum of 1. The factor (1/3) creates
			// sparse but meaningful bridging.
			smaller := min(len(tgA.nodes), len(tgB.nodes))
			nBridges := max(1, smaller/3)

			// Select which nodes to bridge using the seeded PRNG.
			// Different seeds select different bridge nodes —
			// same grammar shape, different wiring realization.
			for range nBridges {
				idxA := rng.IntN(len(tgA.nodes))
				idxB := rng.IntN(len(tgB.nodes))
				l.Connect(tgA.nodes[idxA], tgB.nodes[idxB])
			}
		}
	}

	return l
}

// ExpandLattice adds growBy nodes to an existing lattice, preserving the
// grammar-shaped topology. New nodes are distributed across tags using the
// same proportions as MorphemeDefaultCounts, wired into intra-group rings
// and inter-group bridges. Existing nodes and bonds are untouched.
//
// Returns the number of nodes added.
func ExpandLattice(
	l *lattice.Lattice,
	g *Grammar,
	growBy int,
	seed [32]byte,
	constraintFactory func(string) axiom.Constraint,
	defaultCounts func(int) map[string]int,
) int {
	if growBy <= 0 {
		return 0
	}

	// Get proportional counts for the new batch.
	counts := defaultCounts(growBy)

	// Sort tags for deterministic order.
	tags := make([]string, 0, len(counts))
	for tag := range counts {
		if counts[tag] > 0 {
			tags = append(tags, tag)
		}
	}
	sort.Strings(tags)

	// Collect existing nodes by tag for bridge-wiring.
	existingByTag := make(map[string][]lattice.NodeID)
	for _, n := range l.Nodes() {
		for _, c := range n.Constraints() {
			existingByTag[c.Tag()] = append(existingByTag[c.Tag()], n.ID())
		}
	}

	// Create new nodes grouped by tag.
	type tagGroup struct {
		tag   string
		nodes []*lattice.Node
	}
	groups := make([]tagGroup, 0, len(tags))

	oldSize := l.Size()
	for _, tag := range tags {
		n := counts[tag]
		tg := tagGroup{tag: tag, nodes: make([]*lattice.Node, n)}
		for i := range n {
			node := l.AddNode([]axiom.Constraint{constraintFactory(tag)})
			node.SetEnergy(true)
			tg.nodes[i] = node
		}
		groups = append(groups, tg)
	}

	// Intra-group: ring within each new group.
	for _, tg := range groups {
		if len(tg.nodes) < 2 {
			continue
		}
		for i := range tg.nodes {
			l.Connect(tg.nodes[i], tg.nodes[(i+1)%len(tg.nodes)])
		}
	}

	// Bridge new nodes to existing lattice: connect each new group to
	// existing nodes of every permitted neighbor tag.
	rng := rand.New(rand.NewChaCha8(seed))

	for _, tgNew := range groups {
		for _, rule := range g.Rules {
			if !g.CanNeighbor(tgNew.tag, rule.Tag) {
				continue
			}
			existing := existingByTag[rule.Tag]
			if len(existing) == 0 {
				continue
			}
			// Connect sqrt(new) bridges to existing nodes of this tag.
			nBridges := max(1, len(tgNew.nodes)/3)
			for range nBridges {
				newNode := tgNew.nodes[rng.IntN(len(tgNew.nodes))]
				oldID := existing[rng.IntN(len(existing))]
				oldNode := l.Node(oldID)
				if oldNode != nil {
					l.Connect(newNode, oldNode)
				}
			}
		}
	}

	// Inter-group bridges among the new groups themselves.
	for i, tgA := range groups {
		for j := i + 1; j < len(groups); j++ {
			tgB := groups[j]
			if !g.CanNeighbor(tgA.tag, tgB.tag) && !g.CanNeighbor(tgB.tag, tgA.tag) {
				continue
			}
			smaller := min(len(tgA.nodes), len(tgB.nodes))
			nBridges := max(1, smaller/3)
			for range nBridges {
				l.Connect(
					tgA.nodes[rng.IntN(len(tgA.nodes))],
					tgB.nodes[rng.IntN(len(tgB.nodes))],
				)
			}
		}
	}

	added := l.Size() - oldSize
	log.Printf("expand: added %d nodes (%d -> %d)", added, oldSize, l.Size())
	return added
}