/*
 * 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.util.ArrayList;
import java.util.Arrays;
import java.util.Base64;
import java.util.Base64.Decoder;
import java.util.List;
import java.util.Map;
import java.util.stream.Collectors;

import javax.ejb.Stateless;
import javax.inject.Inject;
import javax.json.JsonObject;

import org.apache.commons.lang3.StringUtils;
import org.apache.commons.lang3.exception.ExceptionUtils;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import com.fasterxml.jackson.core.JsonProcessingException;
import com.fasterxml.jackson.core.type.TypeReference;
import com.scytl.cryptolib.api.exceptions.GeneralCryptoLibException;
import com.scytl.cryptolib.elgamal.bean.ElGamalPublicKey;
import com.scytl.cryptolib.mathematical.groups.activity.GroupElementsCompressor;
import com.scytl.cryptolib.mathematical.groups.impl.ZpGroupElement;
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.ov.commons.beans.EncryptionParameters;
import com.scytl.products.ov.commons.beans.VoteVerificationContextData;
import com.scytl.products.ov.commons.beans.domain.model.vote.Vote;
import com.scytl.products.ov.commons.beans.exceptions.ResourceNotFoundException;
import com.scytl.products.ov.commons.beans.validation.CryptographicOperationException;
import com.scytl.products.ov.commons.dto.ChoiceCodeAndComputeResults;
import com.scytl.products.ov.commons.dto.ChoiceCodesVerificationDecryptResPayload;
import com.scytl.products.ov.commons.dto.ChoiceCodesVerificationInput;
import com.scytl.products.ov.commons.dto.ChoiceCodesVerificationResPayload;
import com.scytl.products.ov.commons.dto.ComputeResults;
import com.scytl.products.ov.commons.dto.PartialCodesInput;
import com.scytl.products.ov.commons.dto.VoteAndComputeResults;
import com.scytl.products.ov.commons.logging.service.AuthTokenHashService;
import com.scytl.products.ov.commons.tracking.TrackIdInstance;
import com.scytl.products.ov.commons.ui.ws.rs.ObjectMappers;
import com.scytl.products.ov.commons.util.JsonUtils;
import com.scytl.products.ov.vv.domain.model.choicecode.CodesComputeResults;
import com.scytl.products.ov.vv.domain.model.choicecode.CodesDecryptionResults;
import com.scytl.products.ov.vv.domain.model.content.ElectoralAuthorityEntity;
import com.scytl.products.ov.vv.domain.model.content.ElectoralAuthorityRepository;
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.verification.Verification;
import com.scytl.products.ov.vv.domain.model.verification.VerificationRepository;
import com.scytl.products.ov.vv.domain.model.verificationset.VerificationSetEntity;
import com.scytl.products.ov.vv.domain.model.verificationset.VerificationSetRepository;
import com.scytl.products.ov.vv.infrastructure.log.SecureLoggerHelper;
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;

/**
 * Choice code generation.
 */
@Stateless
public class ChoiceCodesServiceImpl implements ChoiceCodesService {

    private static final String JSON_PUBLIC_KEY = "publicKey";

    private static final String JSON_CHOICE_CODES_PUBLIC_KEY = "choicesCodesEncryptionPublicKey";

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

    @Inject
    private VerificationSetRepository verificationSetRepository;

    @Inject
    private ElectoralAuthorityRepository electoralAuthorityRepository;

    @Inject
    private VerificationContentRepository verificationContentRepository;

    @Inject
    private VerificationRepository verificationRepository;

    @Inject
    private CodesDecrypterService codesDecrypterService;

    @Inject
    private CodesComputeService codesComputeService;

    @Inject
    private AuthTokenHashService authTokenHashService;

    @Inject
    private LongCodesService longCodesService;

    @Inject
    private ShortCodesService shortCodesService;

    @Inject
    private SecureLoggingWriter secureLoggerWriter;

    @Inject
    private SecureLoggerHelper secureLoggerHelper;

