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

import java.io.IOException;
import java.io.StringWriter;
import java.math.BigInteger;
import java.nio.charset.Charset;
import java.security.KeyPair;
import java.security.PrivateKey;
import java.security.cert.X509Certificate;
import java.time.ZonedDateTime;
import java.time.temporal.ChronoUnit;
import java.util.ArrayList;
import java.util.Base64;
import java.util.Collection;
import java.util.Collections;
import java.util.Date;
import java.util.LinkedList;
import java.util.List;
import java.util.Random;
import java.util.stream.Collectors;

import javax.json.Json;
import javax.json.JsonObject;
import javax.json.JsonObjectBuilder;

import com.fasterxml.jackson.databind.ObjectMapper;
import com.googlecode.jcsv.writer.CSVWriter;
import com.googlecode.jcsv.writer.internal.CSVWriterBuilder;
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.RootCertificateData;
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.certificates.utils.X509CertificateTestDataGenerator;
import com.scytl.cryptolib.elgamal.bean.CiphertextImpl;
import com.scytl.cryptolib.elgamal.bean.ElGamalComputationsValues;
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.commons.beans.Ballot;
import com.scytl.products.ov.commons.beans.Contest;
import com.scytl.products.ov.commons.beans.ElectionOption;
import com.scytl.products.ov.commons.beans.domain.model.messaging.Payload;
import com.scytl.products.ov.commons.beans.domain.model.messaging.PayloadSignature;
import com.scytl.products.ov.commons.beans.domain.model.messaging.PayloadSignatureException;
import com.scytl.products.ov.commons.beans.domain.model.vote.CiphertextEncryptedVote;
import com.scytl.products.ov.commons.beans.domain.model.vote.EncryptedVote;
import com.scytl.products.ov.commons.sign.CryptolibPayloadSigner;
import com.scytl.products.ov.mixnet.commons.io.ElGamalEncryptedBallotEntryConverter;
import com.scytl.products.ov.sdm.domain.model.status.Status;
import com.scytl.products.ov.sdm.infrastructure.JsonConstants;

/**
 * Provide methods to generate ballot boxes for use in SDM tests.
 */
public class SDMTestData {

    public static final String BALLOT_JSON = "input/ballot.json";

    public static final String BALLOT_ID = "1c18f7d21dce4ac1aa1f025599bf2cfd";

    public static final String TENANT_ID = "tenantId";

    public static final String ELECTION_EVENT_ID = "electionEventId";

    public static final String BALLOT_BOX_ID = "ballotBoxId";

    public static final String ELECTORAL_AUTHORITY_ID = "electoralAuthorityId";

    public static final String ADMINISTRATION_AUTHORITY_ID = "certId";

    public static final String WRITE_IN_ALPHABET = "12345";

    private static ElGamalServiceAPI elGamalService;

    private static CertificatesServiceAPI certificateService;

    private static CryptolibPayloadSigner payloadSigner;

    private static Random generator = new Random(1L); // Predictable random
                                                      // values

    /**
     * Generate the ballot box JSON.
     *
     * @param encryptionParameters
     *            the Zp subgroup
     */
    public static String generateBallotBoxJson(String ballotBoxId, String ballotId, String electoralAuthorityId,
            ElGamalPublicKey electoralAuthorityPublicKey, ZpSubgroup encryptionParameters)
            throws GeneralCryptoLibException {
        JsonObjectBuilder ballotBoxBuilder = Json.createObjectBuilder();

        ballotBoxBuilder.add(JsonConstants.JSON_ATTRIBUTE_NAME_ID, ballotBoxId);
        ballotBoxBuilder.add(JsonConstants.JSON_ATTRIBUTE_NAME_STATUS, Status.BB_DOWNLOADED.name());
        ballotBoxBuilder.add(JsonConstants.JSON_ATTRIBUTE_NAME_ELECTORALAUTHORITYID, electoralAuthorityId);

        // Ballot section.
        JsonObjectBuilder ballotBuilder = Json.createObjectBuilder();
        ballotBuilder.add(JsonConstants.JSON_ATTRIBUTE_NAME_ID, ballotId);
        ballotBoxBuilder.add(JsonConstants.JSON_ATTRIBUTE_NAME_BALLOT, ballotBuilder.build());

        // Encryption parameters section.
        JsonObjectBuilder encryptionParametersBuilder = Json.createObjectBuilder();
        encryptionParametersBuilder.add(JsonConstants.JSON_ATTRIBUTE_NAME_P, encryptionParameters.getP().toString());
        encryptionParametersBuilder.add(JsonConstants.JSON_ATTRIBUTE_NAME_Q, encryptionParameters.getQ().toString());
        encryptionParametersBuilder.add(JsonConstants.JSON_ATTRIBUTE_NAME_G, encryptionParameters.getG().toString());
        ballotBoxBuilder.add(JsonConstants.JSON_ATTRIBUTE_NAME_ENCRYPTION_PARAMS, encryptionParametersBuilder.build());

        JsonObject ballotBox = ballotBoxBuilder.build();

        return ballotBox.toString();
    }

