package ring

import (
	"crypto/rand"
	"testing"
)

// ─── vwidth.go tests ─────────────────────────────────────────────────────────

func TestSerializeBoundedUnsigned(t *testing.T) {
	p := SmallGPVParams().Ring
	orig := New(p)
	for i := range orig.Coeffs {
		orig.Coeffs[i] = uint32(i % 16)
	}

	data := SerializeBounded(orig, 4, false)
	if len(data) != (p.N*4+7)/8 {
		t.Fatalf("expected %d bytes, got %d", (p.N*4+7)/8, len(data))
	}

	restored := DeserializeBounded(p, data, 4, false)
	if !Equal(orig, restored) {
		t.Fatal("round-trip failed for unsigned bounded serialization")
	}
}

func TestSerializeBoundedSigned(t *testing.T) {
	p := SmallGPVParams().Ring
	half := p.Q / 2
	orig := New(p)
	for i := range orig.Coeffs {
		// Alternate positive and negative (centered form).
		if i%2 == 0 {
			orig.Coeffs[i] = uint32(i % 100)
		} else {
			orig.Coeffs[i] = p.Q - uint32(i%100) // centered negative
		}
	}

	// Need enough bits for values up to 100.
	data := SerializeBounded(orig, 8, true)
	restored := DeserializeBounded(p, data, 8, true)
	if !Equal(orig, restored) {
		t.Fatal("round-trip failed for signed bounded serialization")
	}

	// Check the centered values are preserved.
	for i := range orig.Coeffs {
		var origCentered uint32
		if orig.Coeffs[i] > half {
			origCentered = p.Q - orig.Coeffs[i]
		} else {
			origCentered = orig.Coeffs[i]
		}
		var restoredCentered uint32
		if restored.Coeffs[i] > half {
			restoredCentered = p.Q - restored.Coeffs[i]
		} else {
			restoredCentered = restored.Coeffs[i]
		}
		if origCentered != restoredCentered {
			t.Fatalf("coefficient %d: centered value %d vs %d", i, origCentered, restoredCentered)
		}
	}
}

// ─── gpv_gadget.go tests ──────────────────────────────────────────────────────

func TestGPVGadgetSignVerify(t *testing.T) {
	gp := SmallGPVParams()
	pk, sk := GPVGadgetKeyGen(gp)

	msg := []byte("hello gadget world")
	sig := GPVGadgetSign(sk, msg)

	if !GPVGadgetVerify(pk, msg, sig) {
		t.Fatal("valid gadget signature rejected")
	}
}

func TestGPVGadgetWrongMessage(t *testing.T) {
	gp := SmallGPVParams()
	pk, sk := GPVGadgetKeyGen(gp)

	msg := []byte("correct message")
	sig := GPVGadgetSign(sk, msg)

	if GPVGadgetVerify(pk, []byte("wrong message"), sig) {
		t.Fatal("gadget signature verified for wrong message")
	}
}

func TestGPVGadgetWrongKey(t *testing.T) {
	gp := SmallGPVParams()
	_, sk1 := GPVGadgetKeyGen(gp)
	pk2, _ := GPVGadgetKeyGen(gp)

	msg := []byte("test message")
	sig := GPVGadgetSign(sk1, msg)

	if GPVGadgetVerify(pk2, msg, sig) {
		t.Fatal("gadget signature verified with wrong public key")
	}
}

func TestGPVGadgetMultipleMessages(t *testing.T) {
	gp := SmallGPVParams()
	pk, sk := GPVGadgetKeyGen(gp)

	messages := []string{
		"",
		"a",
		"hello gadget",
		"lattice-based cryptography with MP12 trapdoor",
	}

	for _, msg := range messages {
		sig := GPVGadgetSign(sk, []byte(msg))
		if !GPVGadgetVerify(pk, []byte(msg), sig) {
			t.Fatalf("gadget verification failed for message %q", msg)
		}
	}
}

func TestGPVGadgetDeterminism(t *testing.T) {
	gp := SmallGPVParams()
	pk, sk := GPVGadgetKeyGen(gp)

	msg := []byte("same message twice")
	sig1 := GPVGadgetSign(sk, msg)
	sig2 := GPVGadgetSign(sk, msg)

	if !GPVGadgetVerify(pk, msg, sig1) || !GPVGadgetVerify(pk, msg, sig2) {
		t.Fatal("valid gadget signatures rejected")
	}

	// Gadget signatures should produce the same E2 (deterministic decomposition)
	// but different E1 (due to fresh perturbation).
	e1Equal := true
	for i := range sig1.E1.Coeffs {
		if sig1.E1.Coeffs[i] != sig2.E1.Coeffs[i] {
			e1Equal = false
			break
		}
	}
	if e1Equal {
		t.Fatal("two gadget signatures have identical E1 — perturbation missing")
	}
}

