package ring

import "testing"

// TestHEEncryptDecrypt verifies single-bit encrypt/decrypt.
func TestHEEncryptDecrypt(t *testing.T) {
	kp := DefaultHEParams()
	pk, sk, _ := HEKeyGen(kp)

	for _, bit := range []int{0, 1} {
		ct := HEEncrypt(pk, bit)
		got := HEDecrypt(sk, ct)
		if got != bit {
			t.Fatalf("bit=%d: decrypted to %d", bit, got)
		}
	}
}

// TestHEAdd verifies homomorphic addition (XOR).
func TestHEAdd(t *testing.T) {
	kp := DefaultHEParams()
	pk, sk, _ := HEKeyGen(kp)

	// Test all XOR combinations.
	for _, a := range []int{0, 1} {
		for _, b := range []int{0, 1} {
			ctA := HEEncrypt(pk, a)
			ctB := HEEncrypt(pk, b)
			ctSum := HEAdd(ctA, ctB)
			got := HEDecrypt(sk, ctSum)
			want := a ^ b
			if got != want {
				t.Fatalf("%d XOR %d: got %d, want %d", a, b, got, want)
			}
		}
	}
}

// TestHENot verifies homomorphic NOT.
func TestHENot(t *testing.T) {
	kp := DefaultHEParams()
	pk, sk, _ := HEKeyGen(kp)

	for _, bit := range []int{0, 1} {
		ct := HEEncrypt(pk, bit)
		ctNot := HENot(ct)
		got := HEDecrypt(sk, ctNot)
		want := 1 - bit
		if got != want {
			t.Fatalf("NOT(%d): got %d, want %d", bit, got, want)
		}
	}
}

// TestHEMul verifies homomorphic multiplication (AND).
func TestHEMul(t *testing.T) {
	kp := DefaultHEParams()
	pk, sk, rlk := HEKeyGen(kp)

	for _, a := range []int{0, 1} {
		for _, b := range []int{0, 1} {
			ctA := HEEncrypt(pk, a)
			ctB := HEEncrypt(pk, b)
			ctProd := HEMul(ctA, ctB, rlk)
			got := HEDecrypt(sk, ctProd)
			want := a & b
			if got != want {
				t.Fatalf("%d AND %d: got %d, want %d", a, b, got, want)
			}
		}
	}
}

// TestHERerandomize verifies that rerandomization preserves the plaintext.
func TestHERerandomize(t *testing.T) {
	kp := DefaultHEParams()
	pk, sk, _ := HEKeyGen(kp)

	for _, bit := range []int{0, 1} {
		ct := HEEncrypt(pk, bit)
		ctRerand := Rerandomize(pk, ct)

		got := HEDecrypt(sk, ctRerand)
		if got != bit {
			t.Fatalf("rerandomize(%d): decrypted to %d", bit, got)
		}

		// Verify ciphertexts are actually different.
		if Equal(ct.U, ctRerand.U) && Equal(ct.V, ctRerand.V) {
			t.Fatal("rerandomized ciphertext is identical — randomness failure")
		}
	}
}

// TestHEChainedAdd verifies multiple additions stay within noise budget.
func TestHEChainedAdd(t *testing.T) {
	kp := DefaultHEParams()
	pk, sk, _ := HEKeyGen(kp)

	// XOR 16 bits together: result should be 0 (even count of 1s) or 1 (odd).
	bits := []int{1, 0, 1, 1, 0, 1, 0, 1, 1, 0, 1, 0, 1, 1, 0, 0}
	expected := 0
	for _, b := range bits {
		expected ^= b
	}

	acc := HEEncrypt(pk, bits[0])
	for _, b := range bits[1:] {
		ct := HEEncrypt(pk, b)
		acc = HEAdd(acc, ct)
	}

	got := HEDecrypt(sk, acc)
	if got != expected {
		t.Fatalf("chained XOR: got %d, want %d (noise estimate: %.1f)", got, expected, acc.NoiseEstimate)
	}
}

// TestHEBGVNoiseParity verifies that BGV encryption produces even noise + message.
func TestHEBGVNoiseParity(t *testing.T) {
	kp := DefaultHEParams()
	pk, sk, _ := HEKeyGen(kp)
	p := kp.Ring

	for _, bit := range []int{0, 1} {
		ct := HEEncrypt(pk, bit)

		// Compute v - s·u
		uNTT := ct.U.Clone()
		NTT(uNTT)
		su := MulPointwise(sk.S, uNTT)
		INTT(su)
		noisy := Sub(ct.V, su)

		c := noisy.Coeffs[0]
		q := p.Q
		if c > q/2 {
			c = q - c
		}
		parity := c % 2

		t.Logf("bit=%d: coeff[0]=%d, centered=%d, parity=%d (want parity=%d)",
			bit, noisy.Coeffs[0], c, parity, bit)

		if int(parity) != bit {
			t.Errorf("bit=%d: wrong parity %d", bit, parity)
		}
	}
}

// TestHEMulDeg2 tests the tensor product without relinearization.
// This isolates whether the issue is in the tensor product or relinearization.
func TestHEMulDeg2(t *testing.T) {
	kp := DefaultHEParams()
	pk, sk, _ := HEKeyGen(kp)
	p := kp.Ring

	for _, a := range []int{0, 1} {
		for _, b := range []int{0, 1} {
			ctA := HEEncrypt(pk, a)
			ctB := HEEncrypt(pk, b)

			// Tensor product: c0 = v_a * v_b, c1 = v_a*u_b + u_a*v_b, c2 = u_a * u_b
			c0 := Mul(ctA.V, ctB.V)
			c1a := Mul(ctA.V, ctB.U)
			c1b := Mul(ctA.U, ctB.V)
			c1 := Add(c1a, c1b)
			c2 := Mul(ctA.U, ctB.U)

			// Degree-2 decryption: c0 - s*c1 + s²*c2
			// s is in NTT form
			sCoeff := sk.S.Clone()

			c1NTT := c1.Clone()
			NTT(c1NTT)
			sc1 := MulPointwise(sCoeff, c1NTT)
			INTT(sc1)

			// s²
			s2 := MulPointwise(sCoeff, sCoeff)
			c2NTT := c2.Clone()
			NTT(c2NTT)
			s2c2 := MulPointwise(s2, c2NTT)
			INTT(s2c2)

			result := Sub(c0, sc1)
			result = Add(result, s2c2)

			// Decode: result[0] mod 2
			q := p.Q
			c := result.Coeffs[0]
			if c > q/2 {
				c = q - c
			}
			got := int(c % 2)
			want := a & b

			t.Logf("%d AND %d: coeff[0]=%d, centered=%d, mod2=%d (want %d)",
				a, b, result.Coeffs[0], c, got, want)

			if got != want {
				t.Errorf("%d AND %d: got %d, want %d", a, b, got, want)
			}
		}
	}
}

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

func BenchmarkHEAdd(b *testing.B) {
	kp := DefaultHEParams()
	pk, _, _ := HEKeyGen(kp)
	ct0 := HEEncrypt(pk, 0)
	ct1 := HEEncrypt(pk, 1)
	b.ResetTimer()
	for range b.N {
		HEAdd(ct0, ct1)
	}
}

func BenchmarkHEMul(b *testing.B) {
	kp := DefaultHEParams()
	pk, _, rlk := HEKeyGen(kp)
	ct0 := HEEncrypt(pk, 1)
	ct1 := HEEncrypt(pk, 1)
	b.ResetTimer()
	for range b.N {
		HEMul(ct0, ct1, rlk)
	}
}