/*
 * 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;

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.List;

import org.apache.commons.io.IOUtils;
import org.junit.Assert;
import org.junit.BeforeClass;
import org.junit.Test;

import com.scytl.cryptolib.api.exceptions.CryptoLibException;
import com.scytl.cryptolib.api.exceptions.GeneralCryptoLibException;
import com.scytl.cryptolib.elgamal.bean.ElGamalPublicKey;
import com.scytl.cryptolib.elgamal.cryptoapi.Ciphertext;
import com.scytl.cryptolib.mathematical.groups.impl.ZpGroupElement;
import com.scytl.cryptolib.mathematical.groups.impl.ZpSubgroup;
import com.scytl.cryptolib.primitives.service.PrimitivesService;
import com.scytl.products.ov.mixnet.commons.ballots.ElGamalEncryptedBallots;
import com.scytl.products.ov.mixnet.commons.beans.ShuffleOutput;
import com.scytl.products.ov.mixnet.commons.beans.proofs.ShuffleProof;
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.locations.LocationsProvider;
import com.scytl.products.ov.mixnet.commons.homomorphic.impl.GjosteenElGamal;
import com.scytl.products.ov.mixnet.commons.io.ElGamalEncryptedBallotsLoader;
import com.scytl.products.ov.mixnet.commons.io.ElGamalEncryptedBallotsWriter;
import com.scytl.products.ov.mixnet.commons.io.ElgamalPublicKeyWriter;
import com.scytl.products.ov.mixnet.commons.io.JSONProofsReader;
import com.scytl.products.ov.mixnet.commons.io.JSONProofsWriter;
import com.scytl.products.ov.mixnet.commons.proofs.bg.commitments.CommitmentParams;
import com.scytl.products.ov.mixnet.commons.proofs.bg.config.ProofsGeneratorConfigurationParams;
import com.scytl.products.ov.mixnet.commons.tools.CiphertextTools;
import com.scytl.products.ov.mixnet.commons.tools.ExponentTools;
import com.scytl.products.ov.mixnet.commons.tools.MatrixArranger;
import com.scytl.products.ov.mixnet.commons.tools.MultiExponentiation;
import com.scytl.products.ov.mixnet.commons.tools.MultiExponentiationImpl;
import com.scytl.products.ov.mixnet.commons.validation.EncryptedBallotsDuplicationValidator;
import com.scytl.products.ov.mixnet.commons.validation.EncryptedBallotsValidator;
import com.scytl.products.ov.mixnet.commons.validation.EncryptedBallotsValidatorManager;
import com.scytl.products.ov.mixnet.proofs.bg.multiexp.ComputeAllE;
import com.scytl.products.ov.mixnet.proofs.bg.multiexp.ParallelComputeAllE;
import com.scytl.products.ov.mixnet.proofs.bg.shuffle.ParallelPrivateAndPublicCommitmentsGenerator;
import com.scytl.products.ov.mixnet.proofs.bg.shuffle.PrivateAndPublicCommitmentsGenerator;
import com.scytl.products.ov.mixnet.proofs.bg.shuffle.ShuffleProofGenerator;
import com.scytl.products.ov.mixnet.proofs.bg.shuffle.ShuffleProofVerifier;
import com.scytl.products.ov.mixnet.shuffler.ElGamalShuffler;
import com.scytl.products.ov.mixnet.shuffler.permutation.PermutationGenerator;
import com.scytl.products.ov.mixnet.shuffler.permutation.Permutator;

/**
 * This abstract test class is the base of other tests which perform ALL of the steps of mixing. It includes concrete
 * methods (implemented here) and abstract methods. The abstract methods are implemented in the test classes which
 * extend this class. This test, and the tests which extend this class, perform the following steps:
 * <P>
 * Prover steps:
 * <ul>
 * <li>Load encrypted ballots</li>
 * <li>Perform a shuffle</li>
 * <li>Generate proofs</li>
 * <li>Store the shuffled encrypted ballots</li>
 * <li>Store the proofs</li>
 * <li>Store the commitment params</li>
 * <li>Store the public key</li>
 * </ul>
 * <P>
 * Verifier steps:
 * <ul>
 * <li>Load original encrypted ballots</li>
 * <li>Load the shuffled and encrypted ballots</li>
 * <li>Verify the proofs</li>
 * </ul>
 */
