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

import static java.util.Collections.singletonList;

import java.io.BufferedReader;
import java.io.IOException;
import java.math.BigInteger;
import java.nio.charset.StandardCharsets;
import java.nio.file.Path;
import java.security.PrivateKey;
import java.security.cert.X509Certificate;
import java.util.ArrayList;
import java.util.Collection;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.stream.Collectors;
import javax.json.JsonObject;

import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.stereotype.Service;

import com.google.common.io.Files;
import com.scytl.cryptolib.api.elgamal.ElGamalServiceAPI;
import com.scytl.cryptolib.api.exceptions.GeneralCryptoLibException;
import com.scytl.cryptolib.api.primitives.PrimitivesServiceAPI;
import com.scytl.cryptolib.elgamal.bean.ElGamalComputationsValues;
import com.scytl.cryptolib.elgamal.bean.ElGamalEncryptionParameters;
import com.scytl.cryptolib.elgamal.bean.ElGamalKeyPair;
import com.scytl.cryptolib.elgamal.bean.ElGamalPublicKey;
import com.scytl.cryptolib.elgamal.cryptoapi.CryptoAPIElGamalEncrypter;
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.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.choices.codes.BallotCastingKeyGenerator;
import com.scytl.products.ov.choices.codes.PartialCodeGenerator;
import com.scytl.products.ov.commons.beans.Ballot;
import com.scytl.products.ov.commons.beans.Contest;
import com.scytl.products.ov.commons.beans.ElectionOption;
import com.scytl.products.ov.commons.beans.domain.model.messaging.PayloadSignatureException;
import com.scytl.products.ov.commons.beans.exceptions.ResourceNotFoundException;
import com.scytl.products.ov.commons.dto.ChoiceCodeGenerationInput;
import com.scytl.products.ov.commons.dto.ChoiceCodeGenerationReqPayload;
import com.scytl.products.ov.commons.dto.ElGamalComputationsValuesCodec;
import com.scytl.products.ov.commons.sign.PayloadSigner;
import com.scytl.products.ov.commons.util.JsonUtils;
import com.scytl.products.ov.config.exceptions.CreateVotingCardSetException;
import com.scytl.products.ov.config.exceptions.specific.GenerateVerificationCardCodesException;
import com.scytl.products.ov.sdm.application.log.SdmSecureLogEvent;
import com.scytl.products.ov.sdm.domain.model.ballot.BallotRepository;
import com.scytl.products.ov.sdm.domain.model.ballotbox.BallotBoxRepository;
import com.scytl.products.ov.sdm.domain.model.status.InvalidStatusTransitionException;
import com.scytl.products.ov.sdm.domain.model.status.Status;
import com.scytl.products.ov.sdm.domain.service.ConfigConstants;
import com.scytl.products.ov.sdm.domain.service.VerificationCardDataGeneratorService;
import com.scytl.products.ov.sdm.infrastructure.JsonConstants;
import com.scytl.products.ov.sdm.infrastructure.cc.ChoiceCodeGenerationRequestPayloadRepository;
import com.scytl.products.ov.sdm.infrastructure.cc.PayloadStorageException;
import com.scytl.products.ov.sdm.infrastructure.service.ConfigurationEntityStatusService;

/**
 * Service that deals with the pre-computation of voting card sets.
 */
@Service
public class VotingCardSetPrecomputationService extends BaseVotingCardSetService {

    private static final Logger LOGGER = LoggerFactory.getLogger(VotingCardSetPrecomputationService.class);

    @Autowired
    private BallotBoxRepository ballotBoxRepository;

    @Autowired
    private ConfigurationEntityStatusService configurationEntityStatusService;

    @Autowired
    private VerificationCardDataGeneratorService verificationCardDataGeneratorService;

    @Autowired
    private SecureLoggingWriter secureLogger;

    @Autowired
    private BallotRepository ballotRepository;

    @Autowired
    private BallotCastingKeyGenerator ballotCastingKeyGenerator;

    @Autowired
    private ElGamalServiceAPI elGamalService;

