package crypto

import (
	"testing"

	"git.mleku.dev/mleku/dendrite/pkg/ratio"
	"git.mleku.dev/mleku/dendrite/pkg/state"
)

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

	sig, err := Sign(&kp.Private, []byte("test message"), params)
	if err != nil {
		t.Fatalf("Sign: %v", err)
	}
	if sig == nil {
		t.Fatal("signature should not be nil")
	}
	if sig.Fingerprint.Hash == "" {
		t.Error("fingerprint hash should not be empty")
	}

	expected := Hash([]byte("test message"))
	if sig.Challenge != expected {
		t.Error("challenge should be Hamadryad hash of message")
	}
}

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

	sig, err := Sign(&kp.Private, []byte("proof test"), params)
	if err != nil {
		t.Fatalf("Sign: %v", err)
	}

	// The lattice already has bonded elements, so proof should have data.
	if len(sig.Proof.LockIns) == 0 {
		t.Error("proof should have lock-in depths")
	}
	if len(sig.Proof.NeighborCounts) == 0 {
		t.Error("proof should have neighbor counts")
	}
	if len(sig.Proof.HexTrace) == 0 {
		t.Error("proof should have hexagram trace")
	}
}

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

	msg := []byte("verify me")
	sig, err := Sign(&kp.Private, msg, params)
	if err != nil {
		t.Fatalf("Sign: %v", err)
	}

	fp := FingerprintFromSpore(kp.Public.Spore)
	if !Verify(fp, msg, sig) {
		t.Error("valid signature should verify")
	}
}

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

	msg := []byte("original")
	sig, err := Sign(&kp.Private, msg, params)
	if err != nil {
		t.Fatalf("Sign: %v", err)
	}

	fp := FingerprintFromSpore(kp.Public.Spore)
	if Verify(fp, []byte("tampered"), sig) {
		t.Error("tampered message should not verify")
	}
}

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

	msg := []byte("test")
	sig, err := Sign(&kp.Private, msg, params)
	if err != nil {
		t.Fatalf("Sign: %v", err)
	}

	// Different fingerprint.
	wrongFP := SporeFingerprint{Hash: "wrong_hash"}
	if Verify(wrongFP, msg, sig) {
		t.Error("wrong fingerprint should not verify")
	}
}

func TestVerifyNilSignature(t *testing.T) {
	fp := SporeFingerprint{Hash: "test"}
	if Verify(fp, []byte("msg"), nil) {
		t.Error("nil signature should not verify")
	}
}

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

	fp := FingerprintFromSpore(kp.Public.Spore)
	if fp.Hash == "" {
		t.Error("fingerprint hash should not be empty")
	}
	if len(fp.TypeSignature) == 0 {
		t.Error("fingerprint should have type signature")
	}
}

func TestSortedTags(t *testing.T) {
	tags := sortedTags(kp().Public.Spore.TypeSignature)
	for i := 1; i < len(tags); i++ {
		if tags[i] < tags[i-1] {
			t.Errorf("tags not sorted: %q after %q", tags[i], tags[i-1])
		}
	}
}

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

func TestSignNilLattice(t *testing.T) {
	privkey := &PrivateKey{Lattice: nil, ConstraintFactory: testFactory}
	_, err := Sign(privkey, []byte("test"), DefaultParams(Security128))
	if err == nil {
		t.Error("expected error with nil lattice")
	}
}

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

	// Record original occupancy.
	origOcc := occupiedCount(l)

	// Sign multiple messages.
	for _, msg := range []string{"msg1", "msg2", "msg3"} {
		_, err := Sign(&kp.Private, []byte(msg), params)
		if err != nil {
			t.Fatalf("Sign(%q): %v", msg, err)
		}
	}

	// Lattice must be unchanged.
	if occupiedCount(l) != origOcc {
		t.Errorf("lattice occupancy changed: was %d, now %d", origOcc, occupiedCount(l))
	}
}

func TestSignMultipleAllVerify(t *testing.T) {
	l := buildMatureLattice()
	params := DefaultParams(Security128)
	kp := GenerateKeyPair(l, params, testFactory)
	fp := FingerprintFromSpore(kp.Public.Spore)

	messages := []string{
		"first", "second", "third", "fourth", "fifth",
		"sixth", "seventh", "eighth", "ninth", "tenth",
	}
	for _, msg := range messages {
		sig, err := Sign(&kp.Private, []byte(msg), params)
		if err != nil {
			t.Fatalf("Sign(%q): %v", msg, err)
		}
		if !Verify(fp, []byte(msg), sig) {
			t.Errorf("signature for %q did not verify", msg)
		}
	}
}

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

	msg := []byte("structural test")
	sig, err := Sign(&kp.Private, msg, params)
	if err != nil {
		t.Fatalf("Sign: %v", err)
	}

	fp := FingerprintFromSpore(kp.Public.Spore)

	// Sanity: valid signature verifies.
	if !Verify(fp, msg, sig) {
		t.Fatal("valid signature should verify")
	}

	// Proof length mismatch: truncate LockIns.
	t.Run("truncated_lockins", func(t *testing.T) {
		bad := copySig(sig)
		if len(bad.Proof.LockIns) > 1 {
			bad.Proof.LockIns = bad.Proof.LockIns[:1]
		}
		if Verify(fp, msg, bad) {
			t.Error("truncated lock-ins should not verify")
		}
	})

	// Proof length mismatch: truncate NeighborCounts.
	t.Run("truncated_neighbors", func(t *testing.T) {
		bad := copySig(sig)
		if len(bad.Proof.NeighborCounts) > 1 {
			bad.Proof.NeighborCounts = bad.Proof.NeighborCounts[:1]
		}
		if Verify(fp, msg, bad) {
			t.Error("truncated neighbor counts should not verify")
		}
	})

	// Zero neighbor count at a bonded site.
	t.Run("zero_neighbor", func(t *testing.T) {
		bad := copySig(sig)
		if len(bad.Proof.NeighborCounts) > 0 {
			bad.Proof.NeighborCounts[0] = 0
		}
		if Verify(fp, msg, bad) {
			t.Error("zero neighbor count should not verify")
		}
	})
}

// copySig returns a shallow copy of a Signature with independent proof slices.
func copySig(s *Signature) *Signature {
	cp := *s
	cp.Proof.LockIns = append([]ratio.Ratio(nil), s.Proof.LockIns...)
	cp.Proof.NeighborCounts = append([]int(nil), s.Proof.NeighborCounts...)
	cp.Proof.HexTrace = append([]state.Hexagram(nil), s.Proof.HexTrace...)
	cp.Response = append([]SiteMark(nil), s.Response...)
	return &cp
}