    @Inject
    private TrackIdInstance trackIdInstance;

    @Inject
    private CorrectnessParser correctnessParser;

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

    /**
     * @throws ResourceNotFoundException
     * @throws IOException
     * @throws GeneralCryptoLibException
     * @throws NumberFormatException
     * @throws CryptographicOperationException
     * @see com.scytl.products.ov.vv.service.ChoiceCodesService#generateChoiceCodes(java.lang.String,
     *      java.lang.String, java.lang.String,
     *      com.scytl.products.ov.commons.beans.domain.model.vote.Vote)
     */
    @Override
    public ChoiceCodeAndComputeResults generateChoiceCodes(final String tenantId, final String electionEventId,
            final String verificationCardId, final VoteAndComputeResults voteAndComputeResults)
            throws ResourceNotFoundException, GeneralCryptoLibException, IOException, CryptographicOperationException {
        LOG.info("Generating choice codes for tenant: {}, election event: {} and verification card: {}.", tenantId,
            electionEventId, verificationCardId);

        Vote vote = voteAndComputeResults.getVote();
        String[] encryptedPartialCodes = vote.getEncryptedPartialChoiceCodes().split(";");
        List<BigInteger> partialCodesElements = new ArrayList<>();
        for (String encryptedPartialCode : encryptedPartialCodes) {
            partialCodesElements.add(new BigInteger(encryptedPartialCode));
        }
        ZpSubgroup zpSubgroup = getZpSubgroupFromVote(vote);

        Verification verification = verificationRepository.findByTenantIdElectionEventIdVerificationCardId(tenantId,
            electionEventId, verificationCardId);
        String verificationCardSetId = verification.getVerificationCardSetId();

        CodesComputeResults computationResults;

        ComputeResults computeAndDecryptResults = voteAndComputeResults.getComputeResults();

        if (computeAndDecryptResults == null || computeAndDecryptResults.getComputationResults() == null
            || computeAndDecryptResults.getComputationResults().isEmpty()) {

            computationResults =
                codesComputeService.computePartialCodes(tenantId, electionEventId, verificationCardSetId,
                    verificationCardId, zpSubgroup, createVerificationPayload(partialCodesElements, vote));

        } else {
            String computationResultsString = computeAndDecryptResults.getComputationResults();
            List<ChoiceCodesVerificationResPayload> ccnContributions = ObjectMappers.fromJson(computationResultsString,
                new TypeReference<List<ChoiceCodesVerificationResPayload>>() {
                });
            computationResults = codesComputeService.computePartialCodesFromComputeResult(zpSubgroup,
                partialCodesElements, ccnContributions);
        }

        Map<BigInteger, ZpGroupElement> optionsToComputedPartialCodes =
            computationResults.getCombinedPartialChoiceCodes();

        String hashAT;
        try {
            hashAT = authTokenHashService.hash(JsonUtils.getJsonObject(vote.getAuthenticationToken()));
        } catch (GeneralCryptoLibException hashException) {
            LOG.error("Could not calculate the hash for the authentication token", hashException);
            hashAT = "ERROR";
        }

        List<ZpGroupElement> computedPartialCodes =
            partialCodesElements.stream().map(optionsToComputedPartialCodes::get).collect(Collectors.toList());

        String hashPcc;
        try {
            hashPcc = SecureLoggerUtil.getHashForLogFromBigIntegerCollection(partialCodesElements);
        } catch (GeneralCryptoLibException hashException) {
            LOG.error("Could not calculate the hash for the partial choice codes to compute", hashException);
            hashPcc = "ERROR";
        }
        List<ZpGroupElement> orderedConsolidatedPartialCodes =
            SecureLoggerUtil.getValuesOrderedByKeyList(partialCodesElements, optionsToComputedPartialCodes);
        String hashPccc;
        try {
            hashPccc =
                SecureLoggerUtil.getHashForLogFromZpGroupElementIntegerCollection(orderedConsolidatedPartialCodes);
        } catch (GeneralCryptoLibException hashException) {
            LOG.error("Could not calculate the hash for the consolidated computed partial choice codes", hashException);
            hashPccc = "ERROR";
        }

        // [SL PCCCON-4]
        // Partial Choice Codes consolidation -- Partial Choice Codes correctly
        // consolidated
        secureLoggerWriter.log(Level.INFO,
            new LogContent.LogContentBuilder().logEvent(VoteVerificationLogEvents.PARTIAL_CHOICE_CODE_CONSOLIDATED)
                .objectId(hashAT).user(verificationCardId).electionEvent(electionEventId)
                .additionalInfo(VoteVerificationLogConstants.INFO_TRACK_ID, trackIdInstance.getTrackId())
                .additionalInfo("#hash_pCC", hashPcc).additionalInfo("#hash_pCCC", hashPccc)
                .additionalInfo("#hashVote", secureLoggerHelper.hash(vote)).createLogInfo());

        CodesDecryptionResults decryptionResult;

        try {
            if (computeAndDecryptResults == null || computeAndDecryptResults.getComputationResults() == null
                || computeAndDecryptResults.getComputationResults().isEmpty()) {

                decryptionResult = codesDecrypterService.decryptPartialCodes(tenantId, electionEventId,
                    verificationCardSetId, verificationCardId, computedPartialCodes);

            } else {
                String decyptionResultsString = computeAndDecryptResults.getDecryptionResults();
                List<ChoiceCodesVerificationDecryptResPayload> ccnContributions = ObjectMappers.fromJson(
                    decyptionResultsString, new TypeReference<List<ChoiceCodesVerificationDecryptResPayload>>() {
                    });
                decryptionResult = codesDecrypterService.decryptPartialCodesFromDecryptResult(electionEventId,
                    verificationCardId, computedPartialCodes, ccnContributions);

            }
        } catch (CryptographicOperationException e) {

            secureLoggerWriter.log(Level.ERROR,
                new LogContent.LogContentBuilder()
                    .logEvent(VoteVerificationLogEvents.ERROR_DECRYPTING_PARTIAL_CHOICE_CODES)
                    .objectId(verificationCardId).user(verificationCardId).electionEvent(electionEventId)
                    .additionalInfo(VoteVerificationLogConstants.INFO_TRACK_ID, trackIdInstance.getTrackId())
                    .additionalInfo(VoteVerificationLogConstants.INFO_HASH_VOTE, secureLoggerHelper.hash(vote))
                    .additionalInfo(VoteVerificationLogConstants.INFO_VERIFICATION_CARD_ID, verificationCardId)
                    .additionalInfo(VoteVerificationLogConstants.INFO_ERR_DESC,
                        "Error decrypting partial choice codes: " + ExceptionUtils.getRootCauseMessage(e))
                    .createLogInfo());
            throw e;
        }

        List<ZpGroupElement> partialCodesList = decryptionResult.getCombinedZpGroupElementLists();

        List<byte[]> longCodes;
        try {
            List<String[]> orderedListOfCorrectnessIds = correctnessParser.parse(vote.getCorrectnessIds());

            longCodes = longCodesService.calculateLongCodes(electionEventId, verificationCardId, partialCodesList,
                orderedListOfCorrectnessIds);
        } catch (CryptographicOperationException e) {
            secureLoggerWriter
                .log(Level.ERROR,
                    new LogContent.LogContentBuilder()
                        .logEvent(VoteVerificationLogEvents.ERROR_COMPUTING_LONG_CHOICE_CODES)
                        .objectId(verificationCardId).user(verificationCardId).electionEvent(electionEventId)
                        .additionalInfo(VoteVerificationLogConstants.INFO_TRACK_ID, trackIdInstance.getTrackId())
                        .additionalInfo(VoteVerificationLogConstants.INFO_HASH_VOTE, secureLoggerHelper.hash(vote))
                        .additionalInfo(VoteVerificationLogConstants.INFO_VERIFICATION_CARD_ID, verificationCardId)
                        .additionalInfo(VoteVerificationLogConstants.INFO_ERR_DESC,
                            "Error computing long choice codes: " + ExceptionUtils.getRootCauseMessage(e))
                        .createLogInfo());
            throw e;
        }

        List<String> shortCodes;
        try {
            shortCodes = shortCodesService.retrieveShortCodes(tenantId, electionEventId, verificationCardId, longCodes);
        } catch (ResourceNotFoundException | CryptographicOperationException e) {
            secureLoggerWriter.log(Level.ERROR,
                new LogContent.LogContentBuilder().logEvent(VoteVerificationLogEvents.ERROR_RETRIEVING_CHOICE_CODES)
                    .objectId(verificationCardId).user(verificationCardId).electionEvent(electionEventId)
                    .additionalInfo(VoteVerificationLogConstants.INFO_TRACK_ID, trackIdInstance.getTrackId())
                    .additionalInfo(VoteVerificationLogConstants.INFO_HASH_VOTE, secureLoggerHelper.hash(vote))
                    .additionalInfo(VoteVerificationLogConstants.INFO_VERIFICATION_CARD_ID, verificationCardId)
                    .additionalInfo(VoteVerificationLogConstants.INFO_ERR_DESC,
                        "Error retrieving choice codes: " + ExceptionUtils.getRootCauseMessage(e))
                    .createLogInfo());
            throw e;
        }

        String serializedShortCodes = StringUtils.join(shortCodes, ';');

        ChoiceCodeAndComputeResults choiceCodes = new ChoiceCodeAndComputeResults();
        choiceCodes.setChoiceCodes(serializedShortCodes);
        try {
            choiceCodes.setComputationResults(ObjectMappers.toJson(computationResults.getComputationResults()));
            choiceCodes.setDecryptionResults(ObjectMappers.toJson(decryptionResult.getDecryptResult()));
        } catch (JsonProcessingException e) {
            LOG.error("Could not serialize choice codes computation results", e);
            throw new IllegalStateException(e);
        }

        secureLoggerWriter.log(Level.INFO,
            new LogContent.LogContentBuilder().logEvent(VoteVerificationLogEvents.CHOICE_CODES_RETRIEVED)
                .objectId(verificationCardId).user(verificationCardId).electionEvent(electionEventId)
                .additionalInfo(VoteVerificationLogConstants.INFO_TRACK_ID, trackIdInstance.getTrackId())
                .additionalInfo(VoteVerificationLogConstants.INFO_HASH_VOTE, secureLoggerHelper.hash(vote))
                .additionalInfo(VoteVerificationLogConstants.INFO_VERIFICATION_CARD_ID, verificationCardId)
                .createLogInfo());
        return choiceCodes;
    }