    @Autowired
    private ElectionEventService electionEventService;

    @Autowired
    private ElGamalComputationsValuesCodec codec;

    @Autowired
    private ChoiceCodeGenerationRequestPayloadRepository choiceCodeGenerationRequestPayloadRepository;

    @Autowired
    private PayloadSigner payloadSigner;

    @Autowired
    private AdminBoardService adminBoardService;

    @Autowired
    private PrimitivesServiceAPI primitivesService;

    private PartialCodeGenerator<ZpGroupElement> partialCodeGenerator = new PartialCodeGenerator<>();

    @Value("${OR_URL}")
    private String orchestratorUrl;

    @Value("${tenantID}")
    private String tenantId;

    @Value("${choiceCodeGenerationChunkSize:100}")
    private int chunkSize;

    @Autowired
    IdleStatusService idleStatusService;

    /**
     * Pre-compute a voting card set.
     *
     * @throws InvalidStatusTransitionException
     *             if the original status does not allow pre-computing
     * @throws IOException
     * @throws GeneralCryptoLibException
     * @throws PayloadStorageException
     * @throws PayloadSignatureException
     */
    public void precompute(final String votingCardSetId, final String electionEventId, final String signingKeyPEM,
            final String adminBoardId)
            throws ResourceNotFoundException, InvalidStatusTransitionException, IOException, GeneralCryptoLibException,
            PayloadStorageException, PayloadSignatureException {

        if (!idleStatusService.isIdReady(votingCardSetId)) {
            return;
        }

        try {

            LOGGER.info("Starting pre-computation of voting card set {}...", votingCardSetId);

            Status fromStatus = Status.LOCKED;
            Status toStatus = Status.PRECOMPUTED;

            validateIds(votingCardSetId, electionEventId);

            checkVotingCardSetStatusTransition(electionEventId, votingCardSetId, fromStatus, toStatus);

            secureLog(SdmSecureLogEvent.PRECOMPUTING_VCS, electionEventId, votingCardSetId);

            // Get the voting card set.
            final JsonObject votingCardSet =
                votingCardSetRepository.getVotingCardSetJson(electionEventId, votingCardSetId);

            // Extract the number of voting cards in the set.
            final int numberOfVotingCardsToGenerate =
                votingCardSet.getInt(JsonConstants.JSON_ATTRIBUTE_NUMBER_OF_VC_TO_GENERATE);
            LOGGER.info("Generating {} voting cards for voting card set {}...", numberOfVotingCardsToGenerate,
                votingCardSetId);

            // Find the matching verification card set identifier.
            String verificationCardSetId = votingCardSetRepository.getVerificationCardSetId(votingCardSetId);
            String ballotBoxId = getReferencedBallotBoxId(votingCardSet);

            // Get the key that will be used to encrypt the representations and
            // the
            // ballot casting keys.
            ElGamalPublicKey encryptionKey = electionEventService.getPrimeEncryptionPublicKey(electionEventId);

            // Create the ElGamal encrypter.
            CryptoAPIElGamalEncrypter encrypter = elGamalService.createEncrypter(encryptionKey);

            // Get the Zp subgroup for the current election event from the
            // encryption key.
            ZpSubgroup encryptionParametersZpSubgroup = encryptionKey.getGroup();

            choiceCodeGenerationRequestPayloadRepository.remove(electionEventId, verificationCardSetId);
            // Build payloads to request the choice code generation.
            // SV-6485 Since the only objective of generating this payload is to
            // store
            // it, the object creation could be skipped in favour of dumping it
            // directly to disk.

            // Build full-sized chunks (i.e. with `chunkSize` elements)
            int fullChunkCount = numberOfVotingCardsToGenerate / chunkSize;
            for (int i = 0; i < fullChunkCount; i++) {
                createAndStorePayload(i, chunkSize, electionEventId, ballotBoxId, votingCardSetId,
                    verificationCardSetId, encrypter, encryptionParametersZpSubgroup, signingKeyPEM, adminBoardId);
            }

            // Build an eventual last chunk with the remaining elements.
            int lastChunkSize = numberOfVotingCardsToGenerate % chunkSize;
            if (lastChunkSize > 0) {
                createAndStorePayload(fullChunkCount, lastChunkSize, electionEventId, ballotBoxId, votingCardSetId,
                    verificationCardSetId, encrypter, encryptionParametersZpSubgroup, signingKeyPEM, adminBoardId);
            }

            // Update the voting card set status to 'pre-computed'.
            configurationEntityStatusService.update(toStatus.name(), votingCardSetId, votingCardSetRepository);

            secureLogger.log(Level.INFO,
                new LogContent.LogContentBuilder().logEvent(SdmSecureLogEvent.VCS_PRECOMPUTED_SUCCESSFULLY)
                    .objectId(votingCardSetId).electionEvent(electionEventId).createLogInfo());
        } finally {
            idleStatusService.freeIdleStatus(votingCardSetId);
        }
    }

