// Package mindsicle implements the frozen lattice — the full graph state
// serialized for persistence and bootstrapping.
//
// A mindsicle is a lossless snapshot of the lattice: every node, edge,
// occupant, candidate, constraint tag, hex state, permutation, projection,
// and age. The spore compresses this to statistics; the mindsicle preserves
// the actual crystal.
//
// Freeze captures a live lattice into a Mindsicle. Thaw reconstitutes
// the lattice from the frozen state. The bootstrapper reads a mindsicle,
// thaws it, emits Go source, compiles, and runs the result.
package mindsicle

import (
	"encoding/json"
	"fmt"
	"io"
	"log"
	"time"

	"git.mleku.dev/mleku/dendrite/pkg/axiom"
	"git.mleku.dev/mleku/dendrite/pkg/lattice"
	"git.mleku.dev/mleku/dendrite/pkg/ratio"
	"git.mleku.dev/mleku/dendrite/pkg/spore"
	"git.mleku.dev/mleku/dendrite/pkg/state"
)

// Version is the current mindsicle format version.
const Version = 1

// Mindsicle is a frozen lattice — the complete graph state serialized.
type Mindsicle struct {
	Spore    *spore.Spore `json:"spore"`
	Nodes    []NodeRecord `json:"nodes"`
	Version  int          `json:"version"`
	FrozenAt time.Time    `json:"frozen_at"`
}

// NodeRecord is the serialized form of a single lattice node.
type NodeRecord struct {
	ID         uint64          `json:"id"`
	Tags       []string        `json:"constraints"`
	Occupant   *ElementRecord  `json:"occupant,omitempty"`
	Candidates []ElementRecord `json:"candidates,omitempty"`
	Neighbors  []uint64        `json:"neighbors"`
	Hex        uint8           `json:"hex"`
	LockIn     ratio.Ratio     `json:"lock_in"`
	BondCount  int             `json:"bond_count"`
	Perm       uint8           `json:"perm"`
	ProjVertex uint8           `json:"proj_vertex"`
	ProjKey    uint8           `json:"proj_key"`
	ProjPath   uint16          `json:"proj_path"`
	Age        uint8           `json:"age"`
}

// ElementRecord is the serialized form of an element.
type ElementRecord struct {
	Type  string `json:"type"`
	Value string `json:"value"`
}

// frozenElement implements axiom.Element for deserialized elements.
type frozenElement struct {
	typ string
	val string
}

func (e frozenElement) Type() string { return e.typ }
func (e frozenElement) Value() any   { return e.val }

// Freeze captures a live lattice into a Mindsicle.
// If sp is nil, no spore is stored (Thaw doesn't need it).
func Freeze(l *lattice.Lattice, sp *spore.Spore) *Mindsicle {
	nodes := l.Nodes()
	records := make([]NodeRecord, len(nodes))

	for i, n := range nodes {
		rec := NodeRecord{
			ID:         uint64(n.ID()),
			Hex:        uint8(n.Hexagram()),
			LockIn:     n.LockIn(),
			BondCount:  n.BondCount(),
			Perm:       n.Permutation(),
			ProjVertex: n.ProjectionVertex(),
			ProjKey:    n.ProjectionKey(),
			ProjPath:   n.ProjectionPath(),
			Age:        n.Age(),
		}

		// Constraint tags.
		for _, c := range n.Constraints() {
			rec.Tags = append(rec.Tags, c.Tag())
		}

		// Occupant.
		if n.Occupied() {
			e := n.Occupant()
			val := ""
			if e.Value() != nil {
				val = fmt.Sprintf("%v", e.Value())
			}
			rec.Occupant = &ElementRecord{
				Type:  e.Type(),
				Value: val,
			}
		}

		// Candidates.
		for _, c := range n.Candidates() {
			val := ""
			if c.Value() != nil {
				val = fmt.Sprintf("%v", c.Value())
			}
			rec.Candidates = append(rec.Candidates, ElementRecord{
				Type:  c.Type(),
				Value: val,
			})
		}

		// Neighbors (as IDs).
		for _, nb := range n.Neighbors() {
			rec.Neighbors = append(rec.Neighbors, uint64(nb.ID()))
		}

		records[i] = rec
	}

	return &Mindsicle{
		Spore:    sp,
		Nodes:    records,
		Version:  Version,
		FrozenAt: time.Now(),
	}
}