public class BaseBGMixnetITest extends BaseBGMixnet {

    private static final String INPUT_FOLDER = "src/test/resources/input/n512_p256_q32";

    private static ElGamalPublicKey pubKey;

    private static ZpSubgroup zp;

    private static ElGamalShuffler shuffler;

    private static int m;

    private static int n;

    private static int mu;

    private static int numiterations;

    private static CommitmentParams commitmentParams;

    private static GjosteenElGamal elgamal;

    private static ProofsGeneratorConfigurationParams config;

    private static LocationsProvider locations;

    private static JSONProofsReader proofsReader;

    private static JSONProofsWriter proofsWriter;

    private static ElGamalEncryptedBallotsLoader elgamalEncryptedBallotsLoader;

    private static MultiExponentiation limMultiExpo;

    private static CiphertextTools ciphertextTools;

    private static ComputeAllE computeAllE;

    private static int concurrencyLevel;

    @BeforeClass
    public static void setUp() throws IOException {

        try {
            final ApplicationConfig mixnetConfig = new ApplicationConfig();
            final MixerConfig mixerConfig = new MixerConfig();
            mixerConfig.setEncryptionParametersInput(
                BaseBGMixnetITest.class.getClassLoader().getResourceAsStream("input/encryption_parameters.json"));
            mixerConfig.setEncryptedBallotsInput(
                new FileInputStream(Paths.get(INPUT_FOLDER + "/encrypted_ballots" + ".csv").toFile()));
            mixerConfig
                .setPublicKeyInput(new FileInputStream(Paths.get(INPUT_FOLDER + "/public_key" + ".json").toFile()));
            mixerConfig.setBallotDecryptionKeyInput(
                new FileInputStream(Paths.get(INPUT_FOLDER + "/private_key" + ".json").toFile()));
            addOutputMixerConfig(mixerConfig);
            mixnetConfig.setMixer(mixerConfig);
            final VerifierConfig verifierConfig = new VerifierConfig();
            mixnetConfig.setVerifier(verifierConfig);

            locations = LocationsProvider.build(mixnetConfig);

            InputStream resourceAsStream =
                BaseBGMixnetITest.class.getClassLoader().getResourceAsStream("input/encryption_parameters.json");
            zp = ZpSubgroup.fromJson(IOUtils.toString(resourceAsStream, StandardCharsets.UTF_8));

            m = 8;

            n = 64;

            mu = 2;

            numiterations = 0;

            concurrencyLevel = Runtime.getRuntime().availableProcessors();

            // Create ElGamal public keys (we dont need the private key)
            final List<ZpGroupElement> publicKeyGroupElements = new ArrayList<>();

            publicKeyGroupElements.add( //
                zp.getGenerator().exponentiate( //
                    ExponentTools.getRandomExponent(zp.getQ())));

            pubKey = new ElGamalPublicKey(publicKeyGroupElements, zp);

            commitmentParams = new CommitmentParams(zp, n);

            elgamal = new GjosteenElGamal(zp, pubKey);

            config = new ProofsGeneratorConfigurationParams(m, n, elgamal, numiterations, mu, zp);

            shuffler = new ElGamalShuffler(new Permutator(), new PermutationGenerator());

            proofsReader = new JSONProofsReader();

            proofsWriter = new JSONProofsWriter();

            final List<EncryptedBallotsValidator> listValidators = new ArrayList<>();
            listValidators.add(new EncryptedBallotsDuplicationValidator());
            EncryptedBallotsValidatorManager validator = new EncryptedBallotsValidatorManager(listValidators);
            elgamalEncryptedBallotsLoader = new ElGamalEncryptedBallotsLoader(validator);

            limMultiExpo = MultiExponentiationImpl.getInstance();

            ciphertextTools = new CiphertextTools(limMultiExpo);

            computeAllE = new ParallelComputeAllE(ciphertextTools);

        } catch (GeneralCryptoLibException ex) {
            throw new CryptoLibException(ex);
        }
    }

