package crypto

import "testing"

func TestBitRev6(t *testing.T) {
	// 6-bit reversal: 0b000001 -> 0b100000 = 32
	if got := bitRev6(1); got != 32 {
		t.Errorf("bitRev6(1) = %d, want 32", got)
	}
	// Identity: 0 -> 0, 63 -> 63
	if got := bitRev6(0); got != 0 {
		t.Errorf("bitRev6(0) = %d, want 0", got)
	}
	if got := bitRev6(63); got != 63 {
		t.Errorf("bitRev6(63) = %d, want 63", got)
	}
	// 0b010101 = 21 -> 0b101010 = 42
	if got := bitRev6(21); got != 42 {
		t.Errorf("bitRev6(21) = %d, want 42", got)
	}
}

func TestNTTRoundTrip(t *testing.T) {
	// Create a polynomial with known coefficients.
	var a [HamN]uint16
	for i := range a {
		a[i] = uint16(i % HamP)
	}
	original := a

	// Forward NTT then inverse should recover original.
	ntt64(&a)
	intt64(&a)

	for i := range a {
		if a[i] != original[i] {
			t.Errorf("round-trip failed at [%d]: got %d, want %d", i, a[i], original[i])
		}
	}
}

func TestNTTZero(t *testing.T) {
	var a [HamN]uint16
	ntt64(&a)
	for i, v := range a {
		if v != 0 {
			t.Errorf("NTT of zero poly: a[%d] = %d, want 0", i, v)
		}
	}
}

func TestNTTConvolution(t *testing.T) {
	// Verify that pointwise multiply in NTT domain corresponds to
	// polynomial multiplication mod (x^64 + 1) mod 257.

	// f(x) = 1 + x
	var f [HamN]uint16
	f[0] = 1
	f[1] = 1

	// g(x) = 1 + x
	var g [HamN]uint16
	g[0] = 1
	g[1] = 1

	// f*g mod (x^64+1) = 1 + 2x + x^2 (since degree < 64, no wraparound).
	ntt64(&f)
	ntt64(&g)

	var h [HamN]uint16
	for i := range h {
		h[i] = mod257(int(f[i]) * int(g[i]))
	}

	intt64(&h)

	// Expected: h[0]=1, h[1]=2, h[2]=1, rest 0.
	if h[0] != 1 {
		t.Errorf("h[0] = %d, want 1", h[0])
	}
	if h[1] != 2 {
		t.Errorf("h[1] = %d, want 2", h[1])
	}
	if h[2] != 1 {
		t.Errorf("h[2] = %d, want 1", h[2])
	}
	for i := 3; i < HamN; i++ {
		if h[i] != 0 {
			t.Errorf("h[%d] = %d, want 0", i, h[i])
		}
	}
}

func TestNTTWraparound(t *testing.T) {
	// Test the negacyclic property: x^64 ≡ -1 mod (x^64 + 1).
	// f(x) = x^63, g(x) = x → f*g = x^64 ≡ -1 ≡ 256 mod 257.

	var f [HamN]uint16
	f[63] = 1

	var g [HamN]uint16
	g[1] = 1

	ntt64(&f)
	ntt64(&g)

	var h [HamN]uint16
	for i := range h {
		h[i] = mod257(int(f[i]) * int(g[i]))
	}

	intt64(&h)

	// Expected: h[0] = 256 (= -1 mod 257), rest 0.
	if h[0] != 256 {
		t.Errorf("h[0] = %d, want 256 (-1 mod 257)", h[0])
	}
	for i := 1; i < HamN; i++ {
		if h[i] != 0 {
			t.Errorf("h[%d] = %d, want 0", i, h[i])
		}
	}
}

func TestPowMod(t *testing.T) {
	// 3 is a generator of Z_257*. 3^256 ≡ 1 mod 257 (Fermat).
	if got := powMod(3, 256); got != 1 {
		t.Errorf("3^256 mod 257 = %d, want 1", got)
	}
	// 3^128 ≡ 256 ≡ -1 mod 257 (since order is 256, half-order gives -1).
	if got := powMod(3, 128); got != 256 {
		t.Errorf("3^128 mod 257 = %d, want 256", got)
	}
	// 9^64 = (3^2)^64 = 3^128 ≡ -1 mod 257.
	if got := powMod(9, 64); got != 256 {
		t.Errorf("9^64 mod 257 = %d, want 256", got)
	}
}