// ─── sigma.go: Phase 2 — interactive sigma protocol ───────────────────────────

func TestSigmaInteractive(t *testing.T) {
	gp := SmallGPVParams()
	pk, sk := GPVKeyGen(gp)

	rng := rand.Reader

	// Retry loop for rejection sampling (can fail ~1-5% of the time).
	var z *Poly
	var ok bool
	for attempt := 0; attempt < 20; attempt++ {
		w, y := Commit(pk, rng)
		c := Challenge(pk.P.Ring, rng)
		z, ok = Respond(sk, y, c)
		if ok {
			if !VerifyInteractive(pk, w, c, z) {
				t.Fatal("valid interactive sigma proof rejected")
			}
			return
		}
	}
	t.Fatal("rejection sampling failed after 20 attempts")
}

func TestSigmaInteractiveWrongKey(t *testing.T) {
	gp := SmallGPVParams()
	pk1, sk1 := GPVKeyGen(gp)
	pk2, _ := GPVKeyGen(gp)
	_ = pk2

	rng := rand.Reader

	w, y := Commit(pk1, rng)
	c := Challenge(pk1.P.Ring, rng)
	z, ok := Respond(sk1, y, c)
	if !ok {
		t.Fatal("rejection sampling failed")
	}

	// Verify with wrong key — should fail because BPoly doesn't match.
	if VerifyInteractive(pk2, w, c, z) {
		t.Fatal("interactive proof verified with wrong key")
	}
}

// ─── sigma.go: Phase 3 — non-interactive (Fiat-Shamir) proof ─────────────────

func TestSigmaNonInteractive(t *testing.T) {
	gp := SmallGPVParams()
	pk, sk := GPVKeyGen(gp)

	proof := Prove(sk, nil)

	if !VerifySigma(pk, nil, proof) {
		t.Fatal("valid non-interactive sigma proof rejected")
	}
}

func TestSigmaNonInteractiveWithContext(t *testing.T) {
	gp := SmallGPVParams()
	pk, sk := GPVKeyGen(gp)

	context := []byte("session-12345")
	proof := Prove(sk, context)

	if !VerifySigma(pk, context, proof) {
		t.Fatal("valid context-bound sigma proof rejected")
	}

	// Wrong context should fail.
	if VerifySigma(pk, []byte("wrong-context"), proof) {
		t.Fatal("proof verified with wrong context")
	}
}

func TestSigmaNonInteractiveWrongKey(t *testing.T) {
	gp := SmallGPVParams()
	_, sk := GPVKeyGen(gp)
	pk2, _ := GPVKeyGen(gp)

	proof := Prove(sk, nil)

	if VerifySigma(pk2, nil, proof) {
		t.Fatal("proof verified with wrong key")
	}
}

func TestSigmaProofMarshal(t *testing.T) {
	gp := SmallGPVParams()
	p := gp.Ring
	_, sk := GPVKeyGen(gp)

	proof := Prove(sk, nil)

	data := MarshalSigmaProof(proof)
	restored := UnmarshalSigmaProof(p, data)

	if restored == nil {
		t.Fatal("unmarshal returned nil")
	}

	if !Equal(proof.Z, restored.Z) {
		t.Fatal("Z not preserved after marshal round-trip")
	}
	if !Equal(proof.C, restored.C) {
		t.Fatal("C not preserved after marshal round-trip")
	}
}

// ─── sigma.go: Phase 4 — randomized GPV signature proof ──────────────────────

func TestRandomizeGPVSignatureRoundTrip(t *testing.T) {
	gp := SmallGPVParams()
	pk, sk := GPVKeyGen(gp)

	msg := []byte("age>=18")
	sig := GPVSign(sk, msg)
	if !GPVVerify(pk, msg, sig) {
		t.Fatal("base GPV signature must verify")
	}

	rng := rand.Reader
	proof := RandomizeGPVSignature(pk, sig, msg, rng)

	if !VerifyRandomizedGPVSignature(pk, msg, proof) {
		t.Fatal("valid randomized signature rejected")
	}
}

func TestRandomizeGPVSignatureWrongMessage(t *testing.T) {
	gp := SmallGPVParams()
	pk, sk := GPVKeyGen(gp)

	msg := []byte("real-claim")
	sig := GPVSign(sk, msg)

	rng := rand.Reader
	proof := RandomizeGPVSignature(pk, sig, msg, rng)

	if VerifyRandomizedGPVSignature(pk, []byte("wrong-claim"), proof) {
		t.Fatal("randomized signature verified with wrong message")
	}
}

