/*
 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.base.Charsets;
import com.google.common.io.ByteSink;
import com.google.common.io.Files;
import org.gnunet.identity.Identity;
import org.gnunet.identity.IdentityCallback;
import org.gnunet.mesh.Mesh;
import org.gnunet.mesh.MeshRunabout;
import org.gnunet.mesh.ChannelEndHandler;
import org.gnunet.secretsharing.ThresholdPublicKey;
import org.gnunet.testbed.CompressedConfig;
import org.gnunet.util.*;
import org.gnunet.util.getopt.Argument;
import org.gnunet.util.getopt.ArgumentAction;
import org.gnunet.voting.messages.*;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.io.File;
import java.io.IOException;
import java.io.InputStream;
import java.util.List;
import java.util.Random;

/**
 * Tool for creating, manipulating and submitting ballot files.
 */
public class BallotTool extends Program {
    private static final Logger logger = LoggerFactory
            .getLogger(BallotTool.class);

    @Argument(
            shortname = "e",
            longname = "ego",
            action = ArgumentAction.STORE_STRING,
            description = "ego to use for the operation")
    String egoName = null;

    @Argument(
            shortname = "q",
            longname = "query",
            action = ArgumentAction.SET,
            description = "query election result")
    boolean query = false;

    @Argument(
            shortname = "s",
            longname = "submit",
            action = ArgumentAction.SET,
            description = "submit the vote to the authorities")
    boolean submit = false;

    @Argument(
            shortname = "r",
            longname = "register",
            action = ArgumentAction.SET,
            description = "register an election with the authorities")
    boolean register = false;

    @Argument(
            shortname = "i",
            longname = "issue",
            action = ArgumentAction.SET,
            description = "sign the ballot as issuer (use with -e)")
    boolean issue = false;

    @Argument(
            shortname = "x",
            longname = "select",
            action = ArgumentAction.STORE_STRING,
            argumentName = "CHOICE",
            description = "select and encrypt a vote option (use with -e)")
    String select = null;

    @Argument(
            shortname = "V",
            longname = "verify",
            action = ArgumentAction.SET,
            description = "verify signatures in the ballot and show information")
    boolean verify = false;

    @Argument(
            shortname = "g",
            longname = "group",
            action = ArgumentAction.STORE_STRING,
            description = "incorporate the group cert into the ballot")
    String groupCertFile = null;

    @Argument(
            shortname = "k",
            longname = "getRequestIdentifier-key",
            action = ArgumentAction.SET,
            description = "getRequestIdentifier the threshold public key from authorities")
    boolean requestKey = false;

    @Argument(
            shortname = "t",
            longname = "template",
            action = ArgumentAction.SET,
            description = "write a template ballot to the give ballot file")
    boolean template = false;

    /**
     * The ego to use for the currently executing action.
     */
    private Identity.Ego ego;

    /**
     * The (possibly modified) ballot originally loaded
     * from 'ballotFilename'
     */
    private Ballot ballot;

    /**
     * Filename to read the ballot from, and write modifications to.
     */
    private String ballotFilename;

    /**
     * Our handle to MESH.
     */
    private Mesh mesh;

    /**
     * A channel to 'currentAuthority' or null.
     */
    private Mesh.Channel channel;

    /**
     * The authority we are currently communicating with.
     */
    private PeerIdentity currentAuthority;

    /**
     * Are we finished with communicating over the mesh channel and don't need to worry about
     * disconnection?
     */
    private boolean tunnelCommunicationFinished;

    private RelativeTime tunnelReconnectBackoff = RelativeTime.STD_BACKOFF;

    public class BallotChannelEndHandler implements ChannelEndHandler {
        @Override
        public void onChannelEnd(final Mesh.Channel channel) {
            // FIXME: just re-running 'doCommands' is a bit of a hack
            BallotTool.this.channel = null;
            if (!tunnelCommunicationFinished) {
                logger.warn("mesh channel disconnected, but operation not finished");
                Scheduler.addDelayed(tunnelReconnectBackoff, new Scheduler.Task() {
                    @Override
                    public void run(Scheduler.RunContext ctx) {
                        doCommands();
                    }
                });
                tunnelReconnectBackoff = tunnelReconnectBackoff.backoff();
            }
        }
    }

    /**
     * Destroy the channel to the authority as well
     * as the mesh handle.
     */
    private void endMesh() {
        tunnelCommunicationFinished = true;
        if (null != channel) {
            channel.destroy();
            channel = null;
        }
        if (null != mesh) {
            mesh.destroy();
            mesh = null;
        }
    }

    public class BallotRegisterReceiver extends MeshRunabout {
        public void visit(BallotRegisterSuccessMessage m) {
            System.out.println("ballot successfully registered");
            ballot.addRegistrationSignature(currentAuthority, m.registrationSignature);
            writeBallot();
            endMesh();
        }

