package meerkat.crypto.dkg.gjkr;

import com.google.protobuf.InvalidProtocolBufferException;
import meerkat.crypto.utils.Arithmetic;
import meerkat.crypto.utils.concrete.Fp;
import meerkat.comm.Channel;
import meerkat.crypto.secretsharing.shamir.Polynomial;
import meerkat.crypto.secretsharing.shamir.SecretSharing;
import meerkat.protobuf.Comm;
import meerkat.protobuf.DKG;

import java.math.BigInteger;
import java.util.ArrayList;

import static meerkat.crypto.dkg.comm.MessageUtils.*;

/**
 * Created by Tzlil on 3/16/2016.
 * <p/>
 * implementation of gjkr protocol user.
 * <p/>
 * this protocol extends joint Feldman protocol by splitting the protocol to commitment stage (stages 1,2,3)
 * and revealing stage (stage 4).
 * <p/>
 * as in joint Feldman, each party in QUAL has his own share of the generated random key.
 * this key can be recover by any subset of QUAL of size at least t + 1.
 */
public class User<T> extends meerkat.crypto.dkg.feldman.User<T> {

    /**
     * All parties participating in key generation.
     * parties[id-1] has my info.
     */
    protected Party<T>[] parties;

    /**
     * gjkr secure protocol object
     */
    protected final Protocol<T> sdkg;

    boolean isStage4;

    /**
     * constructor
     *
     * @param sdkg    gjkr protocol object
     * @param channel channel object
     */
    public User(Protocol<T> sdkg, Channel channel) {
        super(sdkg, channel);
        this.sdkg = sdkg;
        this.parties = sdkg.getParties();
    }

    /**
     * stage1 according to the protocol
     * 1. Pi broadcasts Cik=Aik*Bik for k = 0,...,t.
     * 2. Pi computes the shares Sij,Sij' for j = 1,...,n and sends Sij,Sij' secretly to Pj.
     */
    @Override
    protected void stage1() {
        sdkg.computeAndBroadcastVerificationValues();
        sdkg.sendSecrets();
    }



    @Override
    protected void waitUntilStageOneCompleted() {
        super.waitUntilStageOneCompleted();
        // save the received commitments as verification values
        ArrayList<T> temp;
        for (int i = 0; i < n; i++) {
            temp = parties[i].verifiableValues;
            parties[i].verifiableValues = parties[i].commitments;
            parties[i].commitments = temp;
        }
    }

    /**
     * stage2 according to the protocol
     * Pj verifies all the shares,sharesT he received
     * if check fails for an index i, Pj broadcasts a complaint against Pi.
     * Pj broadcasts done message at the end of this stage
     */
    @Override
    protected void stage2() {
        sdkg.broadcastComplaints();
        //broadcast done message after all complaints
        channel.broadcastMessage(createMessage(DKG.Payload.Type.DONE));
    }

    /**
     * Check if all non-aborting qualified parties have sent commitments.
     * @return
     */
    protected boolean haveAllQualPartiesCommitted() {
        for (int i : QUAL) {
            if (parties[i - 1].aborted)
                continue;
            for (int k = 0; k <= t; k++) {
                if (parties[i - 1].commitments.get(k) == null)
                    return false;
            }
        }
        return true;
    }


    /**
     * Check if all non-aborting qualified parties sent a done message
     * @return
     */
    protected boolean areAllQualPartiesDone() {
        for (int i : QUAL) {
            if (parties[i - 1].aborted)
                continue;
            for (int k = 0; k <= t; k++) {
                if (!parties[i - 1].ysDoneFlag)
                    return false;
            }
        }
        return true;
    }

    /**
     * Check if at least t + 1 secrets were received foreach i in QUAL that aborted
     * @return
     */
    protected boolean haveReceivedEnoughSecretShares() {
        for (int i : QUAL) {
            if (parties[i - 1].aborted && parties[i - 1].recoverSharesSet.size() <= t)
                return false;
        }
        return true;
    }

