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

import static java.util.Objects.requireNonNull;

import java.io.IOException;
import java.io.InputStream;
import java.util.ArrayList;
import java.util.LinkedList;
import java.util.List;
import java.util.concurrent.Callable;
import java.util.concurrent.ExecutionException;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.util.concurrent.ForkJoinPool;
import java.util.concurrent.Future;

import org.apache.commons.io.IOUtils;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import com.scytl.cryptolib.api.exceptions.GeneralCryptoLibException;
import com.scytl.cryptolib.api.primitives.PrimitivesServiceAPI;
import com.scytl.cryptolib.elgamal.bean.ElGamalPrivateKey;
import com.scytl.cryptolib.elgamal.bean.ElGamalPublicKey;
import com.scytl.cryptolib.elgamal.cryptoapi.Ciphertext;
import com.scytl.cryptolib.mathematical.groups.impl.ZpSubgroup;
import com.scytl.products.oscore.logging.api.domain.Level;
import com.scytl.products.oscore.logging.api.domain.LogContent;
import com.scytl.products.oscore.logging.api.writer.SecureLoggingWriter;
import com.scytl.products.oscore.logging.core.provider.TransactionInfoProvider;
import com.scytl.products.ov.mixnet.batches.MixingBatch;
import com.scytl.products.ov.mixnet.batches.MixingBatchFactory;
import com.scytl.products.ov.mixnet.batches.MixingResult;
import com.scytl.products.ov.mixnet.batches.NonMixingBatch;
import com.scytl.products.ov.mixnet.batches.NonMixingBatchFactory;
import com.scytl.products.ov.mixnet.commons.ballots.ElGamalEncryptedBallots;
import com.scytl.products.ov.mixnet.commons.configuration.ApplicationConfig;
import com.scytl.products.ov.mixnet.commons.configuration.BGParams;
import com.scytl.products.ov.mixnet.commons.configuration.BGParamsProvider;
import com.scytl.products.ov.mixnet.commons.configuration.locations.LocationsProvider;
import com.scytl.products.ov.mixnet.commons.configuration.locations.MixerLocationsProvider;
import com.scytl.products.ov.mixnet.commons.exceptions.MixingException;
import com.scytl.products.ov.mixnet.commons.homomorphic.Randomness;
import com.scytl.products.ov.mixnet.commons.homomorphic.impl.GjosteenElGamal;
import com.scytl.products.ov.mixnet.commons.io.ElGamalEncryptedBallotsLoader;
import com.scytl.products.ov.mixnet.commons.io.ElgamalPublicKeyReader;
import com.scytl.products.ov.mixnet.commons.io.JSONProofsWriter;
import com.scytl.products.ov.mixnet.commons.io.ZpGroupReader;
import com.scytl.products.ov.mixnet.commons.preprocess.Partitioner;
import com.scytl.products.ov.mixnet.commons.tools.CiphertextTools;
import com.scytl.products.ov.mixnet.commons.tools.MultiExponentiation;
import com.scytl.products.ov.mixnet.commons.tools.MultiExponentiationImpl;
import com.scytl.products.ov.mixnet.commons.validation.EncryptedBallotsValidator;
import com.scytl.products.ov.mixnet.commons.validation.EncryptedBallotsValidatorManager;
import com.scytl.products.ov.mixnet.infrastructure.log.MixingAndVerifyingLogEvents;
import com.scytl.products.ov.mixnet.proofs.bg.multiexp.ParallelComputeAllE;
import com.scytl.products.ov.mixnet.shuffler.ElGamalShuffler;
import com.scytl.products.ov.mixnet.shuffler.permutation.PermutationGenerator;
import com.scytl.products.ov.mixnet.shuffler.permutation.Permutator;

/**
 * This class is responsible for the concurrent execution of the mixing. It arranges the parallelization according to a
 * specified batch size.
 */
public class BGMixnetLoadBalancer implements MixnetLoadBalancer {

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

    private static final int DEFAULT_BATCH = 1;

    private final BGMixerFactory mixerFactory;

    private final NonMixer nonMixer;

    private final BGParamsProvider bGParamsProvider;

    private final SecureLoggingWriter loggingWriter;

    private final ElGamalEncryptedBallotsLoader elgamalEncryptedBallotsLoader;

    private final PrimitivesServiceAPI primitivesService;