    private static void addOutputMixerConfig(MixerConfig mixerConfig) throws FileNotFoundException {
        int numBatches = 1;
        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);
        Path path = Paths.get("target/basebgMixnetIt/outputs");
        path.toFile().mkdirs();
        commitmentParametersOutputOfBatch.add(new FileOutputStream(path.resolve("commitmentParameters.json").toFile()));
        decryptedVotesOutputOfBatch.add(new FileOutputStream(path.resolve("decryptedVotes.csv").toFile()));
        encryptedBallotsOutputOfBatch.add(new FileOutputStream(path.resolve("encryptedBallots.csv").toFile()));
        encryptionParametersOutputOfBatch.add(new FileOutputStream(path.resolve("encryptionParameters.json").toFile()));
        proofsOutputOfBatch.add(new FileOutputStream(path.resolve("proofs.json").toFile()));
        publicKeyOutputOfBatch.add(new FileOutputStream(path.resolve("publicKey.json").toFile()));
        reEncryptedBallotsOutputOfBatch.add(new FileOutputStream(path.resolve("reencryptedBallots.csv").toFile()));
        votesWithProofOutputOfBatch.add(new FileOutputStream(path.resolve("votesWithProof.csv").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);

    }

    private static void addInputVerifierConfig(VerifierConfig verifierConfig) throws FileNotFoundException {
        List<InputStream> proofsInputOfBatch = new ArrayList<>(1);
        List<InputStream> reEncryptedBallotsInputOfBatch = new ArrayList<>(1);
        Path path = Paths.get("target/basebgMixnetIt/outputs");
        path.toFile().mkdirs();
        proofsInputOfBatch.add(new FileInputStream(path.resolve("proofs.json").toFile()));
        reEncryptedBallotsInputOfBatch.add(new FileInputStream(path.resolve("reencryptedBallots.csv").toFile()));

        verifierConfig.setProofsInputOfBatch(proofsInputOfBatch);
        verifierConfig.setReEncryptedBallotsInputOfBatch(reEncryptedBallotsInputOfBatch);

    }

