package ring

import "testing"

// TestNoiseFloodPreservesPlaintext verifies that noise flooding
// doesn't change the underlying plaintext.
func TestNoiseFloodPreservesPlaintext(t *testing.T) {
	kp := DefaultHEParams()
	pk, sk, _ := HEKeyGen(kp)

	for _, bit := range []int{0, 1} {
		ct := HEEncrypt(pk, bit)
		flooded := NoiseFlood(pk, ct, 6) // moderate flood

		got := HEDecrypt(sk, flooded)
		if got != bit {
			t.Fatalf("bit=%d: after noise flood, decrypted to %d", bit, got)
		}
	}
}

// TestNoiseFloodChangesNoise verifies that flooding changes the noise pattern.
func TestNoiseFloodChangesNoise(t *testing.T) {
	kp := DefaultHEParams()
	pk, _, _ := HEKeyGen(kp)

	ct := HEEncrypt(pk, 1)
	flooded := NoiseFlood(pk, ct, 6)

	// Ciphertext should be different after flooding.
	if Equal(ct.U, flooded.U) && Equal(ct.V, flooded.V) {
		t.Fatal("noise flood didn't change ciphertext")
	}
}

// TestWrapVerifyResult verifies the full anti-malleability stack.
func TestWrapVerifyResult(t *testing.T) {
	kp := DefaultHEParams()
	pk, sk, _ := HEKeyGen(kp)
	gp := SmallGPVParams()
	gpvPK, gpvSK := GPVKeyGen(gp)

	ct := HEEncrypt(pk, 1)
	sessionID := []byte("test-session-123")

	result := WrapResult(pk, ct, sessionID, gpvSK)

	// Verify passes.
	if !VerifyResult(result, sessionID, gpvPK) {
		t.Fatal("valid result rejected")
	}

	// Plaintext preserved through all the defenses.
	got := HEDecrypt(sk, result.Ciphertext)
	if got != 1 {
		t.Fatalf("decrypted to %d after wrapping, want 1", got)
	}
}

// TestWrapVerifyWrongSession verifies session binding.
func TestWrapVerifyWrongSession(t *testing.T) {
	kp := DefaultHEParams()
	pk, _, _ := HEKeyGen(kp)
	gp := SmallGPVParams()
	gpvPK, gpvSK := GPVKeyGen(gp)

	ct := HEEncrypt(pk, 0)
	result := WrapResult(pk, ct, []byte("correct-session"), gpvSK)

	// Wrong session ID should fail.
	if VerifyResult(result, []byte("wrong-session"), gpvPK) {
		t.Fatal("wrong session ID accepted")
	}
}

// TestWrapVerifyTamperedCiphertext verifies tamper detection.
func TestWrapVerifyTamperedCiphertext(t *testing.T) {
	kp := DefaultHEParams()
	pk, _, _ := HEKeyGen(kp)
	gp := SmallGPVParams()
	gpvPK, gpvSK := GPVKeyGen(gp)

	ct := HEEncrypt(pk, 1)
	sessionID := []byte("tamper-test")
	result := WrapResult(pk, ct, sessionID, gpvSK)

	// Tamper with ciphertext.
	result.Ciphertext.U.Coeffs[0] = (result.Ciphertext.U.Coeffs[0] + 1) % kp.Ring.Q

	// Verification should fail (tag won't match).
	if VerifyResult(result, sessionID, gpvPK) {
		t.Fatal("tampered ciphertext accepted")
	}
}

// TestSessionWrapper verifies the session-based API.
func TestSessionWrapper(t *testing.T) {
	kp := DefaultHEParams()
	pk, sk, _ := HEKeyGen(kp)
	gp := SmallGPVParams()
	gpvPK, gpvSK := GPVKeyGen(gp)

	session := NewSession(pk, gpvPK, gpvSK)

	// Encrypt, compute, wrap.
	ct0 := HEEncrypt(pk, 1)
	ct1 := HEEncrypt(pk, 1)
	ctXor := HEAdd(ct0, ct1) // 1 XOR 1 = 0

	result := session.Wrap(ctXor)

	// Verify.
	// Need a matching session for verification.
	verifier := &SessionWrapper{
		SessionID: session.SessionID,
		HEPK:      pk,
		GPVPK:     gpvPK,
	}

	if !verifier.Verify(result) {
		t.Fatal("session verification failed")
	}

	got := HEDecrypt(sk, result.Ciphertext)
	if got != 0 {
		t.Fatalf("expected 0 (1 XOR 1), got %d", got)
	}
}

// TestWrapWithoutSignature verifies wrapping without GPV signature.
func TestWrapWithoutSignature(t *testing.T) {
	kp := DefaultHEParams()
	pk, sk, _ := HEKeyGen(kp)

	ct := HEEncrypt(pk, 1)
	sessionID := []byte("unsigned-session")

	// Wrap without signing key.
	result := WrapResult(pk, ct, sessionID, nil)

	// Verify without signing key.
	if !VerifyResult(result, sessionID, nil) {
		t.Fatal("unsigned result rejected")
	}

	got := HEDecrypt(sk, result.Ciphertext)
	if got != 1 {
		t.Fatalf("decrypted to %d, want 1", got)
	}
}

func BenchmarkNoiseFlood(b *testing.B) {
	kp := DefaultHEParams()
	pk, _, _ := HEKeyGen(kp)
	ct := HEEncrypt(pk, 1)
	b.ResetTimer()
	for range b.N {
		NoiseFlood(pk, ct, 8)
	}
}

func BenchmarkWrapResult(b *testing.B) {
	kp := DefaultHEParams()
	pk, _, _ := HEKeyGen(kp)
	gp := SmallGPVParams()
	_, gpvSK := GPVKeyGen(gp)
	ct := HEEncrypt(pk, 1)
	sessionID := []byte("benchmark")
	b.ResetTimer()
	for range b.N {
		WrapResult(pk, ct, sessionID, gpvSK)
	}
}