/*
 * 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.shuffler.reencryption;

import static org.junit.Assert.assertEquals;

import java.math.BigInteger;
import java.util.ArrayList;
import java.util.List;

import org.junit.Before;
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.CiphertextImpl;
import com.scytl.cryptolib.elgamal.bean.ElGamalPublicKey;
import com.scytl.cryptolib.elgamal.cryptoapi.Ciphertext;
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.mixnet.commons.ballots.ElGamalEncryptedBallots;
import com.scytl.products.ov.mixnet.commons.homomorphic.Randomness;
import com.scytl.products.ov.mixnet.commons.homomorphic.impl.GjosteenElGamal;
import com.scytl.products.ov.mixnet.commons.homomorphic.impl.GjosteenElGamalPlaintext;
import com.scytl.products.ov.mixnet.commons.homomorphic.impl.GjosteenElGamalRandomness;
import com.scytl.products.ov.mixnet.commons.homomorphic.impl.RandomnessGenerator;
import com.scytl.products.ov.mixnet.commons.tools.ExponentTools;
import com.scytl.products.ov.mixnet.commons.tools.GroupTools;

public class ElGamalReEncrypterTest {

    private static final int ONE_PROCESSOR = 1;

    private static int numElements;

    private static ElGamalReEncrypter target;

    private static ElGamalPublicKey publicKey;

    private static Randomness[] rho;

    private static ElGamalEncryptedBallots elGamalEncryptedBallots;

    private static Exponent privateKeyExponent;

    private static ZpSubgroup group;

    private static RandomnessGenerator randomExponentCollectionGenerator;

    private static GjosteenElGamal gjosteenElGamal;

    @BeforeClass
    public static void setUp() {

        try {

            // create group
            final BigInteger p = new BigInteger("15294034768093677312256663166625633354362303");
            final BigInteger q = new BigInteger("7647017384046838656128331583312816677181151");

            group = GroupTools.createGroupWithRandomGenerator(p, q);

            // create public key
            privateKeyExponent = ExponentTools.getRandomExponent(q);
            Exponent[] privateKeyAsExponentArray = new Exponent[1];
            privateKeyAsExponentArray[0] = privateKeyExponent;
            final List<ZpGroupElement> keys = new ArrayList<>();
            keys.add(group.getGenerator().exponentiate(privateKeyExponent));
            publicKey = new ElGamalPublicKey(keys, group);

            numElements = 5;

            rho = new RandomnessGenerator(group.getQ()) //
                .generate(numElements);

            final List<Ciphertext> encryptedBallots = new ArrayList<>();

            // The input preparation should be handled by the PreProcessor
            final String[] gammaString = new String[] {"7606155287466709247547562508296526348227245",
                    "4012556967239296284125919604329103854893595", "11583975953794885768215634543368899452667420",
                    "7155423998998370092051586125176128718489417", "7830708621782553206289139938032543522538535" };

            final ZpGroupElement[] gammas = new ZpGroupElement[numElements];
            for (int i = 0; i < gammaString.length; i++) {
                BigInteger value = new BigInteger(gammaString[i]);
                gammas[i] = new ZpGroupElement(value, p, q);
            }

            final String[] phisString = new String[] {"9894101363857510707841445959199640374054748",
                    "6054959594806155309765068576952771674747761", "2114609311581369961287633946894396492461861",
                    "4070746468346708777122357082197449848916890", "795602734299768285623543899693683842706894" };

            final ZpGroupElement[] phis = new ZpGroupElement[numElements];
            for (int i = 0; i < phisString.length; i++) {
                BigInteger value = new BigInteger(phisString[i]);
                phis[i] = new ZpGroupElement(value, p, q);
            }

            for (int i = 0; i < numElements; i++) {

                final List<ZpGroupElement> ballotPhis = new ArrayList<>();
                ballotPhis.add(phis[i]);
                encryptedBallots.add(new CiphertextImpl(gammas[i], ballotPhis));
            }

            elGamalEncryptedBallots = new ElGamalEncryptedBallots(encryptedBallots);

            randomExponentCollectionGenerator = new RandomnessGenerator(group.getQ());

            gjosteenElGamal = new GjosteenElGamal(group, privateKeyAsExponentArray);

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

    @Before
    public void setUpForEachTest() {
        target = new ElGamalReEncrypter(group, publicKey, ONE_PROCESSOR);
    }

    @Test
    public void givenElGamalEncryptedBallotsWhenReencryptThenExpectedNumberOfBallots() {

        ElGamalEncryptedBallots reencryptedElGamalEncryptedBallots =
            target.reEncrypt(elGamalEncryptedBallots, rho, null);

        int numEncryptedBallots = reencryptedElGamalEncryptedBallots.getBallots().size();

        String errorMsg = "The number of reencrypted ballots was not what was expected";
        assertEquals(errorMsg, numElements, numEncryptedBallots);

        final int availableProcessors = Runtime.getRuntime().availableProcessors();

        System.out.println("Testing concurrency up to " + availableProcessors + " CPUs");

        for (int cpus = 2; cpus < availableProcessors; cpus++) {

            target = new ElGamalReEncrypter(group, publicKey, cpus);

            reencryptedElGamalEncryptedBallots = target.reEncrypt(elGamalEncryptedBallots, rho, null);

            numEncryptedBallots = reencryptedElGamalEncryptedBallots.getBallots().size();

            errorMsg = "The number of reencrypted ballots was not what was expected";
            assertEquals(errorMsg, numElements, numEncryptedBallots);

        }

    }

    /**
     * This test is an end-to-end test of the reencryption process (using simple parameters). It contains the following
     * steps:
     * <ul>
     * <li>Create a ballot - which serves as the plaintext.</li>
     * <li>Encrypt the ballot.</li>
     * <li>ReEncrypt the ballot.</li>
     * <li>Decrypt the ballot.</li>
     * <li>Check if the decrypted ballot matches the original ballot.</li>
     * </ul>
     */
    @Test
    public void givenPlaintextAndReencryptedBallotWhenDecryptThenEqualsPlaintext() {
        try {

            final ZpGroupElement plaintextAsElement = new ZpGroupElement(
                new BigInteger("7606155287466709247547562508296526348227245"), group.getP(), group.getQ());

            final GjosteenElGamalPlaintext plaintextAsGjosteenElGamalPlaintext =
                new GjosteenElGamalPlaintext(new ZpGroupElement[] {plaintextAsElement });

            final GjosteenElGamalRandomness randomness =
                new GjosteenElGamalRandomness(ExponentTools.getRandomExponent(group.getQ()));

            final Ciphertext encryptedBallot = gjosteenElGamal.encrypt(plaintextAsGjosteenElGamalPlaintext, randomness);

            final List<Ciphertext> listElGamalEncryptedBallot = new ArrayList<>();
            listElGamalEncryptedBallot.add(new CiphertextImpl(encryptedBallot.getGamma(), encryptedBallot.getPhis()));

            final Randomness[] randomExponentToBeUsedInEncryption = randomExponentCollectionGenerator.generate(1);

            ElGamalEncryptedBallots reencryptedElGamalEncryptedBallots = target.reEncrypt(
                new ElGamalEncryptedBallots(listElGamalEncryptedBallot), randomExponentToBeUsedInEncryption, null);

            List<Ciphertext> reencryptedBallotList = reencryptedElGamalEncryptedBallots.getBallots();

            ZpGroupElement[] reencryptedBallotListParts = new ZpGroupElement[2];
            reencryptedBallotListParts[0] = reencryptedBallotList.get(0).getGamma();
            reencryptedBallotListParts[1] = reencryptedBallotList.get(0).getPhis().get(0);

            ZpGroupElement decrypted = gjosteenElGamal.decrypt(reencryptedBallotListParts)[0];

            String errorMsg = "The decrypted ballot did not match the plaintext of the ballot";
            assertEquals(errorMsg, plaintextAsElement, decrypted);

            final int availableProcessors = Runtime.getRuntime().availableProcessors();

            System.out.println("Testing concurrency up to " + availableProcessors + " CPUs");

            for (int cpus = 2; cpus < availableProcessors; cpus++) {

                target = new ElGamalReEncrypter(group, publicKey, cpus);

                reencryptedElGamalEncryptedBallots = target.reEncrypt(
                    new ElGamalEncryptedBallots(listElGamalEncryptedBallot), randomExponentToBeUsedInEncryption, null);

                reencryptedBallotList = reencryptedElGamalEncryptedBallots.getBallots();

                reencryptedBallotListParts = new ZpGroupElement[2];
                reencryptedBallotListParts[0] = reencryptedBallotList.get(0).getGamma();
                reencryptedBallotListParts[1] = reencryptedBallotList.get(0).getPhis().get(0);

                decrypted = gjosteenElGamal.decrypt(reencryptedBallotListParts)[0];

                errorMsg = "The decrypted ballot did not match the plaintext of the ballot";
                assertEquals(errorMsg, plaintextAsElement, decrypted);

            }

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