// Multi-party homomorphic encryption session.
//
// MPCSession composes key aggregation, BGV HE, and anti-malleability
// into a single cohesive API for multi-party computation on encrypted data.
//
// Workflow:
//
//  1. All participants agree on a common reference string (seed).
//  2. Each generates a key pair with GenerateSharedA + HEKeyGenWithA.
//  3. Public keys are aggregated via AggregateHEKeys.
//  4. MPCSession is created with the aggregate key.
//  5. Encryption, homomorphic operations, and distributed decryption
//     are performed through the session.
//
// Anti-malleability: every evaluation result is rerandomized, noise-flooded,
// tagged with the session ID, and optionally GPV-signed. Results from
// different sessions cannot be spliced together.

package ring

import "errors"

// MPCSession holds state for a multi-party homomorphic computation.
type MPCSession struct {
	// AggPK is the aggregate public key (sum of all participants' B components).
	AggPK *KEMPublicKey

	// RLK is the relinearization key for multiplicative operations.
	// nil if only additive/XOR operations are needed.
	RLK *RelinearizationKey

	// session wraps results with anti-malleability protections.
	session *SessionWrapper
}

// NewMPCSession creates a multi-party computation session from a set of
// public keys. All keys must share the same A element (from GenerateSharedA).
//
// gpvPK/gpvSK are optional GPV signing keys for result authenticity.
// rlk is optional — only needed for AND/multiplication gates.
func NewMPCSession(pks []*KEMPublicKey, rlk *RelinearizationKey, gpvPK *GPVPublicKey, gpvSK *GPVSecretKey) (*MPCSession, error) {
	aggPK, err := AggregateHEKeys(pks)
	if err != nil {
		return nil, err
	}

	sw := NewSession(aggPK, gpvPK, gpvSK)

	return &MPCSession{
		AggPK:   aggPK,
		RLK:     rlk,
		session: sw,
	}, nil
}

// SessionID returns the unique session identifier.
func (s *MPCSession) SessionID() []byte {
	return s.session.SessionID
}

// Encrypt encrypts a single bit under the aggregate public key.
func (s *MPCSession) Encrypt(bit int) *HECiphertext {
	return HEEncrypt(s.AggPK, bit)
}

// XOR performs homomorphic XOR and wraps the result.
func (s *MPCSession) XOR(a, b *HECiphertext) *SecureHEResult {
	ct := HEXOR(a, b)
	return s.session.Wrap(ct)
}

// AND performs homomorphic AND and wraps the result.
// Requires that the session was created with a relinearization key.
func (s *MPCSession) AND(a, b *HECiphertext) (*SecureHEResult, error) {
	if s.RLK == nil {
		return nil, errors.New("mpc: AND requires relinearization key")
	}
	ct := HEAND(a, b, s.RLK)
	return s.session.Wrap(ct), nil
}

// NOT performs homomorphic NOT and wraps the result.
func (s *MPCSession) NOT(a *HECiphertext) *SecureHEResult {
	ct := HENot(a)
	return s.session.Wrap(ct)
}

// Add performs homomorphic addition and wraps the result.
func (s *MPCSession) Add(a, b *HECiphertext) *SecureHEResult {
	ct := HEAdd(a, b)
	return s.session.Wrap(ct)
}

// Verify checks anti-malleability protections on a received result.
func (s *MPCSession) Verify(result *SecureHEResult) bool {
	return s.session.Verify(result)
}

// Unwrap extracts the ciphertext from a verified result.
// Returns nil if verification fails.
func (s *MPCSession) Unwrap(result *SecureHEResult) *HECiphertext {
	if !s.Verify(result) {
		return nil
	}
	return result.Ciphertext
}

// DecryptDistributed combines partial decryptions to recover the plaintext bit.
// Each party must call PartialDecrypt with their own secret key and contribute
// the result. The partials slice must contain exactly one partial from each
// participant, in any order.
func DecryptDistributed(ct *HECiphertext, partials []*Poly) int {
	return CombinePartialDecryptions(ct, partials)
}