func TestRandomizeGPVSignatureWrongKey(t *testing.T) {
	gp := SmallGPVParams()
	pk1, sk1 := GPVKeyGen(gp)
	pk2, _ := GPVKeyGen(gp)

	msg := []byte("test-message")
	sig := GPVSign(sk1, msg)

	rng := rand.Reader
	proof := RandomizeGPVSignature(pk1, sig, msg, rng)

	// Verify with pk2 should fail.
	if VerifyRandomizedGPVSignature(pk2, msg, proof) {
		t.Fatal("randomized signature verified with wrong public key")
	}
}

func TestRandomizeGPVSignatureMultipleMessages(t *testing.T) {
	gp := SmallGPVParams()
	pk, sk := GPVKeyGen(gp)

	messages := []string{
		"age>=18",
		"country=US",
		"membership-level=gold",
	}

	rng := rand.Reader
	for _, msg := range messages {
		msgBytes := []byte(msg)
		sig := GPVSign(sk, msgBytes)
		proof := RandomizeGPVSignature(pk, sig, msgBytes, rng)
		if !VerifyRandomizedGPVSignature(pk, msgBytes, proof) {
			t.Fatalf("proof failed for message %q", msg)
		}
	}
}

func TestRandomizeGPVSignatureEmptyMessage(t *testing.T) {
	gp := SmallGPVParams()
	pk, sk := GPVKeyGen(gp)

	msg := []byte("")
	sig := GPVSign(sk, msg)
	rng := rand.Reader
	proof := RandomizeGPVSignature(pk, sig, msg, rng)

	if !VerifyRandomizedGPVSignature(pk, msg, proof) {
		t.Fatal("randomized signature failed for empty message")
	}
}

func TestRandomizeGPVSignatureNotLinkable(t *testing.T) {
	gp := SmallGPVParams()
	pk, sk := GPVKeyGen(gp)

	msg := []byte("same-message")
	sig := GPVSign(sk, msg)
	if !GPVVerify(pk, msg, sig) {
		t.Fatal("base sig must verify")
	}

	rng := rand.Reader

	// Two independent randomizations of the same signature.
	proof1 := RandomizeGPVSignature(pk, sig, msg, rng)
	proof2 := RandomizeGPVSignature(pk, sig, msg, rng)

	// Both must verify.
	if !VerifyRandomizedGPVSignature(pk, msg, proof1) {
		t.Fatal("first randomization rejected")
	}
	if !VerifyRandomizedGPVSignature(pk, msg, proof2) {
		t.Fatal("second randomization rejected")
	}

	// The randomized z' must be different (prevents linkability).
	equal := true
	for i := range proof1.ZPrime.Coeffs {
		if proof1.ZPrime.Coeffs[i] != proof2.ZPrime.Coeffs[i] {
			equal = false
			break
		}
	}
	if equal {
		t.Fatal("two randomizations produced identical z' — not randomized")
	}
}

func TestRandomizeGPVSignatureFalcon512(t *testing.T) {
	gp := DefaultGPVParams()
	pk, sk := GPVKeyGen(gp)

	msg := []byte("age>=18")
	sig := GPVSign(sk, msg)
	if !GPVVerify(pk, msg, sig) {
		t.Fatal("Falcon-512 base signature must verify")
	}

	rng := rand.Reader
	proof := RandomizeGPVSignature(pk, sig, msg, rng)

	if !VerifyRandomizedGPVSignature(pk, msg, proof) {
		t.Fatal("Falcon-512 randomized signature rejected")
	}
}

// ─── Benchmarks ───────────────────────────────────────────────────────────────

func BenchmarkGPVGadgetKeyGen(b *testing.B) {
	gp := SmallGPVParams()
	for range b.N {
		GPVGadgetKeyGen(gp)
	}
}

func BenchmarkGPVGadgetSign(b *testing.B) {
	gp := SmallGPVParams()
	_, sk := GPVGadgetKeyGen(gp)
	msg := []byte("benchmark")
	b.ResetTimer()
	for range b.N {
		GPVGadgetSign(sk, msg)
	}
}

func BenchmarkGPVGadgetVerify(b *testing.B) {
	gp := SmallGPVParams()
	pk, sk := GPVGadgetKeyGen(gp)
	msg := []byte("benchmark")
	sig := GPVGadgetSign(sk, msg)
	b.ResetTimer()
	for range b.N {
		GPVGadgetVerify(pk, msg, sig)
	}
}

func BenchmarkSigmaProve(b *testing.B) {
	gp := SmallGPVParams()
	_, sk := GPVKeyGen(gp)
	b.ResetTimer()
	for range b.N {
		Prove(sk, nil)
	}
}

func BenchmarkSigmaVerify(b *testing.B) {
	gp := SmallGPVParams()
	pk, sk := GPVKeyGen(gp)
	proof := Prove(sk, nil)
	b.ResetTimer()
	for range b.N {
		VerifySigma(pk, nil, proof)
	}
}