package meerkat.crypto.dkg.feldman;

import meerkat.comm.Channel;
import meerkat.crypto.secretsharing.feldman.VerifiableSecretSharing;
import meerkat.crypto.secretsharing.shamir.Polynomial;
import com.google.protobuf.ByteString;
import meerkat.protobuf.DKG;
import org.factcenter.qilin.primitives.Group;
import org.factcenter.qilin.util.ByteEncoder;

import java.math.BigInteger;
import java.util.*;

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

/**
 * Created by Tzlil on 3/14/2016.
 *
 * an implementation of JointFeldman distributed key generation protocol.
 *
 * allows set of n parties to generate random key with threshold t.
 */
public class Protocol<T> extends VerifiableSecretSharing<T> {
    public enum ComplaintState {
        /**
         * No complaints, no response required at this point.
         */
        OK,

        /**
         * Party received complaint, waiting for response from party
         */
        Waiting,

        /**
         * Party gave invalid answer to complaint.
         */
        Disqualified,

        /**
         * Party received complaint, gave valid answer.
         */
        NonDisqualified
    }

    /**
     * My share id.
     */
    protected final int id;

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

    /**
     * communication object
     */
    protected Channel channel;

    /**
     * Encode/Decode group elements
     */
    protected final ByteEncoder<T> encoder;

    /**
     * constructor
     * @param q      a large prime.
     * @param t      threshold. Any t+1 share holders can recover the secret,
     *               but any set of at most t share holders cannot
     * @param n      number of share holders
     * @param zi      secret, chosen from Zq
     * @param random use for generate random polynomial
     * @param group
     * @param q a large prime dividing group order.
     * @param g a generator of cyclic group of order q.
     *          the generated group is a subgroup of the given group.
     *          it must be chosen such that computing discrete logarithms is hard in this group.
     * @param encoder Encode/Decode group elements (of type T) to/from byte array
     */
    //TODO: why the use of regular Random? Should it be changed?
    public Protocol(int t, int n, BigInteger zi, Random random, BigInteger q, T g
            , Group<T> group, int id, ByteEncoder<T> encoder)  {
        super(t, n, zi, random, q, g,group);
        this.id = id;
        this.parties = new Party[n];
        for (int i = 1; i <= n ; i++){
            this.parties[i - 1] = new Party(i,n,t);
        }
        this.parties[id - 1].share = getShare(id);
        this.encoder = encoder;
    }

    /**
     * setter
     * @param channel
     */
    public void setChannel(Channel channel){
        this.channel = channel;
    }

    /**
     * setter
     * @param parties
     */
    protected void setParties(Party<T>[] parties){
        this.parties = parties;
    }

    /**
     * getter
     * @return
     */
    protected Party<T>[] getParties(){
        return parties;
    }

    /**
     * stage1.1 according to the protocol
     *  Pi broadcasts Aik for k = 0,...,t.
     */
    public void broadcastCommitments(){
        broadcastCommitments(commitmentsArrayList);
    }

    /**
     * pack commitments as messages and broadcast them
     * @param commitments
     */
    public void broadcastCommitments(ArrayList<T> commitments){
        DKG.CommitmentMessage commitmentMessage;
        for (int k = 0; k <= t ; k++){
            commitmentMessage = DKG.CommitmentMessage.newBuilder()
                    .setCommitment(ByteString.copyFrom(encoder.encode(commitments.get(k))))
                    .setK(k)
                    .build();
            channel.broadcastMessage(createMessage(DKG.Payload.Type.COMMITMENT, commitmentMessage));
        }
    }

    /**
     * Send channel j her secret share (of my polynomial)
     * @param j
     */
    public void sendSecret(int j){
        ByteString secret = ByteString.copyFrom(getShare(j).y.toByteArray());
        channel.sendMessage(j, createMessage(DKG.Payload.Type.SHARE,
                DKG.ShareMessage.newBuilder()
                        .setI(id)
                        .setJ(j)
                        .setShare(secret)
                        ));
    }

    /**
     * stage1.2 according to the protocol
     *  Pi computes the shares Sij for j = 1,...,n and sends Sij secretly to Pj.
     */
    public void sendSecrets(){
        for (int j = 1; j <= n ; j++){
            if(j != id){
                sendSecret(j);
            }
        }
    }

    /**
     *
     * @param i
     * @return computeVerificationValue(j,parties[i - 1].commitments,group) == g ^ parties[i - 1].share mod q
     */
    public boolean isValidShare(int i){
        Party<T> party = parties[i - 1];
        synchronized (parties[i - 1]) {
            return isValidShare(party.share, party.commitments, id);
        }
    }

    /**
     * @param share
     * @param commitments
     * @param j
     * @return computeVerificationValue(j,commitments,group) == g ^ secret.y mod q
     */
    public boolean isValidShare(Polynomial.Point share, ArrayList<T> commitments, int j){
        try{
            T v = computeVerificationValue(j,commitments,group);
            return group.multiply(g,share.y).equals(v);
        }
        catch (NullPointerException e){
            return false;
        }
    }