        public void visit(BallotRegisterFailureMessage m) {
            System.out.println("registering failed: " + m.reason);
            endMesh();
            setReturnValue(1);
        }
    }

    public class QueryReceiver extends MeshRunabout {
        public void visit(ResultQueryResponseMessage m) {
            if (m.results.length != ballot.choices.size()) {
                System.out.println("failure to query result: malformed response");
            } else {
                System.out.println("got results:");
                for (int i = 0; i < m.results.length; i++) {
                    System.out.println("'" + ballot.choices.get(i) + "': " + m.results[i]);
                }
            }
            endMesh();
        }

        public void visit(ResultQueryFailureMessage m) {
            System.out.println("failure to query result: " + m.reason);
            endMesh();
            setReturnValue(1);
        }
    }


    public class PublicKeyReceiver extends MeshRunabout {
        public void visit(KeyQueryResponseMessage m) {
            System.out.println("got threshold public key!");
            ballot.addThresholdPublicKey(currentAuthority, m);
            writeBallot();
            endMesh();
        }
        public void visit(KeyQueryFailureMessage m) {
            System.out.println("failure to query result: " + m.reason);
            endMesh();
            setReturnValue(1);
        }

    }

    public class SubmitReceiver extends MeshRunabout {

        public void visit(SubmitSuccessMessage m) {
            System.out.println("vote successfully submitted");
            ballot.addConfirmation(currentAuthority, m.confirmationSig);
            writeBallot();
            endMesh();
        }

        public void visit(SubmitFailureMessage m) {
            System.out.println("vote not submitted: " + m.reason);
            if (m.signedAuthorityTime != null) {
                // FIXME: verifyRaw
                System.out.println("authority time: " +
                        AbsoluteTime.fromNetwork(m.signedAuthorityTime.time).toFancyString());
            }
            endMesh();
            setReturnValue(1);
        }
    }

    @Override
    protected String makeHelpText() {
        return "gnunet-ballot [OPTIONS]... BALLOT\n" +
                "Create, modify and execute operation on ballots.";
    }

    private void runTemplate() {
        File f = new File(ballotFilename);
        if (f.exists()) {
            System.err.println("file already exists, not overwriting");
            return;
        }
        InputStream is = getClass().getResourceAsStream("ballot-template.espec");
        ByteSink out = Files.asByteSink(f);
        try {
            out.writeFrom(is);
        } catch (IOException e) {
            System.err.println("could not copy template file: " + e.getMessage());
        }
    }

    /**
     * Write the ballot back to disk.
     */
    private void writeBallot() {
        try {
            Files.write(ballot.serialize(), new File(ballotFilename), Charsets.UTF_8);
        } catch (IOException e) {
            System.err.println("could not write ballot file: " + e.getMessage());
        }
    }

