package gnarlring

import (
	"crypto/rand"
	"io"
)

// LWEKeyAgg aggregates multiple LWE public keys sharing the same A element
// into a single group public key. B_agg = Σ B_i.
// All keys must share identical A elements (generated from a common reference).
func LWEKeyAgg(pks []*LWEPublicKey) *LWEPublicKey {
	if len(pks) == 0 {
		return nil
	}
	if len(pks) == 1 {
		return &LWEPublicKey{A: pks[0].A.Clone(), B: pks[0].B.Clone()}
	}

	bAgg := pks[0].B.Clone()
	for i := 1; i < len(pks); i++ {
		bAgg = Add(bAgg, pks[i].B)
	}
	return &LWEPublicKey{A: pks[0].A.Clone(), B: bAgg}
}

// GenerateSharedA generates a common A element from a public seed for key
// aggregation. All participants use the same seed to derive identical A.
func GenerateSharedA(seed []byte) *Poly27 {
	input := append([]byte("gnarl-lwe-shared-a-v1"), seed...)
	return hashBytesToPoly(input)
}

// LWEKeyGenWithA generates a key pair using a pre-specified A element.
// Used for multi-party key aggregation.
func LWEKeyGenWithA(a *Poly27) (*LWEPublicKey, *LWESecretKey) {
	return LWEKeyGenWithAFrom(a, rand.Reader)
}

// LWEKeyGenWithAFrom generates with a given RNG.
func LWEKeyGenWithAFrom(a *Poly27, rng io.Reader) (*LWEPublicKey, *LWESecretKey) {
	if rng == nil {
		rng = rand.Reader
	}
	s := ternaryPoly(rng)

	gs := NewGaussSamplerFrom(DefaultSigma(), rng)
	e := gs.SamplePoly()
	for i := range e.Coeffs {
		if e.Coeffs[i] > Q/2 {
			e.Coeffs[i] = Q - 1
		} else if e.Coeffs[i] != 0 {
			e.Coeffs[i] = 1
		}
	}

	as := Mul(a, s)
	b := Add(as, e)

	pk := &LWEPublicKey{A: a.Clone(), B: b}
	sk := &LWESecretKey{S: s, PK: pk}
	return pk, sk
}

// ShareSecret splits a master secret s into k additive shares.
// Returns k shares where Σ shares[i] = s. The caller distributes each share
// to a different participant.
func ShareSecret(s *Poly27, k int) []*Poly27 {
	shares := make([]*Poly27, k)
	gs := NewGaussSampler(DefaultSigma())

	// Generate k-1 random shares, final share = s - Σ random shares.
	sum := NewPoly27()
	for i := 0; i < k-1; i++ {
		shares[i] = gs.SamplePoly()
		sum = Add(sum, shares[i])
	}
	// Last share: s - sum (mod q).
	shares[k-1] = Sub(s, sum)
	return shares
}

// PartialDecryption computes one child's contribution to distributed
// decryption: d_i = share_i * u. The coordinator sums all partials
// to recover aggregate decryption.
func PartialDecryption(share *Poly27, u *Poly27) *Poly27 {
	return Mul(share, u)
}

// CombinePartials sums partial decryptions and recovers the plaintext bit.
// d_agg = Σ d_i = s_agg * u. Then m = decode(v - d_agg).
func CombinePartials(v *Poly27, partials []*Poly27) int {
	if len(partials) == 0 {
		return 0
	}
	dAgg := partials[0].Clone()
	for i := 1; i < len(partials); i++ {
		dAgg = Add(dAgg, partials[i])
	}
	noisy := Sub(v, dAgg)
	return decodeBit(noisy)
}

// GroupDecryptData holds the per-child shares and the aggregate LWE key
// for distributed decryption.
type GroupDecryptData struct {
	AggPK   *LWEPublicKey // aggregated B (common A)
	Shares  []*Poly27     // one share per child
}

// EncryptGroup encrypts a bit under the aggregated group key. Only used
// by the coordinator to test the distributed decryption path.
func EncryptGroup(aggPK *LWEPublicKey, bit int) *LWECiphertext {
	return LWEEncrypt(aggPK, bit)
}

// DistributedDecrypt performs k-of-n distributed decryption. The caller
// provides partial decryptions from k children. Returns the plaintext bit.
func DistributedDecrypt(ct *LWECiphertext, partials []*Poly27) int {
	return CombinePartials(ct.V, partials)
}