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

import com.scytl.cryptolib.api.exceptions.GeneralCryptoLibException;
import com.scytl.cryptolib.elgamal.bean.CiphertextImpl;
import com.scytl.cryptolib.elgamal.bean.ElGamalComputationsValues;
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.commons.beans.Ballot;
import com.scytl.products.ov.commons.beans.Contest;
import com.scytl.products.ov.commons.beans.ElectionOption;

import java.io.IOException;
import java.math.BigInteger;
import java.util.ArrayList;
import java.util.Collection;
import java.util.Collections;
import java.util.List;
import java.util.Random;
import java.util.stream.Collectors;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

/**
 * Voting data generator for tests.
 */
public class TestVotingDataService {

    private static final Logger logger = LoggerFactory.getLogger(TestVotingDataService.class);

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

    private static final int NANOSECONDS_PER_SECOND = 1_000_000_000;

    private final CryptoAPIElGamalEncrypter encrypter;

    private final ElGamalKeyPair electoralAuthorityKeyPair;

    private final Ballot ballot;

    private final List<List<ZpGroupElement>> unencryptedVotes = new ArrayList<>();

    /**
     * Constructs a voting data provider with a generated electoral authority
     * key pair.
     * 
     * @param ballot
     *            the definition of the ballot
     * @throws GeneralCryptoLibException
     */
    public TestVotingDataService(Ballot ballot) throws GeneralCryptoLibException {
        // Create voting data.
        this(ballot, CryptoTestData.generateElGamalKeyPair(1));
    }

    /**
     * Constructs voting test data with a pre-defined electoral authority key
     * pair.
     * 
     * @param ballot
     *            the definition of the ballot
     * @param electoralAuthorityKeyPair
     *            the pre-defined electoral authority key pair
     * @throws GeneralCryptoLibException
     */
    public TestVotingDataService(Ballot ballot, ElGamalKeyPair electoralAuthorityKeyPair)
            throws GeneralCryptoLibException {
        this.ballot = ballot;
        this.electoralAuthorityKeyPair = electoralAuthorityKeyPair;
        // Create vote encrypter.
        encrypter = new ElGamalService().createEncrypter(electoralAuthorityKeyPair.getPublicKeys());
    }

    /**
     * @return the generated vote encryption key.
     */
    public ElGamalPublicKey getVoteEncryptionKey() {
        return electoralAuthorityKeyPair.getPublicKeys();
    }

    /**
     * Generate votes.
     *
     * @param numVotes
     *            the number of votes to generate
     * @return a list of votes
     * @throws GeneralCryptoLibException
     * @throws IOException
     */
    public List<Ciphertext> generateVotes(int numVotes) throws GeneralCryptoLibException, IOException {
        logger.info("Generating {} votes...", numVotes);

        // Store the Zp subgroup locally.
        ZpSubgroup zpSubgroup = CryptoTestData.getZpSubgroup();

        // Extract the representations from each of the options in each of the
        // contests.
        logger.info("Getting representations...");
        List<String> representations = ballot.getContests().stream().map(Contest::getOptions)
            .flatMap(Collection::stream).map(ElectionOption::getRepresentation).collect(Collectors.toList());
        logger.info("{} representations collected", representations.size());

        logger.info("Voting...");
        long startTime = System.nanoTime();

        // Create a list of votes which is composed of a random choice of
        // representations.
        List<Ciphertext> votes = new ArrayList<>(numVotes);
        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));
            }

            unencryptedVotes.add(vote);

            // Record the vote.
            Ciphertext encryptedVote = encryptVote(electoralAuthorityKeyPair.getPublicKeys(), vote);
            votes.add(encryptedVote);

            if (i % 100 == 0 && i > 1) {
                long elapsedTime = System.nanoTime() - startTime;
                // Extrapolate ETA from current stats.
                double secondsPerItem = (double) elapsedTime / i / NANOSECONDS_PER_SECOND;
                long remainingItems = numVotes - i;
                long eta = Math.round(secondsPerItem * remainingItems);
                long itemsPerSecond = Math.round(1 / secondsPerItem);
                logger.info("{} votes encrypted, {} to go, ~{} votes/s, (ETA ~{} seconds)", i, remainingItems,
                    itemsPerSecond, eta);
            }
        }

        logger.info("{} votes ready", votes.size());

        return votes;
    }

    /**
     * Encrypt a vote.
     *
     * @param electoralAuthorityPublicKey
     *            the public key the votes will be encrypted with
     * @return a list of CSV entries
     */
    private Ciphertext encryptVote(ElGamalPublicKey electoralAuthorityPublicKey, List<ZpGroupElement> vote)
            throws GeneralCryptoLibException {
        GroupElementsCompressor<ZpGroupElement> compressor = new GroupElementsCompressor<>();

        // Cast vote.
        ZpGroupElement compressedVote = compressor.compress(vote);

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

        ElGamalComputationsValues computationsValues =
            encrypter.encryptGroupElements(compressedMessage).getComputationValues();

        // Convert the ElGamal computation values to an encrypted vote.
        List<ZpGroupElement> values = computationsValues.getPhis();
        List<ZpGroupElement> elements = new ArrayList<>(computationsValues.getValues().size());
        for (ZpGroupElement zpGroupElement : values) {
            ZpGroupElement element = new ZpGroupElement(zpGroupElement.getValue(), CryptoTestData.getZpSubgroup());
            elements.add(element);
        }

        return new CiphertextImpl(computationsValues.getGamma(), elements);
    }

    public List<List<ZpGroupElement>> getUnencryptedVotes() {
        return unencryptedVotes;
    }
}
