package crypto

// Track 3: Signature Unforgeability Proofs
//
// These tests demonstrate that forging a dendrite signature without the
// private key (constraint factory) is computationally hard.
//
// EUF-CMA (Existential Unforgeability under Chosen Message Attack):
//   1. The attacker has the public key (spore/fingerprint)
//   2. The attacker can request signatures on chosen messages (signing oracle)
//   3. The attacker must produce a valid signature on a NEW message
//   4. Success means the scheme is broken
//
// We verify unforgeability through:
//   - Random forgery attempts (brute-force over response vectors)
//   - Structural constraint verification (Verify rejects malformed proofs)
//   - EUF-CMA game simulation (sign queries + forgery attempt)
//   - Projection/permutation distribution checks
//   - Cross-message signature independence

import (
	"crypto/rand"
	"encoding/binary"
	"testing"

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

// ---- Random Forgery Resistance ----

func TestRandomForgeryResistance(t *testing.T) {
	// Generate a real keypair and fingerprint.
	l := buildMatureLattice()
	params := DefaultParams(Security128)
	kp := GenerateKeyPair(l, params, testFactory)
	fp := FingerprintFromSpore(kp.Public.Spore)

	msg := []byte("message to forge signature for")
	challenge := Hash(msg)

	// Attempt 10000 random forgeries.
	// Each attempt constructs a random bonding pattern and proof,
	// then checks if Verify accepts it.
	forged := 0
	attempts := 10000

	for range attempts {
		sig := randomSignature(challenge, fp, 10)
		if Verify(fp, msg, sig) {
			forged++
		}
	}

	if forged > 0 {
		t.Errorf("random forgery succeeded %d/%d times — scheme may be weak", forged, attempts)
	}
}

func TestRandomForgeryLargerResponse(t *testing.T) {
	// Try forgeries with responses matching the expected occupancy count.
	l := buildMatureLattice()
	params := DefaultParams(Security128)
	kp := GenerateKeyPair(l, params, testFactory)
	fp := FingerprintFromSpore(kp.Public.Spore)

	msg := []byte("larger response forgery")
	challenge := Hash(msg)

	// Try with responses sized to match what a real signature would have.
	forged := 0
	attempts := 5000
	for range attempts {
		sig := randomSignature(challenge, fp, 20) // ~20 occupied sites
		if Verify(fp, msg, sig) {
			forged++
		}
	}

	if forged > 0 {
		t.Errorf("random forgery (large response) succeeded %d/%d times", forged, attempts)
	}
}

// ---- EUF-CMA Game ----

func TestEUFCMAGame(t *testing.T) {
	// Simulate the EUF-CMA security game:
	// 1. Challenger generates keypair
	// 2. Adversary gets public key
	// 3. Adversary makes Q signing queries (chosen messages)
	// 4. Adversary attempts forgery on a new message

	// Step 1: Generate keypair.
	l := buildMatureLattice()
	params := DefaultParams(Security128)
	kp := GenerateKeyPair(l, params, testFactory)
	fp := FingerprintFromSpore(kp.Public.Spore)

	// Step 2: Adversary sees only public key.
	// (fp is the public key — the fingerprint)

	// Step 3: Signing oracle — adversary requests signatures.
	queryMessages := []string{
		"query-1", "query-2", "query-3",
		"query-4", "query-5",
	}
	querySigs := make([]*Signature, len(queryMessages))
	for i, msg := range queryMessages {
		sig, err := Sign(&kp.Private, []byte(msg), params)
		if err != nil {
			t.Fatalf("Sign(query-%d): %v", i, err)
		}
		querySigs[i] = sig
		// Verify signing oracle produces valid signatures.
		if !Verify(fp, []byte(msg), sig) {
			t.Fatalf("signing oracle signature %d does not verify", i)
		}
	}

	// Step 4: Adversary attempts forgery on a NEW message.
	forgeMsg := []byte("forged-message-not-queried")
	forgeChallenge := Hash(forgeMsg)

	// Strategy A: replay a query signature with new message.
	for i, sig := range querySigs {
		if Verify(fp, forgeMsg, sig) {
			t.Errorf("replay of query-%d signature verifies on new message", i)
		}
	}

	// Strategy B: modify a query signature's challenge.
	for i, sig := range querySigs {
		tampered := copySig(sig)
		tampered.Challenge = forgeChallenge
		if Verify(fp, forgeMsg, tampered) {
			t.Errorf("challenge-swapped query-%d signature verifies", i)
		}
	}

	// Strategy C: random forgery.
	for range 1000 {
		sig := randomSignature(forgeChallenge, fp, 15)
		if Verify(fp, forgeMsg, sig) {
			t.Error("random forgery succeeded in EUF-CMA game")
		}
	}
}

func TestEUFCMANoReuseAcrossMessages(t *testing.T) {
	// Signatures from different messages should not be interchangeable.
	l := buildMatureLattice()
	params := DefaultParams(Security128)
	kp := GenerateKeyPair(l, params, testFactory)
	fp := FingerprintFromSpore(kp.Public.Spore)

	msgs := []string{"alpha", "beta", "gamma", "delta"}
	sigs := make([]*Signature, len(msgs))
	for i, msg := range msgs {
		sig, err := Sign(&kp.Private, []byte(msg), params)
		if err != nil {
			t.Fatalf("Sign(%q): %v", msg, err)
		}
		sigs[i] = sig
	}

	// Each signature should only verify its own message.
	for i, sig := range sigs {
		for j, msg := range msgs {
			result := Verify(fp, []byte(msg), sig)
			if i == j && !result {
				t.Errorf("signature %d should verify message %d", i, j)
			}
			if i != j && result {
				t.Errorf("signature %d should NOT verify message %d", i, j)
			}
		}
	}
}

// ---- Structural Forgery Resistance ----

func TestForgeryWrongPermDist(t *testing.T) {
	// A forged signature with incorrect permutation distribution
	// should be rejected by Verify's permDistCompatible check.
	l := buildMatureLattice()
	params := DefaultParams(Security128)
	kp := GenerateKeyPair(l, params, testFactory)
	fp := FingerprintFromSpore(kp.Public.Spore)

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

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

	// Forge: change all perm values to concentrate in one bucket.
	bad := copySig(sig)
	for i := range bad.Response {
		if bad.Response[i].Occupied {
			bad.Response[i].Perm = 0 // all identity
		}
	}

	// This may or may not fail depending on the original distribution.
	// The point is that an attacker cannot easily match the expected
	// permutation distribution without the factory.
	// We check that at least the distribution changed.
	origDist := [6]int{}
	forgeDist := [6]int{}
	for _, site := range sig.Response {
		if site.Occupied && site.Perm < 6 {
			origDist[site.Perm]++
		}
	}
	for _, site := range bad.Response {
		if site.Occupied && site.Perm < 6 {
			forgeDist[site.Perm]++
		}
	}
	if origDist != forgeDist {
		// Distribution changed — this is detectable.
		t.Logf("original perm dist: %v", origDist)
		t.Logf("forged perm dist:   %v", forgeDist)
	}
}

func TestForgeryEmptyProof(t *testing.T) {
	// A signature with empty proof arrays should be rejected.
	l := buildMatureLattice()
	params := DefaultParams(Security128)
	kp := GenerateKeyPair(l, params, testFactory)
	fp := FingerprintFromSpore(kp.Public.Spore)

	msg := []byte("empty proof")
	challenge := Hash(msg)

	sig := &Signature{
		Fingerprint: fp,
		Challenge:   challenge,
		Response:    []SiteMark{{Index: 0, Occupied: true, TypeTag: "word", Perm: 0}},
		Proof: SporeProof{
			LockIns:        nil,
			NeighborCounts: nil,
			HexTrace:       nil,
		},
	}

	if Verify(fp, msg, sig) {
		t.Error("signature with empty proof should not verify")
	}
}

func TestForgeryMismatchedProofLengths(t *testing.T) {
	// Proof slices with different lengths should be rejected.
	l := buildMatureLattice()
	params := DefaultParams(Security128)
	kp := GenerateKeyPair(l, params, testFactory)
	fp := FingerprintFromSpore(kp.Public.Spore)

	msg := []byte("mismatched proof")
	challenge := Hash(msg)

	sig := &Signature{
		Fingerprint: fp,
		Challenge:   challenge,
		Response:    []SiteMark{{Index: 0, Occupied: true, TypeTag: "word", Perm: 0}},
		Proof: SporeProof{
			LockIns:        []ratio.Ratio{ratio.One},
			NeighborCounts: []int{2, 3}, // wrong length
			HexTrace:       []state.Hexagram{0},
		},
	}

	if Verify(fp, msg, sig) {
		t.Error("mismatched proof lengths should not verify")
	}
}

func TestForgeryZeroLockIn(t *testing.T) {
	// A signature with zero lock-in on all sites should be rejected.
	l := buildMatureLattice()
	params := DefaultParams(Security128)
	kp := GenerateKeyPair(l, params, testFactory)
	fp := FingerprintFromSpore(kp.Public.Spore)

	msg := []byte("zero lock-in")
	challenge := Hash(msg)

	sites := make([]SiteMark, 20)
	lockIns := make([]ratio.Ratio, 20)
	neighborCounts := make([]int, 20)
	hexTrace := make([]state.Hexagram, 20)

	for i := range 20 {
		sites[i] = SiteMark{Index: uint64(i), Occupied: true, TypeTag: "word", Perm: uint8(i % 6)}
		lockIns[i] = ratio.Zero // zero lock-in
		neighborCounts[i] = 2
		hexTrace[i] = state.Hexagram(i % 64)
	}

	sig := &Signature{
		Fingerprint: fp,
		Challenge:   challenge,
		Response:    sites,
		Proof: SporeProof{
			LockIns:        lockIns,
			NeighborCounts: neighborCounts,
			HexTrace:       hexTrace,
		},
	}

	if Verify(fp, msg, sig) {
		t.Error("zero lock-in signature should not verify")
	}
}

func TestForgeryInvalidHexagram(t *testing.T) {
	// Hexagram values > 63 should cause verification to fail.
	l := buildMatureLattice()
	params := DefaultParams(Security128)
	kp := GenerateKeyPair(l, params, testFactory)
	fp := FingerprintFromSpore(kp.Public.Spore)

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

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

	// Forge: inject invalid hexagram value.
	bad := copySig(sig)
	if len(bad.Proof.HexTrace) > 0 {
		bad.Proof.HexTrace[0] = 255 // > 63, invalid
	}
	if Verify(fp, msg, bad) {
		t.Error("signature with invalid hexagram should not verify")
	}
}

func TestForgeryWrongFingerprint(t *testing.T) {
	// Signature verified against wrong fingerprint should fail.
	l := buildMatureLattice()
	params := DefaultParams(Security128)
	kp := GenerateKeyPair(l, params, testFactory)

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

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

func TestForgeryInsufficientOccupancy(t *testing.T) {
	// A signature with too few occupied sites should be rejected.
	// Verify requires: occupied >= max(len(challenge)/4, 1).
	l := buildMatureLattice()
	params := DefaultParams(Security128)
	kp := GenerateKeyPair(l, params, testFactory)
	fp := FingerprintFromSpore(kp.Public.Spore)

	msg := []byte("few sites")
	challenge := Hash(msg)

	// Hamadryad hash is 56 bytes, so min occupied = 56/4 = 14.
	// Create signature with only 1 occupied site.
	sig := &Signature{
		Fingerprint: fp,
		Challenge:   challenge,
		Response:    []SiteMark{{Index: 0, Occupied: true, TypeTag: "word", Perm: 0}},
		Proof: SporeProof{
			LockIns:        []ratio.Ratio{ratio.One},
			NeighborCounts: []int{2},
			HexTrace:       []state.Hexagram{0},
		},
	}

	if Verify(fp, msg, sig) {
		t.Errorf("signature with 1 occupied site should not verify (need >= %d)", len(challenge)/4)
	}
}

// ---- Cross-Key Forgery Resistance ----

func TestCrossKeyForgeryRejected(t *testing.T) {
	// A signature produced by one key should not verify against another key.
	// Use Generate to create genuinely different keypairs (random seeding).
	params := DefaultParams(Security128)
	tags := []string{"word", "punct"}

	kp1, err := Generate(params, tags, testFactory)
	if err != nil {
		t.Fatalf("Generate(1): %v", err)
	}
	fp1 := FingerprintFromSpore(kp1.Public.Spore)

	kp2, err := Generate(params, tags, testFactory)
	if err != nil {
		t.Fatalf("Generate(2): %v", err)
	}
	fp2 := FingerprintFromSpore(kp2.Public.Spore)

	msg := []byte("cross-key test")
	sig1, err := Sign(&kp1.Private, msg, params)
	if err != nil {
		t.Fatalf("Sign(kp1): %v", err)
	}

	// sig1 verifies against fp1 but not fp2.
	if !Verify(fp1, msg, sig1) {
		t.Error("signature should verify against own fingerprint")
	}
	if Verify(fp2, msg, sig1) {
		t.Error("signature should NOT verify against different fingerprint")
	}
}

// ---- Signature Uniqueness ----

func TestSignaturesDifferPerMessage(t *testing.T) {
	// Different messages produce different challenges and therefore
	// different signatures (the bonding pattern is challenge-dependent).
	l := buildMatureLattice()
	params := DefaultParams(Security128)
	kp := GenerateKeyPair(l, params, testFactory)

	msg1 := []byte("message one")
	msg2 := []byte("message two")

	sig1, _ := Sign(&kp.Private, msg1, params)
	sig2, _ := Sign(&kp.Private, msg2, params)

	if sig1.Challenge == sig2.Challenge {
		t.Error("different messages should produce different challenges")
	}
}

// ---- Helpers ----

// randomSignature constructs a random (likely invalid) signature
// with the correct challenge and fingerprint but random response/proof.
func randomSignature(challenge Hamadryad, fp SporeFingerprint, sites int) *Signature {
	response := make([]SiteMark, sites)
	lockIns := make([]ratio.Ratio, sites)
	neighborCounts := make([]int, sites)
	hexTrace := make([]state.Hexagram, sites)

	for i := range sites {
		var rb [4]byte
		rand.Read(rb[:])

		response[i] = SiteMark{
			Index:      uint64(binary.LittleEndian.Uint16(rb[:2])),
			Occupied:   true,
			TypeTag:    randomTag(),
			Projection: rb[2] & 0x3F,
			Perm:       rb[3] % 6,
			LockIn:     ratio.New(int64(rb[0]%100+1), 100),
		}
		lockIns[i] = ratio.New(int64(rb[0]%100+1), 100)
		neighborCounts[i] = int(rb[1]%10) + 1
		hexTrace[i] = state.Hexagram(rb[2] % 64)
	}

	return &Signature{
		Fingerprint: fp,
		Challenge:   challenge,
		Response:    response,
		Proof: SporeProof{
			LockIns:        lockIns,
			NeighborCounts: neighborCounts,
			HexTrace:       hexTrace,
		},
	}
}

func randomTag() string {
	tags := []string{"word", "punct", "space", "number"}
	var b [1]byte
	rand.Read(b[:])
	return tags[int(b[0])%len(tags)]
}