// Searchable encryption via homomorphic pattern matching (the recognizer).
//
// The recognizer evaluates pattern-matching circuits on encrypted data.
// Neither the data nor the query pattern are revealed to the evaluator.
//
// Architecture:
//
//   Data owner:
//     1. Encrypts each bit of the data: c_i = HEEncrypt(pk, data[i])
//     2. Sends encrypted data to the server.
//
//   Query maker:
//     1. Defines a pattern as a Boolean circuit (AND, XOR, NOT gates).
//     2. Encrypts query parameters: q_j = HEEncrypt(pk, pattern[j])
//     3. Sends encrypted query to the server.
//
//   Server (evaluator):
//     1. Evaluates the circuit homomorphically on encrypted data + query.
//     2. Returns encrypted result bit: Enc(match?) to the query maker.
//     3. Never sees plaintext data, query, or result.
//
//   Query maker:
//     1. Decrypts the result: match = HEDecrypt(sk, result)
//     2. Learns only whether the pattern matched, nothing else.
//
// Circuit structure for byte-level matching:
//
//   "Does position i contain byte value v?"
//   For each bit position b ∈ [0,7]:
//     match_b = XNOR(data[i*8+b], v_b)      // bit equality
//             = NOT(XOR(data[i*8+b], v_b))
//   match_i = AND(match_0, match_1, ..., match_7)   // all 8 bits match
//
// The circuit depth is:
//   - XOR: depth 0 (addition, no noise growth)
//   - NOT: depth 0 (constant addition)
//   - AND: depth 1 per gate (multiplication, quadratic noise growth)
//
// For byte matching: 1 level of AND (8-way AND via tree of 2-ANDs = 3 levels).
// With HE64 parameters (1 multiplication level), we can do 1 AND gate per path,
// which suffices for single-bit pattern tests. For full byte matching,
// we'd need a deeper noise budget or modulus switching.
//
// Practical approach: use the 1-level budget for the critical AND gate,
// and batch the XOR/NOT operations (which are free in noise) to reduce
// the circuit to a single AND depth.

package ring

import (
	"crypto/rand"
	"io"
)

// EncryptedBitVector is a sequence of homomorphically encrypted bits.
type EncryptedBitVector struct {
	Bits   []*HECiphertext
	PK     *KEMPublicKey
	params KEMParams
}

// EncryptBits encrypts a byte slice as a bit vector.
// Each bit gets its own ciphertext.
func EncryptBits(pk *KEMPublicKey, data []byte) *EncryptedBitVector {
	return EncryptBitsFrom(pk, data, rand.Reader)
}

// EncryptBitsFrom encrypts using the given randomness source.
func EncryptBitsFrom(pk *KEMPublicKey, data []byte, rng io.Reader) *EncryptedBitVector {
	bits := make([]*HECiphertext, len(data)*8)
	for i, b := range data {
		for j := range 8 {
			bit := int((b >> j) & 1)
			bits[i*8+j] = HEEncryptFrom(pk, bit, rng)
		}
	}
	return &EncryptedBitVector{
		Bits:   bits,
		PK:     pk,
		params: pk.P,
	}
}

// DecryptBits decrypts an encrypted bit vector back to bytes.
func DecryptBits(sk *KEMSecretKey, ebv *EncryptedBitVector) []byte {
	nBytes := (len(ebv.Bits) + 7) / 8
	result := make([]byte, nBytes)
	for i, ct := range ebv.Bits {
		bit := HEDecrypt(sk, ct)
		if bit != 0 {
			result[i/8] |= 1 << (i % 8)
		}
	}
	return result
}

// EncryptedPattern represents an encrypted search pattern.
type EncryptedPattern struct {
	// PatternBits are the encrypted pattern bits.
	PatternBits []*HECiphertext

	// Mask indicates which positions are "don't care" (wildcard).
	// nil means all positions must match.
	Mask []bool
}

// EncryptPattern encrypts a pattern byte for searching.
// Wildcards can be specified via the mask (true = must match, false = wildcard).
func EncryptPattern(pk *KEMPublicKey, pattern []byte, mask []bool) *EncryptedPattern {
	bits := make([]*HECiphertext, len(pattern)*8)
	for i, b := range pattern {
		for j := range 8 {
			bit := int((b >> j) & 1)
			bits[i*8+j] = HEEncrypt(pk, bit)
		}
	}

	// Expand mask to bit level.
	var bitMask []bool
	if mask != nil {
		bitMask = make([]bool, len(pattern)*8)
		for i, m := range mask {
			for j := range 8 {
				bitMask[i*8+j] = m
			}
		}
	}

	return &EncryptedPattern{
		PatternBits: bits,
		Mask:        bitMask,
	}
}

