package crypto

import (
	"sort"

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

// Basis is the public description of a lattice's constraint structure.
// It contains enough information to nucleate a lattice (encrypt)
// but not enough to read the bonding pattern (decrypt).
type Basis struct {
	// Dimension is the number of constraint sites.
	Dimension int

	// Tags enumerates the constraint type tags in sorted order.
	Tags []string

	// Distribution records sites per tag.
	Distribution []spore.TagCount

	// Connectivity records average neighbor count per tag.
	Connectivity []spore.TagRatio

	// PermDist records how nodes distribute across S_3 permutations.
	PermDist [6]int

	// ProjDist records how nodes distribute across 64 projection configs.
	ProjDist [64]int

	// Modulus for any modular arithmetic.
	Modulus ratio.Ratio
}

// FromLattice extracts the public basis from a live lattice.
func FromLattice(l *lattice.Lattice, modulus ratio.Ratio) *Basis {
	s := spore.Extract(l)
	return fromSporeInternal(s, modulus)
}

// FromSpore reconstructs a basis from a spore's fingerprint.
func FromSpore(s *spore.Spore, modulus ratio.Ratio) *Basis {
	return fromSporeInternal(s, modulus)
}

func fromSporeInternal(s *spore.Spore, modulus ratio.Ratio) *Basis {
	b := &Basis{
		Dimension:    s.TotalNodes,
		Distribution: make([]spore.TagCount, len(s.TypeSignature)),
		Connectivity: make([]spore.TagRatio, len(s.Connectivity)),
		PermDist:     s.PermDist,
		ProjDist:     s.ProjDist,
		Modulus:      modulus,
	}

	copy(b.Distribution, s.TypeSignature)
	copy(b.Connectivity, s.Connectivity)

	// Extract and sort tags.
	tags := make([]string, len(s.TypeSignature))
	for i, tc := range s.TypeSignature {
		tags[i] = tc.Tag
	}
	sort.Strings(tags)
	b.Tags = tags

	return b
}

// Equal checks structural equality of two bases.
func (b *Basis) Equal(other *Basis) bool {
	if b.Dimension != other.Dimension {
		return false
	}
	if len(b.Tags) != len(other.Tags) {
		return false
	}
	for i, t := range b.Tags {
		if t != other.Tags[i] {
			return false
		}
	}
	if len(b.Distribution) != len(other.Distribution) {
		return false
	}
	for i, d := range b.Distribution {
		if d.Tag != other.Distribution[i].Tag || d.Count != other.Distribution[i].Count {
			return false
		}
	}
	if b.PermDist != other.PermDist {
		return false
	}
	if b.ProjDist != other.ProjDist {
		return false
	}
	if !b.Modulus.Equal(other.Modulus) {
		return false
	}
	return true
}

// TagCount returns the site count for a given tag, or 0 if not found.
func (b *Basis) TagCount(tag string) int {
	for _, d := range b.Distribution {
		if d.Tag == tag {
			return d.Count
		}
	}
	return 0
}