    /**
     * Constructor. For internal use only.
     * 
     * @param mixerFactory
     * @param nonMixer
     * @param bGParamsProvider
     * @param elgamalEncryptedBallotsLoader
     * @param loggingWriter
     * @param primitivesService
     */
    BGMixnetLoadBalancer(final BGMixerFactory mixerFactory, final NonMixer nonMixer, final BGParamsProvider bGParamsProvider,
            final ElGamalEncryptedBallotsLoader elgamalEncryptedBallotsLoader, SecureLoggingWriter loggingWriter,
            PrimitivesServiceAPI primitivesService) {
        this.mixerFactory = mixerFactory;
        this.nonMixer = nonMixer;
        this.bGParamsProvider = bGParamsProvider;
        this.elgamalEncryptedBallotsLoader = elgamalEncryptedBallotsLoader;
        this.loggingWriter = loggingWriter;
        this.primitivesService = primitivesService;
    }

    @Override
    public void balance(final ApplicationConfig applicationConfig, final LocationsProvider locationsProvider,
            final Randomness[] randomExponents, final List<Ciphertext> preComputations)
            throws MixingException, IOException {

        final ZpSubgroup zpGroup =
            ZpGroupReader.build(locationsProvider.getMixerLocationsProvider().getEncryptionParametersInput());
        ForkJoinPool executor = new ForkJoinPool(applicationConfig.getConcurrencyLevel());
        BGMixer mixer = mixerFactory.newBGMixer(executor);
        try {
            final boolean mixingSuccess = splitInBatchesAndExecute(applicationConfig, locationsProvider,
                randomExponents, preComputations, zpGroup, mixer);
            logMixingBalanceResult(mixingSuccess);
        } catch (GeneralCryptoLibException e) {
            throw new MixingException(e);
        } finally {
            executor.shutdown();
        }

    }

    private void logMixingBalanceResult(final boolean result) {
        if (result) {

            loggingWriter.log(Level.INFO,
                new LogContent.LogContentBuilder()
                    .logEvent(MixingAndVerifyingLogEvents.SHUFREENCR_SUCCESS_GENERATING_MIXED_BALLOT_BOX)
                    .user("adminID").electionEvent("-").additionalInfo("mbb_id", "-").createLogInfo());

            LOG.info("All batches concluded successfully");
        } else {

            loggingWriter.log(Level.ERROR, new LogContent.LogContentBuilder()
                .logEvent(MixingAndVerifyingLogEvents.SHUFREENCR_ERROR_GENERATING_MIXED_BALLOT_BOX).user("adminID")
                .electionEvent("-").additionalInfo("mbb_id", "-")
                .additionalInfo("err_desc",
                    "Not all of the encrypted ballots were mixed successfully, one or more batches " + "had a problem")
                .createLogInfo());

            LOG.error("Some errors occurred during the batches execution");
        }
    }

