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

import static org.junit.Assert.assertEquals;
import static org.junit.Assert.assertNotNull;
import static org.mockito.Mockito.mock;

import java.io.File;
import java.io.FileWriter;
import java.io.IOException;
import java.io.Writer;
import java.math.BigInteger;
import java.nio.file.Files;
import java.nio.file.Path;
import java.security.Security;
import java.security.SignatureException;
import java.util.ArrayList;
import java.util.Collections;
import java.util.List;
import java.util.Random;

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

import com.googlecode.jcsv.writer.CSVWriter;
import com.googlecode.jcsv.writer.internal.CSVWriterBuilder;
import com.scytl.cryptolib.api.elgamal.ElGamalServiceAPI;
import com.scytl.cryptolib.api.exceptions.GeneralCryptoLibException;
import com.scytl.cryptolib.api.primitives.PrimitivesServiceAPI;
import com.scytl.cryptolib.api.proofs.ProofsServiceAPI;
import com.scytl.cryptolib.elgamal.bean.CiphertextImpl;
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.CryptoAPIElGamalEncrypter;
import com.scytl.cryptolib.elgamal.service.ElGamalService;
import com.scytl.cryptolib.mathematical.groups.activity.GroupElementsCompressor;
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.cryptolib.proofs.proof.Proof;
import com.scytl.cryptolib.proofs.service.ProofsService;
import com.scytl.products.oscore.logging.api.writer.SecureLoggingWriter;
import com.scytl.products.oscore.logging.core.provider.TransactionInfoProvider;
import com.scytl.products.ov.mixnet.BallotDecryptionService;
import com.scytl.products.ov.mixnet.SimpleBallotDecryptionService;
import com.scytl.products.ov.mixnet.commons.ballots.ElGamalEncryptedBallots;
import com.scytl.products.ov.mixnet.commons.beans.Ballot;
import com.scytl.products.ov.mixnet.commons.beans.VoteWithProof;
import com.scytl.products.ov.mixnet.commons.io.ElGamalEncryptedBallotEntryConverter;
import com.scytl.products.ov.mixnet.commons.tools.ConfigObjectMapper;

/**
 * Test case to verify that: 1) create two messages (prepareInputData method) 2) compress message values to a only one
 * value (prepareInputData method) 3) encrypt them with a CryptoElGamalEncrypter(prepareInputData method) 4) call to
 * decryptool to decrypt values 5) verify the decryptool output == first initial two messages
 */
public class WriteInDecryptionITest {

    private static final int NUMBER_OF_SUBKEYS = 2;

    private static final int NUMBER_OF_CORES_CONCURRENT = 2;

    private static final int NUMBER_OF_CORES_SEQUENTIAL = 1;

    private static final String SRC_TEST_RESOURCES_BALLOT_JSON = "src/test/resources/ballot.json";

    private static ZpSubgroup _group_g2_q11;

    private static List<List<BigInteger>> _inputList;

    private static ConfigObjectMapper _configObjectMapper;

    private static BallotDecryptionService _target;

    private static ElGamalEncryptionParameters _encParams;

    private static Path _elGamalEncryptedBallotsPath;

    private static ElGamalPrivateKey _privateKey;

    @Rule
    public final TemporaryFolder folder = new TemporaryFolder();

    private static String charset =
        "()0123456789ABCDEFGHIJKLMNOPQRSTUVWXYZ`abcdefghijklmnopqrstuvwxyz\u00A0\u00A1\u00A2\u0160\u0161\u017D\u017E\u0152\u0153\u0178\u00C0\u00C1\u00C2\u00C3\u00C4\u00C5\u00C6\u00C7\u00C8\u00C9\u00CA\u00CB\u00CC\u00CD\u00CE\u00CF\u00D0\u00D1\u00D2\u00D3\u00D4\u00D5\u00D6\u00D8\u00D9\u00DA\u00DB\u00DC\u00DD\u00DE\u00DF\u00E0\u00E1\u00E2\u00E3\u00E4\u00E5\u00E6\u00E7\u00E8\u00E9\u00EA\u00EB\u00EC\u00ED\u00EE\u00EF\u00F0\u00F1\u00F2\u00F3\u00F4\u00F5\u00F6\u00F8\u00F9\u00FA\u00FB\u00FC\u00FD\u00FE\u00FF";

    private static char separator = '#';

    private static BigInteger _p;

    private static ElGamalPublicKey _publicKey;

    private static ElGamalEncryptedBallots _elGamalEncryptedBallots;

