/*
 * 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.ov.mixing.core.elgamal;

import static org.junit.Assert.assertEquals;
import static org.junit.Assert.assertTrue;

import java.math.BigInteger;
import java.security.Security;
import java.util.ArrayList;
import java.util.HashSet;
import java.util.List;
import java.util.Set;

import org.bouncycastle.jce.provider.BouncyCastleProvider;
import org.junit.BeforeClass;
import org.junit.Rule;
import org.junit.Test;
import org.junit.rules.ExpectedException;

import com.scytl.cryptolib.api.elgamal.ElGamalServiceAPI;
import com.scytl.cryptolib.api.exceptions.GeneralCryptoLibException;
import com.scytl.cryptolib.elgamal.bean.ElGamalComputationsValues;
import com.scytl.cryptolib.elgamal.bean.ElGamalEncryptionParameters;
import com.scytl.cryptolib.elgamal.bean.ElGamalKeyPair;
import com.scytl.cryptolib.elgamal.bean.ElGamalPrivateKey;
import com.scytl.cryptolib.elgamal.bean.ElGamalPublicKey;
import com.scytl.cryptolib.elgamal.cryptoapi.Ciphertext;
import com.scytl.cryptolib.elgamal.cryptoapi.CryptoAPIElGamalDecrypter;
import com.scytl.cryptolib.elgamal.cryptoapi.CryptoAPIElGamalEncrypter;
import com.scytl.cryptolib.elgamal.service.ElGamalService;
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.ov.mixing.commons.configuration.MatrixDimensions;
import com.scytl.ov.mixing.core.InvalidJobParametersException;
import com.scytl.ov.mixing.core.InvalidMixablesException;
import com.scytl.ov.mixing.core.JobState;
import com.scytl.ov.mixing.core.JobStateAndProgressListener;
import com.scytl.ov.mixing.core.Mixable;
import com.scytl.ov.mixing.core.MixingException;
import com.scytl.ov.mixing.core.ProofGenerationConfigParams;

public class ElGamalMixAndProveJobTest {

    public static BigInteger p = TestData.P_MEDIUM_SIZE;

    public static BigInteger q = TestData.Q_MEDIUM_SIZE;

    public static BigInteger g = TestData.G_MEDIUM_SIZE;

    public static ZpSubgroup group;

    public static ElGamalEncryptionParameters encryptionParameters;

    public static ElGamalServiceAPI elGamalService;

    public static ElGamalPublicKey publicKey;

    public static ElGamalPrivateKey privateKey;

    public static List<Mixable> mixables;

    public static List<List<ZpGroupElement>> plaintext;

    public static final int NUM_MIXABLES = 3;

    public static final int NUM_ELEMENTS_PER_MIXABLE = 3;

    public static ProofGenerationConfigParams proofGenerationConfigParams;

    @Rule
    public ExpectedException expectedEx = ExpectedException.none();

    @BeforeClass
    public static void setUp() throws GeneralCryptoLibException {

        Security.addProvider(new BouncyCastleProvider());

        group = new ZpSubgroup(g, p, q);

        int m = 1;
        int n = NUM_MIXABLES;
        MatrixDimensions matrixDimensions = new MatrixDimensions(m, n);
        int numIterations = 0;
        int mu = 2;
        proofGenerationConfigParams = new ProofGenerationConfigParams(matrixDimensions, numIterations, mu, group);

        encryptionParameters = new ElGamalEncryptionParameters(p, q, g);

        elGamalService = new ElGamalService();

        /////////////////////////////////////////////////
        //
        // Generate keys
        //
        /////////////////////////////////////////////////

        ElGamalKeyPair keyPair = elGamalService.getElGamalKeyPairGenerator()
            .generateKeys(new ElGamalEncryptionParameters(p, q, g), NUM_ELEMENTS_PER_MIXABLE);

        publicKey = keyPair.getPublicKeys();
        privateKey = keyPair.getPrivateKeys();

        /////////////////////////////////////////////////
        //
        // Generate plaintext data
        //
        /////////////////////////////////////////////////

        plaintext = new ArrayList<>();

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

            List<ZpGroupElement> list = new ArrayList<>();

            for (int j = 0; j < NUM_ELEMENTS_PER_MIXABLE; j++) {
                int semiRandomExponentAsInt = i + j + 10;
                Exponent semiRandomExponent = new Exponent(q, BigInteger.valueOf( + semiRandomExponentAsInt));
                list.add(group.getGenerator().exponentiate(semiRandomExponent));
            }

            plaintext.add(list);
        }

        /////////////////////////////////////////////////
        //
        // Encrypt the data
        //
        /////////////////////////////////////////////////

        mixables = new ArrayList<>();
        CryptoAPIElGamalEncrypter encrypter = elGamalService.createEncrypter(publicKey);

        for (List<ZpGroupElement> singlePlaintext : plaintext) {
            Ciphertext ciphertext = encrypter.encryptGroupElements(singlePlaintext);
            mixables.add(new ElGamalMixableImpl(ciphertext));
        }
    }

    //////////////////////////////////////
    //
    // Jobstate and progress of job tests
    //
    //////////////////////////////////////

    @Test
    public void whenCreatedThenStateAndProgressInitialized() {

        ElGamalMixAndProveJob elGamalMixAndProveJob = new ElGamalMixAndProveJob();

        JobState expectedInitialState = JobState.CREATED;
        int expectedInitialProgress = 0;

        assertEquals(expectedInitialState, elGamalMixAndProveJob.state());
        assertEquals(expectedInitialProgress, elGamalMixAndProveJob.progress());
    }

    @Test
    public void whenFinishedThenStateAndProgressOk() throws InvalidJobParametersException {

        ElGamalMixAndProveJob elGamalMixAndProveJob = new ElGamalMixAndProveJob();

        elGamalMixAndProveJob
            .setParameters(new ElGamalJobParameters(encryptionParameters, publicKey, proofGenerationConfigParams));

        elGamalMixAndProveJob.start(mixables);

        JobState expectedFinishedState = JobState.FINISHED;
        int expectedFinishedProgress = 100;
        assertEquals(expectedFinishedState, elGamalMixAndProveJob.state());
        assertEquals(expectedFinishedProgress, elGamalMixAndProveJob.progress());
    }

    @Test
    public void whenCanceledThenStateAndProgressOk() throws InvalidMixablesException, MixingException {

        ElGamalMixAndProveJob elGamalMixAndProveJob = new ElGamalMixAndProveJob();

        elGamalMixAndProveJob
            .setParameters(new ElGamalJobParameters(encryptionParameters, publicKey, proofGenerationConfigParams));

        elGamalMixAndProveJob.start(mixables);

        elGamalMixAndProveJob.cancel();

        assertEquals(0, elGamalMixAndProveJob.progress());
        assertEquals(JobState.CANCELED, elGamalMixAndProveJob.state());
    }

    //////////////////////////////////////
    //
    // Adding and removing listener tests
    //
    //////////////////////////////////////

    @Test
    public void whenListenerAddedAndRemovedThenOk() {

        ElGamalMixAndProveJob elGamalMixAndProveJob = new ElGamalMixAndProveJob();

        JobStateAndProgressListener listener = new JobStateAndProgressListener();
        elGamalMixAndProveJob.addEventListener(listener);
        elGamalMixAndProveJob.removeEventListener(listener);
    }

    @Test
    public void whenTryToRemoveUnaddedListenerThenNoAffect() {

        ElGamalMixAndProveJob elGamalMixAndProveJob = new ElGamalMixAndProveJob();

        JobStateAndProgressListener listener = new JobStateAndProgressListener();
        elGamalMixAndProveJob.removeEventListener(listener);
    }

    @Test
    public void initialAndFinalInfoFromProgressListenerOk() throws InvalidJobParametersException {

        ElGamalMixAndProveJob elGamalMixAndProveJob = new ElGamalMixAndProveJob();

        elGamalMixAndProveJob
            .setParameters(new ElGamalJobParameters(encryptionParameters, publicKey, proofGenerationConfigParams));

        JobStateAndProgressListener listener = new JobStateAndProgressListener();
        elGamalMixAndProveJob.addEventListener(listener);

        int expectedInitialProgress = 0;
        assertEquals(expectedInitialProgress, listener.progress());

        elGamalMixAndProveJob.start(mixables);

        JobState expectedFinishedState = JobState.FINISHED;
        int expectedFinishedProgress = 100;
        assertEquals(expectedFinishedState, listener.state());
        assertEquals(expectedFinishedProgress, listener.progress());
    }

    @Test
    public void ifListenerAddedAfterJobHasFinishedThenOk()
            throws InvalidJobParametersException, GeneralCryptoLibException {

        ElGamalMixAndProveJob elGamalMixAndProveJob = new ElGamalMixAndProveJob();

        elGamalMixAndProveJob
            .setParameters(new ElGamalJobParameters(encryptionParameters, publicKey, proofGenerationConfigParams));

        elGamalMixAndProveJob.start(mixables);

        JobStateAndProgressListener listener = new JobStateAndProgressListener();
        elGamalMixAndProveJob.addEventListener(listener);

        JobState expectedFinishedState = JobState.FINISHED;
        int expectedFinishedProgress = 100;
        assertEquals(expectedFinishedState, listener.state());
        assertEquals(expectedFinishedProgress, listener.progress());
    }

    //////////////////////////////////////
    //
    // Starting and JobState tests
    //
    //////////////////////////////////////

    @Test
    public void whenStartBeforeSettingParamtersThenException() throws InvalidMixablesException, MixingException {

        expectedEx.expect(IllegalStateException.class);
        expectedEx.expectMessage("Started the job without first setting the job parameters");

        ElGamalMixAndProveJob elGamalMixAndProveJob = new ElGamalMixAndProveJob();

        elGamalMixAndProveJob.start(mixables);
    }

    @Test
    public void whenStartCanceledJobThenException() throws InvalidMixablesException, MixingException {

        expectedEx.expect(IllegalStateException.class);
        expectedEx.expectMessage("Tried to start a job that was canceled");

        ElGamalMixAndProveJob elGamalMixAndProveJob = new ElGamalMixAndProveJob();

        elGamalMixAndProveJob
            .setParameters(new ElGamalJobParameters(encryptionParameters, publicKey, proofGenerationConfigParams));

        elGamalMixAndProveJob.cancel();

        elGamalMixAndProveJob.start(mixables);
    }

    @Test
    public void whenStartFinishedJobThenException() throws InvalidMixablesException, MixingException {

        expectedEx.expect(IllegalStateException.class);
        expectedEx.expectMessage("Tried to start a job that was finished");

        ElGamalMixAndProveJob elGamalMixAndProveJob = new ElGamalMixAndProveJob();

        elGamalMixAndProveJob
            .setParameters(new ElGamalJobParameters(encryptionParameters, publicKey, proofGenerationConfigParams));

        elGamalMixAndProveJob.start(mixables);

        elGamalMixAndProveJob.start(mixables);
    }


    //////////////////////////////////////
    //
    // Problems with input to start
    //
    //////////////////////////////////////

    @Test
    public void whenPassNullThenException()
            throws InvalidMixablesException, MixingException, GeneralCryptoLibException {

        expectedEx.expect(NullPointerException.class);
        expectedEx.expectMessage("The received list of mixables was null");

        ElGamalMixAndProveJob elGamalMixAndProveJob = new ElGamalMixAndProveJob();

        elGamalMixAndProveJob
            .setParameters(new ElGamalJobParameters(encryptionParameters, publicKey, proofGenerationConfigParams));

        elGamalMixAndProveJob.start(null);
    }

    @Test
    public void whenPassEmptyListThenException()
            throws InvalidMixablesException, MixingException, GeneralCryptoLibException {

        expectedEx.expect(IllegalArgumentException.class);
        expectedEx.expectMessage("The received list of mixables was empty");

        ElGamalMixAndProveJob elGamalMixAndProveJob = new ElGamalMixAndProveJob();

        elGamalMixAndProveJob
            .setParameters(new ElGamalJobParameters(encryptionParameters, publicKey, proofGenerationConfigParams));

        List<Mixable> mixables = new ArrayList<>();

        elGamalMixAndProveJob.start(mixables);
    }

    @Test
    public void givenNonElGamalMixableAsInputWhenGetMixedListThenException()
            throws InvalidMixablesException, MixingException, GeneralCryptoLibException {

        expectedEx.expect(InvalidMixablesException.class);
        expectedEx.expectMessage("The list of mixables contained something that is not an ElGamalMixable");

        ElGamalMixAndProveJob elGamalMixAndProveJob = new ElGamalMixAndProveJob();

        elGamalMixAndProveJob
            .setParameters(new ElGamalJobParameters(encryptionParameters, publicKey, proofGenerationConfigParams));

        List<Mixable> mixables = new ArrayList<>();
        mixables.add(new SomeOtherMixable());
        mixables.add(new SomeOtherMixable());
        mixables.add(new SomeOtherMixable());

        elGamalMixAndProveJob.start(mixables);

        elGamalMixAndProveJob.getMix(SomeOtherMixable.class);
    }

    @Test
    public void givenNonElGamalMixableAsInputWhenGetProofThenException()
            throws InvalidMixablesException, MixingException, GeneralCryptoLibException {

        expectedEx.expect(InvalidMixablesException.class);
        expectedEx.expectMessage("The list of mixables contained something that is not an ElGamalMixable");

        ElGamalMixAndProveJob elGamalMixAndProveJob = new ElGamalMixAndProveJob();

        elGamalMixAndProveJob
            .setParameters(new ElGamalJobParameters(encryptionParameters, publicKey, proofGenerationConfigParams));

        List<Mixable> mixables = new ArrayList<>();
        mixables.add(new SomeOtherMixable());
        mixables.add(new SomeOtherMixable());
        mixables.add(new SomeOtherMixable());

        elGamalMixAndProveJob.start(mixables);

        elGamalMixAndProveJob.getProof();
    }

    //////////////////////////////////////
    //
    // Getting results and JobState tests
    //
    //////////////////////////////////////

    @Test
    public void whenGetResultsBeforeStartingThenException() throws InvalidMixablesException, MixingException {

        expectedEx.expect(IllegalStateException.class);
        expectedEx.expectMessage("Tried to retrieve the mixed list before the job was started");

        ElGamalMixAndProveJob elGamalMixAndProveJob = new ElGamalMixAndProveJob();

        elGamalMixAndProveJob
            .setParameters(new ElGamalJobParameters(encryptionParameters, publicKey, proofGenerationConfigParams));

        elGamalMixAndProveJob.getMix(ElGamalMixableImpl.class);
    }

    @Test
    public void whenGetProofBeforeStartingThenException() throws InvalidMixablesException, MixingException {

        expectedEx.expect(IllegalStateException.class);
        expectedEx.expectMessage("Tried to retrieve the proof before the job was started");

        ElGamalMixAndProveJob elGamalMixAndProveJob = new ElGamalMixAndProveJob();

        elGamalMixAndProveJob
            .setParameters(new ElGamalJobParameters(encryptionParameters, publicKey, proofGenerationConfigParams));

        elGamalMixAndProveJob.getProof();
    }

    @Test
    public void whenGetResultsOfCanceledJobThenException() throws InvalidMixablesException, MixingException {

        expectedEx.expect(IllegalStateException.class);
        expectedEx.expectMessage("Tried to retrieve the mixed list from a job that was canceled");

        ElGamalMixAndProveJob elGamalMixAndProveJob = new ElGamalMixAndProveJob();

        elGamalMixAndProveJob
            .setParameters(new ElGamalJobParameters(encryptionParameters, publicKey, proofGenerationConfigParams));

        elGamalMixAndProveJob.start(mixables);

        elGamalMixAndProveJob.cancel();

        elGamalMixAndProveJob.getMix(ElGamalMixableImpl.class);
    }

    @Test
    public void whenGetProofOfCanceledJobThenException() throws InvalidMixablesException, MixingException {

        expectedEx.expect(IllegalStateException.class);
        expectedEx.expectMessage("Tried to retrieve the proof from a job that was canceled");

        ElGamalMixAndProveJob elGamalMixAndProveJob = new ElGamalMixAndProveJob();

        elGamalMixAndProveJob
            .setParameters(new ElGamalJobParameters(encryptionParameters, publicKey, proofGenerationConfigParams));

        elGamalMixAndProveJob.start(mixables);

        elGamalMixAndProveJob.cancel();

        elGamalMixAndProveJob.getProof();
    }

    //////////////////////////////////////////////////////
    //
    // Problems with class specified when retrieving data
    //
    ///////////////////////////////////////////////////////

    @Test
    public void whenNullPassedAsTypeWhenRetrievingThenException() throws InvalidMixablesException, MixingException {

        expectedEx.expect(NullPointerException.class);
        expectedEx.expectMessage("The received class was null");

        ElGamalMixAndProveJob elGamalMixAndProveJob = new ElGamalMixAndProveJob();

        elGamalMixAndProveJob
            .setParameters(new ElGamalJobParameters(encryptionParameters, publicKey, proofGenerationConfigParams));

        elGamalMixAndProveJob.start(mixables);

        elGamalMixAndProveJob.getMix(null);
    }

    @Test
    public void whenWrongTypeWhenRetrievingThenException_subclass_of_ElGamalMixableImpl()
            throws InvalidMixablesException, MixingException {

        expectedEx.expect(InvalidMixablesException.class);
        expectedEx.expectMessage("An element in the mixed list does not match the specified class type");

        ElGamalMixAndProveJob elGamalMixAndProveJob = new ElGamalMixAndProveJob();

        elGamalMixAndProveJob
            .setParameters(new ElGamalJobParameters(encryptionParameters, publicKey, proofGenerationConfigParams));

        elGamalMixAndProveJob.start(mixables);

        elGamalMixAndProveJob.getMix(ClientMixableForTesting.class);
    }

    @Test
    public void whenWrongTypeWhenRetrievingThenException_different_impl_of_Mixable()
            throws InvalidMixablesException, MixingException {

        expectedEx.expect(InvalidMixablesException.class);
        expectedEx.expectMessage("An element in the mixed list does not match the specified class type");

        ElGamalMixAndProveJob elGamalMixAndProveJob = new ElGamalMixAndProveJob();

        elGamalMixAndProveJob
            .setParameters(new ElGamalJobParameters(encryptionParameters, publicKey, proofGenerationConfigParams));

        elGamalMixAndProveJob.start(mixables);

        elGamalMixAndProveJob.getMix(SomeOtherMixable.class);
    }

    //////////////////////////////////////
    //
    // Happypath e2e job tests
    //
    //////////////////////////////////////

    @Test
    public void whenMixAndGetMixThenOk_specifyingInterfaceAsClassType()
            throws InvalidMixablesException, MixingException, GeneralCryptoLibException {

        ElGamalMixAndProveJob elGamalMixAndProveJob = new ElGamalMixAndProveJob();

        elGamalMixAndProveJob
            .setParameters(new ElGamalJobParameters(encryptionParameters, publicKey, proofGenerationConfigParams));

        elGamalMixAndProveJob.start(mixables);

        List<ElGamalMixable> mixed = elGamalMixAndProveJob.getMix(ElGamalMixable.class);

        List<List<ZpGroupElement>> decrypted = decryptListElGamalMixables(mixed);

        // printPlaintext(plaintext);
        // printPlaintext(decrypted);

        Set<List<ZpGroupElement>> plaintextAsSet = new HashSet<>(plaintext);
        Set<List<ZpGroupElement>> decryptedAsSet = new HashSet<>(decrypted);

        assertTrue(plaintextAsSet.equals(decryptedAsSet));
    }

    @Test
    public void whenMixAndGetMixThenOk_specifyingImplementationAsClassType()
            throws InvalidMixablesException, MixingException, GeneralCryptoLibException {

        ElGamalMixAndProveJob elGamalMixAndProveJob = new ElGamalMixAndProveJob();

        elGamalMixAndProveJob
            .setParameters(new ElGamalJobParameters(encryptionParameters, publicKey, proofGenerationConfigParams));

        elGamalMixAndProveJob.start(mixables);

        List<ElGamalMixableImpl> mixed = elGamalMixAndProveJob.getMix(ElGamalMixableImpl.class);

        List<List<ZpGroupElement>> decrypted = decryptListElGamalMixables(getAsListOfInterfaces(mixed));

        // printPlaintext(plaintext);
        // printPlaintext(decrypted);

        Set<List<ZpGroupElement>> plaintextAsSet = new HashSet<>(plaintext);
        Set<List<ZpGroupElement>> decryptedAsSet = new HashSet<>(decrypted);

        assertTrue(plaintextAsSet.equals(decryptedAsSet));
    }

    @Test
    public void whenMixListOfOneMixableThenOk()
            throws InvalidMixablesException, MixingException, GeneralCryptoLibException {

        ElGamalMixAndProveJob elGamalMixAndProveJob = new ElGamalMixAndProveJob();

        int m = 1;
        int n = 1;
        MatrixDimensions matrixDimensions = new MatrixDimensions(m, n);
        int numIterations = 0;
        int mu = 2;
        ProofGenerationConfigParams proofGenerationConfigParams =
            new ProofGenerationConfigParams(matrixDimensions, numIterations, mu, group);

        elGamalMixAndProveJob
            .setParameters(new ElGamalJobParameters(encryptionParameters, publicKey, proofGenerationConfigParams));

        List<Mixable> listMixablesJustOneElement = mixables.subList(0, 1);

        elGamalMixAndProveJob.start(listMixablesJustOneElement);

        List<ElGamalMixableImpl> mixed = elGamalMixAndProveJob.getMix(ElGamalMixableImpl.class);

        List<ZpGroupElement> decrypted = decryptListElGamalMixables(getAsListOfInterfaces(mixed)).get(0);

        assertEquals(1, mixed.size());
        assertEquals(plaintext.get(0), decrypted);
    }

    private List<ElGamalMixable> getAsListOfInterfaces(List<ElGamalMixableImpl> listImplemtations)
            throws GeneralCryptoLibException {

        return new ArrayList<ElGamalMixable>(listImplemtations);
    }

    private List<List<ZpGroupElement>> decryptListElGamalMixables(List<ElGamalMixable> mixed)
            throws GeneralCryptoLibException {

        List<Ciphertext> mixedCiphertext = new ArrayList<>();

        for (Mixable mixable : mixed) {
            ElGamalMixable elGamalMixable = (ElGamalMixable) mixable;
            mixedCiphertext.add(elGamalMixable.getCiphertext());
        }

        return decryptListCiphertexts(mixedCiphertext);
    }

    private List<List<ZpGroupElement>> decryptListCiphertexts(List<Ciphertext> reencrypted)
            throws GeneralCryptoLibException {

        CryptoAPIElGamalDecrypter decrypter = elGamalService.createDecrypter(privateKey);

        List<List<ZpGroupElement>> allDecrypted = new ArrayList<>();

        for (Ciphertext singleCiphertext : reencrypted) {

            ElGamalComputationsValues singleCiphertextAsElGamalComputationsValues =
                new ElGamalComputationsValues(singleCiphertext.getElements());

            List<ZpGroupElement> singleDecrypted =
                decrypter.decrypt(singleCiphertextAsElGamalComputationsValues, false);
            allDecrypted.add(singleDecrypted);
        }

        return allDecrypted;
    }

    @SuppressWarnings("unused")
    private void printPlaintext(List<List<ZpGroupElement>> plaintext) {

        System.out.println("-------- Plaintext ----------");

        for (int i = 0; i < plaintext.size(); i++) {

            for (int j = 0; j < plaintext.get(i).size(); j++) {

                System.out.print(plaintext.get(i).get(j).getValue());
                System.out.print(" ");
            }

            System.out.println();
        }
    }

    class SomeOtherMixable implements Mixable {
    }

}
