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

import static org.junit.Assert.assertEquals;

import java.io.BufferedReader;
import java.io.BufferedWriter;
import java.io.File;
import java.io.FileReader;
import java.io.FileWriter;
import java.io.IOException;
import java.io.PrintWriter;
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.KeyException;
import java.security.KeyPair;
import java.security.Security;
import java.security.cert.X509Certificate;
import java.time.ZoneOffset;
import java.time.ZonedDateTime;
import java.util.ArrayList;
import java.util.Date;
import java.util.HashMap;
import java.util.List;
import java.util.Map;

import org.apache.commons.io.FileUtils;
import org.bouncycastle.jce.provider.BouncyCastleProvider;
import org.junit.Assert;
import org.junit.BeforeClass;
import org.junit.Test;

import com.fasterxml.jackson.core.JsonParseException;
import com.fasterxml.jackson.databind.JsonMappingException;
import com.fasterxml.jackson.databind.ObjectMapper;
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.elgamal.ElGamalServiceAPI;
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.ElGamalEncryptionParameters;
import com.scytl.cryptolib.elgamal.bean.ElGamalKeyPair;
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.products.ov.mixnet.commons.beans.Ballot;
import com.scytl.products.ov.mixnet.commons.configuration.file.FileBasedApplicationConfig;
import com.scytl.products.ov.mixnet.commons.configuration.file.FileBasedMixerConfig;
import com.scytl.products.ov.mixnet.commons.configuration.file.FileBasedVerifierConfig;
import com.scytl.products.ov.mixnet.commons.io.ElGamalEncryptedBallotEntryConverter;

public class ApplicationTest {

    private static ObjectMapper mapper = new ObjectMapper();

    private static ZpSubgroup group;

    private static AsymmetricServiceAPI asymmetricService;

    private static CertificatesServiceAPI certificatesService;

    private static ElGamalServiceAPI elGamalService;

    private static ElGamalEncryptionParameters encryptionParameters;

    private static BigInteger p;

    private static BigInteger q;

    private static BigInteger g;

    private static ElGamalKeyPair elGamalKeyPair;

    @BeforeClass
    public static void setUp() throws GeneralCryptoLibException {

        Security.addProvider(new BouncyCastleProvider());

        asymmetricService = new AsymmetricService();
        certificatesService = new CertificatesService();
        elGamalService = new ElGamalService();

        p = new BigInteger("67555108767603707298738874065797424621214999980483524802737421636930165840299");
        q = new BigInteger("33777554383801853649369437032898712310607499990241762401368710818465082920149");
        g = new BigInteger("3");
        group = new ZpSubgroup(g, p, q);
        encryptionParameters = new ElGamalEncryptionParameters(p, q, g);
    }

    @Test
    public void verify_mixDec() throws IOException, GeneralCryptoLibException, KeyException {

        Path targetPath = Paths.get("target");
        Path outputDirectory = targetPath.resolve("outputs");
        Path inputDirectory = targetPath.resolve("inputs");

        inputDirectory.toFile().mkdirs();
        outputDirectory.toFile().mkdirs();
        Path appConfigLocation = inputDirectory.resolve("applicationConfig.json");

        elGamalKeyPair = createElGamalKeyPair();

        writeEncryptionParameters(encryptionParameters, inputDirectory.resolve("encryptionParameters.json"));

        FileBasedApplicationConfig generatedConfig = createConfig(inputDirectory, outputDirectory);

        mapper.writeValue(appConfigLocation.toFile(), generatedConfig);

        String[] params = new String[] {appConfigLocation.toFile().getAbsolutePath() };

        Application.main(params);

        String expectedVerifierResult = "{\"1\":true}";
        assertEquals(expectedVerifierResult, getFirstLineFromVerificationResultFile());
    }

    private String getFirstLineFromVerificationResultFile() {

        File outputFileOfVerifier = Paths.get("target/outputs/verifierResult.json").toFile();

        String firstLineFromVerifierOutputFile = "";
        try (BufferedReader br = new BufferedReader(new FileReader(outputFileOfVerifier))) {
            firstLineFromVerifierOutputFile = br.readLine();
        } catch (IOException e) {
            System.out.println("Error while trying to read line from verifier output file");
            Assert.fail();
        }

        return firstLineFromVerifierOutputFile;
    }

