package crypto

import (
	"testing"

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

func testFactory(tag string) axiom.Constraint {
	return tagConstraint{tag}
}

func buildMatureLattice() *lattice.Lattice {
	l := lattice.New()
	var nodes []*lattice.Node
	for range 20 {
		n := l.AddNode([]axiom.Constraint{tagConstraint{"word"}})
		nodes = append(nodes, n)
	}
	for range 10 {
		n := l.AddNode([]axiom.Constraint{tagConstraint{"punct"}})
		nodes = append(nodes, n)
	}
	// Ring within first 20.
	for i := 0; i < 20; i++ {
		l.Connect(nodes[i], nodes[(i+1)%20])
	}
	// Ring within last 10.
	for i := 20; i < 30; i++ {
		l.Connect(nodes[i], nodes[20+(i+1-20)%10])
	}
	// Cross-connect.
	l.Connect(nodes[0], nodes[20])
	l.Connect(nodes[5], nodes[25])

	// Bond some elements.
	for i := 0; i < 15; i++ {
		nodes[i].Bond(testElem{"word", string(rune('a' + i))})
	}
	for i := 20; i < 27; i++ {
		nodes[i].Bond(testElem{"punct", "."})
	}
	return l
}

func TestGenerateKeyPair(t *testing.T) {
	l := buildMatureLattice()
	params := DefaultParams(Security128)
	kp := GenerateKeyPair(l, params, testFactory)

	if kp.Public.SporeHash == "" {
		t.Error("spore hash should not be empty")
	}
	if kp.Public.Basis == nil {
		t.Error("basis should not be nil")
	}
	if kp.Public.Spore == nil {
		t.Error("spore should not be nil")
	}
	if kp.Private.Lattice != l {
		t.Error("private key should reference the lattice")
	}
	if kp.Private.ConstraintFactory == nil {
		t.Error("constraint factory should not be nil")
	}
}

func TestKeyPairBasisDimension(t *testing.T) {
	l := buildMatureLattice()
	params := DefaultParams(Security128)
	kp := GenerateKeyPair(l, params, testFactory)

	if kp.Public.Basis.Dimension != l.Size() {
		t.Errorf("basis dimension = %d, want %d", kp.Public.Basis.Dimension, l.Size())
	}
}

func TestKeyPairBasisTags(t *testing.T) {
	l := buildMatureLattice()
	params := DefaultParams(Security128)
	kp := GenerateKeyPair(l, params, testFactory)

	if len(kp.Public.Basis.Tags) != 2 {
		t.Errorf("expected 2 tags, got %d", len(kp.Public.Basis.Tags))
	}
	// Tags should be sorted: punct, word.
	if kp.Public.Basis.Tags[0] != "punct" || kp.Public.Basis.Tags[1] != "word" {
		t.Errorf("tags = %v, want [punct, word]", kp.Public.Basis.Tags)
	}
}

func TestKeyPairFactoryProducesValidConstraints(t *testing.T) {
	l := buildMatureLattice()
	params := DefaultParams(Security128)
	kp := GenerateKeyPair(l, params, testFactory)

	for _, tag := range kp.Public.Basis.Tags {
		c := kp.Private.ConstraintFactory(tag)
		if c.Tag() != tag {
			t.Errorf("factory(%q).Tag() = %q", tag, c.Tag())
		}
		// Constraint should admit matching elements.
		e := testElem{tag, "test"}
		if !c.Admits(e) {
			t.Errorf("factory(%q) does not admit matching element", tag)
		}
		// Should reject mismatched elements.
		e2 := testElem{"other", "test"}
		if c.Admits(e2) {
			t.Errorf("factory(%q) admits mismatched element", tag)
		}
	}
}