// Thaw reconstitutes a live lattice from a frozen Mindsicle.
//
// constraintFactory converts tag strings back into Constraint instances.
// The thaw proceeds in four phases:
//  1. Create nodes with constraints
//  2. Connect neighbors (dedup bidirectional pairs)
//  3. ForceOccupant + candidates
//  4. Restore age, permutation, projection, hex energy/outer
func (m *Mindsicle) Thaw(constraintFactory func(string) axiom.Constraint) *lattice.Lattice {
	l := lattice.New()

	if len(m.Nodes) == 0 {
		return l
	}

	// Phase 1: Create all nodes with their constraints.
	nodeByID := make(map[uint64]*lattice.Node, len(m.Nodes))
	for _, rec := range m.Nodes {
		constraints := make([]axiom.Constraint, len(rec.Tags))
		for j, tag := range rec.Tags {
			constraints[j] = constraintFactory(tag)
		}
		n := l.AddNode(constraints)
		nodeByID[uint64(n.ID())] = n
	}

	// Phase 2: Connect neighbors. Only connect when nbID > rec.ID to
	// deduplicate bidirectional edges without any map allocation.
	for _, rec := range m.Nodes {
		for _, nbID := range rec.Neighbors {
			if nbID > rec.ID {
				na, nb := nodeByID[rec.ID], nodeByID[nbID]
				if na != nil && nb != nil {
					l.Connect(na, nb)
				}
			}
		}
	}

	// Phase 3: Place occupants and candidates. Uses ForceOccupant to
	// bypass constraint checking — the frozen state is pre-validated.
	for _, rec := range m.Nodes {
		n := nodeByID[rec.ID]
		if n == nil {
			continue
		}

		if rec.Occupant != nil {
			elem := frozenElement{typ: rec.Occupant.Type, val: rec.Occupant.Value}
			n.ForceOccupant(elem, rec.BondCount)
		}

		for _, c := range rec.Candidates {
			elem := frozenElement{typ: c.Type, val: c.Value}
			n.AddCandidate(elem)
		}
	}

	// Phase 4: Restore age, permutation, projection, hex state, lock-in.
	for _, rec := range m.Nodes {
		n := nodeByID[rec.ID]
		if n == nil {
			continue
		}

		n.RestoreAge(rec.Age)
		n.SetPermutation(rec.Perm)
		n.SetProjection(rec.ProjVertex, rec.ProjKey, rec.ProjPath)

		// Restore energy bit and outer trigram from the frozen hexagram.
		hex := state.Hexagram(rec.Hex)
		n.SetEnergy(hex.Inner().Energy())
		n.SetOuterTrigram(hex.Outer())

		// Restore the exact trained lock-in ratio. ForceOccupant sets a
		// simple bondCount/1, but the frozen ratio preserves the real value.
		n.RestoreLockIn(rec.LockIn)
	}

	return l
}

// WriteTo serializes the mindsicle to JSON.
func (m *Mindsicle) WriteTo(w io.Writer) (int64, error) {
	data, err := json.MarshalIndent(m, "", "  ")
	if err != nil {
		return 0, err
	}
	n, err := w.Write(data)
	return int64(n), err
}

// ReadMindsicle deserializes a mindsicle from JSON.
func ReadMindsicle(r io.Reader) (*Mindsicle, error) {
	data, err := io.ReadAll(r)
	if err != nil {
		return nil, err
	}
	var m Mindsicle
	if err := json.Unmarshal(data, &m); err != nil {
		return nil, err
	}
	return &m, nil
}

