package ring

import (
	"crypto/rand"
	"io"

	"golang.org/x/crypto/sha3"
)

// Homomorphic evaluation on Ring-LWE ciphertexts.
//
// A ciphertext (u, v) encrypts message m under public key (a, b = a·s + e):
//   u = a·r + e1
//   v = b·r + e2 + encode(m)
//
// Decryption: decode(v - s·u) = m + noise
//
// The noise term grows with operations:
//   - Addition: noise_add = noise1 + noise2 (linear growth)
//   - Multiplication: noise_mul ≈ noise1·noise2 (quadratic growth)
//
// The noise budget determines how many operations can be performed before
// decryption fails. This is the "somewhat homomorphic" (SHE) regime.
//
// For the recognizer use case (searchable encryption), we need:
//   - Homomorphic XOR (= addition mod 2 of encoded bits)
//   - Homomorphic AND (= multiplication of encoded bits)
//   - Bounded depth circuits (pattern matching has fixed depth)

// HEParams holds parameters for homomorphic evaluation.
type HEParams struct {
	KEM KEMParams

	// RelinKey is the relinearization key for converting degree-2
	// ciphertexts back to degree-1 after multiplication.
	// Generated from the secret key.
	RelinKey *RelinearizationKey
}

// HECiphertext is an Ring-LWE ciphertext supporting homomorphic operations.
// It wraps KEMCiphertext with noise tracking.
type HECiphertext struct {
	U *Poly // first component
	V *Poly // second component

	// NoiseEstimate tracks the estimated noise magnitude.
	// When this exceeds q/4, decryption will fail.
	NoiseEstimate float64

	params KEMParams
}

// RelinearizationKey enables conversion of degree-2 ciphertexts
// (produced by multiplication) back to degree-1.
//
// It's an encryption of s² under the public key:
//   rk[i] = (a_i, a_i·s + e_i + (s²)·base^i)
// where base is the decomposition base for noise management.
type RelinearizationKey struct {
	A []*Poly // uniform components
	B []*Poly // a·s + e + s²·base^i
	L int     // number of decomposition levels
}

// DefaultHEParams returns parameters tuned for homomorphic evaluation.
// Uses HE64 ring (n=64, q=10000769) with eta=1 for tight noise.
// Supports depth-1 multiplicative circuits (one AND gate per path).
func DefaultHEParams() KEMParams {
	return KEMParams{
		Ring:         HE64(),
		Eta1:         1,
		Eta2:         1,
		SharedKeyLen: 32,
	}
}

// HEKeyGen generates keys for homomorphic evaluation.
// Uses BGV-structured keygen: b = a·s + 2·e (noise is even).
func HEKeyGen(kp KEMParams) (*KEMPublicKey, *KEMSecretKey, *RelinearizationKey) {
	return HEKeyGenFrom(kp, rand.Reader)
}

// HEKeyGenFrom generates HE keys from the given randomness source.
func HEKeyGenFrom(kp KEMParams, rng io.Reader) (*KEMPublicKey, *KEMSecretKey, *RelinearizationKey) {
	p := kp.Ring

	// a ← uniform
	a := UniformPolyFrom(p, rng)
	NTT(a)

	// s ← ternary secret (small norm)
	s := TernaryPolyFrom(p, rng)
	NTT(s)

	// e ← small noise, scaled by 2 for BGV structure
	e := CBDPolyFrom(p, kp.Eta1, rng)
	e = ScalarMul(e, 2)
	NTT(e)

	// b = a·s + 2·e
	b := MulPointwise(a, s)
	b = Add(b, e)

	// z ← implicit rejection value
	z := make([]byte, kp.SharedKeyLen)
	io.ReadFull(rng, z)

	pk := &KEMPublicKey{A: a, B: b, P: kp}
	sk := &KEMSecretKey{S: s, PK: pk, Z: z}
	rlk := genRelinKey(sk, rng)
	return pk, sk, rlk
}

// genRelinKey generates a relinearization key from the secret key.
// Uses base-decomposition to control noise growth.
func genRelinKey(sk *KEMSecretKey, rng io.Reader) *RelinearizationKey {
	p := sk.PK.P.Ring
	q := p.Q

	// Decomposition base: choose base so that ceil(log_base(q)) levels
	// keeps the added noise manageable. Base = 256 gives ~2 levels for q=12289.
	base := uint32(256)
	levels := 0
	for v := q; v > 0; v /= base {
		levels++
	}

	// s² in NTT form.
	s2 := MulPointwise(sk.S, sk.S)

	rlkA := make([]*Poly, levels)
	rlkB := make([]*Poly, levels)

	power := uint32(1) // base^i
	for i := range levels {
		// a_i ← uniform
		ai := UniformPolyFrom(p, rng)
		NTT(ai)

		// e_i ← small noise, scaled by 2 for BGV structure
		ei := CBDPolyFrom(p, sk.PK.P.Eta1, rng)
		ei = ScalarMul(ei, 2)
		NTT(ei)

		// b_i = a_i·s + e_i + s²·base^i
		bi := MulPointwise(ai, sk.S)
		bi = Add(bi, ei)

		// Add s²·power (in NTT form, scalar multiply is pointwise).
		s2Scaled := ScalarMul(s2, power)
		bi = Add(bi, s2Scaled)

		rlkA[i] = ai
		rlkB[i] = bi

		power = mulMod(power, base, q)
	}

	return &RelinearizationKey{A: rlkA, B: rlkB, L: levels}
}

