/*
 * 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 java.io.ByteArrayInputStream;
import java.io.ByteArrayOutputStream;
import java.io.OutputStream;
import java.nio.charset.StandardCharsets;
import java.security.KeyManagementException;
import java.util.ArrayList;
import java.util.List;
import java.util.stream.Collectors;

import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Service;

import com.scytl.cryptolib.api.exceptions.GeneralCryptoLibException;
import com.scytl.cryptolib.elgamal.bean.ElGamalPrivateKey;
import com.scytl.cryptolib.elgamal.bean.ElGamalPublicKey;
import com.scytl.cryptolib.mathematical.groups.activity.ExponentCompressor;
import com.scytl.cryptolib.mathematical.groups.activity.GroupElementsCompressor;
import com.scytl.cryptolib.mathematical.groups.impl.Exponent;
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.KeyNotFoundException;
import com.scytl.products.ov.commons.dto.MixingDTO;
import com.scytl.products.ov.commons.dto.mixing.MixingPayload;
import com.scytl.products.ov.mixnet.commons.configuration.ApplicationConfig;
import com.scytl.products.ov.mixnet.commons.configuration.MixerConfig;
import com.scytl.products.ov.mixnet.commons.configuration.VerifierConfig;
import com.scytl.products.ov.mixnet.commons.preprocess.Partitioner;

@Service
public class MixingPropertiesProviderImpl implements MixingPropertiesProvider {

    private static final Logger LOG = LoggerFactory.getLogger("std");

    @Autowired
    private KeyRepository keyRepository;

    private MixerConfig fillMixerConfig(ApplicationConfig applicationConfig, MixingDTO mixingDTO)
            throws KeyNotFoundException, KeyManagementException, GeneralCryptoLibException {

        int desiredKeySize = 1;
        if ((mixingDTO.getPayload().getVotes() != null) && (!mixingDTO.getPayload().getVotes().isEmpty())) {
            // all the votes (ciphertexts) in the ballot box will have the same number of phis
            desiredKeySize = mixingDTO.getPayload().getVotes().get(0).getPhis().size();
        } else {
            desiredKeySize = applicationConfig.getNumOptions();
        }
        
        MixerConfig mixer = new MixerConfig();
        mixer.setMix(true);
        mixer.setWarmUpSize(1);

        LOG.info("Preparing mixing config");

        MixingPayload payload = mixingDTO.getPayload();
        // Get the vote decryption key.
        ElGamalPrivateKey ballotDecryptionKey = keyRepository.getBallotDecryptionKey(
                payload.getVoteSetId().getBallotBoxId().getElectionEventId(),
                payload.getElectoralAuthorityId());

        ballotDecryptionKey =
            compressPrivateKeyIfNecessary(ballotDecryptionKey, desiredKeySize);
        mixer.setBallotDecryptionKeyInput(new ByteArrayInputStream(ballotDecryptionKey.toJson().getBytes()));

        ElGamalPublicKey encryptionPublicKey = payload.getVoteEncryptionKey();
        encryptionPublicKey = compressPublicKeyIfNecessary(encryptionPublicKey, desiredKeySize);

        // Convert the vote encryption key to an input stream.
        mixer.setPublicKeyInput(new ByteArrayInputStream(
                encryptionPublicKey.toJson().getBytes(StandardCharsets.UTF_8)));

        // Convert the vote set to an input stream.
        mixer.setEncryptedBallotsInput(new ByteArrayInputStream(
                payload.getVotes().stream().map(vote -> vote.toString())
                        .collect(Collectors.joining(System.lineSeparator()))
                        .getBytes(StandardCharsets.UTF_8)));

        // Convert the encryption parameters to an input stream.
        mixer.setEncryptionParametersInput(new ByteArrayInputStream(
                payload.getEncryptionParameters().toJson().getBytes(StandardCharsets.UTF_8)));

        int numBatches = new Partitioner(applicationConfig.getMinimumBatchSize())
                .generatePartitionsOfMinimumSize(payload.getVotes().size()).size();
        List<OutputStream> commitmentParametersOutputOfBatch = new ArrayList<>(numBatches);
        List<OutputStream> decryptedVotesOutputOfBatch = new ArrayList<>(numBatches);
        List<OutputStream> encryptedBallotsOutputOfBatch = new ArrayList<>(numBatches);
        List<OutputStream> encryptionParametersOutputOfBatch = new ArrayList<>(numBatches);
        List<OutputStream> proofsOutputOfBatch = new ArrayList<>(numBatches);
        List<OutputStream> publicKeyOutputOfBatch = new ArrayList<>(numBatches);
        List<OutputStream> reEncryptedBallotsOutputOfBatch = new ArrayList<>(numBatches);
        List<OutputStream> votesWithProofOutputOfBatch = new ArrayList<>(numBatches);
        for (int i = 0; i < numBatches; i++) {
            commitmentParametersOutputOfBatch.add(new ByteArrayOutputStream());
            decryptedVotesOutputOfBatch.add(new ByteArrayOutputStream());
            encryptedBallotsOutputOfBatch.add(new ByteArrayOutputStream());
            encryptionParametersOutputOfBatch.add(new ByteArrayOutputStream());
            proofsOutputOfBatch.add(new ByteArrayOutputStream());
            publicKeyOutputOfBatch.add(new ByteArrayOutputStream());
            reEncryptedBallotsOutputOfBatch.add(new ByteArrayOutputStream());
            votesWithProofOutputOfBatch.add(new ByteArrayOutputStream());
        }

        mixer.setCommitmentParametersOutputOfBatch(commitmentParametersOutputOfBatch);
        mixer.setDecryptedVotesOutputOfBatch(decryptedVotesOutputOfBatch);
        mixer.setEncryptedBallotsOutputOfBatch(encryptedBallotsOutputOfBatch);
        mixer.setEncryptionParametersOutputOfBatch(encryptionParametersOutputOfBatch);
        mixer.setProofsOutputOfBatch(proofsOutputOfBatch);
        mixer.setPublicKeyOutputOfBatch(publicKeyOutputOfBatch);
        mixer.setReEncryptedBallotsOutputOfBatch(reEncryptedBallotsOutputOfBatch);
        mixer.setVotesWithProofOutputOfBatch(votesWithProofOutputOfBatch);
        return mixer;
    }

    @Override
    public ApplicationConfig createApplicationConfigforBallotBox(MixingDTO mixDecDto)
            throws KeyNotFoundException, KeyManagementException, GeneralCryptoLibException {

        LOG.info("Preparing mixing config");

        ApplicationConfig applicationConfig = new ApplicationConfig();
        applicationConfig.setConcurrencyLevel(16);
        applicationConfig.setMinimumBatchSize(64);
        applicationConfig.setNumberOfParallelBatches(1);
        int defaultNumOptions = 1;
        applicationConfig.setNumOptions(defaultNumOptions);
        applicationConfig.setMixer(fillMixerConfig(applicationConfig, mixDecDto));
        VerifierConfig verifier = new VerifierConfig();
        applicationConfig.setVerifier(verifier);
        return applicationConfig;
    }

    private static ElGamalPrivateKey compressPrivateKeyIfNecessary(final ElGamalPrivateKey elGamalPrivateKey,
            final int numOptions) throws GeneralCryptoLibException {

        LOG.info("Checking if its necessary to compress the ballot decryption key...");

        List<Exponent> keys = elGamalPrivateKey.getKeys();

        if (keys.size() <= numOptions) {
            LOG.info("No need to compress the ballot decryption key");
            return elGamalPrivateKey;
        }
        LOG.info(
            String.format("Compressing the ballot decryption key. Size: %s, required: %s", keys.size(), numOptions));

        ExponentCompressor<ZpSubgroup> compressor = new ExponentCompressor<>(elGamalPrivateKey.getGroup());

        List<Exponent> listWithCompressedFinalElement =
            compressor.buildListWithCompressedFinalElement(numOptions, keys);

        return new ElGamalPrivateKey(listWithCompressedFinalElement, elGamalPrivateKey.getGroup());
    }

    private static ElGamalPublicKey compressPublicKeyIfNecessary(ElGamalPublicKey elGamalPublicKey, int numOptions)
            throws GeneralCryptoLibException {

        LOG.info("Checking if its necessary to compress the mixing public key...");

        List<ZpGroupElement> keys = elGamalPublicKey.getKeys();

        if (keys.size() <= numOptions) {
            LOG.info("No need to compress the mixing public key");
            return elGamalPublicKey;
        }
        LOG.info(String.format("Compressing the mixing public key. Size: %s, required: %s", keys.size(), numOptions));

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

        List<ZpGroupElement> listWithCompressedFinalElement =
            compressor.buildListWithCompressedFinalElement(numOptions, keys);

        return new ElGamalPublicKey(listWithCompressedFinalElement, elGamalPublicKey.getGroup());
    }

}