    /**
     * This method waits for thread completion.
     * @param mixer 
     */
    private boolean splitInBatchesAndExecute(final ApplicationConfig applicationConfig,
            final LocationsProvider locationsProvider, final Randomness[] randomExponents,
            final List<Ciphertext> preComputations, final ZpSubgroup zpGroup, BGMixer mixer)
            throws IOException, GeneralCryptoLibException {

        final ElGamalEncryptedBallots encryptedBallots = loadEncryptedBallots(locationsProvider, zpGroup);
        final int sizeEncryptedBallots = encryptedBallots.getBallots().size();

        final int numberOfThreads = applicationConfig.getNumberOfParallelBatches();

        final List<Callable<MixingResult>> batches;
        if (sizeEncryptedBallots == 0 || sizeEncryptedBallots == 1) {

            LOG.info("Generating partitions...");

            int minBatchSize = applicationConfig.getMinimumBatchSize();

            LOG.info("Using partition size: " + sizeEncryptedBallots);
            LOG.info("Using minimum batch size: " + minBatchSize);

            final List<Integer> partitions =
                new Partitioner(minBatchSize).generatePartitionsOfMinimumSize(sizeEncryptedBallots);

            LOG.info("Generated " + partitions.size() + " partitions " + partitions);

            logPartitionSize(partitions);

            final ElGamalPublicKey publicKey = ElgamalPublicKeyReader
                .readPublicKeyFromStream(locationsProvider.getMixerLocationsProvider().getPublicKeyInput());

            final NonMixingBatchFactory nonMixingBatchFactory = NonMixingBatchFactory.aNew(nonMixer, zpGroup,
                publicKey, getBallotDecryptionKey(locationsProvider.getMixerLocationsProvider()), locationsProvider,
                primitivesService);

            LOG.info("Preparing batch...");

            final int batchNum = 1;
            final int offset = 0;

            batches = new ArrayList<>();

            final int partitionSize = partitions.size();

            final ElGamalEncryptedBallots batchElGamalEncryptedBallots = encryptedBallots;

            final BGParams batchBGParams = bGParamsProvider.getParamsForGiven(partitionSize);

            final NonMixingBatch batch = nonMixingBatchFactory.create(partitionSize, batchElGamalEncryptedBallots,
                batchNum, offset, batchBGParams, DEFAULT_BATCH);

            batches.add(batch);
        } else {

            LOG.info("Generating partitions...");

            int minBatchSize = applicationConfig.getMinimumBatchSize();

            LOG.info("Using partition size: " + sizeEncryptedBallots);
            LOG.info("Using minimum batch size: " + minBatchSize);

            final List<Integer> partitions =
                new Partitioner(minBatchSize).generatePartitionsOfMinimumSize(sizeEncryptedBallots);

            LOG.info("Generated " + partitions.size() + " partitions " + partitions);

            logPartitionSize(partitions);

            final int batchNum = 1;
            final int offset = 0;
            final int concurrencyLevel = applicationConfig.getConcurrencyLevel();

            final ElGamalPublicKey publicKey = ElgamalPublicKeyReader
                .readPublicKeyFromStream(locationsProvider.getMixerLocationsProvider().getPublicKeyInput());

            final GjosteenElGamal gjosteenElGamal = new GjosteenElGamal(zpGroup, publicKey);
            final MultiExponentiation multiExponentiation = MultiExponentiationImpl.getInstance();
            final CiphertextTools ciphertextTools = new CiphertextTools(multiExponentiation);
            final ParallelComputeAllE parallelComputeAllE = new ParallelComputeAllE(ciphertextTools);

            final MixingBatchFactory mixingBatchFactory = MixingBatchFactory.aNew(mixer, zpGroup, gjosteenElGamal,
                concurrencyLevel, publicKey, getBallotDecryptionKey(locationsProvider.getMixerLocationsProvider()),
                locationsProvider, multiExponentiation, parallelComputeAllE, primitivesService);

            LOG.info("Preparing batches...");

            batches = prepareBatches(randomExponents, preComputations, encryptedBallots, partitions, batchNum, offset,
                mixingBatchFactory);

        }
        LOG.info("Executing mixing batches on a thread pool with size " + numberOfThreads);

        return executeBatchesAndAwaitProcessing(numberOfThreads, batches);
    }

    private List<Callable<MixingResult>> prepareBatches(final Randomness[] randomExponents,
            final List<Ciphertext> preComputations, final ElGamalEncryptedBallots encryptedBallots,
            final List<Integer> partitions, int firstBatchNum, int firstOffset, final MixingBatchFactory mixingBatchFactory) {

        final List<Callable<MixingResult>> batches = new ArrayList<>();
        int batchNum = firstBatchNum;
        int offset = firstOffset;
        for (int i = partitions.size() - 1; i >= 0; i--) {

            final int partitionSize = partitions.get(i);

            final ElGamalEncryptedBallots batchElGamalEncryptedBallots =
                getBatchElGamalEncryptedBallots(encryptedBallots, offset, partitionSize);

            final Randomness[] batchRandomness = getBatchRandomness(randomExponents, offset, partitionSize);

            final List<Ciphertext> batchPreComputations =
                getBatchPreComputations(preComputations, offset, partitionSize);

            final BGParams batchBGParams = bGParamsProvider.getParamsForGiven(partitionSize);

            final MixingBatch batch = mixingBatchFactory.create(partitionSize, batchElGamalEncryptedBallots, batchNum,
                offset, batchBGParams, i, batchRandomness, batchPreComputations);

            batches.add(batch);

            offset += partitionSize;
            batchNum++;

        }
        return batches;
    }

    private boolean executeBatchesAndAwaitProcessing(final int numberOfThreads,
            final List<Callable<MixingResult>> batches) {

        final ExecutorService executorService = Executors.newFixedThreadPool(numberOfThreads);
        final List<Future<MixingResult>> threads = new ArrayList<>();

        for (final Callable<MixingResult> current : batches) {
            final Future<MixingResult> mixingThread = executorService.submit(current);
            threads.add(mixingThread);
        }

        executorService.shutdown();
        return awaitForBatchesTermination(threads);
    }