// HEEncrypt encrypts a single bit (0 or 1) for homomorphic evaluation.
//
// Uses BGV-style encoding with plaintext modulus t=2:
//
//	Encrypt(m): u = a·r + 2·e1, v = b·r + 2·e2 + m
//	Decrypt(ct): (v - s·u) mod 2 = m (when noise < q/2)
//
// All noise terms are even (multiplied by t=2), so the message m ∈ {0,1}
// sits in the least significant bit. Addition in the ciphertext domain
// corresponds to XOR in the plaintext domain. Multiplication corresponds
// to AND (product of bits mod 2 = AND).
func HEEncrypt(pk *KEMPublicKey, bit int) *HECiphertext {
	return HEEncryptFrom(pk, bit, rand.Reader)
}

// HEEncryptFrom encrypts a bit using the given randomness source.
func HEEncryptFrom(pk *KEMPublicKey, bit int, rng io.Reader) *HECiphertext {
	p := pk.P.Ring
	q := p.Q

	coins := make([]byte, 32)
	io.ReadFull(rng, coins)

	noiseRng := sha3.NewShake256()
	noiseRng.Write(coins)

	// r ← small (ternary for tighter noise)
	r := TernaryPolyFrom(p, noiseRng)
	NTT(r)

	// e1, e2 ← small noise, then multiply by 2 (BGV encoding)
	e1 := CBDPolyFrom(p, pk.P.Eta1, noiseRng)
	e2 := CBDPolyFrom(p, pk.P.Eta2, noiseRng)
	e1 = ScalarMul(e1, 2)
	e2 = ScalarMul(e2, 2)

	// u = a·r + 2·e1
	u := MulPointwise(pk.A, r)
	INTT(u)
	u = Add(u, e1)

	// v = b·r + 2·e2 + m
	v := MulPointwise(pk.B, r)
	INTT(v)
	v = Add(v, e2)

	if bit != 0 {
		v.Coeffs[0] = addMod(v.Coeffs[0], 1, q)
	}

	return &HECiphertext{
		U:             u,
		V:             v,
		NoiseEstimate: float64(pk.P.Eta1+pk.P.Eta2) * 4,
		params:        pk.P,
	}
}

// HEDecrypt decrypts a homomorphic ciphertext to recover the bit.
// Computes (v - s·u) mod q, then takes result mod 2.
func HEDecrypt(sk *KEMSecretKey, ct *HECiphertext) int {
	uNTT := ct.U.Clone()
	NTT(uNTT)
	su := MulPointwise(sk.S, uNTT)
	INTT(su)

	noisy := Sub(ct.V, su)

	// The result is 2·(noise) + m mod q.
	// Recover m = result mod 2, using centered representation.
	q := ct.params.Ring.Q
	c := noisy.Coeffs[0]

	// Centered lift: if c > q/2, interpret as negative.
	if c > q/2 {
		c = q - c // |c|, but parity is preserved: q is odd, so q-c has same parity as q+c
	}
	return int(c % 2)
}

// HEAdd computes the homomorphic addition of two ciphertexts.
// Decrypts to XOR of the two plaintext bits (addition mod 2 of encoded values).
func HEAdd(a, b *HECiphertext) *HECiphertext {
	return &HECiphertext{
		U:             Add(a.U, b.U),
		V:             Add(a.V, b.V),
		NoiseEstimate: a.NoiseEstimate + b.NoiseEstimate,
		params:        a.params,
	}
}

// HESub computes homomorphic subtraction.
func HESub(a, b *HECiphertext) *HECiphertext {
	return &HECiphertext{
		U:             Sub(a.U, b.U),
		V:             Sub(a.V, b.V),
		NoiseEstimate: a.NoiseEstimate + b.NoiseEstimate,
		params:        a.params,
	}
}