    /**
     * Creates, signs and stores a choice code generation request payload.
     * 
     * @param chunkId
     *            the ID of the chunk this payload represents out of the total
     *            elements of the voting card set.
     * @param chunkSize
     *            the size of the chunk in this payload
     * @param electionEventId
     *            the election event identifier
     * @param ballotBoxId
     *            the ballot box identifier
     * @param votingCardSetId
     *            the voting card set identifier
     * @param verificationCardSetId
     *            the verification card set identifier
     * @param encrypter
     *            the encrypter
     * @param encryptionParametersZpSubgroup
     *            the relevant Zp subgroup
     * @param signingKeyPEM
     *            the singing key in PEM format
     * @param adminBoardId
     *            the administration board identifier
     * @throws IOException
     * @throws ResourceNotFoundException
     * @throws PayloadStorageException
     * @throws PayloadSignatureException
     */
    private void createAndStorePayload(int chunkId, int chunkSize, String electionEventId, String ballotBoxId,
            String votingCardSetId, String verificationCardSetId, CryptoAPIElGamalEncrypter encrypter,
            ZpSubgroup encryptionParametersZpSubgroup, String signingKeyPEM, String adminBoardId)
            throws IOException, ResourceNotFoundException, PayloadStorageException, PayloadSignatureException {
        // Create the payload.
        ChoiceCodeGenerationReqPayload payload = createChoiceCodeGenerationRequestPayload(electionEventId, ballotBoxId,
            verificationCardSetId, chunkId, chunkSize, encrypter, encryptionParametersZpSubgroup);
        LOGGER.info("Payload for voting card set {} with {} voting cards is ready", votingCardSetId, chunkSize);

        // Sign the payload.
        try {
            signPayload(payload, signingKeyPEM, adminBoardId);
        } catch (CertificateManagementException e) {
            throw new PayloadSignatureException(e);
        }

        // Store the payload.
        choiceCodeGenerationRequestPayloadRepository.store(payload);
    }

    private BigInteger hashBallotCastingKey(BigInteger bck, ZpSubgroup encryptionParametersZpSubgroup) {
        try {
            BigInteger bckHash = new BigInteger(primitivesService.getHash(bck.toByteArray()));
            return transformBallotCastingKeyToZpGroupElement(bckHash, encryptionParametersZpSubgroup).getValue();
        } catch (GeneralCryptoLibException e) {
            throw new CreateVotingCardSetException("An error ocurred while hashing the ballot casting key", e);
        }
    }

    /**
     * Signs a choice code generation request payload.
     * 
     * @param payload
     *            the payload to sign
     * @throws PayloadSignatureException
     * @throws CertificateManagementException
     */
    private void signPayload(ChoiceCodeGenerationReqPayload payload, String privateKeyPEM, String adminBoardId)
            throws PayloadSignatureException, CertificateManagementException {
        // Get the admin board's signing key
        PrivateKey signingKey;
        try {
            signingKey = PemUtils.privateKeyFromPem(privateKeyPEM);
        } catch (GeneralCryptoLibException e) {
            throw new PayloadSignatureException(e);
        }

        // Get the admin board's certificate chain
        X509Certificate[] certificateChain = adminBoardService.getCertificateChain(adminBoardId);

        // Sign the payload.
        payload.setSignature(payloadSigner.sign(payload, signingKey, certificateChain));
    }

