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

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

import org.bouncycastle.jce.provider.BouncyCastleProvider;
import org.junit.BeforeClass;
import org.junit.Test;

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.ElGamalEncrypterValues;
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.cryptoapi.CryptoAPIElGamalKeyPairGenerator;
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;

/**
 * This class contains a number of integrations tests that perform the following steps:
 * <ol>
 * <li>Generate test data</li>
 * <li>Encrypt</li>
 * <li>Recrypt (a number of times)</li>
 * <li>Decrypt</li>
 * <li>Confirm that the recovered data matches the original data</li>
 * </ol>
 */
public class ElGamalReencrypterITest {

    public static ElGamalServiceAPI elGamalService;

    public static BigInteger p = TestData.P_MEDIUM_SIZE;

    public static BigInteger q = TestData.Q_MEDIUM_SIZE;

    public static BigInteger g = TestData.G_MEDIUM_SIZE;

    @BeforeClass
    public static void setUp() throws GeneralCryptoLibException {

        Security.addProvider(new BouncyCastleProvider());

        elGamalService = new ElGamalService();
    }

    @Test
    public void encryptReencryptAndDecryptSingleCiphertext_elements1_reencrypt1() throws GeneralCryptoLibException {

        int numElementsPerCiphertext = 1;
        int numTimesToPerformReencrypt = 1;
        encryptReencryptAndDecrypt(numElementsPerCiphertext, numTimesToPerformReencrypt);
    }

    @Test
    public void encryptReencryptAndDecryptSingleCiphertext_elements3_reencrypt3() throws GeneralCryptoLibException {

        int numElementsPerCiphertext = 3;
        int numTimesToPerformReencrypt = 3;
        encryptReencryptAndDecrypt(numElementsPerCiphertext, numTimesToPerformReencrypt);
    }

    @Test
    public void encryptReencryptAndDecryptListCiphertexts_ciphertexts1_elements1_reencrypt1()
            throws GeneralCryptoLibException {

        int numCiphertexts = 1;
        int numElementsPerCiphertext = 1;
        int numTimesToPerformReencrypt = 1;
        encryptReencryptAndDecryptList(numCiphertexts, numElementsPerCiphertext, numTimesToPerformReencrypt);
    }

    @Test
    public void encryptReencryptAndDecryptListCiphertexts_ciphertexts10_elements3_reencrypt1()
            throws GeneralCryptoLibException {

        int numCiphertexts = 10;
        int numElementsPerCiphertext = 3;
        int numTimesToPerformReencrypt = 1;
        encryptReencryptAndDecryptList(numCiphertexts, numElementsPerCiphertext, numTimesToPerformReencrypt);
    }

    @Test
    public void encryptReencryptAndDecryptListCiphertexts_ciphertexts10_elements3_reencrypt3()
            throws GeneralCryptoLibException {

        int numCiphertexts = 10;
        int numElementsPerCiphertext = 3;
        int numTimesToPerformReencrypt = 3;
        encryptReencryptAndDecryptList(numCiphertexts, numElementsPerCiphertext, numTimesToPerformReencrypt);
    }

    public void encryptReencryptAndDecrypt(int numElementsPerCiphertext, int numTimesToPerformReencrypt)
            throws GeneralCryptoLibException {

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

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

        CryptoAPIElGamalKeyPairGenerator keyPairGenerator = elGamalService.getElGamalKeyPairGenerator();

        ElGamalKeyPair keyPair = keyPairGenerator.generateKeys(elGamalEncryptionParameters, numElementsPerCiphertext);

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

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

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

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

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

        // printPlaintext(allPlaintexts);

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

        CryptoAPIElGamalEncrypter encrypter = elGamalService.createEncrypter(keyPair.getPublicKeys());

        Ciphertext singleCiphertext = encrypter.encryptGroupElements(plaintext);

        // printCiphertext(allCiphertexts);

        /////////////////////////////////////////////////
        //
        // Reencrypt the data
        //
        /////////////////////////////////////////////////

        ElGamalReencrypter reencrypter = new ElGamalReencrypter(publicKey);

        Ciphertext reencrypted;

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

            reencrypted = reencrypter.reencrypt(singleCiphertext).getCiphertext();
            singleCiphertext = reencrypted;

            // printCiphertext(reencrypted);
        }

        /////////////////////////////////////////////////
        //
        // Decrypt the data
        //
        /////////////////////////////////////////////////

        CryptoAPIElGamalDecrypter decrypter = elGamalService.createDecrypter(privateKey);

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

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

        // printPlaintext(allDecrypted);

        assertTrue(plaintext.equals(singleDecrypted));
    }

    public void encryptReencryptAndDecryptList(int numCiphertexts, int numElementsPerCiphertext,
            int numTimesToPerformReencrypt) throws GeneralCryptoLibException {

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

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

        CryptoAPIElGamalKeyPairGenerator keyPairGenerator = elGamalService.getElGamalKeyPairGenerator();

        ElGamalKeyPair keyPair = keyPairGenerator.generateKeys(elGamalEncryptionParameters, numElementsPerCiphertext);

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

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

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

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

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

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

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

            allPlaintexts.add(singlePlaintext);
        }

        // printPlaintext(allPlaintexts);

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

        CryptoAPIElGamalEncrypter encrypter = elGamalService.createEncrypter(keyPair.getPublicKeys());

        List<Ciphertext> allCiphertexts = new ArrayList<>();
        for (List<ZpGroupElement> singlePlaintext : allPlaintexts) {

            ElGamalEncrypterValues singleCiphertext = encrypter.encryptGroupElements(singlePlaintext);
            allCiphertexts.add(singleCiphertext);
        }

        // printCiphertext(allCiphertexts);

        /////////////////////////////////////////////////
        //
        // Reencrypt the data
        //
        /////////////////////////////////////////////////

        ElGamalReencrypter reencrypter = new ElGamalReencrypter(publicKey);

        List<ReencryptionOutput> reencrypted = null;
        for (int i = 0; i < numTimesToPerformReencrypt; i++) {

            reencrypted = reencrypter.reencrypt(allCiphertexts);
            allCiphertexts = getListOfAllCiphertexts(reencrypted);

            // printCiphertext(reencrypted);
        }

        /////////////////////////////////////////////////
        //
        // Decrypt the data
        //
        /////////////////////////////////////////////////

        CryptoAPIElGamalDecrypter decrypter = elGamalService.createDecrypter(privateKey);

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

        for (int i = 0; i < reencrypted.size(); i++) {
            Ciphertext singleCiphertext = reencrypted.get(i).getCiphertext();

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

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

        // printPlaintext(allDecrypted);

        assertTrue(allDecrypted.size() == numCiphertexts);
        assertTrue(allPlaintexts.equals(allDecrypted));
    }

    private List<Ciphertext> getListOfAllCiphertexts(List<ReencryptionOutput> reencrypted) {

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

        for (ReencryptionOutput reencryptionOutput : reencrypted) {
            singleCiphertext.add(reencryptionOutput.getCiphertext());

        }
        return singleCiphertext;
    }

    @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();
        }
    }

    @SuppressWarnings("unused")
    private void printCiphertext(List<Ciphertext> allCiphertexts) {

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

        for (Ciphertext c : allCiphertexts) {

            List<ZpGroupElement> e = c.getElements();

            for (ZpGroupElement z : e) {
                System.out.print(z.getValue());
                System.out.print(" ");
            }
            System.out.println();

        }
        System.out.println();
    }
}
