/*
 * Copyright 2018 Scytl Secure Electronic Voting SA
 *
 * All rights reserved
 *
 * See our extended copyright notice in *file 'Copyright.txt' which is part of this source code package
 */
package com.scytl.products.ov.mixnet.webapp.mvc;

import java.io.ByteArrayInputStream;
import java.io.FileInputStream;
import java.io.FileNotFoundException;
import java.io.FileOutputStream;
import java.io.IOException;
import java.io.InputStream;
import java.io.OutputStream;
import java.nio.charset.StandardCharsets;
import java.nio.file.Path;
import java.nio.file.Paths;
import java.util.ArrayList;
import java.util.Base64;
import java.util.List;
import java.util.stream.Collectors;

import org.apache.commons.io.IOUtils;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import com.scytl.cryptolib.api.exceptions.GeneralCryptoLibException;
import com.scytl.cryptolib.elgamal.bean.ElGamalPrivateKey;
import com.scytl.cryptolib.elgamal.bean.ElGamalPublicKey;
import com.scytl.cryptolib.mathematical.groups.activity.ExponentCompressor;
import com.scytl.cryptolib.mathematical.groups.activity.GroupElementsCompressor;
import com.scytl.cryptolib.mathematical.groups.impl.Exponent;
import com.scytl.cryptolib.mathematical.groups.impl.ZpGroupElement;
import com.scytl.cryptolib.mathematical.groups.impl.ZpSubgroup;
import com.scytl.products.ov.commons.beans.domain.model.vote.EncryptedVote;
import com.scytl.products.ov.mixnet.commons.configuration.ApplicationConfig;
import com.scytl.products.ov.mixnet.commons.configuration.MixerConfig;
import com.scytl.products.ov.mixnet.commons.configuration.VerifierConfig;
import com.scytl.products.ov.mixnet.commons.configuration.file.FileBasedMixerConfig;
import com.scytl.products.ov.mixnet.commons.configuration.file.FileBasedVerifierConfig;
import com.scytl.products.ov.mixnet.commons.configuration.locations.DefaultLocationNames;
import com.scytl.products.ov.mixnet.commons.configuration.locations.LocationConstants;
import com.scytl.products.ov.mixnet.commons.constants.Constants;
import com.scytl.products.ov.mixnet.commons.exceptions.MixingException;
import com.scytl.products.ov.mixnet.commons.io.ElgamalPublicKeyReader;
import com.scytl.products.ov.mixnet.commons.preprocess.Partitioner;
import com.scytl.products.ov.validation.ValidatorConfig;

public class MixDecValConfig {

    private static final Logger LOG = LoggerFactory.getLogger("std");

    private Integer concurrencyLevel;

    private Integer numberOfParallelBatches;

    private Integer numOptions;

    private Integer minimumBatchSize;

    private FileBasedMixerConfig mixer;

    private FileBasedVerifierConfig verifier;

    private ValidatorConfig validator;

    public Integer getConcurrencyLevel() {
        return concurrencyLevel;
    }

    public void setConcurrencyLevel(final Integer concurrencyLevel) {
        this.concurrencyLevel = concurrencyLevel;
    }

    public Integer getNumOptions() {
        return numOptions;
    }

    public void setNumOptions(final Integer numOptions) {
        this.numOptions = numOptions;
    }

    public FileBasedMixerConfig getMixer() {
        return mixer;
    }

    public void setMixer(final FileBasedMixerConfig mixer) {
        this.mixer = mixer;
    }

    public FileBasedVerifierConfig getVerifier() {
        return verifier;
    }

    public ValidatorConfig getValidator() {
        return validator;
    }

    public void setVerifier(final FileBasedVerifierConfig verifier) {
        this.verifier = verifier;
    }

    public boolean hasMixerConfig() {
        return this.mixer != null;
    }

    public boolean hasVerifierConfig() {
        return this.verifier != null;
    }