// StreamThawJSON reads a mindsicle JSON file using a streaming decoder,
// processing one NodeRecord at a time. This avoids the double allocation
// of io.ReadAll + json.Unmarshal that makes the non-streaming path OOM
// on large lattices (~4GB JSON → ~20GB peak with full materialization).
//
// Memory profile: ~1.6GB buffer (neighbor IDs + node state for deferred
// phases) + lattice itself. The 3.9GB JSON is never fully in memory.
func StreamThawJSON(r io.Reader, cf func(string) axiom.Constraint) (*lattice.Lattice, error) {
	dec := json.NewDecoder(r)

	// Read opening '{'.
	if _, err := dec.Token(); err != nil {
		return nil, fmt.Errorf("expected opening brace: %w", err)
	}

	// Compact per-node state buffer for deferred phases.
	type nodeState struct {
		neighborIDs []uint64
		lockIn      ratio.Ratio
		hex         uint8
		perm        uint8
		projVertex  uint8
		projKey     uint8
		projPath    uint16
		age         uint8
		bondCount   int
		occType     string
		occValue    string
		candidates  []ElementRecord
	}

	l := lattice.New()
	var states []nodeState

	// Stream through top-level keys.
	for dec.More() {
		tok, err := dec.Token()
		if err != nil {
			return nil, fmt.Errorf("token: %w", err)
		}
		key, ok := tok.(string)
		if !ok {
			continue
		}

		switch key {
		case "nodes":
			// Read opening '[' of nodes array.
			if _, err := dec.Token(); err != nil {
				return nil, fmt.Errorf("nodes array open: %w", err)
			}

			// Phase 1: Decode each NodeRecord one at a time.
			nodeIdx := 0
			for dec.More() {
				var rec NodeRecord
				if err := dec.Decode(&rec); err != nil {
					return nil, fmt.Errorf("node %d decode: %w", nodeIdx, err)
				}

				// Create node with constraints.
				constraints := make([]axiom.Constraint, len(rec.Tags))
				for j, tag := range rec.Tags {
					constraints[j] = cf(tag)
				}
				l.AddNode(constraints)

				// Buffer state for phases 2+3.
				st := nodeState{
					neighborIDs: rec.Neighbors,
					lockIn:      rec.LockIn,
					hex:         rec.Hex,
					perm:        rec.Perm,
					projVertex:  rec.ProjVertex,
					projKey:     rec.ProjKey,
					projPath:    rec.ProjPath,
					age:         rec.Age,
					bondCount:   rec.BondCount,
					candidates:  rec.Candidates,
				}
				if rec.Occupant != nil {
					st.occType = rec.Occupant.Type
					st.occValue = rec.Occupant.Value
				}
				states = append(states, st)

				nodeIdx++
				if nodeIdx%1000000 == 0 {
					log.Printf("  stream-thaw: %dM nodes read", nodeIdx/1000000)
				}
			}

			// Read closing ']'.
			if _, err := dec.Token(); err != nil {
				return nil, fmt.Errorf("nodes array close: %w", err)
			}

		default:
			// Skip non-nodes fields (spore, version, frozen_at).
			var discard json.RawMessage
			if err := dec.Decode(&discard); err != nil {
				return nil, fmt.Errorf("skip %s: %w", key, err)
			}
		}
	}

	// Read closing '}'.
	dec.Token()

	allNodes := l.Nodes()
	log.Printf("  stream-thaw: phase 1 complete, %d nodes created", len(allNodes))

	// Phase 2: Connect neighbors + restore state in a single pass.
	// Edges are bidirectional — connect only when nbID > myID to
	// deduplicate without any map. O(E) time, O(1) extra space.
	// Each state entry is zeroed after use to allow progressive GC.
	for i := range states {
		st := &states[i]
		n := allNodes[i]
		myID := uint64(i)

		// Connect neighbors.
		for _, nbID := range st.neighborIDs {
			if nbID > myID && int(nbID) < len(allNodes) {
				l.Connect(allNodes[myID], allNodes[nbID])
			}
		}

		// Restore occupant.
		if st.occType != "" {
			elem := frozenElement{typ: st.occType, val: st.occValue}
			n.ForceOccupant(elem, st.bondCount)
		}

		// Restore candidates.
		for _, c := range st.candidates {
			elem := frozenElement{typ: c.Type, val: c.Value}
			n.AddCandidate(elem)
		}

		// Restore projection, hex, lock-in.
		n.RestoreAge(st.age)
		n.SetPermutation(st.perm)
		n.SetProjection(st.projVertex, st.projKey, st.projPath)

		hex := state.Hexagram(st.hex)
		n.SetEnergy(hex.Inner().Energy())
		n.SetOuterTrigram(hex.Outer())

		n.RestoreLockIn(st.lockIn)

		// Zero entry to allow GC of strings and slices.
		states[i] = nodeState{}

		if (i+1)%1000000 == 0 {
			log.Printf("  stream-thaw: %dM nodes connected+restored", (i+1)/1000000)
		}
	}
	states = nil // release entire buffer
	log.Printf("  stream-thaw: phase 2 complete, neighbors connected + state restored")

	return l, nil
}