    @Test
    public void givenSmallConfigWhenShuffleThenOK() throws Exception {

        // /////////////////////////////////////////////////
        //
        // Load encrypted ballots
        //
        // /////////////////////////////////////////////////

        final ElGamalEncryptedBallots encryptedBallots =
            elgamalEncryptedBallotsLoader.loadCSV(zp, locations.getMixerLocationsProvider().getEncryptedBallotsInput());

        // /////////////////////////////////////////////////
        //
        // Perform Shuffle
        //
        // /////////////////////////////////////////////////

        final ShuffleOutput<ElGamalEncryptedBallots> shuffleOutput =
            shuffler.sequentialShuffle(zp, pubKey, concurrencyLevel, encryptedBallots, null, null);

        // /////////////////////////////////////////////////
        //
        // Generate proofs
        //
        // /////////////////////////////////////////////////

        final Ciphertext[][] encryptedCiphertexts = MatrixArranger.arrangeInCiphertextMatrix(encryptedBallots, m, n);

        final Ciphertext[][] reencryptedCiphertexts =
            MatrixArranger.arrangeInCiphertextMatrix(shuffleOutput.getReEncryptedBallots(), m, n);

        final PrivateAndPublicCommitmentsGenerator privateAndPublicCommitmentsGenerator =
            createPrivateAndPublicCommitmentsGenerator();

        final ShuffleProofGenerator prover = new ShuffleProofGenerator(config, commitmentParams, encryptedCiphertexts,
            shuffleOutput.getPermutation(), privateAndPublicCommitmentsGenerator, limMultiExpo, concurrencyLevel,
            reencryptedCiphertexts, shuffleOutput.getExponents(), computeAllE);

        final ShuffleProof shuffleProof = prover.generate();

        // /////////////////////////////////////////////////
        //
        // Store shuffled ballots
        //
        // /////////////////////////////////////////////////

        ElGamalEncryptedBallotsWriter.CSV.write(
            locations.getMixerLocationsProvider().getReEncryptedBallotsOutputOfBatch(1),
            shuffleOutput.getReEncryptedBallots(), new PrimitivesService().getRawMessageDigest());

        // /////////////////////////////////////////////////
        //
        // Store Proofs
        //
        // /////////////////////////////////////////////////

        proofsWriter.write(locations.getMixerLocationsProvider().getProofsOutputOfBatch(1), shuffleProof);

        // /////////////////////////////////////////////////
        //
        // Store the commitment parameters
        //
        // /////////////////////////////////////////////////

        commitmentParams
            .serializeToStream(locations.getMixerLocationsProvider().getCommitmentParametersOutputOfBatch(1));

        // /////////////////////////////////////////////////
        //
        // Store the public key
        //
        // /////////////////////////////////////////////////

        ElgamalPublicKeyWriter.serializeToStream(pubKey,
            locations.getMixerLocationsProvider().getPublicKeyOutputOfBatch(1));

        // /////////////////////////////////////////////////
        //
        // Verify the proofs
        //
        // /////////////////////////////////////////////////

        // NOTE: in order to fully test the entire process (all the steps of the prover and verifier), the verifier
        // should load the public key and commitment parameters that the prover serialized to file. We dont do that in
        // this test (we use the copies that exist in memory) because due to I/O slowness, the files might not have been
        // fully written when the test reaches this point, and therefore the test would fail. The functionality of
        // loading the public key and the commitment parameters is tested in the test which performs only the verifier
        // steps.

        final ApplicationConfig mixnetConfig = new ApplicationConfig();
        final MixerConfig mixerConfig = new MixerConfig();
        mixerConfig.setEncryptionParametersInput(
            new FileInputStream(Paths.get("src/test/resources/input/encryption_parameters.json").toFile()));
        mixerConfig.setEncryptedBallotsInput(
            new FileInputStream(Paths.get(INPUT_FOLDER + "/encrypted_ballots" + ".csv").toFile()));
        mixerConfig.setPublicKeyInput(new FileInputStream(Paths.get(INPUT_FOLDER + "/public_key" + ".json").toFile()));
        mixerConfig.setBallotDecryptionKeyInput(
            new FileInputStream(Paths.get(INPUT_FOLDER + "/private_key" + ".json").toFile()));
        mixnetConfig.setMixer(mixerConfig);
        final VerifierConfig verifierConfig = new VerifierConfig();
        mixnetConfig.setVerifier(verifierConfig);
        addInputVerifierConfig(verifierConfig);

        locations = LocationsProvider.build(mixnetConfig);

        final ElGamalEncryptedBallots encryptedBallots_VerifierCopy =
            elgamalEncryptedBallotsLoader.loadCSV(zp, locations.getMixerLocationsProvider().getEncryptedBallotsInput());

        final Ciphertext[][] originalCiphertexts_VerifierCopy =
            MatrixArranger.arrangeInCiphertextMatrix(encryptedBallots_VerifierCopy, m, n);

        final ElGamalEncryptedBallots reencryptedBallots = elgamalEncryptedBallotsLoader.loadCSV(zp,
            locations.getVerifierLocationsProvider().getReEncryptedBallotsInputOfBatch(1));

        final Ciphertext[][] reencryptedCiphertexts_VerifierCopy =
            MatrixArranger.arrangeInCiphertextMatrix(reencryptedBallots, m, n);

        final ShuffleProofVerifier verifier = new ShuffleProofVerifier(zp, elgamal, commitmentParams,
            originalCiphertexts_VerifierCopy, reencryptedCiphertexts_VerifierCopy, m, n, mu, numiterations,
            ciphertextTools, limMultiExpo);

        final ShuffleProof shuffleProof_VerifierCopy =
            proofsReader.read(locations.getVerifierLocationsProvider().getProofsInputOfBatch(1));

        final boolean testResult = verifier.verifyProof(shuffleProof_VerifierCopy.getInitialMessage(),
            shuffleProof_VerifierCopy.getFirstAnswer(), shuffleProof_VerifierCopy.getSecondAnswer());

        final String errorMsg = "proofs failed to validate";
        Assert.assertTrue(errorMsg, testResult);
    }

    protected PrivateAndPublicCommitmentsGenerator createPrivateAndPublicCommitmentsGenerator() {
        return new ParallelPrivateAndPublicCommitmentsGenerator();
    }
}