    public Integer getMinimumBatchSize() {
        return minimumBatchSize;
    }

    public void setMinimumBatchSize(final Integer minimumBatchSize) {
        this.minimumBatchSize = minimumBatchSize;
    }

    public Integer getNumberOfParallelBatches() {
        return numberOfParallelBatches;
    }

    public void setNumberOfParallelBatches(final Integer numberOfParallelBatches) {
        this.numberOfParallelBatches = numberOfParallelBatches;
    }

    public void setValidator(ValidatorConfig validator) {
        this.validator = validator;
    }

    public static ApplicationConfig convertToStreamConfig(MixDecValConfig mdvConfig, List<EncryptedVote> votes)
            throws IOException {
        ApplicationConfig config = new ApplicationConfig();
        config.setConcurrencyLevel(mdvConfig.getConcurrencyLevel());
        config.setMinimumBatchSize(mdvConfig.getMinimumBatchSize());
        config.setNumberOfParallelBatches(mdvConfig.getNumberOfParallelBatches());
        config.setNumOptions(mdvConfig.getNumOptions());
        Integer lines = 0;
        if (mdvConfig.getMixer() != null) {
            lines = votes.size();
            config.setMixer(convertMixer(mdvConfig.getMixer(), mdvConfig.getMinimumBatchSize(), votes));
        }
        if (mdvConfig.getVerifier() != null) {
            config.setVerifier(convertVerifier(mdvConfig.getVerifier(), mdvConfig.getMinimumBatchSize(), lines));
        }
        return config;
    }

