// Key aggregation for multi-party lattice cryptography.
//
// For Ring-LWE keys sharing the same public element a:
//
//	pk_i = (a, b_i = a*s_i + e_i)
//
// The aggregate public key is:
//
//	pk_agg = (a, sum(b_i))
//
// This is a valid Ring-LWE public key for the aggregate secret s_agg = sum(s_i)
// and aggregate error e_agg = sum(e_i). The aggregate error remains short
// because the sum of k independent short vectors has infinity norm bounded
// by k * max(||e_i||_inf), which stays well within the decryption threshold
// for moderate k.
//
// Security: the aggregate key is indistinguishable from a fresh Ring-LWE
// sample. Breaking it requires solving Ring-LWE, which reduces to SVP
// on ideal lattices.
//
// For HE64 (q=10000769, eta=1, n=64): aggregate noise ~ k*eta*sqrt(n) ~ 80
// for k=10 participants, vs decryption threshold q/4 ~ 2.5M. Massive headroom.
//
// Distributed decryption: each party computes d_i = s_i * u, the partial
// decryptions are summed, and the plaintext is recovered from v - d_agg.
// No party reveals its individual secret key.

package ring

import (
	"crypto/rand"
	"errors"

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

// AggregateHEKeys combines multiple Ring-LWE public keys into a single
// aggregate key. All input keys must share the same public element A
// (generated from a common reference string or trusted setup).
//
// Returns a KEMPublicKey whose B component is the sum of all individual
// B components. This aggregate key can be used with standard HEEncrypt
// for multi-party computation.
func AggregateHEKeys(pks []*KEMPublicKey) (*KEMPublicKey, error) {
	if len(pks) == 0 {
		return nil, errors.New("keyagg: no public keys")
	}
	if len(pks) == 1 {
		return pks[0], nil
	}

	// Verify all keys share the same A and parameters.
	ref := pks[0]
	for i := 1; i < len(pks); i++ {
		if !Equal(ref.A, pks[i].A) {
			return nil, errors.New("keyagg: public keys do not share the same A element")
		}
		if ref.P.Ring.N != pks[i].P.Ring.N || ref.P.Ring.Q != pks[i].P.Ring.Q {
			return nil, errors.New("keyagg: incompatible ring parameters")
		}
	}

	// Aggregate: B_agg = sum(B_i).
	bAgg := ref.B.Clone()
	for i := 1; i < len(pks); i++ {
		bAgg = Add(bAgg, pks[i].B)
	}

	return &KEMPublicKey{
		A: ref.A.Clone(),
		B: bAgg,
		P: ref.P,
	}, nil
}

// PartialDecrypt computes a partial decryption of an HE ciphertext.
// Each participant computes d_i = s_i * u using their individual secret key.
// The result is a polynomial that must be combined with other partial
// decryptions to recover the plaintext.
func PartialDecrypt(sk *KEMSecretKey, ct *HECiphertext) *Poly {
	uNTT := ct.U.Clone()
	NTT(uNTT)
	su := MulPointwise(sk.S, uNTT)
	INTT(su)
	return su
}

// CombinePartialDecryptions combines partial decryptions from all
// participants and recovers the plaintext bit.
//
// Each partial decryption d_i = s_i * u. The combination computes:
//
//	d_agg = sum(d_i) = s_agg * u
//	m = decode(v - d_agg) = decode(noise + m_encoded)
//
// The noise term is the sum of all individual error terms, which remains
// short for moderate participant counts.
func CombinePartialDecryptions(ct *HECiphertext, partials []*Poly) int {
	if len(partials) == 0 {
		return 0
	}

	// Sum all partial decryptions.
	dAgg := partials[0].Clone()
	for i := 1; i < len(partials); i++ {
		dAgg = Add(dAgg, partials[i])
	}

	// Recover plaintext: v - d_agg.
	noisy := Sub(ct.V, dAgg)

	// BGV decoding: result mod 2, using centered representation.
	q := ct.params.Ring.Q
	c := noisy.Coeffs[0]
	if c > q/2 {
		c = q - c
	}
	return int(c % 2)
}

// GenerateSharedA generates a common public element A for key aggregation.
// All participants must use the same A when generating their individual keys.
// The seed should be a publicly agreed-upon value (common reference string).
func GenerateSharedA(kp KEMParams, seed []byte) *Poly {
	// Use SHAKE256 as an XOF seeded from the common reference string.
	// All participants with the same seed produce the identical A.
	h := sha3.NewShake256()
	h.Write([]byte("hamadryad-shared-a-v1"))
	h.Write(seed)
	a := UniformPolyFrom(kp.Ring, h)
	NTT(a)
	return a
}

// HEKeyGenWithA generates an HE key pair using a pre-specified A element.
// This is used for multi-party key aggregation where all participants
// share the same A from GenerateSharedA.
func HEKeyGenWithA(kp KEMParams, a *Poly) (*KEMPublicKey, *KEMSecretKey) {
	p := kp.Ring

	// s <- ternary secret.
	s := TernaryPoly(p)
	NTT(s)

	// e <- small noise, scaled by 2 for BGV.
	e := CBDPoly(p, kp.Eta1)
	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)
	if _, err := rand.Read(z); err != nil {
		panic("keyagg: randomness source failed: " + err.Error())
	}

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