package meerkat.mixer.main;

import com.google.protobuf.InvalidProtocolBufferException;
import com.google.protobuf.StringValue;
import joptsimple.OptionParser;
import joptsimple.OptionSet;
import joptsimple.OptionSpec;
import meerkat.crypto.concrete.ECElGamalEncryption;
import meerkat.crypto.concrete.Util;
import meerkat.mixer.MixGenerator;
import meerkat.mixer.MixVerifier;
import meerkat.mixer.MixerOutput;
import meerkat.mixer.proofs.concrete.Mix2nizk;
import meerkat.protobuf.ConcreteCrypto;
import meerkat.protobuf.Crypto;
import org.factcenter.qilin.primitives.concrete.ECElGamal;
import org.factcenter.qilin.primitives.concrete.ECGroup;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.io.*;
import java.math.BigInteger;
import java.security.SecureRandom;
import java.security.spec.InvalidKeySpecException;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.List;
import java.util.Random;

import static java.lang.System.exit;

/**
 * Command-line mixProverVerifier and verifier.
 */
public class Mix {
    final static String DEFAULT_ECGROUP = "secp256r1";
    final static Logger logger = LoggerFactory.getLogger(Mix.class);

    public Random rand;
    public ECGroup group;
    public ECElGamalEncryption enc;
    public ConcreteCrypto.ElGamalPublicKey serializedPk;
    public ECElGamal.SK secretKey;
    public Mix2nizk mixProverVerifier;

    public MixerOutput inMix;

    public Mix() {
        rand = new SecureRandom();
        enc = new ECElGamalEncryption();
        serializedPk = null;
        secretKey = null;
    }

    public void loadMix(File inFile) throws IOException {

        assert(serializedPk != null);
        InputStream in = new FileInputStream(inFile);
        inMix = MixerOutput.parseDelimitedFrom(in);
        in.close();
        logger.info("Loaded mixnet with {} layers and logN={}", inMix.getNumLayers(), inMix.getLogN());
    }

    public boolean verify(boolean strict) {
        try {
            logger.info("Starting verification of {} ciphertexts", 1 << inMix.getLogN());
            long startTime = System.currentTimeMillis();
            boolean ok = MixVerifier.verifyTable(mixProverVerifier, inMix, strict);
            long endTime = System.currentTimeMillis();
            logger.info("Verification took {} seconds", (endTime - startTime) / 1000f);
            return ok;
        } catch (InvalidProtocolBufferException e) {
            logger.error("Badly formatted encryptions", e);
            return false;
        }
    }

    public void mix(File outFile) throws IOException {
        MixGenerator mixer = new MixGenerator(mixProverVerifier, enc);

        List<Crypto.RerandomizableEncryptedMessage> encryptedMessages = Arrays.asList(inMix.getEncryptedMessages()[inMix.getNumLayers()]);

        logger.info("Starting mix of {} ciphertexts", encryptedMessages.size());
        long startTime = System.currentTimeMillis();
        MixerOutput outMix = mixer.mix(encryptedMessages, rand);
        long endTime = System.currentTimeMillis();
        logger.info("Mix done (took {} seconds)", (endTime - startTime) / 1000f);

        OutputStream out = new FileOutputStream(outFile);
        outMix.writeDelimitedTo(out);
        out.close();
    }


    // For testing purposes

    public void setPK(ConcreteCrypto.ElGamalPublicKey serializedPk) {
        this.serializedPk = serializedPk;

        try {
            enc.init(serializedPk);
            group = enc.getGroup();
        } catch (InvalidKeySpecException e) {
            logger.error("Invalid EC-ElGamal public key", e);
            exit(-2);
        }

        mixProverVerifier = new Mix2nizk(rand, enc);
    }

    public void createKeypair(File outFile) throws IOException {
        createKeypair(DEFAULT_ECGROUP, outFile);
    }

    /**
     * Create a new ECElGamal key pair and write it serialized to file.
     *
     * @param outFile
     */
    public void createKeypair(String curveName, File outFile) throws IOException {
        group = new ECGroup(curveName);
        BigInteger sk = ECElGamal.generateSecretKey(group, rand);
        secretKey = new ECElGamal.SK(group, sk);

        serializedPk = Util.encodePK(group, secretKey);
        setPK(serializedPk);
        Crypto.BigInteger serializedSk = Util.encodeBigInteger(sk);

        OutputStream out = new FileOutputStream(outFile);

        serializedPk.writeDelimitedTo(out);
        serializedSk.writeDelimitedTo(out);

        out.close();
    }

    /**
     * Loads a public key and optionally a secret key from a file.
     * @param inFile
     * @throws IOException
     * @throws InvalidKeySpecException
     */
    public void loadKeypair(File inFile) throws IOException, InvalidKeySpecException {
        InputStream in = new FileInputStream(inFile);
        serializedPk = ConcreteCrypto.ElGamalPublicKey.parseDelimitedFrom(in);

        setPK(serializedPk);

        try {
            Crypto.BigInteger serializedSk = Crypto.BigInteger.parseDelimitedFrom(in);
            secretKey = Util.decodeSK(serializedPk, serializedSk);
        } catch (EOFException e) {
            logger.debug("File {} does not have a secret key", inFile);
        } finally {
            in.close();
        }
    }


