/*
 * 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 java.io.IOException;
import java.io.OutputStream;
import java.util.Base64;
import java.util.List;
import java.util.concurrent.ForkJoinPool;
import java.util.concurrent.ForkJoinTask;

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.factory.SecureLoggingFactoryLog4j;
import com.scytl.products.oscore.logging.core.formatter.SplunkFormatter;
import com.scytl.products.oscore.logging.core.provider.TransactionInfoProvider;
import com.scytl.products.ov.mixnet.commons.ballots.ElGamalEncryptedBallots;
import com.scytl.products.ov.mixnet.commons.beans.ShuffleOutput;
import com.scytl.products.ov.mixnet.commons.beans.VoteWithProof;
import com.scytl.products.ov.mixnet.commons.beans.proofs.ShuffleProof;
import com.scytl.products.ov.mixnet.commons.configuration.BGParams;
import com.scytl.products.ov.mixnet.commons.configuration.locations.LocationsProvider;
import com.scytl.products.ov.mixnet.commons.constants.Constants;
import com.scytl.products.ov.mixnet.commons.exceptions.MixingException;
import com.scytl.products.ov.mixnet.commons.exceptions.ShufflerException;
import com.scytl.products.ov.mixnet.commons.homomorphic.Randomness;
import com.scytl.products.ov.mixnet.commons.io.ElGamalEncryptedBallotsWriter;
import com.scytl.products.ov.mixnet.commons.io.ElgamalPublicKeyWriter;
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.proofs.bg.commitments.CommitmentParams;
import com.scytl.products.ov.mixnet.commons.proofs.bg.config.ProofsGeneratorConfigurationParams;
import com.scytl.products.ov.mixnet.commons.tools.MultiExponentiation;
import com.scytl.products.ov.mixnet.infrastructure.log.MixingAndVerifyingLogEvents;
import com.scytl.products.ov.mixnet.io.VotesWithProofWriter;
import com.scytl.products.ov.mixnet.proofs.bg.multiexp.ParallelComputeAllE;
import com.scytl.products.ov.mixnet.shuffler.ElGamalShuffler;

/**
 * Bayer-Groth mixer
 */
public class BGMixer {

    private final TransactionInfoProvider transactionInfoProvider = new TransactionInfoProvider();

    private final SecureLoggingWriter loggingWriter;

    private final JSONProofsWriter proofsWriter;

    private final ElGamalShuffler shuffler;

    private final BallotDecryptionService ballotDecryptionService;

    private final ShuffleProofService shuffleProofService;
    
    private final ForkJoinPool executor;

    /**
     * Constructor. For internal use only.
     * 
     * @param shuffler
     * @param proofsWriter
     * @param ballotDecryptionService
     * @param shuffleProofService
     * @param executor
     */
    BGMixer(final ElGamalShuffler shuffler, final JSONProofsWriter proofsWriter,
            BallotDecryptionService ballotDecryptionService, ShuffleProofService shuffleProofService, 
            ForkJoinPool executor) {

        this.proofsWriter = proofsWriter;
        this.shuffler = shuffler;
        this.loggingWriter = initializeLoggingWritter();

        this.ballotDecryptionService = ballotDecryptionService;
        this.shuffleProofService = shuffleProofService;
        this.executor = executor;
    }

    public void mix(final ElGamalEncryptedBallots encryptedBallots, final BGParams bGparams,
            final CommitmentParams commitmentParams,
            final ProofsGeneratorConfigurationParams proofsGeneratorConfigurationParams, final Integer batchNum,
            final int partitionSize, final int offset, final int concurrencyLevel,
            final ElGamalPublicKey publicKey, final ElGamalPrivateKey ballotDecryptionKey, final ZpSubgroup group,
            final LocationsProvider locationsProvider, final MultiExponentiation multiExponentiation,
            final ParallelComputeAllE parallelComputeAllE, final Randomness[] batchRandomness,
            final List<Ciphertext> batchPreComputations, final PrimitivesServiceAPI primitivesService)
            throws MixingException {
        Runnable task = () -> doMix(encryptedBallots, bGparams, commitmentParams,
            proofsGeneratorConfigurationParams, batchNum, partitionSize,
            offset, concurrencyLevel, publicKey, ballotDecryptionKey,
            group, locationsProvider, multiExponentiation,
            parallelComputeAllE, batchRandomness, batchPreComputations,
            primitivesService);
        executor.submit(task).join();
    }

