/*
 * 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.ui.ws.rs.application;

import static org.junit.Assert.assertEquals;

import java.io.IOException;
import java.nio.charset.Charset;
import java.nio.file.Files;
import java.nio.file.Path;
import java.security.KeyPair;
import java.security.PrivateKey;
import java.security.Security;
import java.security.cert.X509Certificate;
import java.util.Base64;
import java.util.Collections;
import java.util.List;
import java.util.Map;
import java.util.function.Function;
import java.util.stream.Collectors;

import org.bouncycastle.jce.provider.BouncyCastleProvider;
import org.junit.BeforeClass;
import org.junit.Test;
import org.junit.runner.RunWith;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.http.HttpStatus;
import org.springframework.http.ResponseEntity;
import org.springframework.test.context.ContextConfiguration;
import org.springframework.test.context.TestPropertySource;
import org.springframework.test.context.junit4.SpringJUnit4ClassRunner;

import com.scytl.cryptolib.api.exceptions.GeneralCryptoLibException;
import com.scytl.cryptolib.elgamal.bean.ElGamalKeyPair;
import com.scytl.cryptolib.elgamal.bean.ElGamalPublicKey;
import com.scytl.cryptolib.mathematical.groups.impl.ZpGroupElement;
import com.scytl.cryptolib.mathematical.groups.impl.ZpSubgroup;
import com.scytl.cryptolib.primitives.primes.utils.PemUtils;
import com.scytl.products.ov.commons.beans.domain.model.ballotbox.BallotBoxId;
import com.scytl.products.ov.commons.beans.domain.model.ballotbox.BallotBoxIdImpl;
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.EncryptedVote;
import com.scytl.products.ov.commons.beans.domain.model.vote.VoteSetId;
import com.scytl.products.ov.commons.beans.domain.model.vote.VoteSetIdImpl;
import com.scytl.products.ov.commons.dto.mixing.MixingPayload;
import com.scytl.products.ov.commons.dto.mixing.MixingPayloadImpl;
import com.scytl.products.ov.commons.path.PathResolver;
import com.scytl.products.ov.commons.ui.ws.rs.ObjectMappers;
import com.scytl.products.ov.sdm.SDMTestData;
import com.scytl.products.ov.sdm.application.service.CertificateManagementException;
import com.scytl.products.ov.sdm.application.service.MixDecValService;
import com.scytl.products.ov.sdm.application.service.PlatformRootCAService;
import com.scytl.products.ov.sdm.domain.common.IdleState;
import com.scytl.products.ov.sdm.domain.model.mixing.MixingData;
import com.scytl.products.ov.sdm.domain.service.ConfigConstants;
import com.scytl.products.ov.validation.services.ValidationServiceImpl;

@RunWith(SpringJUnit4ClassRunner.class)
@ContextConfiguration(classes = {BaseTestConfig.class, PathResolverTestConfig.class, KeyTestConfig.class,
        CryptoServiceTestConfig.class, BackendTestConfig.class })
@TestPropertySource(properties = {"tenantID=tenantId" })
public class MixDecValResourceIntegrationTest {
    public static final int NUM_VOTES = 10;

    @Autowired
    private ElGamalKeyPair electoralAuthorityKeyPair;

    @Autowired
    private KeyPair signingKeyPair;

    @Autowired
    private MixDecValResource sut = new MixDecValResource();

    @Autowired
    private PathResolver pathResolver;

    @Autowired
    private PlatformRootCAService platformRootCAService;

    @Autowired
    private MixDecValService mixDecValService;

    @BeforeClass
    public static void setup() {
        Security.addProvider(new BouncyCastleProvider());
    }

    /**
     * Perform a full mix-decrypt-validate cycle from the SDM and check that the
     * results match expectations.
     * 
     * @throws PayloadSignatureException
     * @throws CertificateManagementException 
     */
    @Test
    public void testMixDecVal()
            throws GeneralCryptoLibException, IOException, PayloadSignatureException, CertificateManagementException {
        // Create a directory structure with all the elements the mix-dec-val
        // service expects.
        prepareTestElectoralAuthority(electoralAuthorityKeyPair);
        List<List<ZpGroupElement>> votes = SDMTestData.generateVotes(NUM_VOTES);
        prepareTestEncryptedBallots(electoralAuthorityKeyPair, votes);
        prepareTestMixingPayloads(electoralAuthorityKeyPair, votes, SDMTestData.TENANT_ID,
            SDMTestData.ELECTION_EVENT_ID, SDMTestData.BALLOT_BOX_ID, SDMTestData.ELECTORAL_AUTHORITY_ID);

        // Prepare the input data.
        String electoralAuthorityPrivateKeyString = Base64.getEncoder()
            .encodeToString(electoralAuthorityKeyPair.getPrivateKeys().toJson().getBytes(Charset.forName("UTF-8")));
        MixingData mixingData =
            new MixingData(getPrivateKeyPEM(signingKeyPair.getPrivate()), electoralAuthorityPrivateKeyString);

        // Run the process.
        ResponseEntity<IdleState> response =
            sut.getBallotBox(SDMTestData.ELECTION_EVENT_ID, SDMTestData.BALLOT_BOX_ID, mixingData);

        // Examine the results.
        assertEquals(HttpStatus.OK, response.getStatusCode());

        // The decompressed votes must match the representations of the votes.
        // For that, the decompressed vote list is obtained.
        Path decompressedVotesFile = ValidationServiceImpl.getDecompressedVotesFilePath(getBallotBoxPath());
        // Each vote type is tallied.
        Map<String, Integer> outputVoteTally = tally(Files.readAllLines(decompressedVotesFile));
        // Then the same is done for the original votes. For each vote, get its
        // vote elements' values joined by a semicolon, in the same format as
        // the decompressed votes.
        List<String> voteAliases =
            votes
                .stream().map(voteElements -> voteElements.stream()
                    .map(voteElement -> voteElement.getValue().toString()).collect(Collectors.joining(";")))
                .collect(Collectors.toList());
        Map<String, Integer> inputVoteTally = tally(voteAliases);
        // Finally, both tallies are compared.
        outputVoteTally.forEach((option, voteCount) -> assertEquals(voteCount, inputVoteTally.get(option)));

        // There must not be auditable votes.
        Path auditableVotesFile = ValidationServiceImpl.getAuditableVotesFilePath(getBallotBoxPath());
        assertEquals(0, Files.readAllLines(auditableVotesFile).size());
    }

    /**
     * Return a list of occurrences of each of the strings.
     */
    private Map<String, Integer> tally(List<String> strings) {
        return strings.stream().collect(Collectors.groupingBy(Function.identity(), Collectors.summingInt(e -> 1)));
    }

    private void prepareTestEncryptedBallots(ElGamalKeyPair electoralAuthorityKeyPair, List<List<ZpGroupElement>> votes)
            throws IOException, GeneralCryptoLibException {
        Path ballotBoxPath = getBallotBoxPath();
        // Get the encrypted ballots JSON.
        String encryptedBallots =
            SDMTestData.generateEncryptedBallots(electoralAuthorityKeyPair.getPublicKeys(), votes);
        // Ensure the expected location exists.
        Files.createDirectories(ballotBoxPath);
        // Write the encrypted ballots JSON to its expected location.
        Files.write(ballotBoxPath.resolve(ConfigConstants.CONFIG_FILE_NAME_ENCRYPTED_BALLOT_BOX + ConfigConstants.CSV),
            encryptedBallots.getBytes(Charset.defaultCharset()));
    }

    private void prepareTestMixingPayloads(ElGamalKeyPair electoralAuthorityKeyPair, List<List<ZpGroupElement>> votes,
            String tenantId, String electionEventId, String ballotBoxId, String electoralAuthorityId)
            throws IOException, GeneralCryptoLibException, PayloadSignatureException, CertificateManagementException {
        Path ballotBoxPath = getBallotBoxPath();

        // The public part of the votes decryption keypair (used to encrypt the
        // votes)
        ElGamalPublicKey decryptionPublicKey = electoralAuthorityKeyPair.getPublicKeys();

        // Mock key with which the votes would have been encrypted before going
        // through the mixing
        ElGamalPublicKey combinedPublicKey = decryptionPublicKey.multiply(decryptionPublicKey);

        // Get the previous encrypted ballots.
        List<EncryptedVote> previousEncryptedBallots = SDMTestData.generateEncryptedVotes(combinedPublicKey, votes);
        // Get the encrypted ballots.
        List<EncryptedVote> encryptedBallots = SDMTestData.generateEncryptedVotes(decryptionPublicKey, votes);
        // Ensure the expected location exists.
        Files.createDirectories(ballotBoxPath);

        ZpSubgroup encryptionParameters = electoralAuthorityKeyPair.getPublicKeys().getGroup();

        int voteSetIndex = 0;
        BallotBoxId ballotBoxIdObject = new BallotBoxIdImpl(tenantId, electionEventId, ballotBoxId);
        VoteSetId voteSetId = new VoteSetIdImpl(ballotBoxIdObject, voteSetIndex);
        MixingPayload originalPayload = new MixingPayloadImpl(combinedPublicKey, voteSetId, encryptedBallots,
            electoralAuthorityId, encryptionParameters);

        MixingPayload mixedPayload = new MixingPayloadImpl(originalPayload, decryptionPublicKey, encryptedBallots,
            Collections.emptyList(), previousEncryptedBallots, "", Collections.emptyList());

        PayloadSignature mixedPayloadSignature = SDMTestData.signPayload(mixedPayload, signingKeyPair);
        mixedPayload.setSignature(mixedPayloadSignature);

        X509Certificate[] certChain = mixedPayloadSignature.getCertificateChain();

        // Save the last cert chain element as if it was the platformroot so it
        // can be retrieved for validation later
        platformRootCAService.save(certChain[certChain.length - 1]);

        String previousNodeName = "previous_node_name";

        // Write the encrypted ballots JSON to its expected location.
        ObjectMappers.toJson(ballotBoxPath.resolve(String.format("%s-%s.json", voteSetId.toString(), previousNodeName)),
            mixedPayload);

    }

    private void prepareTestElectoralAuthority(ElGamalKeyPair electoralAuthorityKeyPair)
            throws GeneralCryptoLibException, IOException {
        Path electoralAuthorityPath = pathResolver.resolve(ConfigConstants.CONFIG_FILES_BASE_DIR,
            SDMTestData.ELECTION_EVENT_ID, ConfigConstants.CONFIG_DIR_NAME_ONLINE,
            ConfigConstants.CONFIG_DIR_NAME_ELECTORAL_AUTHORITY, SDMTestData.ELECTORAL_AUTHORITY_ID);
        // Get the electoral authority JSON.
        String electoralAuthorityJson =
            SDMTestData.generateElectoralAuthorityJson(electoralAuthorityKeyPair.getPublicKeys());
        // Ensure the expected location exists.
        Files.createDirectories(electoralAuthorityPath);
        // Write the electoral authority JSON to its expected location.
        Files.write(electoralAuthorityPath.resolve(ConfigConstants.CONFIG_FILE_NAME_DECRYPTION_KEY_JSON),
            electoralAuthorityJson.getBytes(Charset.defaultCharset()));
    }

    private Path getBallotBoxPath() {
        return mixDecValService.getBallotBoxPath(SDMTestData.BALLOT_ID, SDMTestData.BALLOT_BOX_ID,
            SDMTestData.ELECTION_EVENT_ID);
    }

    /**
     * @param privateKey
     *            a private key
     * @return the private key in PEM format
     */
    private String getPrivateKeyPEM(PrivateKey privateKey) throws GeneralCryptoLibException {
        return PemUtils.privateKeyToPem(privateKey);
    }
}