// HEMul computes homomorphic multiplication of two ciphertexts.
// Requires a relinearization key to bring the result back to degree 1.
//
// In BGV with t=2, decryption gives v - s·u = 2·e + m (mod q).
// The tensor product of two such ciphertexts produces a degree-2
// ciphertext (c0, c1, c2) where:
//
//	decrypt_deg2 = c0 - c1·s + c2·s²
//	             = (2·e_a + m_a)·(2·e_b + m_b)
//	             = 4·e_a·e_b + 2·(m_a·e_b + m_b·e_a) + m_a·m_b
//	             = 2·(stuff) + m_a·m_b    (mod 2)
//
// The product preserves the BGV structure: noise is still even,
// message sits in the LSB. No BFV-style rescaling needed.
//
// Relinearization converts (c0, c1, c2) → (u', v') using the
// relinearization key (encryption of s²).
func HEMul(a, b *HECiphertext, rlk *RelinearizationKey) *HECiphertext {
	p := a.params.Ring

	// Tensor product components.
	// For ciphertext (u, v), decryption is v - s·u.
	// For the product: (v_a - s·u_a)(v_b - s·u_b)
	//   = v_a·v_b - s·(v_a·u_b + u_a·v_b) + s²·u_a·u_b
	//   = c0 - s·c1 + s²·c2
	c0 := Mul(a.V, b.V)         // constant term
	c1a := Mul(a.V, b.U)        // linear terms
	c1b := Mul(a.U, b.V)        //
	c1 := Add(c1a, c1b)         // c1 = v_a·u_b + u_a·v_b
	c2 := Mul(a.U, b.U)         // quadratic term

	// Relinearize c2: replace s²·c2 with an equivalent degree-1 term
	// using the RLK which encrypts s².
	// RLK[i] = (a_i, b_i = a_i·s + 2·e_i + s²·base^i)
	// So: sum(digit_i · b_i) = sum(digit_i · a_i)·s + 2·(noise) + s²·c2
	// This gives us: v' = c0 + sum(digit_i · b_i), u' = c1 + sum(digit_i · a_i)
	base := uint32(256)
	relinU := New(p)
	relinV := New(p)

	if c2.isNTT {
		INTT(c2)
	}

	for level := 0; level < rlk.L; level++ {
		digit := New(p)
		for j := range digit.Coeffs {
			d := c2.Coeffs[j]
			for k := 0; k < level; k++ {
				d /= base
			}
			digit.Coeffs[j] = d % base
		}

		NTT(digit)

		uPart := MulPointwise(digit, rlk.A[level])
		vPart := MulPointwise(digit, rlk.B[level])
		INTT(uPart)
		INTT(vPart)

		relinU = Add(relinU, uPart)
		relinV = Add(relinV, vPart)
	}

	return &HECiphertext{
		U:             Add(c1, relinU),
		V:             Add(c0, relinV),
		NoiseEstimate: a.NoiseEstimate * b.NoiseEstimate * 4,
		params:        a.params,
	}
}

// HENot computes the homomorphic NOT (bit flip) of a ciphertext.
// NOT(x) = 1 - x. In BGV encoding: the constant term has the message in its LSB.
// Adding 1 to v flips the LSB. u stays the same (NOT doesn't change the s term).
func HENot(a *HECiphertext) *HECiphertext {
	p := a.params.Ring

	one := New(p)
	one.Coeffs[0] = 1

	// v' = 1 + v (flip LSB), u stays (noise unchanged by constant add)
	// Actually NOT(x) = 1 XOR x = 1 + x mod 2, which in the ciphertext domain
	// means: add encrypt(1) to the ciphertext. But we can do it cheaper:
	// just add 1 to v's constant coefficient.
	newV := a.V.Clone()
	newV.Coeffs[0] = addMod(newV.Coeffs[0], 1, p.Q)

	return &HECiphertext{
		U:             a.U.Clone(),
		V:             newV,
		NoiseEstimate: a.NoiseEstimate,
		params:        a.params,
	}
}

// HEXOR computes XOR via addition (same as HEAdd for binary plaintexts).
func HEXOR(a, b *HECiphertext) *HECiphertext {
	return HEAdd(a, b)
}

// HEAND computes AND via multiplication.
func HEAND(a, b *HECiphertext, rlk *RelinearizationKey) *HECiphertext {
	return HEMul(a, b, rlk)
}

// Rerandomize adds a fresh encryption of zero to break linkability.
// This is the key anti-malleability defense: same plaintext, new ciphertext.
func Rerandomize(pk *KEMPublicKey, ct *HECiphertext) *HECiphertext {
	return RerandomizeFrom(pk, ct, rand.Reader)
}

// RerandomizeFrom uses the given randomness source.
func RerandomizeFrom(pk *KEMPublicKey, ct *HECiphertext, rng io.Reader) *HECiphertext {
	// Encrypt zero and add to ciphertext.
	zero := HEEncryptFrom(pk, 0, rng)
	return HEAdd(ct, zero)
}