    private ZpSubgroup getZpSubgroupFromVote(Vote vote)
            throws ResourceNotFoundException, GeneralCryptoLibException, NumberFormatException, IOException {

        VerificationContentEntity verificationContentEntity;

        try {
            Verification verification = verificationRepository.findByTenantIdElectionEventIdVerificationCardId(
                vote.getTenantId(), vote.getElectionEventId(), vote.getVerificationCardId());
            verificationContentEntity =
                verificationContentRepository.findByTenantIdElectionEventIdVerificationCardSetId(vote.getTenantId(),
                    vote.getElectionEventId(), verification.getVerificationCardSetId());
        } catch (ResourceNotFoundException e) {
            LOG.error("Error creating ZpSubgroup from Vote.");
            throw e;
        }

        ZpSubgroup mathematicalGroup;
        try {
            VoteVerificationContextData voteVerificationContextData =
                ObjectMappers.fromJson(verificationContentEntity.getJson(), VoteVerificationContextData.class);
            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);
            mathematicalGroup =
                new ZpSubgroup(generatorEncryptParamBigInteger, pEncryptParamBigInteger, qEncryptParamBigInteger);
        } catch (GeneralCryptoLibException | NumberFormatException | IOException e) {
            LOG.error("Error creating ZpSubgroup from Vote.");
            throw e;
        }

