package crypto

import (
	"git.mleku.dev/mleku/dendrite/pkg/axiom"
	"git.mleku.dev/mleku/dendrite/pkg/dissolve"
	"git.mleku.dev/mleku/dendrite/pkg/lattice"
	"git.mleku.dev/mleku/dendrite/pkg/ratio"
)

// SiteMark records one lattice site's state in a ciphertext or sample.
type SiteMark struct {
	Index      uint64      // position in the lattice
	Occupied   bool        // element bonded here?
	TypeTag    string      // element type tag (if occupied)
	ValueHash  Hamadryad   // Hamadryad hash of element value (if occupied)
	Projection uint8       // 8-bit: [age(2)|key(3)|vertex(3)]
	ProjPath   uint16      // rendering path index (token stream position)
	Age        uint8       // 2-bit ADSR phase (0=Attack, 1=Decay, 2=Sustain, 3=Release)
	Perm       uint8       // S_3 permutation index
	LockIn     ratio.Ratio // bond strength
}

// NoiseSample records a dissolution event during encryption.
type NoiseSample struct {
	Index  uint64      // which site dissolved
	TypeTag string     // what was there
	LockIn ratio.Ratio // how strong the bond was before dissolution
}

// GaussianSampler produces discrete Gaussian samples using the
// lattice's own dissolution dynamics. Elements that survive
// dissolution are samples: their probability of survival at
// a site is governed by ContextualLockIn relative to the threshold.
type GaussianSampler struct {
	Lattice   *lattice.Lattice
	Width     ratio.Ratio // sigma parameter (controls spread)
	Threshold ratio.Ratio // dissolution cutoff = sampling boundary
}

// NewSampler creates a sampler from an existing lattice.
// The threshold determines the dissolution cutoff: elements with
// ContextualLockIn below threshold are dissolved (become noise).
func NewSampler(l *lattice.Lattice, threshold ratio.Ratio) *GaussianSampler {
	return &GaussianSampler{
		Lattice:   l,
		Width:     threshold,
		Threshold: threshold,
	}
}

// Sample runs one dissolution scan and returns the surviving
// bonding pattern as a sample from the discrete Gaussian.
func (gs *GaussianSampler) Sample() []SiteMark {
	// Run dissolution — elements below threshold are removed.
	dissolved := make(chan axiom.Element, gs.Lattice.Size())
	events := make(chan dissolve.Event, gs.Lattice.Size())

	cfg := dissolve.Config{
		Threshold: gs.Threshold,
	}
	dissolve.ScanOnce(gs.Lattice, cfg, dissolved, events)

	// Drain channels.
	close(dissolved)
	close(events)
	for range dissolved {
	}
	for range events {
	}

	// Collect surviving pattern.
	return snapshot(gs.Lattice)
}

// NoiseVector produces a noise vector for LWE encryption by
// running dissolution and recording which sites dissolved.
// Length = lattice dimension. Values are lock-in depths of dissolved sites.
func (gs *GaussianSampler) NoiseVector() ([]ratio.Ratio, []NoiseSample) {
	dissolved := make(chan axiom.Element, gs.Lattice.Size())
	events := make(chan dissolve.Event, gs.Lattice.Size())

	cfg := dissolve.Config{
		Threshold: gs.Threshold,
	}
	dissolve.ScanOnce(gs.Lattice, cfg, dissolved, events)

	close(dissolved)
	close(events)

	// Drain dissolved elements.
	for range dissolved {
	}

	// Collect noise samples from dissolution events.
	var noise []NoiseSample
	noiseVec := make([]ratio.Ratio, gs.Lattice.Size())
	for ev := range events {
		idx := uint64(ev.NodeID)
		tag := ""
		if ev.Element != nil {
			tag = ev.Element.Type()
		}
		noise = append(noise, NoiseSample{
			Index:  idx,
			TypeTag: tag,
			LockIn: ev.LockIn,
		})
		if int(idx) < len(noiseVec) {
			noiseVec[idx] = ev.LockIn
		}
	}

	return noiseVec, noise
}

// snapshot captures the current bonding pattern of the lattice.
func snapshot(l *lattice.Lattice) []SiteMark {
	nodes := l.Nodes()
	marks := make([]SiteMark, len(nodes))
	for i, n := range nodes {
		marks[i] = SiteMark{
			Index:      uint64(n.ID()),
			Occupied:   n.Occupied(),
			Projection: n.ProjectionByte(), // [age(2)|key(3)|vertex(3)]
			ProjPath:   n.ProjectionPath(),
			Age:        n.Age(),
			Perm:       n.Permutation(),
			LockIn:     n.LockIn(),
		}
		if n.Occupied() {
			occ := n.Occupant()
			marks[i].TypeTag = occ.Type()
			marks[i].ValueHash = hashValue(occ.Value())
		}
	}
	return marks
}