    private static MixerConfig convertMixer(FileBasedMixerConfig fileBasedMixerConfig, Integer partitionSize,
            List<EncryptedVote> votes) throws IOException {

        LOG.info("Obtaining mixer config...");

        MixerConfig mixerConfig = new MixerConfig();
        mixerConfig.setMix(fileBasedMixerConfig.isMix());
        mixerConfig.setWarmUpSize(fileBasedMixerConfig.getWarmUpSize());
        String ballotDecryptionKey =
            fileBasedMixerConfig.getInput().get(LocationConstants.BALLOT_DECRYPTION_KEY_JSON_TAG);

        int desiredKeySize = 1;
        if ((votes != null) && (!votes.isEmpty())) {
            // all the votes (ciphertexts) in the ballot box will have the same number of phis
            desiredKeySize = votes.get(0).getPhis().size();
        }

        if (ballotDecryptionKey != null) {

            InputStream is = new ByteArrayInputStream(Base64.getDecoder().decode(ballotDecryptionKey));
            try {

                ElGamalPrivateKey privateKey = ElGamalPrivateKey.fromJson(IOUtils.toString(is));
                privateKey = compressPrivateKeyIfNecessary(privateKey, desiredKeySize);
                mixerConfig.setBallotDecryptionKeyInput(IOUtils.toInputStream(privateKey.toJson()));

            } catch (GeneralCryptoLibException e) {
                LOG.error("Exception when trying to obtain the ballot decryption key: " + e.getMessage());
                throw new MixingException(e);
            }
        }
        if (votes != null) {
            mixerConfig.setEncryptedBallotsInput(new ByteArrayInputStream(votes.stream().map(vote -> vote.toString())
                .collect(Collectors.joining(System.lineSeparator())).getBytes(StandardCharsets.UTF_8)));
        }
        if (fileBasedMixerConfig.getInput().get(LocationConstants.ENCRYPTION_PARAMETERS_JSON_TAG) != null) {
            mixerConfig.setEncryptionParametersInput(new FileInputStream(
                Paths.get(fileBasedMixerConfig.getInput().get(LocationConstants.ENCRYPTION_PARAMETERS_JSON_TAG)
                    + Constants.JSON_FILE_EXTENSION).toFile()));
        }
        if (fileBasedMixerConfig.getInput().get(LocationConstants.PUBLIC_KEY_JSON_TAG) != null) {

            FileInputStream fileInputStream =
                new FileInputStream(Paths.get(fileBasedMixerConfig.getInput().get(LocationConstants.PUBLIC_KEY_JSON_TAG)
                    + Constants.JSON_FILE_EXTENSION).toFile());
            try {

                ElGamalPublicKey publicKey = ElgamalPublicKeyReader.readPublicKeyFromStream(fileInputStream);
                publicKey = compressPublicKeyIfNecessary(publicKey, desiredKeySize);
                mixerConfig.setPublicKeyInput(IOUtils.toInputStream(publicKey.toJson()));

            } catch (GeneralCryptoLibException e) {
                LOG.error("Exception when trying to obtain the mixing public key: " + e.getMessage());
                throw new MixingException(e);
            }

            LOG.info("Successfully obtained mixer config");
        }

        if (votes != null) {
            String outputDirectory = fileBasedMixerConfig.getOutputDirectory();
            if (outputDirectory != null) {
                Path path = Paths.get(outputDirectory);
                path.toFile().mkdirs();
                int numBatches = new Partitioner(partitionSize).generatePartitionsOfMinimumSize(votes.size()).size();
                List<OutputStream> commitmentParametersOutputOfBatch = new ArrayList<>(numBatches);
                List<OutputStream> decryptedVotesOutputOfBatch = new ArrayList<>(numBatches);
                List<OutputStream> encryptedBallotsOutputOfBatch = new ArrayList<>(numBatches);
                List<OutputStream> encryptionParametersOutputOfBatch = new ArrayList<>(numBatches);
                List<OutputStream> proofsOutputOfBatch = new ArrayList<>(numBatches);
                List<OutputStream> publicKeyOutputOfBatch = new ArrayList<>(numBatches);
                List<OutputStream> reEncryptedBallotsOutputOfBatch = new ArrayList<>(numBatches);
                List<OutputStream> votesWithProofOutputOfBatch = new ArrayList<>(numBatches);
                for (int i = 0; i < numBatches; i++) {
                    String batchNum = Integer.toString(i);
                    path.resolve(batchNum).toFile().mkdirs();
                    commitmentParametersOutputOfBatch.add(new FileOutputStream(Paths
                        .get(outputDirectory, batchNum,
                            DefaultLocationNames.COMMITMENT_PARAMETERS_OUTPUT_FILE_NAME + Constants.JSON_FILE_EXTENSION)
                        .toFile()));
                    decryptedVotesOutputOfBatch.add(new FileOutputStream(Paths
                        .get(outputDirectory, batchNum,
                            DefaultLocationNames.DECRYPTED_VOTES_OUTPUT_FILE_NAME + Constants.CSV_FILE_EXTENSION)
                        .toFile()));
                    encryptedBallotsOutputOfBatch.add(new FileOutputStream(Paths
                        .get(outputDirectory, batchNum,
                            DefaultLocationNames.ENCRYPTED_BALLOTS_OUTPUT_FILE_NAME + Constants.CSV_FILE_EXTENSION)
                        .toFile()));
                    encryptionParametersOutputOfBatch.add(new FileOutputStream(Paths
                        .get(outputDirectory, batchNum,
                            DefaultLocationNames.ENCRYPTION_PARAMETERS_OUTPUT_FILE_NAME + Constants.JSON_FILE_EXTENSION)
                        .toFile()));
                    proofsOutputOfBatch
                        .add(new FileOutputStream(Paths
                            .get(outputDirectory, batchNum,
                                DefaultLocationNames.PROOFS_OUTPUT_FILE_NAME + Constants.JSON_FILE_EXTENSION)
                            .toFile()));
                    publicKeyOutputOfBatch.add(new FileOutputStream(Paths
                        .get(outputDirectory, batchNum,
                            DefaultLocationNames.PUBLIC_KEY_OUTPUT_FILE_NAME + Constants.JSON_FILE_EXTENSION)
                        .toFile()));
                    reEncryptedBallotsOutputOfBatch.add(new FileOutputStream(Paths
                        .get(outputDirectory, batchNum,
                            DefaultLocationNames.REENCRYPTED_BALLOTS_OUTPUT_FILE_NAME + Constants.CSV_FILE_EXTENSION)
                        .toFile()));
                    votesWithProofOutputOfBatch
                        .add(new FileOutputStream(Paths
                            .get(outputDirectory, batchNum,
                                DefaultLocationNames.VOTES_WITH_PROOF_FILE_NAME + Constants.CSV_FILE_EXTENSION)
                            .toFile()));
                }

                mixerConfig.setCommitmentParametersOutputOfBatch(commitmentParametersOutputOfBatch);
                mixerConfig.setDecryptedVotesOutputOfBatch(decryptedVotesOutputOfBatch);
                mixerConfig.setEncryptedBallotsOutputOfBatch(encryptedBallotsOutputOfBatch);
                mixerConfig.setEncryptionParametersOutputOfBatch(encryptionParametersOutputOfBatch);
                mixerConfig.setProofsOutputOfBatch(proofsOutputOfBatch);
                mixerConfig.setPublicKeyOutputOfBatch(publicKeyOutputOfBatch);
                mixerConfig.setReEncryptedBallotsOutputOfBatch(reEncryptedBallotsOutputOfBatch);
                mixerConfig.setVotesWithProofOutputOfBatch(votesWithProofOutputOfBatch);
            }
        }

        return mixerConfig;
    }