        return mathematicalGroup;
    }

    private PartialCodesInput createVerificationPayload(List<BigInteger> partialCodesElements, Vote vote)
            throws ResourceNotFoundException {

        try {
            Verification verification = verificationRepository.findByTenantIdElectionEventIdVerificationCardId(
                vote.getTenantId(), vote.getElectionEventId(), vote.getVerificationCardId());

            VerificationContentEntity verificationContent =
                verificationContentRepository.findByTenantIdElectionEventIdVerificationCardSetId(vote.getTenantId(),
                    vote.getElectionEventId(), verification.getVerificationCardSetId());

            VerificationSetEntity verificationSetEntity =
                verificationSetRepository.findByTenantIdElectionEventIdVerificationCardSetId(vote.getTenantId(),
                    vote.getElectionEventId(), verification.getVerificationCardSetId());

            VoteVerificationContextData voteVerificationContextData =
                ObjectMappers.fromJson(verificationContent.getJson(), VoteVerificationContextData.class);

            JsonObject verificationSetJSON = JsonUtils.getJsonObject(verificationSetEntity.getJSON());

            String electoralAuthorityId = voteVerificationContextData.getElectoralAuthorityId();

            ElectoralAuthorityEntity electoralAuthority =
                electoralAuthorityRepository.findByTenantIdElectionEventIdElectoralAuthorityId(vote.getTenantId(),
                    vote.getElectionEventId(), electoralAuthorityId);

            String electoralPublicKeyString =
                JsonUtils.getJsonObject(electoralAuthority.getJson()).getString(JSON_PUBLIC_KEY);

            EncryptionParameters encryptionParameters = voteVerificationContextData.getEncryptionParameters();
            String generatorEncryptParam = encryptionParameters.getG();
            String pEncryptParam = encryptionParameters.getP();
            String qEncryptParam = encryptionParameters.getQ();
            ZpSubgroup mathematicalGroup = new ZpSubgroup(new BigInteger(generatorEncryptParam),
                new BigInteger(pEncryptParam), new BigInteger(qEncryptParam));
            Decoder decoder = Base64.getDecoder();
            String choiceCodesEncryptionPKString =
                new String(decoder.decode(verificationSetJSON.getString(JSON_CHOICE_CODES_PUBLIC_KEY)));
            ElGamalPublicKey elGamalChoiceCodesPublicKey = ElGamalPublicKey.fromJson(choiceCodesEncryptionPKString);
            ZpGroupElement compressedChoiceCodesPKGroupElement =
                compressor.compress(elGamalChoiceCodesPublicKey.getKeys());
            ElGamalPublicKey compressesedChoiceCodesEncryptionPK =
                new ElGamalPublicKey(Arrays.asList(compressedChoiceCodesPKGroupElement), mathematicalGroup);

            ChoiceCodesVerificationInput choiceCodesVerificationInput = new ChoiceCodesVerificationInput();
            choiceCodesVerificationInput.setVote(ObjectMappers.toJson(vote));
            choiceCodesVerificationInput.setElGamalElectoralPublicKey(
                new String(decoder.decode(electoralPublicKeyString), StandardCharsets.UTF_8));
            choiceCodesVerificationInput
                .setCompressesedChoiceCodesEncryptionPK(compressesedChoiceCodesEncryptionPK.toJson());
            PartialCodesInput partialCodesInput = new PartialCodesInput();
            partialCodesInput.setChoiceCodesVerificationInput(choiceCodesVerificationInput);
            partialCodesInput.setPartialCodesElements(partialCodesElements);
            return partialCodesInput;
        } catch (IOException | GeneralCryptoLibException e) {
            LOG.error("", e);
            throw new ResourceNotFoundException(e.getMessage(), e);
        }

    }
}