    /**
     * Creates a payload to request the generation of a voting card set's choice
     * codes.
     * 
     * @param electionEventId
     *            the election event the voting card set belongs to
     * @param votingCardSetId
     *            the identifier of the voting card set to generate the payload
     *            for
     * @param encryptedRepresentations
     *            the voting options' representations, encrypted
     * @param encryptedBallotCastingKeys
     *            the ballot casting keys, encrypted
     * @return a choice code generation request payload
     * @throws ResourceNotFoundException
     * @throws IOException
     */
    private ChoiceCodeGenerationReqPayload createChoiceCodeGenerationRequestPayload(String electionEventId,
            String ballotBoxId, String verificationCardSetId, int chunkId, int numberOfVotingCardsToGenerate,
            CryptoAPIElGamalEncrypter encrypter, ZpSubgroup encryptionParametersZpSubgroup)
            throws IOException, ResourceNotFoundException {

        List<String[]> verificationCardKeyPairs = new ArrayList<>();

        // Generate, encrypt and encode the representations.
        String encodedEncryptedRepresentations = getEncodedEncryptedRepresentationsForBallotBox(ballotBoxId,
            electionEventId, encrypter, encryptionParametersZpSubgroup);

        ChoiceCodeGenerationReqPayload payload;

        // Obtain the verification card IDs.
        try (BufferedReader br = verificationCardDataGeneratorService.precompute(electionEventId, verificationCardSetId,
            numberOfVotingCardsToGenerate)) {

            // Create a new payload.
            payload = new ChoiceCodeGenerationReqPayload(tenantId, electionEventId, verificationCardSetId, chunkId);
            payload
                .setChoiceCodeGenerationInputList(
                    br.lines()
                        .map(vcid -> createChoiceCodeGenerationEntry(encrypter, encryptionParametersZpSubgroup,
                            verificationCardKeyPairs, encodedEncryptedRepresentations, vcid))
                        .collect(Collectors.toList()));
        }

        // Save the generated verification card key pairs to be used in the
        // generation phase
        saveVerificationCardKeyPairs(electionEventId, verificationCardSetId, verificationCardKeyPairs);

        return payload;
    }

    private ChoiceCodeGenerationInput createChoiceCodeGenerationEntry(CryptoAPIElGamalEncrypter encrypter,
            ZpSubgroup encryptionParametersZpSubgroup, List<String[]> verificationCardKeyPairs,
            String encodedEncryptedRepresentations, String vcid) {
        // Generate Verification card key pair
        ElGamalKeyPair keyPair =
            generateVerificationCardKeyPair(vcid, encryptionParametersZpSubgroup, verificationCardKeyPairs);
        // Create a ballot casting key.
        ZpGroupElement elGamalBck = createBallotCastingKey(encryptionParametersZpSubgroup);
        // Exponentiate the BCK to VC key and hash it to prevent the "fake vote
        // confirmation" attack (SV-6259)
        BigInteger exponentiatedBck =
            partialCodeGenerator.generatePartialCode(elGamalBck, keyPair.getPrivateKeys().getKeys().get(0)).getValue();
        BigInteger bckHash = hashBallotCastingKey(exponentiatedBck, encryptionParametersZpSubgroup);
        // Encrypt it.
        ElGamalComputationsValues encryptedBCK =
            encryptValue(elGamalBck.getValue(), encrypter, encryptionParametersZpSubgroup);
        ElGamalComputationsValues encryptedHashedBCK = encryptValue(bckHash, encrypter, encryptionParametersZpSubgroup);
        // Encode it.
        String encodedBallotCastingKey = codec.encodeSingle(encryptedBCK);
        String encodedHashedBallotCastingKey = codec.encodeSingle(encryptedHashedBCK);

        return new ChoiceCodeGenerationInput(vcid, encodedBallotCastingKey, encodedHashedBallotCastingKey,
            encodedEncryptedRepresentations);
    }