    /**
     * broadcast commitments and recover parties information if necessary
     */
    private void resolveQualifyingPublicKey() {
        sdkg.broadcastCommitments();
        // wait until all parties in QUAL broadcast their commitments or aborted
        while (!stop && !haveAllQualPartiesCommitted())
            waitAndHandleReceivedMessages();

        if (stop)
            return;

        sdkg.computeAndBroadcastComplaints(QUAL);

        //broadcast done message after all complaints
        channel.broadcastMessage(createMessage(DKG.Payload.Type.DONE));

        // wait until all parties in QUAL done or aborted
        while (!stop && !areAllQualPartiesDone())
            waitAndHandleReceivedMessages();

        if (stop)
            return;

        // broadcast i private secret foreach i in QUAL that aborted
        for (int i : QUAL) {
            if (parties[i - 1].aborted) {
                sdkg.broadcastAnswer(parties[i - 1].share, parties[i - 1].shareT, i);
            }
        }
        // wait until at least t + 1 secrets will received foreach i in QUAL that aborted
        while (!stop && !haveReceivedEnoughSecretShares())
            waitAndHandleReceivedMessages();

        if (stop)
            return;

        Arithmetic<BigInteger> arithmetic = new Fp(sdkg.getQ());
        // restore necessary information
        for (int i = 0; i < n; i++) {
            if (parties[i].recoverSharesSet.isEmpty()) {
                continue;
            }
            Polynomial.Point[] shares = new Polynomial.Point[t + 1];
            int j = 0;
            for (Polynomial.Point share : parties[i].recoverSharesSet) {
                shares[j++] = share;
                if (j >= shares.length) {
                    break;
                }
            }
            Polynomial polynomial = SecretSharing.recoverPolynomial(shares, arithmetic);
            BigInteger[] coefficients = polynomial.getCoefficients();
            for (int k = 0; k <= t; k++) {
                parties[i].commitments.add(k, group.multiply(g, coefficients[k]));
            }
            parties[i].share = new Polynomial.Point(BigInteger.valueOf(id), polynomial);
        }
    }

    /**
     * notifies message handler and message handler that stage 4 was started
     */
    protected void setStage4() {
        isStage4 = true;
    }

    @Override
    protected void stage4() {
        setStage4();
        resolveQualifyingPublicKey();
        if (stop) return;
        super.stage4();
    }


    /**
     * if !isStage4  as super, with extension to double secret message
     * else answer message is valid if:
     * 1. it was received in broadcast chanel
     * 2. secret.j == sender
     * 3. QUAL contains i and j
     */
    protected boolean isValidAnswerMessage(int sender, boolean isBroadcast, DKG.ShareMessage doubleSecretMessage) {
        if (!isStage4) {
            return super.isValidAnswerMessage(sender, isBroadcast, doubleSecretMessage);
        } else {
            int i = doubleSecretMessage.getI();
            int j = doubleSecretMessage.getJ();
            return isBroadcast && j == sender && parties[i - 1].aborted && !parties[j - 1].aborted
                    && QUAL.contains(i) && QUAL.contains(j);
        }
    }


    /**
     * as in super with respect to protocol stage
     */
    @Override
    protected boolean isValidDoneMessage(int sender, boolean isBroadcast) {
        if (!isStage4) {
            return super.isValidDoneMessage(sender, isBroadcast);
        } else {
            return isBroadcast && !parties[sender - 1].ysDoneFlag;
        }
    }

    /**
     * use only in stage4
     * complaint message is valid if:
     * 1. it was received in broadcast chanel
     * 2. secret.j == sender
     * 3. QUAL contains i and j
     */
    protected boolean isValidComplaintMessage(int sender, boolean isBroadcast,
                                              DKG.ShareMessage complaintMessage) {
        int i = complaintMessage.getI();
        int j = complaintMessage.getJ();
        return isBroadcast && j == sender && QUAL.contains(i) && QUAL.contains(j);
    }