    @BeforeClass
    public static void setUp() throws GeneralCryptoLibException, IOException {

        Security.addProvider(new BouncyCastleProvider());

        _p = new BigInteger(
            "16370518994319586760319791526293535327576438646782139419846004180837103527129035954742043590609421369665944746587885814920851694546456891767644945459124422553763416586515339978014154452159687109161090635367600349264934924141746082060353483306855352192358732451955232000593777554431798981574529854314651092086488426390776811367125009551346089319315111509277347117467107914073639456805159094562593954195960531136052208019343392906816001017488051366518122404819967204601427304267380238263913892658950281593755894747339126531018026798982785331079065126375455293409065540731646939808640273393855256230820509217411510058759");

        BigInteger q = new BigInteger(
            "8185259497159793380159895763146767663788219323391069709923002090418551763564517977371021795304710684832972373293942907460425847273228445883822472729562211276881708293257669989007077226079843554580545317683800174632467462070873041030176741653427676096179366225977616000296888777215899490787264927157325546043244213195388405683562504775673044659657555754638673558733553957036819728402579547281296977097980265568026104009671696453408000508744025683259061202409983602300713652133690119131956946329475140796877947373669563265509013399491392665539532563187727646704532770365823469904320136696927628115410254608705755029379");

        BigInteger g = new BigInteger("3");

        _encParams = new ElGamalEncryptionParameters(_p, q, g);

        _group_g2_q11 = new ZpSubgroup(g, _p, q);

        _configObjectMapper = new ConfigObjectMapper();

        TransactionInfoProvider transactionInfoProvider = new TransactionInfoProvider();
        PrimitivesServiceAPI primitivesService = new PrimitivesService();
        ElGamalServiceAPI elGamalService = new ElGamalService();
        ProofsServiceAPI proofsService = new ProofsService();
        _target = new SimpleBallotDecryptionService(transactionInfoProvider, primitivesService, elGamalService,
            proofsService, mock(SecureLoggingWriter.class));

        _inputList = new ArrayList<>();

        generate100EncryptedBallots();

    }

    @AfterClass
    public static void removeBouncyCastleProvider() {
        Security.removeProvider("BC");

    }

    private static void generate100EncryptedBallots() throws GeneralCryptoLibException, IOException {

        // please adjust the removed values if the test fails because of the
        // change of the ballot file that is referenced here
        Ballot ballot = _configObjectMapper.fromJSONFileToJava(new File(SRC_TEST_RESOURCES_BALLOT_JSON), Ballot.class);

        List<String> representations = new ArrayList<>();
        ballot.getContests().forEach(
            contest -> contest.getOptions().forEach(option -> representations.add(option.getRepresentation())));
        // these numbers are not compatible with the hardcoded p value, so
        // should not be used in the test
        representations.remove("101");
        representations.remove("41");
        representations.remove("103");
        representations.remove("19");
        ElGamalService elGamalService = new ElGamalService();
        ElGamalKeyPair elGamalKeyPair =
            elGamalService.getElGamalKeyPairGenerator().generateKeys(_encParams, NUMBER_OF_SUBKEYS);
        _privateKey = elGamalKeyPair.getPrivateKeys();
        _publicKey = elGamalKeyPair.getPublicKeys();
        CryptoAPIElGamalEncrypter encrypter = elGamalService.createEncrypter(_publicKey);
        List<ElGamalComputationsValues> ciphertexts = new ArrayList<>();
        GroupElementsCompressor<ZpGroupElement> compressor = new GroupElementsCompressor<>();
        char[] chars = charset.toCharArray();
        String dictionary = new StringBuilder().append(charset).append(separator).toString();
        Random random = new Random();
        for (int i = 0; i < 100; i++) {
            Collections.shuffle(representations);
            List<String> values = new ArrayList<>(3 + NUMBER_OF_SUBKEYS - 1);
            List<ZpGroupElement> zpGroupElements =
                prepareMessage(representations.get(0), representations.get(1), representations.get(2));
            values.add(representations.get(0));
            values.add(representations.get(1));
            values.add(representations.get(2));

            ZpGroupElement compressedElement = compressor.compress(zpGroupElements);
            List<ZpGroupElement> compressedMessage = new ArrayList<>();
            compressedMessage.add(compressedElement);
            for (int j = 1; j < NUMBER_OF_SUBKEYS; j++) {
                String createdEncodedWriteIn = createEncodedWriteIn(random, chars, dictionary);
                List<ZpGroupElement> preparedMessage = prepareMessage(createdEncodedWriteIn);
                ZpGroupElement preparedMessageAsGroup = compressor.compress(preparedMessage);
                values.add(createdEncodedWriteIn);
                compressedMessage.add(preparedMessageAsGroup);
            }
            insertToInputList(values);

            ElGamalComputationsValues ciphertext =
                encrypter.encryptGroupElements(compressedMessage).getComputationValues();
            ciphertexts.add(ciphertext);
        }

        List<Ciphertext> elGamalEncryptedBallots = fromElGamalComputationsValuesToElGamalEncryptedBallot(ciphertexts);

        _elGamalEncryptedBallots = new ElGamalEncryptedBallots(elGamalEncryptedBallots);

        _elGamalEncryptedBallotsPath = Files.createTempFile("encryptedBallots.csv", null);
        _elGamalEncryptedBallotsPath.toFile().deleteOnExit();

        try (final Writer out = new FileWriter(_elGamalEncryptedBallotsPath.toFile());
                final CSVWriter<Ciphertext> csvWriter = new CSVWriterBuilder<Ciphertext>(out)
                    .entryConverter(new ElGamalEncryptedBallotEntryConverter()).build()) {
            csvWriter.writeAll(elGamalEncryptedBallots);
            csvWriter.flush();
        }

    }

