/*
 * 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.channel.cc.mixing.service;

import static org.hamcrest.CoreMatchers.equalTo;
import static org.hamcrest.CoreMatchers.is;
import static org.hamcrest.CoreMatchers.not;
import static org.hamcrest.core.IsCollectionContaining.hasItem;
import static org.junit.Assert.assertThat;
import static org.junit.Assert.assertTrue;
import static org.mockito.Mockito.mock;
import static org.mockito.Mockito.when;

import java.io.IOException;
import java.math.BigInteger;
import java.security.KeyManagementException;
import java.util.List;
import java.util.stream.Collectors;
import java.util.stream.Stream;

import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import com.fasterxml.jackson.databind.ObjectMapper;
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.products.ov.channel.cc.commons.keymanagement.KeyManager;
import com.scytl.products.ov.commons.beans.Ballot;
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.vote.CiphertextEncryptedVote;
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.MixingDTO;
import com.scytl.products.ov.commons.dto.MixingDTOImpl;
import com.scytl.products.ov.commons.dto.mixing.MixingPayload;
import com.scytl.products.ov.commons.dto.mixing.MixingPayloadImpl;
import com.scytl.products.ov.commons.sign.CryptoTestData;
import com.scytl.products.ov.commons.sign.TestVotingDataService;

public class TestUtil {
    private static final Logger logger = LoggerFactory.getLogger(TestUtil.class);

    private final InternalMixDecryptService sut;

    private final String queueName;

    private final KeyManager keyManager;

    private final ZpSubgroup encryptionParameters;

    private final TestVotingDataService votingTestData;

    TestUtil(InternalMixDecryptService sut, String queueName, KeyManager keyManager,
            ElGamalKeyPair electoralAuthorityKeyPair)
            throws GeneralCryptoLibException, IOException {
        this.sut = sut;
        this.queueName = queueName;
        this.keyManager = keyManager;
        // Generate the cryptographic data.
        encryptionParameters = CryptoTestData.getEncryptionParameters();

        // Generate the votes and encrypt them.
        votingTestData = new TestVotingDataService(getBallotDefinition(), electoralAuthorityKeyPair);
    }

    MixingDTO createMixingDTO(int numberOfVotes) throws KeyManagementException, GeneralCryptoLibException, IOException {
        ElGamalPublicKey voteEncryptionKey =
            keyManager.getMixingPublicKey(TestConfig.ELECTION_EVENT_ID, TestConfig.ELECTORAL_AUTHORITY_ID);
        List<EncryptedVote> votes = votingTestData.generateVotes(numberOfVotes).stream()
            .map(CiphertextEncryptedVote::new).collect(Collectors.toList());

        BallotBoxId ballotBoxId =
            new BallotBoxIdImpl(TestConfig.TENANT_ID_VALUE, TestConfig.ELECTION_EVENT_ID, TestConfig.BALLOT_BOX_ID);
        VoteSetId voteSetId = new VoteSetIdImpl(ballotBoxId, TestConfig.VOTESET_INDEX);
        // Create an input DTO.
        return createTestMixingDTO(encryptionParameters, voteEncryptionKey, voteSetId, votes);
    }

    void executeMixing(MixingDTO dto) {

        logger.info("Started mixing on thread {}", Thread.currentThread().getName());

        // Make sure the input DTO is as expected.
        assertThat(dto.getQueuesToVisit(), hasItem(queueName));

        // Run the mixing and decryption and return the results.
        sut.executeMixing(dto);
    }

    void testResults(MixingDTO outputMixingDTO, MixingPayload previousPayload) {

        // The node the DTO has has just been through should no longer be in the
        // list of nodes to visit.
        assertThat(outputMixingDTO.getQueuesToVisit(), not(hasItem(queueName)));
        // The last visited node should be the only one there is.
        assertThat(outputMixingDTO.getLastQueueName(), is(queueName));
        // The encryption key should have been altered.
        assertThat(previousPayload.getVoteEncryptionKey(),
            not(equalTo(outputMixingDTO.getPayload().getVoteEncryptionKey())));
        // The resulting encryption key should be the unitary key.
        ElGamalPublicKey resultingVoteEncryptionKey = outputMixingDTO.getPayload().getVoteEncryptionKey();
        assertThat(resultingVoteEncryptionKey.getKeys().get(0).getValue(), is(BigInteger.ONE));
        // The previous votes in the output match the input votes.
        previousPayload.getVotes()
            .forEach(vote -> assertTrue(outputMixingDTO.getPayload().getPreviousVotes().contains(vote)));
        // The previous vote encryption key in the output matches the input's.
        assertThat(outputMixingDTO.getPayload().getPreviousVoteEncryptionKey(),
            equalTo(previousPayload.getVoteEncryptionKey()));

        // Collect the original votes.
        List<BigInteger> unencryptedVotes = votingTestData.getUnencryptedVotes().stream().map(
            // For each vote, get its elements
            vote -> vote.stream()
                // Get each of the values
                .map(ZpGroupElement::getValue)
                // and multiply them
                .reduce(BigInteger.ONE, BigInteger::multiply))
            // and create a list of the results
            .collect(Collectors.toList());

        // Collect the phis.
        List<BigInteger> phis = outputMixingDTO.getPayload().getVotes().stream().map(
            // Get each vote's phis
            vote -> vote.getPhis().stream()
                // Multiply them
                .reduce(BigInteger.ONE, BigInteger::multiply))
            // and create a list of the results
            .collect(Collectors.toList());
        // Finally, make sure the decrypted phis match the pre-encryption votes.
        assertTrue(phis.containsAll(unencryptedVotes));
    }

    /**
     * Create a mixing DTO with test data to be used as the input of the mixing
     * and decryption service.
     *
     * @param encryptionParameters
     *            the encryption parameters
     * @param voteEncryptionKey
     *            the vote encryption key
     * @param voteSetId
     *            a reference to the original set of votes
     * @param votes
     *            the set of votes to mix
     * @return an input mixing DTO
     */
    private MixingDTO createTestMixingDTO(ZpSubgroup encryptionParameters, ElGamalPublicKey voteEncryptionKey,
            VoteSetId voteSetId, List<EncryptedVote> votes) {
        MixingPayload payload = new MixingPayloadImpl(voteEncryptionKey, voteSetId, votes,
            TestConfig.ELECTORAL_AUTHORITY_ID, encryptionParameters);
        PayloadSignature mockSignature = mock(PayloadSignature.class);
        when(mockSignature.getSignatureContents()).thenReturn(TestConfig.SIGNATURE_CONTENTS);
        payload.setSignature(mockSignature);

        return new MixingDTOImpl(TestConfig.REQUEST_ID, Stream.of(queueName).collect(Collectors.toSet()), payload);
    }

    /**
     * Get a ballot definition from a JSON file.
     *
     * @return the ballot definition
     */
    private static Ballot getBallotDefinition() throws IOException {
        return new ObjectMapper().readValue(
            TestVotingDataService.class.getClassLoader().getResourceAsStream(TestConfig.BALLOT_JSON), Ballot.class);
    }
}
