package meerkat.crypto.dkg.feldman;

import com.google.protobuf.InvalidProtocolBufferException;
import com.google.protobuf.TextFormat;
import meerkat.comm.Channel;
import meerkat.crypto.dkg.comm.MessageHandler;
import meerkat.crypto.secretsharing.shamir.Polynomial;
import com.google.protobuf.ByteString;
import meerkat.protobuf.Comm;
import meerkat.protobuf.DKG;
import org.factcenter.qilin.primitives.Group;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.math.BigInteger;
import java.util.ArrayList;
import java.util.Set;
import java.util.concurrent.ArrayBlockingQueue;
import java.util.concurrent.BlockingQueue;
import java.util.concurrent.LinkedBlockingDeque;

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

/**
 * Created by Tzlil on 3/14/2016.
 *
 * implementation of joint feldman protocol user.
 *
 * according to the protocol, each user run feldman verifiable secret sharing
 * as a dealer.
 *
 * by the end of run(), 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> implements Runnable {
    final Logger logger = LoggerFactory.getLogger(getClass());

    /**
     * joint feldman protocol object
     */
    protected final Protocol<T> dkg;

    /**
     * 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.
     */
    protected final T g;

    /**
     * cyclic group contains g.
     */
    protected final Group<T> group;

    /**
     * user id
     */
    protected final int id;

    /**
     * threshold
     */
    protected final int t;

    /**
     * number of shares
     */
    protected final int n;


    /**
     * channel object
     */
    protected final Channel channel; //

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

    /**
     * set of all non-disqualified parties
     */
    protected Set<Integer> QUAL;

    /**
     * my own share of the generated random key.
     */
    protected Polynomial.Point share;

    /**
     * public verification values
     */
    protected ArrayList<T> commitments;

    /**
     * public value,
     * y = g ^ key
     */
    protected T y;

    protected BlockingQueue<Comm.BroadcastMessage> receiveQueue;

    /**
     * constructor
     * @param dkg joint feldman protocol object
     * @param channel channel object
     */
    public User(Protocol<T> dkg, Channel channel) {
        this.dkg = dkg;

        this.g = dkg.getGenerator();
        this.group = dkg.getGroup();
        this.n = dkg.getN();
        this.t = dkg.getT();
        this.id = dkg.getId();

        this.channel = channel;
        dkg.setChannel(channel);
        registerReceiverCallback();

        this.parties = dkg.getParties();
        this.QUAL = null;
        this.commitments = null;
        this.share = null;
        this.y = null;

        this.receiveQueue = new LinkedBlockingDeque<>();

    }

    /**
     * create MessageHandler and register it as ReceiverCallback
     */
    protected void registerReceiverCallback() {
        channel.registerReceiverCallback(new meerkat.crypto.dkg.comm.MessageHandler() {
            @Override
            public void handleMessage(Comm.BroadcastMessage envelope) throws InvalidProtocolBufferException {
                receiveQueue.add(envelope);
            }
        });

//        this.messageHandler = new MessageHandler();
//        channel.registerReceiverCallback(messageHandler);
    }

    /**
     * Wait for at least one message to arrive, then handle any messages currently in the queue
     */
    protected void waitAndHandleReceivedMessages()  {
        Comm.BroadcastMessage msg = null;
        while (!stop && msg == null) {
            try {
                msg = receiveQueue.take();
            } catch (InterruptedException e) {
                // Possibly stop
            }
        }
        while (!stop && msg != null) {
            try {
                handleMessage(msg);
            } catch (InvalidProtocolBufferException e) {
                logger.warn("Received invalid message: {}", TextFormat.printToString(msg));
            }
            msg = receiveQueue.poll();
        }
    }

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

    /**
     * Check if all shares and commitments have arrived from other parties
     */
    protected boolean isStageOneCompleted() {
        for (int i = 0 ; i < n ; i++) {
            if (!parties[i].aborted) {
                if (parties[i].share == null)
                    return false;
                for (int k = 0 ; k <=  t ; k++) {
                    if (parties[i].commitments.get(k) == null)
                        return false;
                }
            }
        }
        return true;
    }

    protected void waitUntilStageOneCompleted() {
        while (!stop && !isStageOneCompleted())
            waitAndHandleReceivedMessages();
    }

    protected  boolean isStageTwoCompleted() {
        for (int i = 0 ; i < n ; i++) {
            if (!parties[i].aborted && !parties[i].doneFlag)
                return false;
        }
        return true;
    }

    /**
     * stage2 according to the protocol
     *  Pj verifies all the shares 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
     */
    protected void stage2() {
        dkg.broadcastComplaints();
        //broadcast done message after all complaints
        channel.broadcastMessage(createMessage(DKG.Payload.Type.DONE));
    }

    /**
     * wait until all other parties done complaining by receiving done message
     */
    protected void waitUntilStageTwoCompleted(){
        while (!stop && !isStageTwoCompleted())
            waitAndHandleReceivedMessages();
    }


    protected boolean haveReceivedAllStage3ComplaintAnswers() {
        for (int i = 0; i < n; i++) {
            if (parties[i].aborted)
                continue;
            for (int j = 0; j < n; j++) {
                if (parties[i].complaints[j].equals(Protocol.ComplaintState.Waiting))
                    return false;

            }
        }
        return true;
    }

    /**
     * stage3 according to the protocol
     *  1. if more than t players complain against a player Pi he is disqualified.
     *     otherwise Pi broadcasts the share Sij for each complaining player Pj.
     *  2. if any of the revealed shares fails the verification test, player Pi is disqualified.
     *     set QUAL to be the set of non-disqualified players.
     */
    protected void stage3(){
        dkg.answerAllComplainingPlayers();

        // wait until there is no complaint waiting for answer
        while (!stop && !haveReceivedAllStage3ComplaintAnswers())
            waitAndHandleReceivedMessages();

        this.QUAL = dkg.calcQUAL();
    }

    /**
     * stage4 according to the protocol
     *  1. public value y is computed as y = multiplication of yi mod p for i in QUAL
     *  2. public verification values are computed as Ak = multiplication of Aik mod p for i in QUAL for k = 0,...,t
     *  3. Pj sets is share of the secret as xj = sum of Sij mod q for i in QUAL
     */
    protected void stage4(){
        this.y = dkg.calcY(QUAL);
        this.commitments = dkg.calcCommitments(QUAL);
        this.share = dkg.calcShare(QUAL);
    }

    @Override
    public void run() {
        this.runThread = Thread.currentThread();
        // For debugging
        String previousName = runThread.getName();
        runThread.setName(getClass().getName() +":" + getID());

        try {
            stage1();
            waitUntilStageOneCompleted();
            if (stop) return;
            stage2();
            waitUntilStageTwoCompleted();
            if (stop) return;
            stage3();
            if (stop) return;
            stage4();
        } finally {
            runThread.setName(previousName);
        }
    }

    /**
     * current thread in the main loop
     */
    protected Thread runThread;

    /**
     * flag indicates if there was request to stop the current run of the protocol
     */
    protected boolean stop = false;

    /**
     * Request the current run loop to exit gracefully
     */
    public void stop() {
        try {
            stop = true;
            runThread.interrupt();
        }catch (Exception e){
            //do nothing
        }

    }

    /**
     * getter
     * @return commitments
     */
    public ArrayList<T> getCommitments() {
        return commitments;
    }

    /**
     * getter
     * @return g
     */
    public T getGenerator() {
        return g;
    }

    /**
     * getter
     * @return group
     */
    public Group<T> getGroup() {
        return group;
    }

    /**
     * getter
     * @return share
     */
    public Polynomial.Point getShare() {
        return share;
    }

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

    /**
     * getter
     * @return n
     */
    public int getN() {
        return n;
    }

    /**
     * getter
     * @return t
     */
    public int getT() {
        return t;
    }

    /**
     * getter
     * @return y
     */
    public T getPublicValue() {
        return y;
    }

    /**
     * getter
     * @return QUAL
     */
    public Set<Integer> getQUAL() {
        return QUAL;
    }

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

    /**
     * commitment message is valid if:
     *  1. it was received in broadcast chanel
     *  2. the sender didn't sent this commitment before
     */
    protected boolean isValidCommitmentMessage(int sender, boolean isBroadcast, DKG.CommitmentMessage commitmentMessage){
        int i = sender - 1;
        int k = commitmentMessage.getK();
        return isBroadcast && parties[i].commitments.get(k) == null;
    }

    /**
     * secret message is valid if:
     *  1. it was received in private chanel
     *  2. the sender didn't sent secret message before
     *  3. secret.i == i
     *  4. secret.j == id
     */
    protected boolean isValidSecretMessage(int sender, boolean isBroadcast, DKG.ShareMessage secretMessage){
        int i = secretMessage.getI();
        int j = secretMessage.getJ();
        if(sender != i  || isBroadcast)
            return false;
        else
            return parties[i - 1].share == null && j == id;

    }

    /**
     * done message is valid if:
     *  1. it was received in broadcast chanel
     *  2. the sender didn't sent done message before
     */
    protected boolean isValidDoneMessage(int sender, boolean isBroadcast){
        return isBroadcast && !parties[sender - 1].doneFlag;
    }


    /**
     * complaint message is valid if:
     *  1. it was received in broadcast chanel
     *  2. the sender didn't complain against id before
     */
    protected boolean isValidComplaintMessage(int sender, boolean isBroadcast, DKG.IDMessage complaintMessage){
        int i = sender;
        int j = complaintMessage.getId();

        assert(i > 0);
        assert(j > 0);
        assert(i <= parties.length);
        assert(j <= parties[i-1].complaints.length);

        return isBroadcast && parties[i - 1].complaints[j - 1].equals( Protocol.ComplaintState.OK);
    }

    /**
     * answer message is valid if:
     *  1. it was received in broadcast chanel
     *  2. secret.i == i
     *  3. 1 <= secret.j <= n
     *  4. it is marked that j complained against i and i didn't received
     */
    protected boolean isValidAnswerMessage(int sender, boolean isBroadcast, DKG.ShareMessage secretMessage){
        int i = secretMessage.getI();
        int j = secretMessage.getJ();
        if(sender != i  || !isBroadcast)
            return false;
        else
            return j >= 1 && j <= n && parties[i - 1].complaints[j - 1].equals(Protocol.ComplaintState.Waiting);
    }


    public void handleMessage(Comm.BroadcastMessage envelope) throws InvalidProtocolBufferException {
        int sender = envelope.getSender();
        boolean isBroadcast = !envelope.getIsPrivate();
        DKG.Payload msg = DKG.Payload.parseFrom(envelope.getPayload());

        logger.debug("handling Message: Dst={}, Src={}, [{}]",
                envelope.getDestination(), envelope.getSender(), TextFormat.printToString(msg));

        switch (msg.getType()) {
            case COMMITMENT:
                /**
                 * saves the commitment
                 */
                assert msg.getPayloadDataCase() == DKG.Payload.PayloadDataCase.COMMITMENT;
                DKG.CommitmentMessage commitmentMessage = msg.getCommitment();
                if (isValidCommitmentMessage(sender, isBroadcast, commitmentMessage)) {
                    int i = sender - 1;
                    int k = commitmentMessage.getK();

                    parties[i].commitments.set(k, extractCommitment(commitmentMessage));
                }
                break;


            case SHARE:
                /**
                 * saves the secret
                 */
                assert msg.getPayloadDataCase() == DKG.Payload.PayloadDataCase.SHARE;
                DKG.ShareMessage secretMessage = msg.getShare();
                if(isValidSecretMessage(sender,isBroadcast,secretMessage)) {
                    int i = secretMessage.getI();
                    Polynomial.Point secret = extractShare(id,secretMessage.getShare());
                    parties[i - 1].share = secret;
                }
                break;

            case DONE:

                /**
                 * marks that the sender was finished sending all his complaints
                 */
                if(isValidDoneMessage(sender,isBroadcast)) {
                    parties[sender - 1].doneFlag = true;
                }
                break;

            case COMPLAINT:
                /**
                 * marks that the sender was complained against id
                 */
                if (msg.getPayloadDataCase() != DKG.Payload.PayloadDataCase.ID) {
                    logger.error("User {} Expecting ID message, got from SRC={} msg {}", getID(), envelope.getSender(), TextFormat.printToString(msg));
                    assert (msg.getPayloadDataCase() == DKG.Payload.PayloadDataCase.ID);
                }

                DKG.IDMessage complaintMessage = msg.getId();
                if(isValidComplaintMessage(sender,isBroadcast,complaintMessage)){
                    int i = sender;
                    int j = complaintMessage.getId();
                    parties[j - 1].complaints[i - 1] = Protocol.ComplaintState.Waiting;
                }
                break;
            case ANSWER:
                /**
                 * if the secret is valid, marks the complaint as NonDisqualified
                 * else marks it as Disqualified
                 * in case that the complainer is id ( j == id ), saves the secret
                 */
                assert msg.getPayloadDataCase() == DKG.Payload.PayloadDataCase.SHARE;
                secretMessage = msg.getShare();
                if(isValidAnswerMessage(sender,isBroadcast,secretMessage)) {
                    int i = secretMessage.getI();
                    int j = secretMessage.getJ();
                    Polynomial.Point secret = extractShare(j,secretMessage.getShare());
                    if (dkg.isValidShare(secret, parties[i - 1].commitments, j)) {
                        parties[i - 1].complaints[j - 1] = Protocol.ComplaintState.NonDisqualified;
                    } else {
                        parties[i - 1].complaints[j - 1] = Protocol.ComplaintState.Disqualified;
                    }
                    if (j == id) {
                        parties[i - 1].share = secret;
                    }
                }
                break;
            case ABORT:
                /**
                 * marks that the sender was aborted
                 */
                parties[sender - 1].aborted = true;
                break;
            default:
                logger.error("Bad message: SRC={}, DST={}, Payload={}", envelope.getSender(), envelope.getDestination(), TextFormat.printToString(msg));
                break;

        }
    }

    /**
     * extract share value from ByteString
     * @param i
     * @param share
     * @return new Point (i,share)
     */
    public Polynomial.Point extractShare(int i, ByteString share){
        BigInteger x = BigInteger.valueOf(i);
        BigInteger y = new BigInteger(share.toByteArray());
        return new Polynomial.Point(x,y);
    }

    /**
     *
     * @param commitmentMessage
     * @return
     */
    public T extractCommitment(DKG.CommitmentMessage commitmentMessage){
        return dkg.decodeCommitment(commitmentMessage.getCommitment().toByteArray());
    }
}