    private static boolean awaitForBatchesTermination(final List<Future<MixingResult>> threads) {
        boolean result = true;

        for (final Future<MixingResult> mixingThread : threads) {
            try {
                final MixingResult mixingResult = mixingThread.get();
                final boolean batchResultPositive = mixingResult.getResult();
                if (!batchResultPositive) {
                    final int rangeInit = mixingResult.getOffset();
                    final int rangeEnd = mixingResult.getOffset() + mixingResult.getPartitionSize();
                    LOG.error("An error occurred during the execution of batch number " + mixingResult.getBatchNum()
                        + " which was mixing the range[" + rangeInit + ", " + rangeEnd + "]");
                }

                result &= mixingResult.getResult();

            } catch (ExecutionException e) {

                Throwable cause = e.getCause();
                String errorMsg = "An error occurred during the execution of a mixing thread";
                LOG.error(errorMsg + ". Exception message: " + cause.getMessage(), e);
                throw new MixingException(errorMsg, cause);
            } catch (InterruptedException e) {
                final String errorMsg = "Got interrupted while mixing.";
                LOG.warn(errorMsg);
                Thread.currentThread().interrupt();
                throw new MixingException(errorMsg, e);
            }
        }
        return result;
    }

    private static ElGamalEncryptedBallots getBatchElGamalEncryptedBallots(final ElGamalEncryptedBallots encryptedBallots,
            final int offset, final Integer partitionSize) {
        final List<Ciphertext> batchListEncryptedBallots =
            encryptedBallots.getBallots().subList(offset, offset + partitionSize);

        return new ElGamalEncryptedBallots(batchListEncryptedBallots);
    }

    private Randomness[] getBatchRandomness(final Randomness[] randomExponents, final int offset,
            final Integer partitionSize) {

        Randomness[] returnValue = null;
        if (randomExponents != null && randomExponents.length > offset) {

            int remainingAvailableRandomExponents = getRemainingAvailableRandomExponents(randomExponents, offset);

            if (remainingAvailableRandomExponents >= partitionSize) {
                returnValue = new Randomness[partitionSize];
                System.arraycopy(randomExponents, offset, returnValue, 0, partitionSize);
            } else {
                returnValue = new Randomness[remainingAvailableRandomExponents];
                System.arraycopy(randomExponents, offset, returnValue, 0, remainingAvailableRandomExponents);
            }
        }

        return returnValue;
    }

    private static int getRemainingAvailableRandomExponents(final Randomness[] randomExponents, final int offset) {
        return randomExponents.length - offset;
    }

    private static List<Ciphertext> getBatchPreComputations(final List<Ciphertext> preComputations, final int offset,
            final Integer partitionSize) {

        List<Ciphertext> batchPreComputations = null;
        if (preComputations != null && preComputations.size() > offset) {
            // Check that there are enough precomputations to take them from the
            // offset.
            if ((preComputations.size() - offset) >= partitionSize) {
                batchPreComputations = preComputations.subList(offset, offset + partitionSize);
            } else {
                batchPreComputations = preComputations.subList(offset, preComputations.size());
            }
        }

        return batchPreComputations;
    }

    private ElGamalEncryptedBallots loadEncryptedBallots(final LocationsProvider locationsProvider,
            final ZpSubgroup zpGroupParams) {
        ElGamalEncryptedBallots encryptedBallots;

        try {
            encryptedBallots = elgamalEncryptedBallotsLoader.loadCSV(zpGroupParams,
                locationsProvider.getMixerLocationsProvider().getEncryptedBallotsInput());
        } catch (IOException e) {
            LOG.error("An error occurred while loading the encrypted ballots");
            throw new IllegalArgumentException("An error occurred while loading the encrypted ballots", e);
        }
        return encryptedBallots;
    }

    private void logPartitionSize(final List<Integer> partitions) {
        final int numBatches = partitions.size();

        if (numBatches == 1) {
            log1batch();
        } else {
            logManyBatches(numBatches);
        }
    }

    private void logManyBatches(final int numBatches) {
        loggingWriter.log(Level.INFO,
            new LogContent.LogContentBuilder()
                .logEvent(MixingAndVerifyingLogEvents.SHUFREENCR_STARTING_MIXING_MULTIPLE_BATCHES).user("adminID")
                .electionEvent("-").additionalInfo("mbb_id", "-").additionalInfo("num_batches", Integer.toString(numBatches))
                .createLogInfo());
    }

    private void log1batch() {
        loggingWriter.log(Level.INFO,
            new LogContent.LogContentBuilder()
                .logEvent(MixingAndVerifyingLogEvents.SHUFREENCR_STARTING_MIXING_SINGLE_BATCH).user("adminID")
                .electionEvent("-").additionalInfo("mbb_id", "-").createLogInfo());
    }