// MatchBit evaluates whether a single encrypted data bit equals
// a single encrypted pattern bit. Returns Enc(data_bit XNOR pattern_bit).
//
// XNOR(a, b) = NOT(XOR(a, b)) = NOT(a + b) = 1 + a + b (mod 2).
// In the ciphertext domain: HEAdd(HEAdd(enc_a, enc_b), Enc(1))... but
// that's HEAdd(HEAdd(a, b), one). Actually NOT(XOR(a,b)):
// = HENot(HEAdd(a, b)) = HENot(HEXOR(a, b))
//
// This is depth-0 (no multiplication), so it's free in noise budget.
func MatchBit(dataBit, patternBit *HECiphertext) *HECiphertext {
	xored := HEXOR(dataBit, patternBit) // 1 if bits differ
	return HENot(xored)                 // 1 if bits equal
}

// MatchByte evaluates whether 8 consecutive data bits equal the
// pattern byte. Returns Enc(all_8_bits_match).
//
// This requires AND-ing 8 bit-equality results together.
// With depth-1 budget, we can do ONE AND gate total.
// For 8 bits, we need log2(8) = 3 levels of AND — too deep.
//
// Solution: reduce to a single AND using batch XOR.
// The trick: instead of checking each bit individually,
// compute the Hamming distance homomorphically.
//
// Alternative single-depth approach:
// Σ (data_b XOR pattern_b) = 0 iff all bits match.
// In BGV mod 2: sum of XORs is just XOR of all bit-differences.
// If any bit differs, the sum is 1 (odd count = nonzero).
// But sum mod 2 doesn't detect even numbers of differences.
//
// Better: use the identity for AND of XNOR results.
// For a single pair, just return the XNOR (depth 0).
// The caller decides how many pairs to AND together based on budget.
func MatchByte(dataBits, patternBits []*HECiphertext, rlk *RelinearizationKey) *HECiphertext {
	if len(dataBits) != 8 || len(patternBits) != 8 {
		panic("MatchByte requires exactly 8 bits each")
	}

	// Compute XNOR for each bit position.
	matches := make([]*HECiphertext, 8)
	for i := range 8 {
		matches[i] = MatchBit(dataBits[i], patternBits[i])
	}

	// AND tree: depth log2(8) = 3.
	// Level 1: AND pairs → 4 results (depth 1).
	// Level 2: AND pairs → 2 results (depth 2).
	// Level 3: AND pair → 1 result (depth 3).
	//
	// With depth-1 budget, we can only do level 1.
	// Return the AND of the first pair as a partial match indicator.
	//
	// For full byte matching, use MatchByteSingle which reduces
	// to a single multiplication.
	return andTree(matches, rlk)
}

// MatchByteSingle evaluates byte equality using a single AND gate.
//
// Trick: XOR all 8 bit-differences together, then check if result is 0.
// XOR(d0⊕p0, d1⊕p1, ..., d7⊕p7) = 0 iff ALL bits match... no, it's 0
// iff an EVEN number of bits differ. That's not what we want.
//
// Alternative trick using the OR-NOT equivalence:
// All bits match = NOT(any bit differs)
// Any bit differs = OR(d0⊕p0, d1⊕p1, ..., d7⊕p7)
// OR = NOT(AND(NOT(x), NOT(y), ...))
//
// So: allMatch = NOT(OR(diffs)) = AND(NOT(d0⊕p0), NOT(d1⊕p1), ...)
//             = AND(XNOR(d0,p0), XNOR(d1,p1), ...)
// This is the same 8-way AND, still depth 3.
//
// With depth-1 limitation, the practical approach is:
// Split into two 4-bit groups, AND each group (1 level), then
// the caller can AND the two groups in a second pass with fresh keys.
//
// Or: use the "indicator polynomial" trick:
// f(x) = 1 - x evaluates to 1 at x=0 and 0 at x=1.
// For the sum s = Σ (d_i ⊕ p_i), we want f(s) = 1 iff s = 0.
// But s is the XOR (mod 2), not the integer sum, so this doesn't work directly.
//
// Practical depth-1 solution:
// AND the first pair, returning Enc(bit0_matches AND bit1_matches).
// The server can report partial match results for each pair of bits,
// and the client combines them.
func MatchByteSingle(dataBits, patternBits []*HECiphertext, rlk *RelinearizationKey) []*HECiphertext {
	if len(dataBits) != 8 || len(patternBits) != 8 {
		panic("MatchByteSingle requires exactly 8 bits each")
	}

	// Compute XNOR (bit equality) for each position.
	xnors := make([]*HECiphertext, 8)
	for i := range 8 {
		xnors[i] = MatchBit(dataBits[i], patternBits[i])
	}

	// AND adjacent pairs (depth 1).
	results := make([]*HECiphertext, 4)
	for i := range 4 {
		results[i] = HEAND(xnors[2*i], xnors[2*i+1], rlk)
	}

	return results
}

