package crypto

import (
	"testing"

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

func testKeyPair() *KeyPair {
	l := buildMatureLattice()
	return GenerateKeyPair(l, DefaultParams(Security128), testFactory)
}

func TestEncryptProducesCiphertext(t *testing.T) {
	kp := testKeyPair()
	msg := []byte("hello world")
	params := DefaultParams(Security128)

	ct, err := Encrypt(&kp.Public, msg, params)
	if err != nil {
		t.Fatalf("Encrypt: %v", err)
	}
	if ct == nil {
		t.Fatal("ciphertext should not be nil")
	}
	if len(ct.Sites) == 0 {
		t.Error("ciphertext should have sites")
	}
	if ct.Basis == nil {
		t.Error("ciphertext should have basis")
	}
}

func TestEncryptOccupiesSites(t *testing.T) {
	kp := testKeyPair()
	msg := []byte("test")
	params := DefaultParams(Security128)

	ct, err := Encrypt(&kp.Public, msg, params)
	if err != nil {
		t.Fatalf("Encrypt: %v", err)
	}

	occupied := 0
	for _, s := range ct.Sites {
		if s.Occupied {
			occupied++
		}
	}
	if occupied == 0 {
		t.Error("encryption should produce occupied sites")
	}
}

func TestEncryptInvalidParams(t *testing.T) {
	kp := testKeyPair()
	msg := []byte("hello")
	bad := Params{} // all zeros — invalid

	_, err := Encrypt(&kp.Public, msg, bad)
	if err == nil {
		t.Error("expected error with invalid params")
	}
}

func TestEncryptNilSpore(t *testing.T) {
	pub := &PublicKey{
		Basis: &Basis{Tags: []string{"word"}},
		Spore: nil,
	}
	_, err := Encrypt(pub, []byte("x"), DefaultParams(Security128))
	if err == nil {
		t.Error("expected error with nil spore")
	}
}

func TestEncryptDecryptRoundTrip(t *testing.T) {
	// Build a lattice and keypair where we control the constraints.
	l := buildMatureLattice()
	params := DefaultParams(Security128)
	kp := GenerateKeyPair(l, params, testFactory)

	msg := []byte("hi")
	ct, err := Encrypt(&kp.Public, msg, params)
	if err != nil {
		t.Fatalf("Encrypt: %v", err)
	}

	// Decrypt should recover at least some bytes.
	// Full round-trip correctness depends on the lattice being
	// deterministic with matching constraints. With public constraints
	// on both sides, the decrypt should find the bytes via hash matching.
	result, err := Decrypt(&kp.Private, ct)
	if err != nil {
		// Decryption may fail if dissolution removed all message elements.
		// This is expected with certain parameter combinations.
		t.Logf("Decrypt returned error (may be expected): %v", err)
		return
	}

	// Check that at least some bytes were recovered.
	if len(result) == 0 {
		t.Error("expected non-empty decryption result")
	}
}

func TestDecryptNilCiphertext(t *testing.T) {
	kp := testKeyPair()
	_, err := Decrypt(&kp.Private, nil)
	if err == nil {
		t.Error("expected error with nil ciphertext")
	}
}

func TestDecryptNoFactory(t *testing.T) {
	ct := &Ciphertext{
		Sites:  []SiteMark{{Occupied: true, TypeTag: "word"}},
		Params: DefaultParams(Security128),
		Basis:  &Basis{Tags: []string{"word"}},
	}
	privkey := &PrivateKey{ConstraintFactory: nil}
	_, err := Decrypt(privkey, ct)
	if err == nil {
		t.Error("expected error with nil factory")
	}
}

func TestMessageElementInterface(t *testing.T) {
	me := MessageElement{Index: 0, Byte: 'A', TypeTag: "word"}
	var e axiom.Element = me
	if e.Type() != "word" {
		t.Errorf("Type() = %q, want word", e.Type())
	}
	if e.Value().(byte) != 'A' {
		t.Errorf("Value() = %v, want 'A'", e.Value())
	}
}

func TestPublicConstraint(t *testing.T) {
	factory := publicConstraintFactory(&Basis{Tags: []string{"word", "punct"}})
	c := factory("word")
	if c.Tag() != "word" {
		t.Errorf("Tag() = %q, want word", c.Tag())
	}
	if !c.Admits(testElem{"word", "x"}) {
		t.Error("should admit matching element")
	}
	if c.Admits(testElem{"punct", "."}) {
		t.Error("should reject mismatched element")
	}
}