    /**
     * Load the ballot decryption key from the appropriate locations provider.
     * 
     * @param mixerLocationsProvider
     *            the mixer's locations provider.
     * @return an ElGamal private key bean
     */
    public ElGamalPrivateKey getBallotDecryptionKey(MixerLocationsProvider mixerLocationsProvider) {
        try (InputStream is = mixerLocationsProvider.getBallotDecryptionKeyInput()) {
            return ElGamalPrivateKey.fromJson(IOUtils.toString(is));
        } catch (IOException | GeneralCryptoLibException e) {
            throw new MixingException("Failed to load the ballot decryption key", e);
        }
    }
    
    /**
     * Builder for creating {@link BGMixnetLoadBalancer} instances.
     */
    public static final class Builder {
        private final List<EncryptedBallotsValidator> ballotsValidators = new LinkedList<>();
        private PrimitivesServiceAPI primitivesService;
        private BallotDecryptionService ballotDecryptionService;
        private TransactionInfoProvider transactionInfoProvider; 
        private SecureLoggingWriter secureLoggingWriter;
        
        /**
         * Builds a new {@link BGMixnetLoadBalancer} instance.
         * @return a new instance.
         */
        public BGMixnetLoadBalancer build() {
            requireNonNull(primitivesService, "Primitives service is null.");
            requireNonNull(ballotDecryptionService, "Ballot decryption service is null.");
            requireNonNull(transactionInfoProvider, "Transaction info provider is null.");
            requireNonNull(secureLoggingWriter, "Secure logging writer is null.");
            
            BGMixerFactory mixerFactory = newBGMixerFactory();
            NonMixer nonMixer = new NonMixer(ballotDecryptionService);
            BGParamsProvider bGParamsProvider = new BGParamsProvider();
            ElGamalEncryptedBallotsLoader elgamalEncryptedBallotsLoader = 
                    newElGamalEncryptedBallotsLoader();
            return new BGMixnetLoadBalancer(mixerFactory, nonMixer, bGParamsProvider, 
                elgamalEncryptedBallotsLoader, secureLoggingWriter, primitivesService);
        }
        
        /**
         * Adds a given ballots validator.
         * 
         * @param ballotsValidator the validator
         * @return this instance.
         */
        public Builder addBallotsValidator(EncryptedBallotsValidator ballotsValidator) {
            requireNonNull(ballotsValidator, "Ballots validator is null.");
            ballotsValidators.add(ballotsValidator);
            return this;
        }
        
        /**
         * Sets the primitives service.
         * 
         * @param primitivesService the service
         * @return this instance.
         */
        public Builder setPrimitivesService(
                PrimitivesServiceAPI primitivesService) {
            this.primitivesService = primitivesService;
            return this;
        }
        
        /**
         * Sets the ballot decryption service.
         * 
         * @param ballotDecryptionService the service
         * @return this instance.
         */
        public Builder setBallotDecryptionService(
                BallotDecryptionService ballotDecryptionService) {
            this.ballotDecryptionService = ballotDecryptionService;
            return this;
        }
        
        /**
         * Sets the transaction info provider.
         * 
         * @param transactionInfoProvider the provider
         * @return this instance.
         */
        public Builder setTransactionInfoProvider(
                TransactionInfoProvider transactionInfoProvider) {
            this.transactionInfoProvider = transactionInfoProvider;
            return this;
        }
        
        /**
         * Sets the logging writer.
         * 
         * @param loggingWriter the writer
         * @return this instance.
         */
        public Builder setLoggingWriter(SecureLoggingWriter loggingWriter) {
            this.secureLoggingWriter = loggingWriter;
            return this;
        }

        private ElGamalEncryptedBallotsLoader newElGamalEncryptedBallotsLoader() {
            EncryptedBallotsValidator validator = 
                    new EncryptedBallotsValidatorManager(new ArrayList<>(ballotsValidators));
            return new ElGamalEncryptedBallotsLoader(validator);
        }

        private BGMixerFactory newBGMixerFactory() {
            ElGamalShuffler shuffler = newElGamalShuffler();
            JSONProofsWriter proofsWriter = new JSONProofsWriter();
            ShuffleProofService shuffleProofService = 
                    new SimpleShuffleProofService(transactionInfoProvider);
            return new BGMixerFactoryImpl(shuffler, proofsWriter, ballotDecryptionService, 
                shuffleProofService);
        }

        private static ElGamalShuffler newElGamalShuffler() {
            Permutator permutator = new Permutator();
            PermutationGenerator permutationGenerator = new PermutationGenerator();
            return new ElGamalShuffler(permutator, permutationGenerator);
        }
    }
}
