/*
 * 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.sdm.config.spring.batch;

import static java.nio.file.Files.deleteIfExists;
import static java.nio.file.Files.write;

import java.io.IOException;
import java.nio.file.Path;
import java.nio.file.StandardOpenOption;
import java.security.cert.X509Certificate;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.Queue;
import java.util.Set;

import org.springframework.batch.core.StepContribution;
import org.springframework.batch.core.scope.context.ChunkContext;
import org.springframework.batch.core.step.tasklet.Tasklet;
import org.springframework.batch.repeat.RepeatStatus;

import com.fasterxml.jackson.core.JsonProcessingException;
import com.fasterxml.jackson.core.type.TypeReference;
import com.fasterxml.jackson.databind.ObjectMapper;
import com.scytl.cryptolib.api.exceptions.GeneralCryptoLibException;
import com.scytl.cryptolib.elgamal.bean.ElGamalComputationsValues;
import com.scytl.cryptolib.mathematical.groups.impl.ZpGroupElement;
import com.scytl.products.ov.commons.beans.EncryptionParameters;
import com.scytl.products.ov.commons.beans.domain.model.messaging.InvalidSignatureException;
import com.scytl.products.ov.commons.beans.domain.model.messaging.PayloadVerificationException;
import com.scytl.products.ov.commons.dto.ChoiceCodeGenerationDTO;
import com.scytl.products.ov.commons.dto.ChoiceCodeGenerationOutput;
import com.scytl.products.ov.commons.dto.ChoiceCodeGenerationResPayload;
import com.scytl.products.ov.commons.dto.ElGamalComputationsValuesCodec;
import com.scytl.products.ov.commons.verify.PayloadVerifier;
import com.scytl.products.ov.config.spring.batch.CommonBatchInfrastructure;
import com.scytl.products.ov.datapacks.beans.VerificationCardIdComputedValues;

public class NodeContributionsVerificationCombinationTasklet implements Tasklet {

    private CommonBatchInfrastructure commonBatchInfrastructure;

    private String basePath;

    private Queue<VerificationCardIdComputedValues> outputQueue;

    private EncryptionParameters encryptionParameters;

    private ElGamalComputationsValuesCodec codec;

    private ObjectMapper objectMapper;

    private PayloadVerifier payloadVerifier;

    private X509Certificate trustedCertificate;

    public NodeContributionsVerificationCombinationTasklet(String basePath,
            CommonBatchInfrastructure commonBatchInfrastructure, EncryptionParameters encryptionParameters,
            Queue<VerificationCardIdComputedValues> outputQueue, ObjectMapper objectMapper,
            PayloadVerifier payloadVerifier, X509Certificate trustedCertificate, ElGamalComputationsValuesCodec codec) {
        this.basePath = basePath;
        this.commonBatchInfrastructure = commonBatchInfrastructure;
        this.outputQueue = outputQueue;
        this.encryptionParameters = encryptionParameters;
        this.objectMapper = objectMapper;
        this.codec = codec;
        this.trustedCertificate = trustedCertificate;
        this.payloadVerifier = payloadVerifier;
    }

    /**
     * This tasklet execution involves the following: - Obtaining the nodes
     * computations of choice codes / ballot casting key from the file system
     * (downloaded by the SDM) - Verify the signatures of each of the nodes
     * contributions - Combine the cc and bck of the nodes, and extract the
     * derived keys commitment also from the nodes - Save the derived keys
     * commitments to file system (later to be synchronized to Vote
     * Verification) - Save the combined cc and bck to the output buffer (will
     * be used for credential generation)
     */
    @Override
    public RepeatStatus execute(StepContribution contribution, ChunkContext chunkContext) throws Exception {

        VotingCardGenerationJobExecutionContext jobExecutionContext = new VotingCardGenerationJobExecutionContext(
            chunkContext.getStepContext().getStepExecution().getJobExecution().getExecutionContext());

        List<Path> nodeContributionsFiles =
            commonBatchInfrastructure.getDataSerializationProvider(basePath, jobExecutionContext.getVotingCardSetId(),
                jobExecutionContext.getVerificationCardSetId()).getNodeContributions();

        Path derivedKeysPath = commonBatchInfrastructure.getDataSerializationProvider(basePath,
            jobExecutionContext.getVotingCardSetId(), jobExecutionContext.getVerificationCardSetId()).getDerivedKeys();
        deleteIfExists(derivedKeysPath);
        
        for (Path nodeContributionsFile : nodeContributionsFiles) {
            verifyAndCombineNodeContributions(nodeContributionsFile, derivedKeysPath);
        }

        return RepeatStatus.FINISHED;
    }

    private void verifyAndCombineNodeContributions(Path nodeContributionsFile, Path derivedKeysPath)
            throws IOException, PayloadVerificationException, GeneralCryptoLibException, InvalidSignatureException {
        List<ChoiceCodeGenerationDTO<ChoiceCodeGenerationResPayload>> nodeContributions =
                objectMapper.readValue(nodeContributionsFile.toFile(),
                    new TypeReference<List<ChoiceCodeGenerationDTO<ChoiceCodeGenerationResPayload>>>() {
                    });

        verifySignatures(nodeContributions);

        List<VerificationCardIdComputedValues> computedValues = combineChoiceCodeNodesComputeContributions(
            nodeContributions);

        writeDerivedKeyCommitments(derivedKeysPath, computedValues);

        outputQueue.addAll(computedValues);
    }

    private void verifySignatures(List<ChoiceCodeGenerationDTO<ChoiceCodeGenerationResPayload>> nodeContributions)
            throws PayloadVerificationException, InvalidSignatureException {

        for (ChoiceCodeGenerationDTO<ChoiceCodeGenerationResPayload> choiceCodeNodeContribution : nodeContributions) {
            if (!payloadVerifier.isValid(choiceCodeNodeContribution.getPayload(), trustedCertificate)) {
                throw new InvalidSignatureException("verification card set",
                    choiceCodeNodeContribution.getPayload().getVerificationCardSetId());
            }
        }
    }

    private void writeDerivedKeyCommitments(Path derivedKeysPath, List<VerificationCardIdComputedValues> computedValues)
            throws IOException {
        List<String> derivedKeysLines = new ArrayList<>();

        for (VerificationCardIdComputedValues computedValuesEntry : computedValues) {
            derivedKeysLines
                .add(getDerivedKeysLine(computedValuesEntry.getId(), computedValuesEntry.getChoiceCodesKeyCommitments(),
                    computedValuesEntry.getBallotCastingKeyExponentCommitments()));
        }

        write(derivedKeysPath, derivedKeysLines, StandardOpenOption.CREATE, StandardOpenOption.APPEND);
    }

    private String getDerivedKeysLine(String verificationCardId, List<String> choiceCodeDerivedKeyCommitments, List<String> ballotCastKeyDerivedExponentCommitments)
            throws JsonProcessingException {
        return String.join(";", verificationCardId, objectMapper.writeValueAsString(choiceCodeDerivedKeyCommitments),
            objectMapper.writeValueAsString(ballotCastKeyDerivedExponentCommitments));
    }

    private List<VerificationCardIdComputedValues> combineChoiceCodeNodesComputeContributions(
            List<ChoiceCodeGenerationDTO<ChoiceCodeGenerationResPayload>> nodesContributions)
            throws GeneralCryptoLibException {

        List<VerificationCardIdComputedValues> combinedContributions = new ArrayList<>();

        Map<String, String> vcIdEncryptedBallotCastingKeys = new HashMap<>();
        Map<String, List<String>> vcIdComputedBallotCastingKeys = new HashMap<>();
        Map<String, List<String>> vcIdComputedRepresentations = new HashMap<>();
        Map<String, List<String>> vcIdChoiceCodesDerivedKeyCommitments = new HashMap<>();
        Map<String, List<String>> vcIdballotCastingKeyExponentCommitments = new HashMap<>();

        for (ChoiceCodeGenerationOutput generationOutput : nodesContributions.get(0).getPayload()
            .getChoiceCodeGenerationOutputList()) {
            String verificationCardId = generationOutput.getVerificationCardId();
            vcIdEncryptedBallotCastingKeys.put(verificationCardId, generationOutput.getEncryptedBallotCastingKey());
            vcIdComputedBallotCastingKeys.put(verificationCardId, new ArrayList<>());
            vcIdComputedRepresentations.put(verificationCardId, new ArrayList<>());
            vcIdChoiceCodesDerivedKeyCommitments.put(verificationCardId, new ArrayList<>());
            vcIdballotCastingKeyExponentCommitments.put(verificationCardId, new ArrayList<>());
        }

        for (ChoiceCodeGenerationDTO<ChoiceCodeGenerationResPayload> nodeContributions : nodesContributions) {
            ChoiceCodeGenerationResPayload partialChoiceCodesResult = nodeContributions.getPayload();

            for (ChoiceCodeGenerationOutput generationOutput : partialChoiceCodesResult
                .getChoiceCodeGenerationOutputList()) {

                final String verificationCardId = generationOutput.getVerificationCardId();

                vcIdComputedBallotCastingKeys.get(verificationCardId)
                    .add(generationOutput.getComputedBallotCastingKey());

                vcIdComputedRepresentations.get(verificationCardId).add(generationOutput.getComputedRepresentations());

                vcIdChoiceCodesDerivedKeyCommitments.get(verificationCardId).add(generationOutput.getChoiceCodesKeyCommitmentJson());
                
                vcIdballotCastingKeyExponentCommitments.get(verificationCardId).add(generationOutput.getBallotCastingKeyCommitmentJson());
            }
        }

        for (Map.Entry<String, List<String>> entry : vcIdComputedRepresentations.entrySet()) {
			String combinedOptionRepresentation = combineOptionRepresentationsContributions(entry.getValue());
        	
			String combinedBallotCastingKey = combineBallotCastingKeyContributions(
					vcIdComputedBallotCastingKeys.get(entry.getKey()));

			combinedContributions.add(new VerificationCardIdComputedValues(entry.getKey(), combinedOptionRepresentation,
					vcIdEncryptedBallotCastingKeys.get(entry.getKey()), combinedBallotCastingKey,
					vcIdChoiceCodesDerivedKeyCommitments.get(entry.getKey()),
			        vcIdballotCastingKeyExponentCommitments.get(entry.getKey())));
        }
        
        return combinedContributions;
    }

    private String combineOptionRepresentationsContributions(List<String> contributions)
            throws GeneralCryptoLibException {
        Map<ElGamalComputationsValues, ElGamalComputationsValues> combination = new HashMap<>();
        Set<ElGamalComputationsValues> options = null;
        for (String contribution : contributions) {
            Map<ElGamalComputationsValues, ElGamalComputationsValues> part = codec.decodeMap(contribution);
            if (options == null) {
                options = part.keySet();
            } else if (options.size() != part.size() || !options.containsAll(part.keySet())) {
                throw new IllegalArgumentException("All node contributions must have the same options.");
            }
            for (Map.Entry<ElGamalComputationsValues, ElGamalComputationsValues> entry : part.entrySet()) {
                ElGamalComputationsValues values = combination.get(entry.getKey());
                if (values == null) {
                    values = entry.getValue();
                } else {
                    values = combineElGamalComputationsValues(values, entry.getValue());
                }
                combination.put(entry.getKey(), values);
            }
        }
        return codec.encodeMap(combination);
    }

    private ElGamalComputationsValues combineElGamalComputationsValues(ElGamalComputationsValues values1,
            ElGamalComputationsValues values2) throws GeneralCryptoLibException {
        List<ZpGroupElement> ciphertext1 = values1.getValues();
        List<ZpGroupElement> ciphertext2 = values2.getValues();
        if (ciphertext1.size() != ciphertext2.size()) {
            throw new IllegalArgumentException("ElGamal computations values must have the same length.");
        }
        List<ZpGroupElement> combination = new ArrayList<>(ciphertext1.size());
        for (int i = 0; i < ciphertext1.size(); i++) {
            combination.add(ciphertext1.get(i).multiply(ciphertext2.get(i)));
        }
        return new ElGamalComputationsValues(combination);
    }

    private String combineBallotCastingKeyContributions(List<String> contributions) throws GeneralCryptoLibException {
        ElGamalComputationsValues combination = null;
        for (String contribution : contributions) {
            ElGamalComputationsValues values = codec.decodeSingle(contribution);
            if (combination == null) {
                combination = values;
            } else {
                combination = combineElGamalComputationsValues(combination, values);
            }
        }
        return codec.encodeSingle(combination);
    }

}
