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

import java.io.IOException;
import java.math.BigInteger;
import java.nio.charset.StandardCharsets;
import java.security.cert.X509Certificate;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.Base64;
import java.util.Collections;
import java.util.List;
import java.util.Optional;
import java.util.stream.Collectors;

import javax.ejb.Stateless;
import javax.inject.Inject;
import javax.ws.rs.core.MediaType;

import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import com.fasterxml.jackson.core.type.TypeReference;
import com.fasterxml.jackson.databind.ObjectMapper;
import com.scytl.cryptolib.api.exceptions.GeneralCryptoLibException;
import com.scytl.cryptolib.api.proofs.ProofsServiceAPI;
import com.scytl.cryptolib.elgamal.bean.ElGamalPublicKey;
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.cryptolib.primitives.primes.utils.PemUtils;
import com.scytl.cryptolib.proofs.cryptoapi.ProofVerifierAPI;
import com.scytl.cryptolib.proofs.proof.Proof;
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.ov.commons.beans.EncryptionParameters;
import com.scytl.products.ov.commons.beans.VoteVerificationContextData;
import com.scytl.products.ov.commons.beans.domain.model.messaging.PayloadVerificationException;
import com.scytl.products.ov.commons.beans.exceptions.ResourceNotFoundException;
import com.scytl.products.ov.commons.beans.validation.CryptographicOperationException;
import com.scytl.products.ov.commons.domain.model.platform.PlatformCARepository;
import com.scytl.products.ov.commons.dto.ChoiceCodesVerificationDecryptResPayload;
import com.scytl.products.ov.commons.infrastructure.remote.client.RetrofitConsumer;
import com.scytl.products.ov.commons.tracking.TrackIdInstance;
import com.scytl.products.ov.commons.ui.ws.rs.ObjectMappers;
import com.scytl.products.ov.commons.verify.PayloadVerifier;
import com.scytl.products.ov.vv.domain.model.choicecode.CodesDecryptionResults;
import com.scytl.products.ov.vv.domain.model.content.VerificationContentEntity;
import com.scytl.products.ov.vv.domain.model.content.VerificationContentRepository;
import com.scytl.products.ov.vv.domain.model.platform.VvPlatformCARepository;
import com.scytl.products.ov.vv.infrastructure.log.SecureLoggerUtil;
import com.scytl.products.ov.vv.infrastructure.log.VoteVerificationLogConstants;
import com.scytl.products.ov.vv.infrastructure.log.VoteVerificationLogEvents;
import com.scytl.products.ov.vv.infrastructure.remote.OrchestratorClient;

import okhttp3.RequestBody;
import okhttp3.ResponseBody;

/**
 * Decrypt encrypted partial choice codes
 */
@Stateless
public class CodesDecrypterService {

    @Inject
    private OrchestratorClient ccOrchestratorClient;

    @Inject
    private TrackIdInstance trackId;

    private static final Logger LOG = LoggerFactory.getLogger(CodesDecrypterService.class);

    @Inject
    private ProofsServiceAPI proofsService;

    @Inject
    private VerificationContentRepository verificationContentRepository;

    @Inject
    private GroupElementsCompressor<ZpGroupElement> groupElementsCompressor;

    @Inject
    private SecureLoggingWriter secureLoggerWriter;

    @Inject
    private PayloadVerifier payloadVerifier;

    private ObjectMapper objectMapper = new ObjectMapper();

    @Inject
    @VvPlatformCARepository
    private PlatformCARepository platformCARepository;

    static final String CC_ORCHESTRATOR_PATH = "choicecodes";

    /**
     * Removes encryption from the encrypted choice codes received.
     * 
     * @param tenantId
     *            - tenant identifier.
     * @param eeId
     *            - election event identifier.
     * @param verificationCardSetId
     *            - verification card set id
     * @param verificationCardId
     *            - verification card id
     * @param encryptedPartialCodes
     *            - encryptionPartial codes of the vote
     * @return combined decrypted partial codes and all decryption results
     * @throws ResourceNotFoundException
     * @throws CryptographicOperationException
     */
    public CodesDecryptionResults decryptPartialCodes(String tenantId, String eeId, String verificationCardSetId,
            String verificationCardId, List<ZpGroupElement> encryptedPartialCodes)
            throws ResourceNotFoundException, CryptographicOperationException {

        try {

            LOG.info("Obtaining the necessary data to decrypt the choice codes for tenant: {}, election event: {} "
                + "and verification card: {}.", tenantId, eeId, verificationCardId);

            ZpGroupElement gamma = encryptedPartialCodes.get(0);

            LOG.info("Requesting the choice code nodes decryption contributions");

            List<ChoiceCodesVerificationDecryptResPayload> decryptResult =
                collectChoiceCodeNodesDecryptionContributions(tenantId, eeId, verificationCardSetId, verificationCardId,
                    gamma);

            return decryptPartialCodesFromDecryptResult(eeId, verificationCardId, encryptedPartialCodes, decryptResult);
        } catch (GeneralCryptoLibException e) {
            LOG.error("Error decrypting partial choice codes: " + e.getMessage());
            throw new CryptographicOperationException("Error decrypting partial choice codes:", e);
        } catch (PayloadVerificationException e) {
            LOG.error("Error verifying partial choice codes signatures: " + e.getMessage());
            throw new IllegalStateException("Error verifying partial choice codes signatures:", e);
        } catch (IOException e) {
            LOG.error("Error retrieving decrypted partial choice codes: " + e.getMessage());
            throw new IllegalStateException("Error retrieving decrypted partial choice codes:", e);
        }
    }