    private static String createEncodedWriteIn(Random random, char[] chars, String dictionary) {
        String firstName = createRandomString(random, chars);
        String lastName = createRandomString(random, chars);
        String municipality = createRandomString(random, chars);
        String encodedValue = encode(firstName, lastName, municipality, dictionary);
        BigInteger number = new BigInteger(encodedValue);
        String calculatedValue = number.pow(2).mod(_p).toString();
        return calculatedValue;
    }

    private static String encode(String firstName, String lastName, String municipality, String dictionary) {
        String concatenated = new StringBuilder().append(firstName).append(separator).append(lastName).append(separator)
            .append(municipality).toString();
        char[] charArray = concatenated.toCharArray();
        StringBuilder sb = new StringBuilder();
        for (char c : charArray) {
            sb.append(String.format("%03d", dictionary.indexOf(c) + 1));
        }
        return sb.toString();
    }

    private static String createRandomString(Random random, char[] chars) {
        StringBuilder sb = new StringBuilder();
        int size = random.nextInt(30);
        for (int i = 0; i < size; i++) {
            char c = chars[random.nextInt(chars.length)];
            sb.append(c);
        }
        return sb.toString();
    }

    private static void insertToInputList(List<String> values) {

        List<BigInteger> list = new ArrayList<>();
        for (String value : values) {
            list.add(new BigInteger(value));
        }
        _inputList.add(list);
    }

    private static List<ZpGroupElement> prepareMessage(String... numbers) throws GeneralCryptoLibException {

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

        for (String number : numbers) {
            ZpGroupElement element = new ZpGroupElement(new BigInteger(number), _group_g2_q11);
            messages.add(element);
        }

        return messages;
    }

    private static List<Ciphertext> fromElGamalComputationsValuesToElGamalEncryptedBallot(
            final List<ElGamalComputationsValues> ciphertexts) throws GeneralCryptoLibException {

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

        for (ElGamalComputationsValues ciphertext : ciphertexts) {
            Ciphertext elGamalEncryptedBallot = new CiphertextImpl(ciphertext.getGamma(), ciphertext.getPhis());
            encryptedBallots.add(elGamalEncryptedBallot);

        }
        return encryptedBallots;
    }

    @Test
    public void decrypt_a_valid_input_sequentially() throws GeneralCryptoLibException, IOException, SignatureException {
        decryptGivenNumberOfCores(NUMBER_OF_CORES_SEQUENTIAL);
    }

    @Test
    public void decrypt_a_valid_input_concurrently() throws GeneralCryptoLibException, IOException, SignatureException {
        decryptGivenNumberOfCores(NUMBER_OF_CORES_CONCURRENT);
    }

    private void decryptGivenNumberOfCores(int numberOfCores)
            throws GeneralCryptoLibException, IOException, SignatureException {

        List<Ciphertext> encryptedBallots = _elGamalEncryptedBallots.getBallots();
        ElGamalPrivateKey privateKey = _privateKey;
        Integer batchNum = 1;
        int partitionSize = 1;
        int offset = 0;
        List<VoteWithProof> decrypt =
            _target.decrypt(encryptedBallots, privateKey, batchNum, partitionSize, offset);
        checkDecryptionProofs(decrypt);
        checkDecryptedOptions(decrypt);
        assertEquals(_inputList.size(), decrypt.size());
    }

    private void checkDecryptionProofs(final List<VoteWithProof> proofs) throws GeneralCryptoLibException {

        for (VoteWithProof proof : proofs) {
            assertNotNull(Proof.fromJson(proof.getProof()));
        }
    }

    private void checkDecryptedOptions(final List<VoteWithProof> proofs) {

        for (int i = 0; i < _inputList.size(); i++) {
            int finalValueOfDecryptedValues = 1;
            for (int j = 0; j < proofs.get(i).getDecrypted().size(); j++) {
                finalValueOfDecryptedValues *= proofs.get(i).getDecrypted().get(j).intValue();
            }
            int finalValueOfOriginalValues = 1;
            for (int j = 0; j < _inputList.get(i).size(); j++) {
                finalValueOfOriginalValues *= _inputList.get(i).get(j).intValue();
            }
            assertEquals(finalValueOfDecryptedValues, finalValueOfOriginalValues);
        }
    }

}
