// Multi-party homomorphic encryption session.
package ring

import "errors"

type MPCSession struct {
	AggPK   *KEMPublicKey
	RLK     *RelinearizationKey
	session *SessionWrapper
}

func NewMPCSession(pks []*KEMPublicKey, rlk *RelinearizationKey, gpvPK *GPVPublicKey, gpvSK *GPVSecretKey) (s *MPCSession, err 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
}

func (s *MPCSession) SessionID() (id []byte) {
	return s.session.SessionID
}

func (s *MPCSession) Encrypt(bit int32) (ct *HECiphertext) {
	return HEEncrypt(s.AggPK, bit)
}

func (s *MPCSession) XOR(a, b *HECiphertext) (result *SecureHEResult) {
	ct := HEXOR(a, b)
	return s.session.Wrap(ct)
}

func (s *MPCSession) AND(a, b *HECiphertext) (result *SecureHEResult, err 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
}

func (s *MPCSession) NOT(a *HECiphertext) (result *SecureHEResult) {
	ct := HENot(a)
	return s.session.Wrap(ct)
}

func (s *MPCSession) MPCAdd(a, b *HECiphertext) (result *SecureHEResult) {
	ct := HEAdd(a, b)
	return s.session.Wrap(ct)
}

func (s *MPCSession) Verify(result *SecureHEResult) (ok bool) {
	return s.session.Verify(result)
}

func (s *MPCSession) Unwrap(result *SecureHEResult) (ct *HECiphertext) {
	if !s.Verify(result) {
		return nil
	}
	return result.Ciphertext
}

func DecryptDistributed(ct *HECiphertext, partials []*Poly) (result int32) {
	return CombinePartialDecryptions(ct, partials)
}