    /**
     * Removes encryption from the encrypted choice codes received.
     * 
     * @param eeId
     *            - election event identifier.
     * @param verificationCardId
     *            - verification card id
     * @param encryptedPartialCodes
     *            - encryptionPartial codes of the vote
     * @param decryptResult
     *            - the result of the partial choice codes nodes decryption
     * @return combined decrypted partial codes and all decryption results
     * @throws ResourceNotFoundException
     * @throws CryptographicOperationException
     */
    public CodesDecryptionResults decryptPartialCodesFromDecryptResult(String eeId, String verificationCardId,
            List<ZpGroupElement> encryptedPartialCodes, List<ChoiceCodesVerificationDecryptResPayload> decryptResult)
            throws CryptographicOperationException {
        try {
            List<ZpGroupElement> phis = encryptedPartialCodes.subList(1, encryptedPartialCodes.size());

            List<List<ZpGroupElement>> ccnContributions = getCcnContributions(decryptResult);

            List<ZpGroupElement> gammasWithCCNContributions =
                combineChoiceCodeNodesDecryptionContributions(ccnContributions, phis.size());
            String hashPdcCc =
                SecureLoggerUtil.getHashForLogFromZpGroupElementIntegerCollection(gammasWithCCNContributions);
            // Partial Decryption -- Partial Decryptions correctly consolidated
            secureLoggerWriter.log(Level.INFO,
                new LogContent.LogContentBuilder()
                    .logEvent(VoteVerificationLogEvents.PARTIAL_DECRYPTION_SUCESSFULLY_CONSOLIDATED)
                    .user(verificationCardId).electionEvent(eeId)
                    .additionalInfo(VoteVerificationLogConstants.INFO_TRACK_ID, trackId.getTrackId())
                    .additionalInfo("#hash_pDC_CC", hashPdcCc).createLogInfo());

            LOG.info("Decrypting partial choice codes");

            List<ZpGroupElement> gammasforDecryption =
                gammasWithCCNContributions.stream().map(ZpGroupElement::invert).collect(Collectors.toList());

            List<ZpGroupElement> combinedZpGroupElementLists = combineZpGroupElementLists(gammasforDecryption, phis);
            return new CodesDecryptionResults(combinedZpGroupElementLists, decryptResult);
        } catch (GeneralCryptoLibException e) {
            LOG.error("Error decrypting partial choice codes using the decrypt contributions: " + e.getMessage());
            throw new CryptographicOperationException(
                "Error decrypting partial choice codes using the decrypt contributions:", e);
        }
    }

    private List<List<ZpGroupElement>> getCcnContributions(
            List<ChoiceCodesVerificationDecryptResPayload> decryptResult) {
        List<List<ZpGroupElement>> ccnContributions = new ArrayList<>();
        decryptResult.forEach(decryptResPayload -> {
            List<ZpGroupElement> castedGammas =
                decryptResPayload.getDecryptContributionResult().stream().map(gammaStr -> {
                    try {
                        return ZpGroupElement.fromJson(gammaStr);
                    } catch (GeneralCryptoLibException e) {
                        LOG.error("Gamma in json format could not be casted to ZpGroupElement: " + e.getMessage());
                        throw new IllegalStateException(
                            "Gamma in json format could not be casted to ZpGroupElement: " + e.getMessage());
                    }
                }).collect(Collectors.toList());
            ccnContributions.add(castedGammas);
        });

        return ccnContributions;
    }

