// Searchable encryption via homomorphic pattern matching.
package ring

import (
	"crypto/rand"
	"io"
)

type EncryptedBitVector struct {
	Bits   []*HECiphertext
	PK     *KEMPublicKey
	params KEMParams
}

func EncryptBits(pk *KEMPublicKey, data []byte) (ebv *EncryptedBitVector) {
	return EncryptBitsFrom(pk, data, rand.Reader)
}

func EncryptBitsFrom(pk *KEMPublicKey, data []byte, rng io.Reader) (ebv *EncryptedBitVector) {
	bits := []*HECiphertext{:int32(len(data)) * 8}
	for i := int32(0); i < int32(len(data)); i++ {
		b := data[i]
		for j := int32(0); j < 8; j++ {
			bit := int32((b >> uint32(j)) & 1)
			bits[i*8+j] = HEEncryptFrom(pk, bit, rng)
		}
	}
	return &EncryptedBitVector{
		Bits:   bits,
		PK:     pk,
		params: pk.P,
	}
}

func DecryptBits(sk *KEMSecretKey, ebv *EncryptedBitVector) (result []byte) {
	nBytes := (int32(len(ebv.Bits)) + 7) / 8
	result = []byte{:nBytes}
	for i := int32(0); i < int32(len(ebv.Bits)); i++ {
		bit := HEDecrypt(sk, ebv.Bits[i])
		if bit != 0 {
			result[i/8] |= 1 << uint32(i%8)
		}
	}
	return result
}

type EncryptedPattern struct {
	PatternBits []*HECiphertext
	Mask        []bool
}

func EncryptPattern(pk *KEMPublicKey, pattern []byte, mask []bool) (ep *EncryptedPattern) {
	bits := []*HECiphertext{:int32(len(pattern)) * 8}
	for i := int32(0); i < int32(len(pattern)); i++ {
		b := pattern[i]
		for j := int32(0); j < 8; j++ {
			bit := int32((b >> uint32(j)) & 1)
			bits[i*8+j] = HEEncrypt(pk, bit)
		}
	}

	var bitMask []bool
	if mask != nil {
		bitMask = []bool{:int32(len(pattern)) * 8}
		for i := int32(0); i < int32(len(mask)); i++ {
			for j := int32(0); j < 8; j++ {
				bitMask[i*8+j] = mask[i]
			}
		}
	}

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

func MatchBit(dataBit, patternBit *HECiphertext) (result *HECiphertext) {
	xored := HEXOR(dataBit, patternBit)
	return HENot(xored)
}

func MatchByte(dataBits, patternBits []*HECiphertext, rlk *RelinearizationKey) (result *HECiphertext) {
	if int32(len(dataBits)) != 8 || int32(len(patternBits)) != 8 {
		panic("MatchByte requires exactly 8 bits each")
	}

	matches := []*HECiphertext{:8}
	for i := int32(0); i < 8; i++ {
		matches[i] = MatchBit(dataBits[i], patternBits[i])
	}

	return andTree(matches, rlk)
}

func MatchByteSingle(dataBits, patternBits []*HECiphertext, rlk *RelinearizationKey) (results []*HECiphertext) {
	if int32(len(dataBits)) != 8 || int32(len(patternBits)) != 8 {
		panic("MatchByteSingle requires exactly 8 bits each")
	}

	xnors := []*HECiphertext{:8}
	for i := int32(0); i < 8; i++ {
		xnors[i] = MatchBit(dataBits[i], patternBits[i])
	}

	results = []*HECiphertext{:4}
	for i := int32(0); i < 4; i++ {
		results[i] = HEAND(xnors[2*i], xnors[2*i+1], rlk)
	}
	return results
}

func SearchAt(data *EncryptedBitVector, pattern *EncryptedPattern, pos int32, rlk *RelinearizationKey) (allResults []*HECiphertext) {
	patLen := int32(len(pattern.PatternBits)) / 8
	if pos*8+patLen*8 > int32(len(data.Bits)) {
		return nil
	}

	for byteIdx := int32(0); byteIdx < patLen; byteIdx++ {
		dataStart := (pos + byteIdx) * 8
		patStart := byteIdx * 8

		if pattern.Mask != nil {
			isWild := true
			for b := int32(0); b < 8; b++ {
				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
}

func andTree(inputs []*HECiphertext, rlk *RelinearizationKey) (result *HECiphertext) {
	if int32(len(inputs)) == 0 {
		return nil
	}
	if int32(len(inputs)) == 1 {
		return inputs[0]
	}

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

type RecognizerResult struct {
	MatchPairs []*HECiphertext
	Position   int32
}

func DecryptResult(sk *KEMSecretKey, result *RecognizerResult) (ok bool) {
	for i := int32(0); i < int32(len(result.MatchPairs)); i++ {
		if HEDecrypt(sk, result.MatchPairs[i]) != 1 {
			return false
		}
	}
	return true
}

func Recognize(pk *KEMPublicKey, sk *KEMSecretKey, rlk *RelinearizationKey,
	data []byte, pattern []byte) (matches []int32) {

	encData := EncryptBits(pk, data)
	encPattern := EncryptPattern(pk, pattern, nil)

	patLen := int32(len(pattern))

	for pos := int32(0); pos <= int32(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
}