    /**
     * Actually execute the action the user requested.
     */
    void doCommands() {
        if (null != groupCertFile) {
            Configuration groupCertConfig = new Configuration();
            groupCertConfig.parse(groupCertFile);
            GroupCert groupCert = GroupCert.fromGroupCertConfig(groupCertConfig);
            ballot.encodeGroup(groupCert);
            writeBallot();
            return;
        }
        if (register) {
            List<PeerIdentity> remainingAuthorities = ballot.getRemainingRegisterAuthorities();
            if (remainingAuthorities.isEmpty()) {
                System.err.println("all authorities already received the ballot");
                return;
            }
            Random r = new Random();
            currentAuthority = remainingAuthorities.get(r.nextInt(remainingAuthorities.size()));
            System.out.println("registering ballot with authority " + currentAuthority.toString());
            mesh = new Mesh(getConfiguration(), new BallotChannelEndHandler(), new BallotRegisterReceiver());
            channel = mesh.createChannel(currentAuthority, TallyAuthorityDaemon.MESH_PORT, true, true);
            BallotRegisterRequestMessage m = new BallotRegisterRequestMessage();
            CompressedConfig ccfg = new CompressedConfig(ballot.toConfiguration());
            m.compressedBallotConfig = ccfg.compressedData;
            channel.send(m);
            return;
        }
        if (issue) {
            if (null == ego) {
                System.err.println("no ego given");
                setReturnValue(1);
                return;
            }
            ballot.issue(ego.getPrivateKey());
            writeBallot();
            return;
        }
        if (select != null) {
            if (null == ego) {
                System.err.println("no ego given");
                setReturnValue(1);
                return;
            }
            ThresholdPublicKey thresholdPublicKey = ballot.getMajorityThresholdPublicKey();
            if (null == thresholdPublicKey) {
                System.err.println(String.format("no majority threshold public key in ballot (got keys of %s authorities)",
                        (ballot.thresholdPublicKeys == null) ? 0 : ballot.thresholdPublicKeys.size()));
                setReturnValue(1);
                return;
            }
            ballot.encodeChoice(select, thresholdPublicKey, ego.getPrivateKey());
            writeBallot();
            return;
        }
        if (submit) {
            List<PeerIdentity> remainingAuthorities = ballot.getRemainingSubmitAuthorities();
            if (remainingAuthorities.isEmpty()) {
                System.err.println("all authorities already received the ballot");
                return;
            }
            Random r = new Random();
            PeerIdentity authority = remainingAuthorities.get(r.nextInt(remainingAuthorities.size()));
            System.out.println("submitting to authority " + authority.toString());
            currentAuthority = authority;
            mesh = new Mesh(cfg, new BallotChannelEndHandler(), new SubmitReceiver());
            channel = mesh.createChannel(authority, TallyAuthorityDaemon.MESH_PORT, true, true);
            SubmitMessage m = new SubmitMessage();
            if (ballot.voterPub == null) {
                throw new InvalidBallotException("no voter in ballot");
            }
            m.voterPub = ballot.voterPub;
            if (ballot.groupCert == null) {
                throw new InvalidBallotException("no group cert in ballot");
            }
            m.groupCertExpiration = ballot.groupCert.getExpiration().asMessage();
            m.groupCert = ballot.groupCert.getSignature();
            m.ballotGuid = ballot.getBallotGuid();
            if (null == ballot.encryptedVote) {
                throw new InvalidBallotException("no encrypted vote in ballot");
            }
            m.encryptedVote = ballot.encryptedVote;
            channel.send(m);
            return;
        }
        if (verify) {
            System.out.print(ballot.describe());
            return;
        }
        if (query) {
            List<PeerIdentity> remainingAuthorities = ballot.getAuthorities();
            if (remainingAuthorities.isEmpty()) {
                System.err.println("no authorities available");
                setReturnValue(2);
                return;
            }
            Random r = new Random();
            currentAuthority = remainingAuthorities.get(r.nextInt(remainingAuthorities.size()));
            System.out.println("querying authority " + currentAuthority.toString());
            mesh = new Mesh(cfg, new BallotChannelEndHandler(), new QueryReceiver());
            channel = mesh.createChannel(currentAuthority, TallyAuthorityDaemon.MESH_PORT, true, true);
            ResultQueryMessage m = new ResultQueryMessage();
            m.ballotGuid = ballot.getBallotGuid();
            channel.send(m);
            return;
        }
        if (requestKey) {
            List<PeerIdentity> remainingAuthorities = ballot.getRemainingKeyAuthorities();
            if (remainingAuthorities.isEmpty()) {
                System.err.println("all authorities already signed group key");
                return;
            }
            Random r = new Random();
            currentAuthority = remainingAuthorities.get(r.nextInt(remainingAuthorities.size()));
            System.out.println("asking authority for key " + currentAuthority.toString());
            mesh = new Mesh(cfg, new BallotChannelEndHandler(), new PublicKeyReceiver());
            channel = mesh.createChannel(currentAuthority, TallyAuthorityDaemon.MESH_PORT, true, true);
            KeyQueryMessage m = new KeyQueryMessage();
            m.ballotGuid = ballot.getBallotGuid();
            channel.send(m);
            return;
        }
        setReturnValue(1);
        System.err.println("no action specified");
    }

    @Override
    public void run() {
        if (unprocessedArgs.length != 1) {
            System.err.println("no ballot file specified");
            setReturnValue(1);
            return;
        }
        ballotFilename = unprocessedArgs[0];
        if (template)  {
            runTemplate();
            return;
        }
        // all other commands need a loaded ballot file
        File bf = new File(ballotFilename);
        if (!bf.exists()) {
            System.err.println("ballot file does not exist");
            return;
        }
        try {
            ballot = new Ballot(ballotFilename);
        } catch (InvalidBallotException e) {
            System.err.println("Invalid or incomplete ballot:");
            System.err.println(e.getMessage());
            setReturnValue(1);
            return;
        }
        // if there's an ego name, look it up ...
        if (null != egoName) {
            Identity.lookup(getConfiguration(), egoName, new IdentityCallback() {
                @Override
                public void onEgo(Identity.Ego ego) {
                    BallotTool.this.ego = ego;
                    doCommands();
                }
                @Override
                public void onError(String errorMessage) {
                    System.err.println("can't retrieve ego: " + errorMessage);
                    setReturnValue(1);
                }
            });
        } else {
            doCommands();
        }
    }
    public static void main(String args[]) {
        Program tool = new BallotTool();
        int ret = tool.start(args);
        System.exit(ret);
    }
}