    private List<ChoiceCodesVerificationDecryptResPayload> collectChoiceCodeNodesDecryptionContributions(
            String tenantId, String eeId, String verificationCardSetId, String verificationCardId, ZpGroupElement gamma)
            throws GeneralCryptoLibException, ResourceNotFoundException, PayloadVerificationException, IOException {

        RequestBody requestBody =
            RequestBody.create(okhttp3.MediaType.parse(MediaType.APPLICATION_JSON), gamma.toJson());

        try (ResponseBody response = RetrofitConsumer
            .processResponse(ccOrchestratorClient.getChoiceCodeNodesDecryptContributions(trackId.getTrackId(),
                CC_ORCHESTRATOR_PATH, tenantId, eeId, verificationCardSetId, verificationCardId, requestBody))) {

            List<ChoiceCodesVerificationDecryptResPayload> decryptResult = objectMapper.readValue(response.byteStream(),
                new TypeReference<List<ChoiceCodesVerificationDecryptResPayload>>() {
                });

            String hashPdCc =
                SecureLoggerUtil.getHashForLogFromZpGroupElementIntegerCollection(Collections.singletonList(gamma));
            // [SL PDREQ-5]
            // Partial Decryption request -- Partial decryption correctly
            // received

            secureLoggerWriter.log(Level.INFO,
                new LogContent.LogContentBuilder()
                    .logEvent(VoteVerificationLogEvents.PARTIAL_DECRYPTION_CORRECTLY_RECEIVED).user(verificationCardId)
                    .additionalInfo("#request_id", trackId.getTrackId()).additionalInfo("#hash_pD_CC", hashPdCc)
                    .electionEvent(eeId).createLogInfo());

            verifySignatures(tenantId, eeId, verificationCardId, decryptResult);

            verifyProofs(tenantId, eeId, verificationCardSetId, verificationCardId, gamma, decryptResult);

            List<List<ZpGroupElement>> ccnContributions = new ArrayList<>();
            decryptResult.forEach(decryptResPayload -> {
                List<ZpGroupElement> castedGammas =
                    decryptResPayload.getDecryptContributionResult().stream().map(gammaStr -> {
                        try {
                            return ZpGroupElement.fromJson(gammaStr);
                        } catch (GeneralCryptoLibException e) {
                            LOG.error("Gamma in json format could not be casted to ZpGroupElement: " + e.getMessage());
                            throw new IllegalStateException(
                                "Gamma in json format could not be casted to ZpGroupElement: " + e.getMessage(), e);
                        }
                    }).collect(Collectors.toList());
                ccnContributions.add(castedGammas);
            });

            return decryptResult;
        }
    }

    private void verifySignatures(String tenantId, String eeId, String verificationCardId,
            List<ChoiceCodesVerificationDecryptResPayload> response)
            throws GeneralCryptoLibException, ResourceNotFoundException, PayloadVerificationException {

        X509Certificate rootCertificate = (X509Certificate) PemUtils
            .certificateFromPem(platformCARepository.getRootCACertificate().getCertificateContent());

        for (ChoiceCodesVerificationDecryptResPayload choiceCodesComputationResult : response) {
            if (!payloadVerifier.isValid(choiceCodesComputationResult, rootCertificate)) {
                throw new IllegalStateException("Signature invalid for tenantId " + tenantId + " eeId " + eeId
                    + " verificationCardId " + verificationCardId);
            }
        }
    }