    /**
     * Decrypt the output of a mixnet
     * Outputs to a file, one line per decoded element.
     *
     * @throws IOException
     */
    public void decrypt(File outFile) throws IOException {
        PrintStream out = new PrintStream(outFile);

        assert(inMix != null);
        assert(inMix.getEncryptedMessages() != null);

        logger.info("Decrypting {} ciphertexts", inMix.getEncryptedMessages()[inMix.getNumLayers()].length);
        for (Crypto.RerandomizableEncryptedMessage ciphertext : inMix.getEncryptedMessages()[inMix.getNumLayers()]) {
            StringValue plaintext = Util.decrypt(StringValue.class, secretKey, group, ciphertext);
            out.println(plaintext.getValue());
        }

        out.close();
    }

    public void encrypt(File inFile, File outFile) throws IOException {
        BufferedReader in = new BufferedReader(new InputStreamReader(new FileInputStream(inFile)));

        ArrayList<String> lines = new ArrayList<>();
        String line;
        while ((line = in.readLine()) != null) {
            lines.add(line);
        }
        in.close();

        int n = lines.size();
        int logN = 32 - Integer.numberOfLeadingZeros(n - 1);
        n = 1 << logN; // We need the smallest power of 2 greater than the size


        Crypto.RerandomizableEncryptedMessage ciphers[][] = new Crypto.RerandomizableEncryptedMessage[1][n];

        logger.info("Encrypting {} plaintexts (padded to {})", lines.size(), n);
        for (int i = 0; i < lines.size(); ++i) {
            ciphers[0][i] = enc.encrypt(StringValue.newBuilder().setValue(lines.get(i)).build(), enc.generateRandomness(rand));
        }

        // Pad with empty values
        StringValue empty = StringValue.newBuilder().setValue("").build();
        for (int i = lines.size(); i < n; ++i) {
            ciphers[0][i] = enc.encrypt(empty, enc.generateRandomness(rand));
        }

        inMix = new MixerOutput(logN, null, ciphers);

        OutputStream out = new FileOutputStream(outFile);
        inMix.writeDelimitedTo(out);
        out.close();
        logger.info("Wrote mixnet with {} layers and logN={}", inMix.getNumLayers(), inMix.getLogN());
    }


    static void printHelp(OptionParser parser) {
        printHelp(null, parser);
    }

    static void printHelp(String msg, OptionParser parser) {
        if (msg != null)
            System.out.println(msg);
        try {
            parser.printHelpOn(System.out);
        } catch (IOException e) {
            // should never happen
        }
    }

    public static void main(String[] args) {
        OptionParser parser = new OptionParser();
        final OptionSpec<Void> OPT_HELP = parser.accepts("help", "Print help");
        final OptionSpec<String> OPT_CURVENAME = parser.accepts("curve-name", "Use the given named elliptic curve").withRequiredArg().ofType(String.class)
                .defaultsTo(DEFAULT_ECGROUP);
        final OptionSpec<Void> OPT_GENKEY = parser.accepts("genkey", "Generate a key-pair (write into key file)");
        final OptionSpec<Void> OPT_DECRYPT = parser.accepts("decrypt", "Decrypt using given keypair");
        final OptionSpec<File> OPT_KEYFILE = parser.accepts("keys", "File containing public key (or keypair for decryption)").withRequiredArg().ofType(File.class);
        final OptionSpec<Void> OPT_ENCRYPT = parser.accepts("encrypt", "Decrypt using given keypair");
        final OptionSpec<Void> OPT_STRICT = parser.accepts("strict", "Use strict verification (verify that network matches our Benes implementation)");
        final OptionSpec<File> OPT_INPUTFILE = parser.accepts("infile", "Input file (if mixing, should contain ciphertext or mixProverVerifier output; if verifying, mixProverVerifier output)").withRequiredArg().ofType(File.class);
        final OptionSpec<File> OPT_OUTPUTFILE = parser.accepts("outfile", "Output file. If given, operate in Mix mode; otherwise in verify mode.").withRequiredArg().ofType(File.class);

        OptionSet options = parser.parse(args);
        if (options.has(OPT_HELP)) {
            printHelp(parser);
            return;
        }

        File keyFile = options.valueOf(OPT_KEYFILE);
        if (keyFile == null) {
            printHelp("Must specify key file", parser);
            return;
        }

        File inFile = options.valueOf(OPT_INPUTFILE);
        if ((inFile == null || !inFile.canRead()) && !options.has(OPT_GENKEY)) {
            printHelp("Must specify input file except for key generation" + (inFile.canRead() ? "" : " (can't read inFile)"), parser);
            return;
        }


        File outFile = options.valueOf(OPT_OUTPUTFILE);

        Mix mix = new Mix();

        try {
            if (options.has(OPT_GENKEY)) {
                mix.createKeypair(options.valueOf(OPT_CURVENAME), keyFile);
            } else {
                mix.loadKeypair(keyFile);
                if (options.has(OPT_ENCRYPT)) {
                    mix.encrypt(inFile, outFile);
                } else if (options.has(OPT_DECRYPT)) {
                    mix.loadMix(inFile);
                    mix.decrypt(outFile);
                } else if (options.has(OPT_OUTPUTFILE)) {
                    // Mix mode
                    mix.loadMix(inFile);
                    mix.mix(outFile);
                } else {
                    mix.loadMix(inFile);
                    boolean ok = mix.verify(options.has(OPT_STRICT));
                    if (ok) {
                        logger.info("Verification successful");
                    } else {
                        logger.error("Verification failed!");
                        exit(-1);
                    }
                }
            }
        } catch (FileNotFoundException e) {
            e.printStackTrace();
        } catch (IOException e) {
            e.printStackTrace();
        } catch (InvalidKeySpecException e) {
            e.printStackTrace();
        }

    }
}