    private void doMix(final ElGamalEncryptedBallots encryptedBallots,
            final BGParams bGparams,
            final CommitmentParams commitmentParams,
            final ProofsGeneratorConfigurationParams proofsGeneratorConfigurationParams,
            final int batchNum, final int partitionSize,
            final int offset, final int concurrencyLevel,
            final ElGamalPublicKey publicKey,
            final ElGamalPrivateKey ballotDecryptionKey,
            final ZpSubgroup group,
            final LocationsProvider locationsProvider,
            final MultiExponentiation multiExponentiation,
            final ParallelComputeAllE parallelComputeAllE,
            final Randomness[] batchRandomness,
            final List<Ciphertext> batchPreComputations,
            final PrimitivesServiceAPI primitivesService) {
        // It's a new thread, so we need to create a new trackId
        transactionInfoProvider.generate(Constants.LOGGING_TENANT, Constants.LOGGING_LOCALHOST,
            Constants.LOGGING_LOCALHOST);

        int m = bGparams.getM();
        int n = bGparams.getN();

        ShuffleOutput<ElGamalEncryptedBallots> shuffleOutput =
            shuffle(group, publicKey, concurrencyLevel, encryptedBallots, batchRandomness,
                batchPreComputations, batchNum, partitionSize, offset);

        // Decryption and proof generation can be run concurrently. Start them
        // both
        // and then wait to
        // collect the results.

        // Start the proof generation task.
        ForkJoinTask<ShuffleProof> proofGenerationTask =
            ForkJoinTask.adapt(() -> shuffleProofService.generate(encryptedBallots, shuffleOutput, m, n,
                commitmentParams, proofsGeneratorConfigurationParams, concurrencyLevel, multiExponentiation,
                parallelComputeAllE, batchNum, partitionSize, offset));
        proofGenerationTask.fork();

        // Do the decryption task.
        List<Ciphertext> encryptedBallotStream = shuffleOutput.getReEncryptedBallots().getBallots();
        List<VoteWithProof> votesWithProof = ballotDecryptionService.decrypt(encryptedBallotStream, ballotDecryptionKey, 
            batchNum, partitionSize, offset);
        
        ShuffleProof shuffleProof;
        if (proofGenerationTask.tryUnfork()) {
            shuffleProof = proofGenerationTask.invoke();
        } else {
            shuffleProof = proofGenerationTask.join();
        }
        
        try {
            storeData(encryptedBallots, votesWithProof, shuffleOutput, shuffleProof, commitmentParams, batchNum, group,
                publicKey, locationsProvider, primitivesService);
        } catch (Exception e) {
            throw new MixingException("An error trying to store the outputs of the mixing", e);
        }
    }

    private SecureLoggingWriter initializeLoggingWritter() {
        SplunkFormatter splunkFormatter =
            new SplunkFormatter(Constants.LOGGING_APPLICATION, Constants.LOGGING_APPLICATION, transactionInfoProvider);
        final SecureLoggingFactoryLog4j loggerFactory = new SecureLoggingFactoryLog4j(splunkFormatter);
        return loggerFactory.getLogger(Constants.SECURE_LOGGER_NAME);
    }

    private ShuffleOutput<ElGamalEncryptedBallots> shuffle(final ZpSubgroup group, final ElGamalPublicKey publicKey,
            final int concurrencyLevel, final ElGamalEncryptedBallots encryptedBallots,
            final Randomness[] batchRandomness, final List<Ciphertext> batchPreComputations, final Integer batchNum,
            final int partitionSize, final int offset) {

        try {
            ShuffleOutput<ElGamalEncryptedBallots> shuffleOutput = shuffler.sequentialShuffle(group, publicKey,
                concurrencyLevel, encryptedBallots, batchRandomness, batchPreComputations);

            loggingWriter.log(Level.INFO, new LogContent.LogContentBuilder()
                .logEvent(
                    MixingAndVerifyingLogEvents.SHUFREENCR_SUCCESS_SHUFFLING_AND_REENCRYPTING_THE_ENCRYPTED_BALLOTS)
                .user("adminID").electionEvent("-").additionalInfo("mbb_id", "-")
                .additionalInfo("batch_num", batchNum.toString()).additionalInfo("partition_size", Integer.toString(partitionSize))
                .additionalInfo("offset", Integer.toString(offset)).createLogInfo());

            return shuffleOutput;
        } catch (Exception e) {
            loggingWriter.log(Level.ERROR, new LogContent.LogContentBuilder()
                .logEvent(MixingAndVerifyingLogEvents.SHUFREENCR_ERROR_SHUFFLING_AND_REENCRYPTING_THE_ENCRYPTED_BALLOTS)
                .user("adminID").electionEvent("-").additionalInfo("mbb_id", "-")
                .additionalInfo("batch_num", batchNum.toString()).additionalInfo("partition_size", Integer.toString(partitionSize))
                .additionalInfo("offset", Integer.toString(offset)).additionalInfo("error_msg", e.getMessage()).createLogInfo());
            throw new ShufflerException(
                "An error occurred while trying to shuffle and re-encrypt the encrypted ballots", e);
        }
    }