    private void verifyProofs(String tenantId, String eeId, String verificationCardSetId, String verificationCardId,
            ZpGroupElement gamma, List<ChoiceCodesVerificationDecryptResPayload> decryptResponses)
            throws GeneralCryptoLibException, ResourceNotFoundException {

        VoteVerificationContextData voteVerificationContextData =
            getVoteVerificationContextData(tenantId, eeId, verificationCardSetId);

        List<ElGamalPublicKey> commitmentPublicKeys = getCommitmentPublicKeys(voteVerificationContextData);

        for (ChoiceCodesVerificationDecryptResPayload choiceCodesDecryptResult : decryptResponses) {

            ElGamalPublicKey resultPublicKey = ElGamalPublicKey.fromJson(choiceCodesDecryptResult.getPublicKeyJson());

            if (!commitmentPublicKeys.remove(resultPublicKey)) {
                throw new IllegalStateException("Resulting proof public key was not valid for tenantId " + tenantId
                    + " eeId " + eeId + " verificationCardSetId " + verificationCardSetId + " verificationCardId "
                    + verificationCardId);
            }

            List<String> exponentiatedGammaStrings = choiceCodesDecryptResult.getDecryptContributionResult();
            List<ZpGroupElement> exponentiatedGammas = new ArrayList<>(exponentiatedGammaStrings.size());
            for (String exponentiatedGammaString : exponentiatedGammaStrings) {
                ZpGroupElement exponentiatedGamma = ZpGroupElement.fromJson(exponentiatedGammaString);
                exponentiatedGammas.add(exponentiatedGamma);
            }

            ZpGroupElement compressedPublicKeys = groupElementsCompressor.compress(resultPublicKey.getKeys());
            ZpGroupElement compressedExponentiatedGamma = groupElementsCompressor.compress(exponentiatedGammas);

            ZpSubgroup encryptionParametersGroup = getEncryptionParametersGroup(voteVerificationContextData);

            List<ZpGroupElement> exponentiatedElements =
                Arrays.asList(compressedPublicKeys, compressedExponentiatedGamma);
            List<ZpGroupElement> baseElements = Arrays.asList(encryptionParametersGroup.getGenerator(), gamma);

            Proof proof = Proof.fromJson(choiceCodesDecryptResult.getExponentiationProofJson());

            ProofVerifierAPI proofsVerifier = proofsService.createProofVerifierAPI(encryptionParametersGroup);

            if (!proofsVerifier.verifyExponentiationProof(exponentiatedElements, baseElements, proof)) {
                // [SL] hash of the consolidated partial choice codes
                String hashPccc = SecureLoggerUtil.getHashForLogFromZpGroupElementIntegerCollection(baseElements);
                // [SL] hash of the partial decryption
                String hashPdCc =
                    SecureLoggerUtil.getHashForLogFromZpGroupElementIntegerCollection(exponentiatedGammas);
                // [SL] hash of the partial decryption proof

                ArrayList<Exponent> exponentiatonProofExponents = new ArrayList<>();
                exponentiatonProofExponents.add(proof.getHashValue());
                exponentiatonProofExponents.addAll(proof.getValuesList());
                String decProofCc = SecureLoggerUtil.getHashForLogFromExponentCollection(exponentiatonProofExponents);

                // [SL PDCON-3] Medium
                // Partial Decryption consolidation -- Proof not valid
                secureLoggerWriter.log(Level.INFO,
                    new LogContent.LogContentBuilder().logEvent(VoteVerificationLogEvents.PROOF_NOT_VALID_PDC)
                        .additionalInfo(VoteVerificationLogConstants.INFO_ERR_DESC, "Proof not valid")
                        .additionalInfo("#hash_pCCC", hashPccc).additionalInfo("#hash_pD_CC", hashPdCc)
                        .additionalInfo("#dec_proof_CC", decProofCc).createLogInfo());
                throw new IllegalStateException("Proof invalid for tenantId " + tenantId + " eeId " + eeId
                    + " verificationCardSetId " + verificationCardSetId + " verificationCardId " + verificationCardId);
            }
        }

        if (!commitmentPublicKeys.isEmpty()) {
            throw new IllegalStateException("Resulting proof set did not contain all public keys expected for tenantId "
                + tenantId + " eeId " + eeId + " verificationCardSetId " + verificationCardSetId
                + " verificationCardId " + verificationCardId);
        }
    }

    private List<ElGamalPublicKey> getCommitmentPublicKeys(VoteVerificationContextData voteVerificationContextData)
            throws GeneralCryptoLibException {
        String[] commitmentPublicKeyBase64Strings =
            voteVerificationContextData.getNonCombinedChoiceCodesEncryptionPublicKeys().split(";");
        List<ElGamalPublicKey> commitmentPublicKeys = new ArrayList<>(commitmentPublicKeyBase64Strings.length);
        for (String commitmentPublicKeyBase64String : commitmentPublicKeyBase64Strings) {
            String commitmentPublicKeyString =
                new String(Base64.getDecoder().decode(commitmentPublicKeyBase64String.getBytes(StandardCharsets.UTF_8)),
                    StandardCharsets.UTF_8);
            ElGamalPublicKey commitmentPublicKey = ElGamalPublicKey.fromJson(commitmentPublicKeyString);
            commitmentPublicKeys.add(commitmentPublicKey);
        }
        return commitmentPublicKeys;
    }

    private VoteVerificationContextData getVoteVerificationContextData(String tenantId, String eeId,
            String verificationCardSetId) throws ResourceNotFoundException {
        VerificationContentEntity verificationContent = verificationContentRepository
            .findByTenantIdElectionEventIdVerificationCardSetId(tenantId, eeId, verificationCardSetId);

        VoteVerificationContextData voteVerificationContextData;
        try {
            voteVerificationContextData =
                ObjectMappers.fromJson(verificationContent.getJson(), VoteVerificationContextData.class);
        } catch (IOException e) {
            throw new IllegalStateException("Verification context data was not valid for tenantId " + tenantId
                + " eeId " + eeId + " verificationCardSetId " + verificationCardSetId, e);
        }
        return voteVerificationContextData;
    }

