/*
 This file is part of GNUnet.
  (C) 2012, 2013 Christian Grothoff (and other contributing authors)

  GNUnet is free software; you can redistribute it and/or modify
  it under the terms of the GNU General Public License as published
  by the Free Software Foundation; either version 3, or (at your
  option) any later version.

  GNUnet is distributed in the hope that it will be useful, but
  WITHOUT ANY WARRANTY; without even the implied warranty of
  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
  General Public License for more details.

  You should have received a copy of the GNU General Public License
  along with GNUnet; see the file COPYING.  If not, write to the
  Free Software Foundation, Inc., 59 Temple Place - Suite 330,
  Boston, MA 02111-1307, USA.
 */
package org.gnunet.voting;


import com.google.common.collect.Maps;
import org.gnunet.consensus.Consensus;
import org.gnunet.consensus.ConsensusCallback;
import org.gnunet.consensus.ConsensusElement;
import org.gnunet.construct.Construct;
import org.gnunet.mesh.Mesh;
import org.gnunet.mesh.MeshRunabout;
import org.gnunet.mesh.ChannelEndHandler;
import org.gnunet.secretsharing.*;
import org.gnunet.testbed.CompressedConfig;
import org.gnunet.util.*;
import org.gnunet.util.crypto.EcdsaPublicKey;
import org.gnunet.util.crypto.EddsaPrivateKey;
import org.gnunet.util.crypto.EddsaPublicKey;
import org.gnunet.util.crypto.EddsaSignature;
import org.gnunet.voting.messages.*;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.util.HashMap;
import java.util.HashSet;
import java.util.Set;


/**
 * Daemon that is responsible for accepting and counting votes.
 */
public class TallyAuthorityDaemon extends Program {
    private static final Logger logger = LoggerFactory
            .getLogger(TallyAuthorityDaemon.class);

    /**
     * Mesh port used to connect to to the tally authority daemon.
     */
    public static final int MESH_PORT = 1002;

    /**
     * Mesh handle.
     */
    private Mesh mesh;

    /**
     * Private key of the local peer.
     */
    private EddsaPrivateKey authorityPrivateKey;

    /**
     * Public key of the local peer.
     */
    private EddsaPublicKey authorityPublicKey;

    /**
     * All elections known to this authority
     */
    private HashMap<HashCode, ElectionState> elections = Maps.newHashMap();

    /**
     * State of one election.
     */
    class ElectionState {
        /**
         * The ballot that describes this election.
         */
        Ballot ballot;

        /**
         * The threshold crypto share, null if the key has not yet been
         * established.
         */
        Share share;

        /**
         * A voter is in this set if its vote has been in the consensus.
         */
        Set<EcdsaPublicKey> countedVoters = new HashSet<EcdsaPublicKey>();

        /**
         * Key generation session.
         */
        KeyGeneration keyGeneration;

        /**
         * Consensus with the other authorities on the set of ballots.
         */
        Consensus consensus;

        /**
         * Are we done with the vote consensus?
         */
        boolean consensusDone;

        /**
         * Product of all encrypted votes (mod q), used to compute the final tally.
         */
        Ciphertext voteProduct = Ciphertext.identity();

        /**
         * Maping from choice to number of votes for that choice.
         * In our currently simplified implementation, tally.length is always 2.
         * If the tally has not been counted yet, 'tally' is null.
         */
        long[] tally;

        /**
         * The decrypt session.
         */
        Decryption decryption;
    }

    /**
     * Callbacks for the vote consensus.
     */
    class ElectionConsensusConclude implements ConsensusCallback {
        private final ElectionState electionState;

        public ElectionConsensusConclude(ElectionState electionState) {
            this.electionState = electionState;
        }
        @Override
        public void onElement(ConsensusElement element) {
            System.out.println("got element from consensus");
            EncryptedVote vote = Construct.parseAs(element.data, EncryptedVote.class);
            System.out.println("got vote from consensus, ciphertext: " + vote.v.toString());
            if (electionState.countedVoters.contains(vote.voterPublicKey)) {
                // Complain.  FIXME: keep lexically largest vote, so ballot is unambigous
                logger.error("voter {} voted twice", vote.voterPublicKey);
                return;
            }
            electionState.voteProduct = electionState.voteProduct.multiply(vote.v);
            electionState.countedVoters.add(vote.voterPublicKey);

            System.out.println("threshold key (of this authority): " + electionState.share.publicKey.toString());
        }