    private static ElGamalPrivateKey compressPrivateKeyIfNecessary(final ElGamalPrivateKey elGamalPrivateKey,
            final int numOptions) throws GeneralCryptoLibException {

        LOG.info("Checking if its necessary to compress the ballot decryption key...");

        List<Exponent> keys = elGamalPrivateKey.getKeys();

        if (keys.size() <= numOptions) {
            LOG.info("No need to compress the ballot decryption key");
            return elGamalPrivateKey;
        }
        LOG.info(
            String.format("Compressing the ballot decryption key. Size: %s, required: %s", keys.size(), numOptions));

        ExponentCompressor<ZpSubgroup> compressor = new ExponentCompressor<>(elGamalPrivateKey.getGroup());

        List<Exponent> listWithCompressedFinalElement =
            compressor.buildListWithCompressedFinalElement(numOptions, keys);

        return new ElGamalPrivateKey(listWithCompressedFinalElement, elGamalPrivateKey.getGroup());
    }

    private static ElGamalPublicKey compressPublicKeyIfNecessary(ElGamalPublicKey elGamalPublicKey, int numOptions)
            throws GeneralCryptoLibException {

        LOG.info("Checking if its necessary to compress the mixing public key...");

        List<ZpGroupElement> keys = elGamalPublicKey.getKeys();

        if (keys.size() <= numOptions) {
            LOG.info("No need to compress the mixing public key");
            return elGamalPublicKey;
        }
        LOG.info(String.format("Compressing the mixing public key. Size: %s, required: %s", keys.size(), numOptions));

        GroupElementsCompressor<ZpGroupElement> compressor = new GroupElementsCompressor<>();

        List<ZpGroupElement> listWithCompressedFinalElement =
            compressor.buildListWithCompressedFinalElement(numOptions, keys);

        return new ElGamalPublicKey(listWithCompressedFinalElement, elGamalPublicKey.getGroup());
    }