    private FileBasedApplicationConfig createConfig(Path inputDirectory, Path outputDirectory)
            throws IOException, GeneralCryptoLibException {

        Path publicKeyLocation = writePublicKey(inputDirectory.resolve("publicKey.json"), elGamalKeyPair);
        Path privateKeyLocation = writePrivateKey(inputDirectory.resolve("ballotDecryptionKey.json"), elGamalKeyPair);
        Path ballotBoxToBeMixedLocation = createBallotBoxToBeMixed(inputDirectory);

        FileBasedApplicationConfig generatedConfig = new FileBasedApplicationConfig();
        generatedConfig.setConcurrencyLevel(1);
        generatedConfig.setNumberOfParallelBatches(1);
        generatedConfig.setNumOptions(Integer.valueOf(1));
        generatedConfig.setMinimumBatchSize(64);
        FileBasedMixerConfig mixer = new FileBasedMixerConfig();
        mixer.setMix(true);
        Map<String, String> inputMap = new HashMap<>();
        inputMap.put("ballotDecryptionKey", privateKeyLocation.toFile().getAbsolutePath().replaceAll(".json", ""));
        inputMap.put("publicKey", publicKeyLocation.toFile().getAbsolutePath().replaceAll(".json", ""));
        inputMap.put("encryptedBallots", ballotBoxToBeMixedLocation.toFile().getAbsolutePath().replaceAll(".csv", ""));
        inputMap.put("encryptionParameters",
            inputDirectory.resolve("encryptionParameters.json").toFile().getAbsolutePath().replaceAll(".json", ""));
        mixer.setInput(inputMap);
        String outputAbsPath = outputDirectory.toFile().getAbsolutePath();
        mixer.setOutputDirectory(outputAbsPath);
        generatedConfig.setMixer(mixer);
        FileBasedVerifierConfig verifier = new FileBasedVerifierConfig();
        verifier.setVerify(true);
        verifier.setInputDirectory(outputAbsPath);
        verifier.setOutputDirectory(outputAbsPath);
        generatedConfig.setVerifier(verifier);

        return generatedConfig;
    }

    private void writeEncryptionParameters(ElGamalEncryptionParameters encryptionParameters2, Path resolve)
            throws IOException, GeneralCryptoLibException {
        final List<String> linesToBeWritten = new ArrayList<>();
        linesToBeWritten.add(((ZpSubgroup) encryptionParameters2.getGroup()).toJson());
        FileUtils.writeLines(resolve.toFile(), linesToBeWritten);

    }

    private ElGamalKeyPair createElGamalKeyPair() throws KeyException {
        ElGamalKeyPair generatedKeyPair;
        try {
            generatedKeyPair = elGamalService.getElGamalKeyPairGenerator().generateKeys(encryptionParameters, 1);
        } catch (GeneralCryptoLibException e) {
            throw new KeyException("An error occurred while generating the ElGamal key pair", e);
        }
        return generatedKeyPair;

    }

    private Path createBallotBoxToBeMixed(Path elGamalEncryptedBallotsPath)
            throws GeneralCryptoLibException, JsonParseException, JsonMappingException, IOException {

        Ballot ballot = mapper.readValue(new File("src/test/resources/input/ballot.json"), Ballot.class);

        List<String> representations = new ArrayList<>();
        ballot.getContests().forEach(
            contest -> contest.getOptions().forEach(option -> representations.add(option.getRepresentation())));

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

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

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

        int numberMixables = 5;
        int numberOptionsInBallot = 3;
        for (int i = 0; i < numberMixables; i++) {

            // way of generating encrypted ballots with a range of values
            int indexOfOptionToChoose = Math.floorMod(i, numberOptionsInBallot);
            List<ZpGroupElement> zpGroupElements = prepareMessage(representations.get(indexOfOptionToChoose));

            ZpGroupElement compressedElement = compressor.compress(zpGroupElements);

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

            elGamalEncryptedBallots.add(encrypter.encryptGroupElements(compressedMessage));
        }

        Path ballots = elGamalEncryptedBallotsPath.resolve("encryptedBallots.csv");
        try (final Writer out = new FileWriter(ballots.toFile());
                final CSVWriter<Ciphertext> csvWriter = new CSVWriterBuilder<Ciphertext>(out)
                    .entryConverter(new ElGamalEncryptedBallotEntryConverter()).build()) {
            csvWriter.writeAll(elGamalEncryptedBallots);
            csvWriter.flush();
        }

        generateTestCertAndSign(elGamalEncryptedBallotsPath);
        return ballots;
    }

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

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

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

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

        return messages;
    }

    private Path writePublicKey(Path inputFilePath, ElGamalKeyPair keyPair)
            throws IOException, GeneralCryptoLibException {
        try (FileWriter fw = new FileWriter(inputFilePath.toFile()); PrintWriter pw = new PrintWriter(fw)) {
            ElGamalPublicKey publicKeys = keyPair.getPublicKeys();
            pw.write(publicKeys.toJson());
        }
        return inputFilePath;
    }

    private Path writePrivateKey(Path inputFilePath, ElGamalKeyPair keyPair)
            throws IOException, GeneralCryptoLibException {
        try (FileWriter fw = new FileWriter(inputFilePath.toFile()); PrintWriter pw = new PrintWriter(fw)) {
            pw.write(keyPair.getPrivateKeys().toJson());
        }
        return inputFilePath;
    }
}