    private ZpSubgroup getEncryptionParametersGroup(VoteVerificationContextData voteVerificationContextData)
            throws ResourceNotFoundException, GeneralCryptoLibException {
        EncryptionParameters encryptionParameters = voteVerificationContextData.getEncryptionParameters();
        String generatorEncryptParam = encryptionParameters.getG();
        String pEncryptParam = encryptionParameters.getP();
        String qEncryptParam = encryptionParameters.getQ();

        BigInteger generatorEncryptParamBigInteger = new BigInteger(generatorEncryptParam);
        BigInteger pEncryptParamBigInteger = new BigInteger(pEncryptParam);
        BigInteger qEncryptParamBigInteger = new BigInteger(qEncryptParam);
        return new ZpSubgroup(generatorEncryptParamBigInteger, pEncryptParamBigInteger, qEncryptParamBigInteger);
    }

    private List<ZpGroupElement> combineChoiceCodeNodesDecryptionContributions(
            List<List<ZpGroupElement>> ccnContributions, int numRequiredElements) throws GeneralCryptoLibException {

        LOG.info("Combining choice code nodes contributions for the decryption");

        Optional<List<ZpGroupElement>> result = ccnContributions.stream().reduce(this::combineZpGroupElementLists);

        if (!result.isPresent()) {
            secureLoggerWriter.log(Level.ERROR,
                new LogContent.LogContentBuilder()
                    .logEvent(VoteVerificationLogEvents.PARTIAL_DECRYPTIONS_CANNOT_BE_CONSOLIDATED)
                    .additionalInfo(VoteVerificationLogConstants.INFO_TRACK_ID, trackId.getTrackId())
                    .additionalInfo(VoteVerificationLogConstants.INFO_ERR_DESC,
                        "Partial Decryptions cannot be consolidated")
                    .createLogInfo());
            LOG.error("Error combining the Choice Code Nodes contributions");
            throw new IllegalStateException("Choice Code Nodes did not return their contributions");
        }
        return compressZpGroupElementListIfNecessary(result.get(), numRequiredElements);
    }

    /**
     * Combines (multiplies its elements) two lists of ZpGroupElement. It is
     * required that both input ZpGroupElement lists have the same length,
     * otherwise an IllegalStateException is thrown
     *
     * @param zpgel1
     * @param zpgel2
     * @return the list of combined (multiplied) ZpGroupElements
     */
    private List<ZpGroupElement> combineZpGroupElementLists(List<ZpGroupElement> zpgel1, List<ZpGroupElement> zpgel2) {

        LOG.info("Combining ZpGroupElement lists");

        if (zpgel1.size() != zpgel2.size()) {
            LOG.error("ZpGroupElement lists have different number of elements. "
                + "Both lists need to have same size to combine them.");
            throw new IllegalStateException("ZpGroupElement lists have different number of elements. "
                + "Both lists need to have same size to combine them.");
        }

        List<ZpGroupElement> result = new ArrayList<>(zpgel1.size());
        try {
            for (int i = 0; i < zpgel1.size(); i++) {
                result.add(zpgel1.get(i).multiply(zpgel2.get(i)));
            }
            return result;
        } catch (GeneralCryptoLibException e) {
            LOG.error("ZpGroupElement lists could not be combined. " + e.getMessage());
            throw new IllegalStateException("ZpGroupElement lists could not be combined. " + e.getMessage(), e);
        }
    }

    /**
     * Compress a list of ZpGroup elements if the size of the list is longer
     * than the number of required messages. Otherwise return the same list.
     *
     * @param inputList
     *            The list of group elements to be compressed
     * @param numRequiredMessages
     *            The number of the required messages
     * @return The result of the compression process.
     * @throws GeneralCryptoLibException
     */
    private List<ZpGroupElement> compressZpGroupElementListIfNecessary(final List<ZpGroupElement> inputList,
            final int numRequiredMessages) throws GeneralCryptoLibException {

        if (inputList.size() <= numRequiredMessages) {
            return inputList;
        }

        List<ZpGroupElement> untouchedElements = inputList.subList(0, numRequiredMessages - 1);

        List<ZpGroupElement> elementsToBeCombined = inputList.subList(numRequiredMessages - 1, inputList.size());

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

        ZpGroupElement compressedElements = grpElementsCompressor.compress(elementsToBeCombined);

        untouchedElements.add(compressedElements);

        return untouchedElements;
    }
}