// SearchAt evaluates whether encrypted data starting at position pos
// matches the encrypted pattern. Returns per-pair AND results.
func SearchAt(data *EncryptedBitVector, pattern *EncryptedPattern, pos int, rlk *RelinearizationKey) []*HECiphertext {
	patLen := len(pattern.PatternBits) / 8 // pattern length in bytes
	if pos*8+patLen*8 > len(data.Bits) {
		return nil // pattern would extend beyond data
	}

	// For each byte of the pattern, compute pair-AND results.
	var allResults []*HECiphertext
	for byteIdx := range patLen {
		dataStart := (pos + byteIdx) * 8
		patStart := byteIdx * 8

		// Skip wildcard bytes.
		if pattern.Mask != nil {
			isWild := true
			for b := range 8 {
				if pattern.Mask[patStart+b] {
					isWild = false
					break
				}
			}
			if isWild {
				continue
			}
		}

		pairs := MatchByteSingle(
			data.Bits[dataStart:dataStart+8],
			pattern.PatternBits[patStart:patStart+8],
			rlk,
		)
		allResults = append(allResults, pairs...)
	}

	return allResults
}

// andTree computes a tree of AND operations.
// If the depth exceeds the noise budget, the result may not decrypt correctly.
func andTree(inputs []*HECiphertext, rlk *RelinearizationKey) *HECiphertext {
	if len(inputs) == 0 {
		return nil
	}
	if len(inputs) == 1 {
		return inputs[0]
	}

	// Reduce pairwise.
	for len(inputs) > 1 {
		var next []*HECiphertext
		for i := 0; i+1 < len(inputs); i += 2 {
			next = append(next, HEAND(inputs[i], inputs[i+1], rlk))
		}
		if len(inputs)%2 == 1 {
			next = append(next, inputs[len(inputs)-1])
		}
		inputs = next
	}
	return inputs[0]
}

// RecognizerResult holds the outcome of a homomorphic pattern match.
type RecognizerResult struct {
	// MatchPairs are the per-pair AND results.
	// With depth-1 budget, each result covers 2 bit positions.
	// ALL must decrypt to 1 for a full match.
	MatchPairs []*HECiphertext

	// Position is the data position that was tested.
	Position int
}

// DecryptResult decrypts the recognizer output.
// Returns true if ALL pair results are 1 (full match).
func DecryptResult(sk *KEMSecretKey, result *RecognizerResult) bool {
	for _, ct := range result.MatchPairs {
		if HEDecrypt(sk, ct) != 1 {
			return false
		}
	}
	return true
}

// Recognize runs the full searchable encryption pipeline.
//
// 1. Encrypts the data and pattern.
// 2. Homomorphically evaluates matches at each position.
// 3. Returns encrypted results for each position.
//
// The server sees only ciphertexts. The data owner decrypts the results
// to learn which positions matched.
func Recognize(pk *KEMPublicKey, sk *KEMSecretKey, rlk *RelinearizationKey,
	data []byte, pattern []byte) []int {

	// Encrypt data and pattern.
	encData := EncryptBits(pk, data)
	encPattern := EncryptPattern(pk, pattern, nil)

	patLen := len(pattern)
	var matches []int

	// Try each position.
	for pos := 0; pos <= len(data)-patLen; pos++ {
		pairs := SearchAt(encData, encPattern, pos, rlk)
		if pairs == nil {
			continue
		}

		result := &RecognizerResult{
			MatchPairs: pairs,
			Position:   pos,
		}

		if DecryptResult(sk, result) {
			matches = append(matches, pos)
		}
	}

	return matches
}