package ring

import "testing"

// TestMatchBitEquality verifies single-bit equality test.
func TestMatchBitEquality(t *testing.T) {
	kp := DefaultHEParams()
	pk, sk, _ := HEKeyGen(kp)

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

			want := 0
			if a == b {
				want = 1
			}
			if got != want {
				t.Fatalf("MatchBit(%d, %d): got %d, want %d", a, b, got, want)
			}
		}
	}
}

// TestEncryptDecryptBits verifies bit vector round-trip.
func TestEncryptDecryptBits(t *testing.T) {
	kp := DefaultHEParams()
	pk, sk, _ := HEKeyGen(kp)

	data := []byte{0xAB, 0xCD}
	enc := EncryptBits(pk, data)
	dec := DecryptBits(sk, enc)

	if len(dec) != len(data) {
		t.Fatalf("length mismatch: got %d, want %d", len(dec), len(data))
	}
	for i := range data {
		if dec[i] != data[i] {
			t.Fatalf("byte[%d]: got %02x, want %02x", i, dec[i], data[i])
		}
	}
}

// TestMatchByteSingle verifies per-pair AND for byte matching.
func TestMatchByteSingle(t *testing.T) {
	kp := DefaultHEParams()
	pk, sk, rlk := HEKeyGen(kp)

	// Matching bytes: all pair-ANDs should be 1.
	dataByte := byte(0xA5)
	patByte := byte(0xA5)

	dataBits := make([]*HECiphertext, 8)
	patBits := make([]*HECiphertext, 8)
	for i := range 8 {
		dataBits[i] = HEEncrypt(pk, int((dataByte>>i)&1))
		patBits[i] = HEEncrypt(pk, int((patByte>>i)&1))
	}

	pairs := MatchByteSingle(dataBits, patBits, rlk)
	if len(pairs) != 4 {
		t.Fatalf("expected 4 pair results, got %d", len(pairs))
	}

	for i, ct := range pairs {
		got := HEDecrypt(sk, ct)
		if got != 1 {
			t.Fatalf("pair[%d]: got %d, want 1 (matching bytes)", i, got)
		}
	}
}

// TestMatchByteSingleMismatch verifies that a different byte produces
// at least one pair with value 0.
func TestMatchByteSingleMismatch(t *testing.T) {
	kp := DefaultHEParams()
	pk, sk, rlk := HEKeyGen(kp)

	dataByte := byte(0xA5)
	patByte := byte(0x5A) // every bit is different

	dataBits := make([]*HECiphertext, 8)
	patBits := make([]*HECiphertext, 8)
	for i := range 8 {
		dataBits[i] = HEEncrypt(pk, int((dataByte>>i)&1))
		patBits[i] = HEEncrypt(pk, int((patByte>>i)&1))
	}

	pairs := MatchByteSingle(dataBits, patBits, rlk)

	anyZero := false
	for _, ct := range pairs {
		if HEDecrypt(sk, ct) == 0 {
			anyZero = true
			break
		}
	}
	if !anyZero {
		t.Fatal("all pairs returned 1 for mismatched bytes")
	}
}

// TestRecognize verifies the full searchable encryption pipeline.
func TestRecognize(t *testing.T) {
	kp := DefaultHEParams()
	pk, sk, rlk := HEKeyGen(kp)

	// Search for pattern in data.
	data := []byte("hello")
	pattern := []byte("ll")

	matches := Recognize(pk, sk, rlk, data, pattern)

	t.Logf("data=%q, pattern=%q, matches=%v", data, pattern, matches)

	// "ll" should appear at position 2 in "hello".
	found := false
	for _, pos := range matches {
		if pos == 2 {
			found = true
		}
	}
	if !found {
		t.Fatalf("expected match at position 2, got %v", matches)
	}
}

// TestRecognizeNoMatch verifies no false positives.
func TestRecognizeNoMatch(t *testing.T) {
	kp := DefaultHEParams()
	pk, sk, rlk := HEKeyGen(kp)

	data := []byte("abc")
	pattern := []byte("xy")

	matches := Recognize(pk, sk, rlk, data, pattern)

	if len(matches) > 0 {
		t.Fatalf("expected no matches, got %v", matches)
	}
}

// TestRecognizeSingleByte verifies single-byte pattern search.
func TestRecognizeSingleByte(t *testing.T) {
	kp := DefaultHEParams()
	pk, sk, rlk := HEKeyGen(kp)

	data := []byte{0x42, 0x43, 0x42, 0x44}
	pattern := []byte{0x42}

	matches := Recognize(pk, sk, rlk, data, pattern)

	t.Logf("pattern=0x42, matches=%v", matches)

	// 0x42 appears at positions 0 and 2.
	if len(matches) != 2 {
		t.Fatalf("expected 2 matches, got %d: %v", len(matches), matches)
	}
}

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

func BenchmarkMatchByteSingle(b *testing.B) {
	kp := DefaultHEParams()
	pk, _, rlk := HEKeyGen(kp)
	dataBits := make([]*HECiphertext, 8)
	patBits := make([]*HECiphertext, 8)
	for i := range 8 {
		dataBits[i] = HEEncrypt(pk, i%2)
		patBits[i] = HEEncrypt(pk, i%2)
	}
	b.ResetTimer()
	for range b.N {
		MatchByteSingle(dataBits, patBits, rlk)
	}
}