    @Override
    public void handleMessage(Comm.BroadcastMessage envelope) throws InvalidProtocolBufferException {
        int sender = envelope.getSender();
        boolean isBroadcast = !envelope.getIsPrivate();
        DKG.Payload msg = DKG.Payload.parseFrom(envelope.getPayload());
        switch (msg.getType()) {
            case SHARE:
                /**
                 * as in super, with extension to double secret message
                 */
                DKG.ShareMessage doubleSecretMessage = msg.getShare();
                if (isValidSecretMessage(sender, isBroadcast, doubleSecretMessage)) {
                    int i = doubleSecretMessage.getI();
                    synchronized (parties[i - 1]) {
                        parties[i - 1].share = extractShare(id, doubleSecretMessage.getShare());
                        parties[i - 1].shareT = extractShare(id, doubleSecretMessage.getShareT());
                        parties[i - 1].notify();
                    }
                }
                break;
            case ANSWER:
                /**
                 * if !isStage4 as super, with extension to double secret message
                 * else saves secret
                 */
                assert msg.getPayloadDataCase() == DKG.Payload.PayloadDataCase.SHARE;
                doubleSecretMessage = msg.getShare();
                if (isValidAnswerMessage(sender, isBroadcast, doubleSecretMessage)) {
                    int i = doubleSecretMessage.getI();
                    int j = doubleSecretMessage.getJ();
                    Polynomial.Point secret = extractShare(j, doubleSecretMessage.getShare());
                    Polynomial.Point secretT = extractShare(j, doubleSecretMessage.getShareT());
                    synchronized (parties[i - 1]) {
                        if (!isStage4) {
                            if (sdkg.isValidShare(secret, secretT, parties[j - 1].verifiableValues, i)) {
                                parties[i - 1].complaints[j - 1] = meerkat.crypto.dkg.feldman.Protocol.ComplaintState.NonDisqualified;

                            } else {
                                parties[i - 1].complaints[j - 1] = meerkat.crypto.dkg.feldman.Protocol.ComplaintState.Disqualified;
                            }
                            if (j == id) {
                                parties[i - 1].share = secret;
                                parties[i - 1].shareT = secretT;
                            }
                        } else if (sdkg.isValidShare(secret, secretT, parties[i - 1].verifiableValues, j)) {
                            parties[i - 1].recoverSharesSet.add(secret);
                        }
                        parties[i - 1].notify();
                    }
                }
                break;
            case DONE:
                /**
                 * as in super with respect to protocol state
                 */
                if (!isStage4)
                    super.handleMessage(envelope);
                else {
                    if (isValidDoneMessage(sender, isBroadcast)) {
                        synchronized (parties[sender - 1]) {
                            parties[sender - 1].ysDoneFlag = true;
                            parties[sender - 1].notify();
                        }
                    }
                }
                break;
            case COMPLAINT:
                /**
                 * if !isStage4 as in super
                 * else if secret,secretT are valid with respect to verifiableValues but
                 *      secret is not valid with respect to commitments then
                 *      marks i as aborted
                 */
                if (!isStage4) {
                    super.handleMessage(envelope);
                } else {
                    assert (msg.getPayloadDataCase() == DKG.Payload.PayloadDataCase.SHARE);
                    DKG.ShareMessage ysComplaintMessage = msg.getShare();
                    if (isValidComplaintMessage(sender, isBroadcast, ysComplaintMessage)) {
                        int i = ysComplaintMessage.getI();
                        int j = ysComplaintMessage.getJ();
                        Polynomial.Point secret = extractShare(i, ysComplaintMessage.getShare());
                        Polynomial.Point secretT = extractShare(i, ysComplaintMessage.getShareT());
                        if (sdkg.isValidShare(secret, secretT, parties[i - 1].verifiableValues, j)
                                && !dkg.isValidShare(secret, parties[i - 1].commitments, j)) {
                            synchronized (parties[i - 1]) {
                                parties[i - 1].aborted = true;
                                parties[i - 1].notify();
                            }
                        }
                    }
                }
                break;
            default:
                super.handleMessage(envelope);
                break;
        }
    }

}