package crypto

import (
	"context"
	"crypto/rand"
	"encoding/binary"
	"errors"
	"fmt"

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

// Generate produces a self-contained keypair. It builds a lattice from
// scratch, seeds it with random elements, runs dissolution to establish
// equilibrium, and extracts the keypair.
//
// tags defines the constraint type layers. factory maps each tag to
// a constraint implementation — this becomes the private key (trapdoor).
func Generate(params Params, tags []string, factory func(string) axiom.Constraint) (*KeyPair, error) {
	if !params.Valid() {
		return nil, errors.New("crypto: invalid parameters")
	}
	if len(tags) == 0 {
		return nil, errors.New("crypto: no type tags")
	}
	if factory == nil {
		return nil, errors.New("crypto: nil constraint factory")
	}

	// 1. Build lattice with N nodes distributed across tags.
	l := lattice.New()
	nodesPerTag := params.N / len(tags)
	if nodesPerTag < 1 {
		nodesPerTag = 1
	}

	type tagGroup struct {
		tag   string
		nodes []*lattice.Node
	}
	groups := make([]tagGroup, len(tags))

	for i, tag := range tags {
		groups[i].tag = tag
		for range nodesPerTag {
			n := l.AddNode([]axiom.Constraint{factory(tag)})
			n.SetEnergy(true)
			groups[i].nodes = append(groups[i].nodes, n)
		}
	}

	// 2. Ring topology within each type + cross-type bridges.
	for _, g := range groups {
		for j := range g.nodes {
			l.Connect(g.nodes[j], g.nodes[(j+1)%len(g.nodes)])
		}
	}
	for i := 0; i < len(groups); i++ {
		for j := i + 1; j < len(groups); j++ {
			a, b := groups[i].nodes, groups[j].nodes
			step := max(1, min(len(a), len(b))/3)
			for k := 0; k < min(len(a), len(b)); k += step {
				l.Connect(a[k], b[k])
			}
		}
	}

	// 3. Seed with random elements via Brownian walk growth.
	// Each seed element carries random permutation and projection data
	// so the resulting lattice has a unique structural fingerprint.
	seedCount := l.Size() / 2 // aim for ~50% occupancy
	solution := make(chan axiom.Element, seedCount)
	for i := range seedCount {
		tag := tags[i%len(tags)]
		val := randomByte()
		rb := randomByte()
		rb2 := randomByte()
		solution <- seedElement{
			tag:        tag,
			val:        val,
			perm:       rb % 6,
			projVertex: rb2 & 0x07,
			projKey:    (rb2 >> 3) & 0x07,
			projPath:   randomUint16(),
		}
	}
	close(solution)

	events := make(chan grow.Event, seedCount*2)
	ctx := context.Background()
	grow.Run(ctx, l, solution, grow.Config{
		MaxSteps: params.MaxWalkSteps,
		Workers:  4,
	}, events)
	close(events)
	for range events {
	}

	// 4. Dissolution passes to establish equilibrium.
	for range params.DissolutionPasses {
		dissolved := make(chan axiom.Element, l.Size())
		dissEvents := make(chan dissolve.Event, l.Size())
		dissolve.ScanOnce(l, dissolve.Config{
			Threshold: params.SmoothingParam,
		}, dissolved, dissEvents)
		close(dissolved)
		close(dissEvents)
		for range dissolved {
		}
		for range dissEvents {
		}
	}

	// 5. Extract keypair.
	return GenerateKeyPair(l, params, factory), nil
}

// seedElement is a random element used during key generation.
// It carries random permutation and projection data so that
// each generated keypair produces a unique structural fingerprint
// (PermDist, ProjDist) in the spore. Without these, all lattices
// built from the same tags/params would be structurally identical.
type seedElement struct {
	tag        string
	val        byte
	perm       uint8  // random S_3 permutation (0-5)
	projVertex uint8  // random projection vertex (0-7)
	projKey    uint8  // random projection key (0-7)
	projPath   uint16 // random rendering path
}

func (e seedElement) Type() string          { return e.tag }
func (e seedElement) Value() any            { return e.val }
func (e seedElement) Permutation() uint8    { return e.perm }
func (e seedElement) ProjectionVertex() uint8 { return e.projVertex }
func (e seedElement) ProjectionKey() uint8    { return e.projKey }
func (e seedElement) ProjectionPath() uint16  { return e.projPath }

// randomByte returns a cryptographically random byte.
func randomByte() byte {
	var b [1]byte
	if _, err := rand.Read(b[:]); err != nil {
		panic(fmt.Sprintf("crypto/rand: %v", err))
	}
	return b[0]
}

// randomUint16 returns a cryptographically random uint16.
func randomUint16() uint16 {
	var b [2]byte
	if _, err := rand.Read(b[:]); err != nil {
		panic(fmt.Sprintf("crypto/rand: %v", err))
	}
	return binary.LittleEndian.Uint16(b[:])
}

// cloneLattice creates a new lattice with the same topology and occupancy.
// The clone is independent — mutations don't affect the original.
func cloneLattice(src *lattice.Lattice, factory func(string) axiom.Constraint) *lattice.Lattice {
	dst := lattice.New()
	srcNodes := src.Nodes()

	// Recreate all nodes with same constraints.
	dstNodes := make([]*lattice.Node, len(srcNodes))
	for i, n := range srcNodes {
		constraints := n.Constraints()
		clonedConstraints := make([]axiom.Constraint, len(constraints))
		for j, c := range constraints {
			if factory != nil {
				clonedConstraints[j] = factory(c.Tag())
			} else {
				clonedConstraints[j] = c
			}
		}
		dstNodes[i] = dst.AddNode(clonedConstraints)

		// Copy energy state.
		h := n.Hexagram()
		dstNodes[i].SetEnergy(h.Inner().Energy())

		// Copy projection and permutation.
		dstNodes[i].SetPermutation(n.Permutation())
		dstNodes[i].SetProjection(n.ProjectionVertex(), n.ProjectionKey(), n.ProjectionPath())
	}

	// Recreate neighbor connections.
	for i, n := range srcNodes {
		for _, nb := range n.Neighbors() {
			nbID := int(nb.ID())
			if nbID > i { // only connect once per pair
				dst.Connect(dstNodes[i], dstNodes[nbID])
			}
		}
	}

	// Re-bond occupied sites.
	for i, n := range srcNodes {
		if n.Occupied() {
			occ := n.Occupant()
			dstNodes[i].Bond(occ)
		}
	}

	return dst
}

// occupiedCount returns the number of occupied nodes in a lattice.
func occupiedCount(l *lattice.Lattice) int {
	count := 0
	for _, n := range l.Nodes() {
		if n.Occupied() {
			count++
		}
	}
	return count
}

// Unused import guard for ratio — needed by dissolution Config.
var _ = ratio.Zero