    private static VerifierConfig convertVerifier(FileBasedVerifierConfig fileBasedVerifierConfig,
            Integer partitionSize, Integer lines) throws FileNotFoundException {
        VerifierConfig verifierConfig = new VerifierConfig();
        verifierConfig.setVerify(fileBasedVerifierConfig.isVerify());
        String outputDirectory = fileBasedVerifierConfig.getOutputDirectory();
        String inputDirectory = fileBasedVerifierConfig.getInputDirectory();
        if (outputDirectory != null) {
            Path path = Paths.get(outputDirectory);
            path.toFile().mkdirs();
            verifierConfig.setOutput(new FileOutputStream(Paths
                .get(outputDirectory, DefaultLocationNames.VERIFIER_OUTPUT_FILE_NAME + Constants.JSON_FILE_EXTENSION)
                .toFile()));
        }
        if (lines != null) {
            int numBatches = new Partitioner(partitionSize).generatePartitionsOfMinimumSize(lines).size();
            List<InputStream> commitmentParametersInputOfBatch = new ArrayList<>(numBatches);
            List<InputStream> decryptedVotesInputOfBatch = new ArrayList<>(numBatches);
            List<InputStream> encryptedBallotsInputOfBatch = new ArrayList<>(numBatches);
            List<InputStream> encryptionParametersInputOfBatch = new ArrayList<>(numBatches);
            List<InputStream> proofsInputOfBatch = new ArrayList<>(numBatches);
            List<InputStream> publicKeyInputOfBatch = new ArrayList<>(numBatches);
            List<InputStream> reEncryptedBallotsInputOfBatch = new ArrayList<>(numBatches);
            List<InputStream> votesWithProofInputOfBatch = new ArrayList<>(numBatches);
            for (int i = 0; i < numBatches; i++) {
                String batchNum = Integer.toString(i);
                commitmentParametersInputOfBatch.add(new FileInputStream(Paths
                    .get(inputDirectory, batchNum,
                        DefaultLocationNames.COMMITMENT_PARAMETERS_OUTPUT_FILE_NAME + Constants.JSON_FILE_EXTENSION)
                    .toFile()));
                encryptedBallotsInputOfBatch.add(new FileInputStream(Paths
                    .get(inputDirectory, batchNum,
                        DefaultLocationNames.ENCRYPTED_BALLOTS_OUTPUT_FILE_NAME + Constants.CSV_FILE_EXTENSION)
                    .toFile()));
                encryptionParametersInputOfBatch.add(new FileInputStream(Paths
                    .get(inputDirectory, batchNum,
                        DefaultLocationNames.ENCRYPTION_PARAMETERS_OUTPUT_FILE_NAME + Constants.JSON_FILE_EXTENSION)
                    .toFile()));
                proofsInputOfBatch
                    .add(
                        new FileInputStream(Paths
                            .get(inputDirectory, batchNum,
                                DefaultLocationNames.PROOFS_OUTPUT_FILE_NAME + Constants.JSON_FILE_EXTENSION)
                            .toFile()));
                publicKeyInputOfBatch
                    .add(new FileInputStream(Paths
                        .get(inputDirectory, batchNum,
                            DefaultLocationNames.PUBLIC_KEY_OUTPUT_FILE_NAME + Constants.JSON_FILE_EXTENSION)
                        .toFile()));
                reEncryptedBallotsInputOfBatch.add(new FileInputStream(Paths
                    .get(inputDirectory, batchNum,
                        DefaultLocationNames.REENCRYPTED_BALLOTS_OUTPUT_FILE_NAME + Constants.CSV_FILE_EXTENSION)
                    .toFile()));
            }

            verifierConfig.setCommitmentParametersInputOfBatch(commitmentParametersInputOfBatch);
            verifierConfig.setDecryptedVotesInputOfBatch(decryptedVotesInputOfBatch);
            verifierConfig.setEncryptedBallotsInputOfBatch(encryptedBallotsInputOfBatch);
            verifierConfig.setEncryptionParametersInputOfBatch(encryptionParametersInputOfBatch);
            verifierConfig.setProofsInputOfBatch(proofsInputOfBatch);
            verifierConfig.setPublicKeyInputOfBatch(publicKeyInputOfBatch);
            verifierConfig.setReEncryptedBallotsInputOfBatch(reEncryptedBallotsInputOfBatch);
            verifierConfig.setVotesWithProofInputOfBatch(votesWithProofInputOfBatch);
        }
        return verifierConfig;
    }
}