    private void saveVerificationCardKeyPairs(String electionEventId, String verificationCardSetId,
            List<String[]> verificationCardKeyPairs) {
        Path verificationCardSetKeyPairsPath = pathResolver.resolve(ConfigConstants.CONFIG_FILES_BASE_DIR)
            .resolve(electionEventId).resolve(ConfigConstants.CONFIG_DIR_NAME_OFFLINE)
            .resolve(ConfigConstants.CONFIG_VERIFICATION_CARDS_KEY_PAIR_DIRECTORY).resolve(verificationCardSetId);

        verificationCardSetKeyPairsPath.toFile().mkdirs();

        for (String[] verificationCardKeyPair : verificationCardKeyPairs) {
            String verificationCardId = verificationCardKeyPair[0];
            String privateKeyPair = verificationCardKeyPair[1];
            String publicKeyPair = verificationCardKeyPair[2];
            Path verificationCardKeyPairFilePath =
                verificationCardSetKeyPairsPath.resolve(verificationCardId + ConfigConstants.KEY);

            try {
                Files.write((privateKeyPair + System.lineSeparator() + publicKeyPair).getBytes(StandardCharsets.UTF_8),
                    verificationCardKeyPairFilePath.toFile());
            } catch (IOException e) {
                throw new IllegalStateException("Error saving verification cards key pairs", e);
            }
        }
    }

    private ElGamalKeyPair generateVerificationCardKeyPair(String vcid, ZpSubgroup encryptionParameters,
            List<String[]> verificationCardKeyPairs) {
        try {
            ElGamalEncryptionParameters elGamalEncryptionParameters = new ElGamalEncryptionParameters(
                encryptionParameters.getP(), encryptionParameters.getQ(), encryptionParameters.getG());

            ElGamalKeyPair keyPair = elGamalService.getElGamalKeyPairGenerator()
                .generateKeys(elGamalEncryptionParameters, ConfigConstants.NUM_KEYS_VERIFICATION_CARD_KEYPAIR);

            verificationCardKeyPairs
                .add(new String[] {vcid, keyPair.getPrivateKeys().toJson(), keyPair.getPublicKeys().toJson() });

            return keyPair;
        } catch (GeneralCryptoLibException e) {
            throw new IllegalStateException("Error while generating verification card key pairs", e);
        }
    }

    private String getReferencedBallotBoxId(JsonObject votingCardSetJson) {
        JsonObject ballotBoxReferenceJson =
            votingCardSetJson.getJsonObject(JsonConstants.JSON_ATTRIBUTE_NAME_BALLOT_BOX);
        return ballotBoxReferenceJson.getString(JsonConstants.JSON_ATTRIBUTE_NAME_ID);
    }

    private Ballot getBallot(String ballotBoxId, String electionEventId) throws ResourceNotFoundException, IOException {
        String ballotId = getBallotIdForBallotBox(ballotBoxId, electionEventId);
        Map<String, Object> attributeValueMap = new HashMap<>();
        attributeValueMap.put(JsonConstants.JSON_STRUCTURE_NAME_ELECTION_EVENT_ID, electionEventId);
        attributeValueMap.put(JsonConstants.JSON_ATTRIBUTE_NAME_ID, ballotId);
        String ballotAsJson = ballotRepository.find(attributeValueMap);
        return objectMapper.readValue(ballotAsJson, Ballot.class);
    }

    private String getBallotIdForBallotBox(String ballotBoxId, String electionEventId) throws IOException {
        JsonObject ballotBox = getBallotBoxObject(ballotBoxId, electionEventId);
        JsonObject ballotObject = ballotBox.getJsonObject(JsonConstants.JSON_ATTRIBUTE_NAME_BALLOT);

        return ballotObject.getString(JsonConstants.JSON_ATTRIBUTE_NAME_ID);
    }