    public static List<List<ZpGroupElement>> generateVotes(long numVotes)
            throws GeneralCryptoLibException, IOException {
        // Store the Zp subgroup locally.
        ZpSubgroup zpSubgroup = CryptoTestData.getZpSubgroup();
        // Load the ballot from a test fixture.
        Ballot ballot = new ObjectMapper()
            .readValue(SDMTestData.class.getClassLoader().getResourceAsStream(BALLOT_JSON), Ballot.class);

        // Extract the representations from each of the options in each of the
        // contests.
        List<String> representations = ballot.getContests().stream().map(Contest::getOptions)
            .flatMap(Collection::stream).map(ElectionOption::getRepresentation).collect(Collectors.toList());

        // Create a list of 'random' votes.
        List<List<ZpGroupElement>> votes = new LinkedList<>();
        // Create a list of votes which is composed of a random choice of
        // representations.
        for (int i = 0; i < numVotes; i++) {
            // Loop over all representations for this vote.
            List<ZpGroupElement> vote = new ArrayList<>();
            for (String representation : representations) {
                // Whether this representation is chosen.
                if (generator.nextBoolean()) {
                    vote.add(new ZpGroupElement(new BigInteger(representation), zpSubgroup));
                }
            }
            // Blank votes are not allowed. Choose the first representation if
            // that's the case.
            if (vote.isEmpty()) {
                vote =
                    Collections.singletonList(new ZpGroupElement(new BigInteger(representations.get(0)), zpSubgroup));
            }
            // Record the vote.
            votes.add(vote);
        }

        return votes;
    }

