// Key aggregation for multi-party lattice cryptography.
package ring

import (
	"crypto/rand"
	"errors"

	"crypto/sha3"
)

func AggregateHEKeys(pks []*KEMPublicKey) (aggPK *KEMPublicKey, err error) {
	if int32(len(pks)) == 0 {
		return nil, errors.New("keyagg: no public keys")
	}
	if int32(len(pks)) == 1 {
		return pks[0], nil
	}

	ref := pks[0]
	for i := int32(1); i < int32(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")
		}
	}

	bAgg := ref.B.Clone()
	for i := int32(1); i < int32(len(pks)); i++ {
		bAgg = Add(bAgg, pks[i].B)
	}

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

func PartialDecrypt(sk *KEMSecretKey, ct *HECiphertext) (su *Poly) {
	uNTT := ct.U.Clone()
	NTT(uNTT)
	su = MulPointwise(sk.S, uNTT)
	INTT(su)
	return su
}

func CombinePartialDecryptions(ct *HECiphertext, partials []*Poly) (result int32) {
	if int32(len(partials)) == 0 {
		return 0
	}

	dAgg := partials[0].Clone()
	for i := int32(1); i < int32(len(partials)); i++ {
		dAgg = Add(dAgg, partials[i])
	}

	noisy := Sub(ct.V, dAgg)

	q := ct.params.Ring.Q
	c := noisy.Coeffs[0]
	if c > q/2 {
		c = q - c
	}
	return int32(c % 2)
}

func GenerateSharedA(kp KEMParams, seed []byte) (a *Poly) {
	h := sha3.NewSHAKE256()
	h.Write([]byte("hamadryad-shared-a-v1"))
	h.Write(seed)
	a = UniformPolyFrom(kp.Ring, h)
	NTT(a)
	return a
}

func HEKeyGenWithA(kp KEMParams, a *Poly) (pk *KEMPublicKey, sk *KEMSecretKey) {
	p := kp.Ring

	s := TernaryPoly(p)
	NTT(s)

	e := CBDPoly(p, kp.Eta1)
	e = ScalarMul(e, 2)
	NTT(e)

	b := MulPointwise(a, s)
	b = Add(b, e)

	z := []byte{:kp.SharedKeyLen}
	_, err := rand.Read(z)
	if 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
}