        @Override
        public void onDone() {
            System.out.println("consensus concluded");
            electionState.consensusDone = true;
            electionState.consensus.destroy();
            electionState.consensus = null;

            electionState.decryption = new Decryption(
                    getConfiguration(),
                    electionState.share,
                    electionState.voteProduct,
                    electionState.ballot.concludeTime,
                    electionState.ballot.queryTime,
                    new DecryptCallback() {
                        @Override
                        public void onResult(Plaintext plaintext) {
                            logger.info("got decypt result");
                            long l = electionState.countedVoters.size();
                            long t[] = plaintext.bruteForceDiscreteLog(l, electionState.ballot.generators);
                            if (null == t) {
                                logger.warn("could not brute-force result");
                            } else {
                                logger.info("brute-forced result");
                                electionState.tally = t;
                            }
                        }
                    });
        }
    }

    class ConsensusConcludeTask implements Scheduler.Task {
        /**
         * Which election on this authority is the consensus conclude for?
         */
        private final ElectionState electionState;

        public ConsensusConcludeTask(ElectionState electionState) {
            this.electionState = electionState;
        }
        @Override
        public void run(Scheduler.RunContext ctx) {
            electionState.consensus.conclude(new ElectionConsensusConclude(electionState));
        }
    }

    static class SecretReady implements SecretReadyCallback {
        private final ElectionState electionState;

        public SecretReady(ElectionState electionState) {
            this.electionState = electionState;
        }

        @Override
        public void onSecretReady(Share share) {
            electionState.keyGeneration = null;
            electionState.share = share;
        }
    }

    private SubmitFailureMessage.SignedAuthorityTime getTimeSigMessage() {
        SubmitFailureMessage.SignedAuthorityTime tm = new SubmitFailureMessage.SignedAuthorityTime();
        // FIXME!
        tm.purpose = 0;
        tm.time = AbsoluteTime.now().asMessage();
        tm.signature = authorityPrivateKey.sign(Construct.toBinary(tm.time), tm.purpose,authorityPublicKey);
        return tm;
    }

    private class TallyMeshReceiver extends MeshRunabout {
        public void visit(SubmitMessage m) {
            logger.debug("got submit message");
            ElectionState electionState = elections.get(m.ballotGuid);
            if (null == electionState) {
                SubmitFailureMessage fm = new SubmitFailureMessage();
                fm.reason = "no matching ballot found";
                getSender().send(fm);
            } else if (!electionState.ballot.startTime.isDue()) {
                SubmitFailureMessage fm = new SubmitFailureMessage();
                fm.reason = "too early to submit vote";
                fm.signedAuthorityTime = getTimeSigMessage();
                getSender().send(fm);
            } else if (electionState.ballot.closingTime.isDue()) {
                SubmitFailureMessage fm = new SubmitFailureMessage();
                fm.reason = "too late to submit vote";
                fm.signedAuthorityTime = getTimeSigMessage();
                getSender().send(fm);
            }
            // FIXME: check signatures of voter and CA
            else {
                // we do *not* check for duplicate votes here,
                // as consensus takes care of this, and there is no harm in sending
                // exact duplicates
                byte[] elem = Construct.toBinary(m.encryptedVote);
                electionState.consensus.insertElement(new ConsensusElement(elem, 0));
                SubmitSuccessMessage sm = new SubmitSuccessMessage();
                sm.confirmationSig = EddsaSignature.randomGarbage();
                getSender().send(sm);
            }

            getSender().receiveDone();
        }

        public void visit(BallotRegisterRequestMessage m) {
            logger.info("ballot register requested");
            CompressedConfig ccfg = new CompressedConfig(m.compressedBallotConfig);
            Ballot b;
            HashCode guid;
            try {
                b = new Ballot(ccfg.decompress());
                guid = b.getBallotGuid();
            } catch (InvalidBallotException e) {
                BallotRegisterFailureMessage fm = new BallotRegisterFailureMessage();
                fm.reason = "invalid ballot (" + e.getMessage() + ")";
                getSender().send(fm);
                getSender().receiveDone();
                return;
            }
            if (elections.containsKey(guid)) {
                BallotRegisterFailureMessage fm = new BallotRegisterFailureMessage();
                fm.reason = "ballot with same GUID already registered";
                getSender().send(fm);
                return;
            }
            ElectionState electionState = new ElectionState();
            electionState.ballot = b;
            PeerIdentity[] ids = new PeerIdentity[b.getAuthorities().size()];
            ids = b.getAuthorities().toArray(ids);
            electionState.consensus = new Consensus(
                    getConfiguration(),
                    ids,
                    b.getBallotGuid(),
                    electionState.ballot.closingTime,
                    electionState.ballot.concludeTime);

            ConsensusConcludeTask t = new ConsensusConcludeTask(electionState);
            if (b.concludeTime.isDue()) {
                logger.info("concluding now");
                Scheduler.add(t);
            } else {
                logger.info("concluding in {}", b.closingTime.getRemaining().getSeconds());
                Scheduler.addDelayed(b.closingTime.getRemaining(), t);
            }
            System.out.println("authority threshold: " + electionState.ballot.threshold);
            System.out.println("authority num_peers: " + electionState.ballot.authorities.size());
            // we hash the GUID a second time, so that there's no
            // collision with the consensus (as secretsharing also uses consensus internally)
            electionState.keyGeneration = new KeyGeneration(
                    getConfiguration(),
                    ids,
                    HashCode.hash(b.getBallotGuid().data),
                    electionState.ballot.keygenStartTime,
                    electionState.ballot.keygenEndTime,
                    electionState.ballot.threshold, new SecretReady(electionState));
            elections.put(guid, electionState);

            BallotRegisterSuccessMessage rm = new BallotRegisterSuccessMessage();
            rm.registrationSignature = EddsaSignature.randomGarbage();
            getSender().send(rm);
        }