    private void storeData(final ElGamalEncryptedBallots encryptedBallots, final List<VoteWithProof> votesWithProof,
            final ShuffleOutput<ElGamalEncryptedBallots> shuffleOutput, final ShuffleProof shuffleProof,
            final CommitmentParams commitmentParams, final int batchNum, final ZpSubgroup groupConfig,
            final ElGamalPublicKey publicKey, final LocationsProvider locationsProvider,
            PrimitivesServiceAPI primitivesService) throws IOException, GeneralCryptoLibException {

        ElGamalEncryptedBallotsWriter.CSV.write(
            locationsProvider.getMixerLocationsProvider().getEncryptedBallotsOutputOfBatch(batchNum), encryptedBallots,
            primitivesService.getRawMessageDigest());

        VotesWithProofWriter.write(
            locationsProvider.getMixerLocationsProvider().getVotesWithProofOutputOfBatch(batchNum),
            locationsProvider.getMixerLocationsProvider().getDecryptedVotesOutputOfBatch(batchNum), votesWithProof);

        ZpGroupReader.serializeToStream(groupConfig,
            locationsProvider.getMixerLocationsProvider().getEncryptionParametersOutputOfBatch(batchNum));

        OutputStream pathThisBatchMixedBallot =
            locationsProvider.getMixerLocationsProvider().getReEncryptedBallotsOutputOfBatch(batchNum);

        byte[] hash = ElGamalEncryptedBallotsWriter.CSV.write(pathThisBatchMixedBallot,
            shuffleOutput.getReEncryptedBallots(), primitivesService.getRawMessageDigest());
        String hashAsString = Base64.getEncoder().encodeToString(hash);
        logHashOfWrittenBallotBox(batchNum, hashAsString);

        proofsWriter.write(locationsProvider.getMixerLocationsProvider().getProofsOutputOfBatch(batchNum),
            shuffleProof);

        commitmentParams.serializeToStream(
            locationsProvider.getMixerLocationsProvider().getCommitmentParametersOutputOfBatch(batchNum));

        ElgamalPublicKeyWriter.serializeToStream(publicKey,
            locationsProvider.getMixerLocationsProvider().getPublicKeyOutputOfBatch(batchNum));

    }

    private void logHashOfWrittenBallotBox(final int batchNum, String hashAsString)
            throws GeneralCryptoLibException, IOException {

        LogContent.LogContentBuilder logContent =
            new LogContent.LogContentBuilder().user("adminID").electionEvent("-").additionalInfo("mbb_id", "-")
                .additionalInfo("batch_num", Integer.toString(batchNum)).additionalInfo("mbb_hash", hashAsString);

        try {
            loggingWriter.log(Level.INFO,
                logContent
                    .logEvent(
                        MixingAndVerifyingLogEvents.SHUFREENCR_SUCCESS_WRITING_THE_SHUFFLED_AND_REENCRYPTING_BALLOT_BOX)
                    .createLogInfo());
        } catch (Exception e) {
            loggingWriter.log(Level.ERROR, logContent
                .logEvent(MixingAndVerifyingLogEvents.SHUFREENCR_ERROR_WRITING_THE_SHUFFLED_AND_REENCRYPTING_BALLOT_BOX)
                .additionalInfo("error_msg", e.getMessage()).createLogInfo());
            throw new ShufflerException("An error occurred while writing the shuffled ballots box", e);
        }
    }

}
