/*
 * 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.BufferedWriter;
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.nio.file.Paths;
import java.security.KeyPair;
import java.security.Security;
import java.security.SignatureException;
import java.security.cert.X509Certificate;
import java.time.ZoneOffset;
import java.time.ZonedDateTime;
import java.util.ArrayList;
import java.util.Collections;
import java.util.Date;
import java.util.List;

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 org.junit.runner.RunWith;

import com.googlecode.jcsv.writer.CSVWriter;
import com.googlecode.jcsv.writer.internal.CSVWriterBuilder;
import com.scytl.cryptolib.api.asymmetric.AsymmetricServiceAPI;
import com.scytl.cryptolib.api.certificates.CertificatesServiceAPI;
import com.scytl.cryptolib.api.exceptions.GeneralCryptoLibException;
import com.scytl.cryptolib.asymmetric.service.AsymmetricService;
import com.scytl.cryptolib.certificates.bean.CertificateData;
import com.scytl.cryptolib.certificates.bean.ValidityDates;
import com.scytl.cryptolib.certificates.bean.X509DistinguishedName;
import com.scytl.cryptolib.certificates.cryptoapi.CryptoAPIX509Certificate;
import com.scytl.cryptolib.certificates.service.CertificatesService;
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;

import mockit.integration.junit4.JMockit;

/**
 * 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
 */
@RunWith(JMockit.class)
public class DecryptionITest {

    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;

    private static AsymmetricServiceAPI _asymmetricService;

    private static CertificatesServiceAPI _certificatesService;

    private static ElGamalEncryptedBallots _elGamalEncryptedBallots;

    private static ElGamalPublicKey _publicKey;

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

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

        Security.addProvider(new BouncyCastleProvider());

        BigInteger p = new BigInteger("67555108767603707298738874065797424621214999980483524802737421636930165840299");

        BigInteger q = new BigInteger("33777554383801853649369437032898712310607499990241762401368710818465082920149");

        BigInteger g = new BigInteger("3");

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

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

        _configObjectMapper = new ConfigObjectMapper();

        _target = new SimpleBallotDecryptionService(new TransactionInfoProvider(), new PrimitivesService(),
            new ElGamalService(), new ProofsService(), mock(SecureLoggingWriter.class));

        _inputList = new ArrayList<>();

        _asymmetricService = new AsymmetricService();

        _certificatesService = new CertificatesService();

        generate100EncryptedBallots();

    }

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

    }

    private static void generate100EncryptedBallots() throws GeneralCryptoLibException, IOException {

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

        ElGamalService elGamalService = new ElGamalService();
        ElGamalKeyPair elGamalKeyPair = elGamalService.getElGamalKeyPairGenerator().generateKeys(_encParams, 1);
        _privateKey = elGamalKeyPair.getPrivateKeys();
        _publicKey = elGamalKeyPair.getPublicKeys();
        CryptoAPIElGamalEncrypter encrypter = elGamalService.createEncrypter(_publicKey);
        List<ElGamalComputationsValues> ciphertexts = new ArrayList<>();

        GroupElementsCompressor<ZpGroupElement> compressor = new GroupElementsCompressor<>();

        for (int i = 0; i < 100; i++) {
            Collections.shuffle(representations);
            List<ZpGroupElement> zpGroupElements = prepareMessage(representations.get(0), representations.get(1));

            insertToInputList(representations.get(0), representations.get(1));

            ZpGroupElement compressedElement = compressor.compress(zpGroupElements);
            List<ZpGroupElement> compressedMessage = new ArrayList<>();
            compressedMessage.add(compressedElement);

            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 void insertToInputList(final String value1, final String value2) {

        List<BigInteger> list = new ArrayList<>();
        list.add(new BigInteger(value1));
        list.add(new BigInteger(value2));

        _inputList.add(list);
    }

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

        // numKeys = 4
        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) {
            CiphertextImpl 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(final 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);
        }
    }

    @SuppressWarnings("unused")
    private static X509Certificate generateTestCertAndSign(final Path... paths)
            throws GeneralCryptoLibException, IOException {
        KeyPair keyPair = _asymmetricService.getKeyPairForSigning();

        ZonedDateTime now = ZonedDateTime.now(ZoneOffset.UTC);
        ZonedDateTime end = now.plusYears(1);
        ValidityDates validityDates = new ValidityDates(Date.from(now.toInstant()), Date.from(end.toInstant()));

        CertificateData certificateData = new CertificateData();
        certificateData.setSubjectPublicKey(keyPair.getPublic());
        certificateData.setValidityDates(validityDates);
        X509DistinguishedName distinguishedName = new X509DistinguishedName.Builder("certId", "ES").build();
        certificateData.setSubjectDn(distinguishedName);
        certificateData.setIssuerDn(distinguishedName);

        CryptoAPIX509Certificate cert =
            _certificatesService.createSignX509Certificate(certificateData, keyPair.getPrivate());

        for (Path path : paths) {
            Path signaturePath = Paths.get(path.toString() + ".metadata");
            String signatureB64 = new String();
            Files.deleteIfExists(signaturePath);
            try (BufferedWriter writer = new BufferedWriter(new FileWriter(signaturePath.toFile()))) {
                writer.write(signatureB64);
            }
        }
        return cert.getCertificate();
    }
}