        public void visit(ResultQueryMessage m) {
            logger.debug("got result query message");
            ElectionState electionState = elections.get(m.ballotGuid);
            if (null == electionState) {
                ResultQueryFailureMessage rm = new ResultQueryFailureMessage();
                rm.reason = "no matching ballot found";
                getSender().send(rm);
            } else {
                if (!electionState.ballot.queryTime.isDue()) {
                    ResultQueryFailureMessage rm = new ResultQueryFailureMessage();
                    rm.reason = "result query not allowed yet";
                    getSender().send(rm);
                }
                else if (null == electionState.tally) {
                    ResultQueryFailureMessage rm = new ResultQueryFailureMessage();
                    rm.reason = "tally not yet available";
                    getSender().send(rm);
                }
                else {
                    ResultQueryResponseMessage rm = new ResultQueryResponseMessage();
                    rm.results = electionState.tally;
                    getSender().send(rm);
                }
            }
            getSender().receiveDone();
        }

        public void visit(KeyQueryMessage m) {
            logger.debug("got key query message");
            getSender().receiveDone();
            ElectionState electionState = elections.get(m.ballotGuid);
            if (null == electionState) {
                KeyQueryFailureMessage rm = new KeyQueryFailureMessage();
                rm.reason = "no matching ballot found";
                getSender().send(rm);
                return;
            }
            if (!electionState.ballot.keygenEndTime.isDue()) {
                KeyQueryFailureMessage rm = new KeyQueryFailureMessage();
                rm.reason = "key query not allowed yet";
                getSender().send(rm);
                return;
            }
            if (null == electionState.share) {
                KeyQueryFailureMessage rm = new KeyQueryFailureMessage();
                rm.reason = "key not yet established";
                getSender().send(rm);
                return;
            }
            KeyQueryResponseMessage.BallotPublicKey ballotPublicKey = new KeyQueryResponseMessage.BallotPublicKey();
            ballotPublicKey.ballotGuid = electionState.ballot.getBallotGuid();
            ballotPublicKey.publicKey = electionState.share.publicKey;

            KeyQueryResponseMessage rm = new KeyQueryResponseMessage();
            rm.signedGuidKey = ballotPublicKey;
            // FIXME!
            rm.purpose = 0;
            rm.signature = authorityPrivateKey.sign(Construct.toBinary(rm.signedGuidKey),
                    rm.purpose, authorityPublicKey);
            getSender().send(rm);
        }
    }

    public TallyAuthorityDaemon() {
        authorityPrivateKey = EddsaPrivateKey.createRandom();
        authorityPublicKey = authorityPrivateKey.getPublicKey();
    }

    public static void main(String[] args) {
        TallyAuthorityDaemon daemon = new TallyAuthorityDaemon();
        int ret = daemon.start(args);
        System.exit(ret);
    }

    @Override
    public void run() {
        logger.info("running tally daemon");
        mesh = new Mesh(getConfiguration(), null, new ChannelEndHandler() {
            @Override
            public void onChannelEnd(Mesh.Channel channel) {
                logger.warn("on channel end");
            }
        }, new TallyMeshReceiver(), MESH_PORT);

        Scheduler.addDelayed(RelativeTime.FOREVER, new Scheduler.Task() {
            @Override
            public void run(Scheduler.RunContext ctx) {
                if (null != mesh) {
                    mesh.destroy();
                    mesh = null;
                }
            }
        });
    }
}