    private JsonObject getBallotBoxObject(String ballotBoxId, String electionEventId) {
        Map<String, Object> attributeValueMap = new HashMap<>();
        attributeValueMap.put(JsonConstants.JSON_STRUCTURE_NAME_ELECTION_EVENT_ID, electionEventId);
        attributeValueMap.put(JsonConstants.JSON_ATTRIBUTE_NAME_ID, ballotBoxId);
        String ballotBoxAsJson = ballotBoxRepository.find(attributeValueMap);

        return JsonUtils.getJsonObject(ballotBoxAsJson);
    }

    /**
     * Gets a ballot box's voting option representations as BigIntegers.
     *
     * @param ballotBoxId
     *            the ballot box with the voting options
     * @param electionEventId
     *            the election event the ballot box belongs to
     * @param encrypter
     *            the encryption engine for representations
     * @param encryptionParametersZpSubgroup
     *            the encryption parameters for the engine
     * @return a stream of encrypted representations
     * @throws IOException
     * @throws ResourceNotFoundException
     */
    private String getEncodedEncryptedRepresentationsForBallotBox(String ballotBoxId, String electionEventId,
            CryptoAPIElGamalEncrypter encrypter, ZpSubgroup encryptionParametersZpSubgroup)
            throws IOException, ResourceNotFoundException {
        return codec.encodeList(
            // From the specified ballot...
            getBallot(ballotBoxId, electionEventId)
                // ... get its contests...
                .getContests().stream()
                // ... get the contest's options...
                .map(Contest::getOptions)
                // ... from each option, get its representation...
                .flatMap(Collection::stream).map(ElectionOption::getRepresentation)
                // ... as a BigInteger...
                .map(BigInteger::new)
                // ... encrypt it...
                .map(value -> encryptValue(value, encrypter, encryptionParametersZpSubgroup))
                // ... and encode it all
                .collect(Collectors.toList()));
    }

    private ZpGroupElement createBallotCastingKey(ZpSubgroup zpSubgroup) {

        String ballotCastingKeyStr;
        try {
            ballotCastingKeyStr = ballotCastingKeyGenerator.generateBallotCastingKey();
        } catch (Exception e) {
            LOGGER.error("Error generating the ballot casting key: ", e);
            throw new GenerateVerificationCardCodesException("Error generating the ballot casting key", e);
        }

        ZpGroupElement ballotCastingKey =
            transformBallotCastingKeyToZpGroupElement(new BigInteger(ballotCastingKeyStr), zpSubgroup);

        return ballotCastingKey;
    }

    private ZpGroupElement transformBallotCastingKeyToZpGroupElement(final BigInteger ballotCastingKey,
            final ZpSubgroup zpsubgroup) {
        try {
            BigInteger ballotCastingKeySquared = ballotCastingKey.multiply(ballotCastingKey);
            return new ZpGroupElement(ballotCastingKeySquared, zpsubgroup);
        } catch (GeneralCryptoLibException e) {
            throw new CreateVotingCardSetException("An error occurred when encoding the Ballot Casting Key.", e);
        }
    }

    /**
     * Encrypts a ballot casting keys.
     *
     * @param electionEventId
     *            the election event the voting card set belongs to
     * @param verificationCardSetId
     *            the identifier of the voting card set
     * @param ballotCastingKeys
     *            the ballot casting keys to encrypt
     * @return the encrypted ballot casting keys
     * @throws ResourceNotFoundException
     * @throws IOException
     */
    private ElGamalComputationsValues encryptValue(BigInteger value, CryptoAPIElGamalEncrypter encrypter,
            ZpSubgroup encryptionKeyGroup) {
        try {
            ZpGroupElement element = new ZpGroupElement(value, encryptionKeyGroup);
            return encrypter.encryptGroupElements(singletonList(element)).getComputationValues();
        } catch (GeneralCryptoLibException e) {
            throw new IllegalStateException("Failed to encrypt ballot casting keys or voting option representation", e);
        }
    }
}