    /**
     * Generate encrypted ballots ready for mixing.
     *
     * @param electoralAuthorityPublicKey
     *            the public key the votes will be encrypted with
     * @return a string with the encrypted ballots as CSV.
     */
    public static String generateEncryptedBallots(ElGamalPublicKey electoralAuthorityPublicKey,
            List<List<ZpGroupElement>> votes) throws IOException, GeneralCryptoLibException {
        // Encrypt the votes and place them in the ballot box.
        CryptoAPIElGamalEncrypter encrypter = getElGamalService().createEncrypter(electoralAuthorityPublicKey);
        List<ElGamalComputationsValues> ciphertexts = new ArrayList<>();
        GroupElementsCompressor<ZpGroupElement> compressor = new GroupElementsCompressor<>();

        // Cast votes.
        for (List<ZpGroupElement> vote : votes) {
            ZpGroupElement compressedVote = compressor.compress(vote);

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

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

        List<Ciphertext> elGamalEncryptedBallots = fromElGamalComputationsValuesToElGamalEncryptedBallots(ciphertexts);

        // Write the encrypted votes to a CSV string.
        try (final StringWriter writer = new StringWriter();
                final CSVWriter<Ciphertext> csvWriter = new CSVWriterBuilder<Ciphertext>(writer)
                    .entryConverter(new ElGamalEncryptedBallotEntryConverter()).build()) {
            csvWriter.writeAll(elGamalEncryptedBallots);
            csvWriter.flush();

            return writer.toString();
        }
    }

    /**
     * Generate encrypted votes ready for mixing.
     *
     * @param electoralAuthorityPublicKey
     *            the public key the votes will be encrypted with
     * @return the list of encrypted votes.
     */
    public static List<EncryptedVote> generateEncryptedVotes(ElGamalPublicKey electoralAuthorityPublicKey,
            List<List<ZpGroupElement>> votes) throws IOException, GeneralCryptoLibException {
        // Encrypt the votes and place them in the ballot box.
        CryptoAPIElGamalEncrypter encrypter = getElGamalService().createEncrypter(electoralAuthorityPublicKey);
        List<ElGamalComputationsValues> ciphertexts = new ArrayList<>();
        GroupElementsCompressor<ZpGroupElement> compressor = new GroupElementsCompressor<>();

        // Cast votes.
        for (List<ZpGroupElement> vote : votes) {
            ZpGroupElement compressedVote = compressor.compress(vote);

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

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

        List<Ciphertext> elGamalEncryptedBallots = fromElGamalComputationsValuesToElGamalEncryptedBallots(ciphertexts);

        return elGamalEncryptedBallots.stream().map(CiphertextEncryptedVote::new).collect(Collectors.toList());
    }

    /**
     * Generate the electoral authority JSON.
     *
     * @param electoralAuthorityPublicKey
     *            the electoral authority public key
     */
    public static String generateElectoralAuthorityJson(ElGamalPublicKey electoralAuthorityPublicKey)
            throws GeneralCryptoLibException {
        JsonObjectBuilder electoralAuthorityBuilder = Json.createObjectBuilder();

        electoralAuthorityBuilder.add(JsonConstants.JSON_ATTRIBUTE_NAME_PUBLIC_KEY, Base64.getEncoder()
            .encodeToString(electoralAuthorityPublicKey.toJson().getBytes(Charset.defaultCharset())));

        return electoralAuthorityBuilder.build().toString();
    }

    /**
     * Generate the election event JSON.
     */
    public static String generateElectionEventJson() {
        JsonObjectBuilder electionEventBuilder = Json.createObjectBuilder();

        JsonObjectBuilder administrationAuthorityBuilder = Json.createObjectBuilder();
        administrationAuthorityBuilder.add(JsonConstants.JSON_ATTRIBUTE_NAME_ID, ADMINISTRATION_AUTHORITY_ID);
        electionEventBuilder.add(JsonConstants.JSON_ATTRIBUTE_NAME_ADMINISTRATION_AUTHORITY,
            administrationAuthorityBuilder);

        JsonObjectBuilder settingsBuilder = Json.createObjectBuilder();
        settingsBuilder.add(JsonConstants.JSON_PARAM_NAME_WRITE_IN_ALPHABET,
            Base64.getEncoder().encodeToString(WRITE_IN_ALPHABET.getBytes(Charset.defaultCharset())));
        electionEventBuilder.add(JsonConstants.JSON_ATTRIBUTE_NAME_SETTINGS, settingsBuilder);

        return electionEventBuilder.build().toString();
    }

    /**
     * Convert El Gamal computation values to encrypted ballots
     *
     * @return the encrypted ballots
     */
    private static List<Ciphertext> fromElGamalComputationsValuesToElGamalEncryptedBallots(
            final List<ElGamalComputationsValues> elGamalComputationsValues) throws GeneralCryptoLibException {

        List<Ciphertext> encryptedBallots = new ArrayList<>();
        for (ElGamalComputationsValues ciphertext : elGamalComputationsValues) {
            List<ZpGroupElement> values = ciphertext.getPhis();
            List<ZpGroupElement> elements = new ArrayList<>(ciphertext.getValues().size());
            for (ZpGroupElement zpGroupElement : values) {
                ZpGroupElement element = new ZpGroupElement(zpGroupElement.getValue(), CryptoTestData.getZpSubgroup());
                elements.add(element);
            }
            Ciphertext elGamalEncryptedBallot = new CiphertextImpl(ciphertext.getGamma(), elements);
            encryptedBallots.add(elGamalEncryptedBallot);
        }

        return encryptedBallots;
    }

    private static ElGamalServiceAPI getElGamalService() throws GeneralCryptoLibException {
        if (elGamalService == null) {
            elGamalService = new ElGamalService();
        }

        return elGamalService;
    }

    private static CertificatesServiceAPI getCertificateService() throws GeneralCryptoLibException {
        if (certificateService == null) {
            certificateService = new CertificatesService();
        }

        return certificateService;
    }

    private static CryptolibPayloadSigner getPayloadSigner() throws GeneralCryptoLibException {
        if (payloadSigner == null) {
            payloadSigner = new CryptolibPayloadSigner(new AsymmetricService());
        }

        return payloadSigner;
    }

    /**
     * Create a certificate chain with a configurable number of intermediate
     * CAs.
     * 
     * @param intermediateCAs
     *            the number of intermediate CAs to build
     * @return a certificate chain
     * @throws GeneralCryptoLibException
     */
    private static X509Certificate[] createCertificateChain(int intermediateCAs, KeyPair signingKeyPair)
            throws GeneralCryptoLibException {
        ZonedDateTime startDate = ZonedDateTime.now();
        ZonedDateTime endDate = startDate.plus(1, ChronoUnit.HOURS);
        ValidityDates validityDates =
            new ValidityDates(Date.from(startDate.toInstant()), Date.from(endDate.toInstant()));

        List<CryptoAPIX509Certificate> certificateChainBuilder = new ArrayList<>();
        // Create the self-signed root CA.
        X509DistinguishedName rootDN = getDistinguishedName("Root CA");
        KeyPair rootKeyPair = signingKeyPair;
        RootCertificateData rootCertificateData = X509CertificateTestDataGenerator.getRootCertificateData(rootKeyPair);
        rootCertificateData.setSubjectDn(rootDN);
        rootCertificateData.setValidityDates(validityDates);
        CryptoAPIX509Certificate rootCertificate =
            getCertificateService().createRootAuthorityX509Certificate(rootCertificateData, rootKeyPair.getPrivate());

        certificateChainBuilder.add(rootCertificate);

        // Create the required intermediate certificates.
        PrivateKey caSigningKey = rootKeyPair.getPrivate();
        X509DistinguishedName caIssuerDN = rootDN;
        for (int i = 0; i < intermediateCAs; i++) {
            // Build the intermediate CA.
            KeyPair intermediateKeyPair = signingKeyPair;
            X509DistinguishedName intermediateDN = getDistinguishedName("Intermediate CA #" + i);
            CertificateData intermediateCertificateData =
                X509CertificateTestDataGenerator.getCertificateData(intermediateKeyPair);
            intermediateCertificateData.setSubjectDn(intermediateDN);
            intermediateCertificateData.setIssuerDn(caIssuerDN);
            intermediateCertificateData.setValidityDates(validityDates);
            CryptoAPIX509Certificate intermediateCertificate = getCertificateService()
                .createIntermediateAuthorityX509Certificate(intermediateCertificateData, caSigningKey);
            // Add the intermediate CA certificate to the certificate chain
            // builder.
            certificateChainBuilder.add(intermediateCertificate);
            // Keep this CA's private key to sign the next certificate.
            caSigningKey = intermediateKeyPair.getPrivate();
            // Keep this CA's subject DN to assign to the next certificate as
            // issuer.
            caIssuerDN = intermediateDN;
        }

        // Create the leaf certificate.
        KeyPair leafKeyPair = signingKeyPair;
        CertificateData leafCertificateData = X509CertificateTestDataGenerator.getCertificateData(leafKeyPair);
        leafCertificateData.setSubjectDn(getDistinguishedName("Leaf certificate"));
        leafCertificateData.setIssuerDn(caIssuerDN);
        leafCertificateData.setValidityDates(validityDates);
        CryptoAPIX509Certificate leafCertificate =
            getCertificateService().createSignX509Certificate(leafCertificateData, caSigningKey);
        certificateChainBuilder.add(leafCertificate);

        // Create a certificate array from the reversed certificate chain
        // builder.
        Collections.reverse(certificateChainBuilder);
        return certificateChainBuilder.stream().map(CryptoAPIX509Certificate::getCertificate)
            .toArray(X509Certificate[]::new);
    }

    /**
     * Builds an X509 distinguished name fit for test certificates.
     * 
     * @param commonName
     *            the common name
     * @return an X509 distinguished name
     * @throws GeneralCryptoLibException
     */
    private static X509DistinguishedName getDistinguishedName(String commonName) throws GeneralCryptoLibException {
        return new X509DistinguishedName.Builder(commonName, "XX").addOrganization("Scytl").addLocality("Barcelona")
            .build();
    }

    public static PayloadSignature signPayload(Payload payload, KeyPair signingKeyPair)
            throws PayloadSignatureException, GeneralCryptoLibException {
        X509Certificate[] certificateChain = createCertificateChain(1, signingKeyPair);

        return getPayloadSigner().sign(payload, signingKeyPair.getPrivate(), certificateChain);

    }
}