    /**
     * stage2 according to the protocol
     *  Pj verifies all the shares he received (using isValidShare)
     *  if check fails for an index i, Pj broadcasts a complaint against Pi.
     */
    public void broadcastComplaints(){
        for (int i = 1; i <= n ; i++ ){
            if(i != id && !isValidShare(i)) {
                broadcastComplaint(i);
            }
        }
    }

    /**
     * create a complaint message against i and broadcast it
     * @param i
     */
    private void broadcastComplaint(int i){
        //message = new Message(Type.Complaint, j)
        DKG.IDMessage complaint = DKG.IDMessage.newBuilder()
                .setId(i)
                .build();
        channel.broadcastMessage(createMessage(DKG.Payload.Type.COMPLAINT, complaint));
    }

    /**
     * create an answer message for j and broadcast it
     * @param j
     */
    public void broadcastComplaintAnswer(int j){
        channel.broadcastMessage(createMessage(DKG.Payload.Type.ANSWER, DKG.ShareMessage.newBuilder()
                .setI(id)
                .setJ(j)
                .setShare(ByteString.copyFrom(getShare(j).y.toByteArray()))));
    }

    /**
     * stage3.1 according to the protocol
     *   if more than t players complain against a player Pi he is disqualified.
     */
    public void answerAllComplainingPlayers(){
        ComplaintState[] complaints = parties[id - 1].complaints;
        for (int i = 1; i <= n; i++) {
            switch (complaints[i - 1]) {
                case Waiting:
                    broadcastComplaintAnswer(i);
                    break;
                default:
                    break;
            }
        }

    }

    /**
     * stage3.2 according to the protocol
     *  if any of the revealed shares fails the verification test, player Pi is disqualified.
     *  set QUAL to be the set of non-disqualified players.
     */
    public Set<Integer> calcQUAL(){
        Set<Integer> QUAL = new HashSet<Integer>();
        boolean nonDisqualified;
        int counter;
        for (int i = 1; i <= n; i++) {
            synchronized (parties[i - 1]) {
                ComplaintState[] complaints = parties[i - 1].complaints;
                nonDisqualified = true;
                counter = 0;
                for (int j = 1; j <= n; j++) {
                    switch (complaints[j - 1]) {
                        case OK:
                            break;
                        case NonDisqualified:
                            counter++;
                            break;
                        default:
                            nonDisqualified = false;
                            break;
                    }
                    if (!nonDisqualified)
                        break;
                }
                if (nonDisqualified && counter <= t) {
                    QUAL.add(i);
                }
            }
        }
        return QUAL;
    }

    /**
     * compute Y, the commitment to the final public key (includes only qualifying set)
     * stage4.1 according to the protocol
     *   public value y is computed as y = multiplication of yi mod p for i in QUAL
     */
    public T calcY(Set<Integer> QUAL){
        T y = group.zero();
        for (int i : QUAL) {
            synchronized (parties[i - 1]) {
                y = group.add(y, parties[i - 1].commitments.get(0));
            }
        }
        return y;
    }

    /**
     * stage4.2 according to the protocol
     *  public verification values are computed as Ak = multiplication
     *  of Aik mod p for i in QUAL for k = 0,...,t
     */
    public ArrayList<T> calcCommitments(Set<Integer> QUAL){
        ArrayList<T> commitments = new ArrayList<T>(t+1);
        T value;
        for (int k = 0; k <= t; k++){
            value = group.zero();
            for (int i : QUAL) {
                synchronized (parties[i - 1]) {
                    value = group.add(value, parties[i - 1].commitments.get(k));
                }
            }
            commitments.add(k,value);
        }
        return commitments;
    }

    /**
     * stage4.3 according to the protocol
     *  Pj sets is share of the share as xj = sum of Sij mod q for i in QUAL
     */
    public Polynomial.Point calcShare(Set<Integer> QUAL){
        BigInteger xj = BigInteger.ZERO;
        for (int i : QUAL) {
            synchronized (parties[i - 1]) {
                xj = xj.add(parties[i - 1].share.y);
            }
        }
        return new Polynomial.Point(BigInteger.valueOf(id) , xj.mod(q));
    }

    /**
     * decode commitment from arr
     * @param arr
     * @return
     */
    public T decodeCommitment(byte[] arr){
        return encoder.decode(arr);
    }

    /**
     * getter
     * @return id
     */
    public int getId() {
        return id;
    }

    /**
     * getter
     * @return channel
     */
    public Channel getChannel() {
        return channel;
    }


    /**
     * getter
     * @return encoder
     */
    public ByteEncoder<T> getEncoder() {
        return encoder;
    }

}
