/*
 * 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.ei.services.domain.service.ballotbox;

import com.scytl.cryptolib.api.exceptions.CryptoLibException;
import com.fasterxml.jackson.databind.ObjectMapper;
import com.scytl.cryptolib.api.exceptions.GeneralCryptoLibException;
import com.scytl.cryptolib.elgamal.bean.ElGamalPublicKey;
import com.scytl.cryptolib.mathematical.groups.impl.ZpSubgroup;
import com.scytl.products.ov.commons.beans.domain.model.ballotbox.BallotBoxId;
import com.scytl.products.ov.commons.beans.domain.model.messaging.PayloadSignatureException;
import com.scytl.products.ov.commons.beans.domain.model.vote.EncryptedVote;
import java.io.IOException;
import java.io.OutputStream;
import java.nio.charset.StandardCharsets;
import java.security.InvalidKeyException;
import java.security.PrivateKey;
import java.security.Signature;
import java.security.SignatureException;
import java.time.ZoneOffset;
import java.time.ZonedDateTime;
import java.util.Base64;
import java.util.List;

import javax.ejb.TransactionAttribute;
import javax.ejb.TransactionAttributeType;
import javax.inject.Inject;

import org.apache.commons.io.output.CloseShieldOutputStream;
import org.apache.commons.lang3.StringUtils;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import com.scytl.products.ov.commons.beans.domain.model.vote.Vote;
import com.scytl.products.ov.commons.beans.domain.model.vote.VoteSetIdImpl;
import com.scytl.products.ov.commons.beans.exceptions.CleansedBallotBoxRepositoryException;
import com.scytl.products.ov.commons.beans.exceptions.CleansedBallotBoxServiceException;
import com.scytl.products.ov.commons.beans.exceptions.DuplicateEntryException;
import com.scytl.products.ov.commons.beans.exceptions.ResourceNotFoundException;
import com.scytl.products.ov.commons.beans.validation.CryptographicOperationException;
import com.scytl.products.ov.commons.crypto.PrivateKeyForObjectRepository;
import com.scytl.products.ov.commons.dto.mixing.MixingPayload;
import com.scytl.products.ov.commons.dto.mixing.MixingPayloadImpl;
import com.scytl.products.ov.commons.sign.PayloadSigner;
import com.scytl.products.ov.commons.signature.SignatureFactory;
import com.scytl.products.ov.commons.signature.SignatureOutputStream;
import com.scytl.products.ov.commons.util.JsonUtils;
import com.scytl.products.ov.ei.services.domain.common.csv.CleansedExportedBallotBoxItemWriter;
import com.scytl.products.ov.ei.services.domain.model.ballotbox.BallotBoxInformationRepository;
import com.scytl.products.ov.ei.services.domain.model.ballotbox.CleansedBallotBox;
import com.scytl.products.ov.ei.services.domain.model.ballotbox.CleansedBallotBoxRepository;
import com.scytl.products.ov.ei.services.domain.model.ballotbox.CleansedExportedBallotBoxItem;
import com.scytl.products.ov.ei.services.domain.model.content.ElectoralAuthorityEntity;
import com.scytl.products.ov.ei.services.domain.model.content.ElectoralAuthorityRepository;
import com.scytl.products.ov.ei.services.domain.model.tenant.EiTenantSystemKeys;
import com.scytl.products.ov.ei.services.domain.model.vote.SuccessfulVote;
import com.scytl.products.ov.ei.services.domain.service.ballotbox.PartitionCalculator.PartitionDetails;
import com.scytl.products.ov.ei.services.infrastructure.persistence.CleansedBallotBoxAccess;
import com.scytl.products.ov.ei.services.infrastructure.persistence.SuccessfulVotesAccess;
import java.math.BigInteger;
import java.security.cert.X509Certificate;
import java.util.stream.Collectors;
import java.util.stream.Stream;
import javax.ejb.EJB;
import javax.json.JsonObject;

public class CleansedBallotBoxServiceImpl implements CleansedBallotBoxService {

    private static final int PAGE_SIZE = 500;

    // The keystore alias
    private static final String KEYSTORE_ALIAS = "privatekey";

    private static final byte[] LINE_SEPARATOR = "\n".getBytes(StandardCharsets.UTF_8);

    @Inject
    private SignatureFactory signatureFactory;

    @Inject
    private CleansedBallotBoxRepository cleansedBallotBoxRepository;

    @Inject
    private BallotBoxInformationRepository ballotBoxInformationRepository;

    @Inject
    private CleansedBallotBoxAccess cleansedBallotBoxAccess;

    @Inject
    private SuccessfulVotesAccess successfulVotesAccess;

    private static final Logger LOG = LoggerFactory.getLogger(CleansedBallotBoxServiceImpl.class);
    
    @Inject
    private PrivateKeyForObjectRepository privateKeyRepository;

    @Inject
    private PayloadSigner payloadSigner;

    @EJB
    private EiTenantSystemKeys eiTenantSystemKeys;

    @EJB
    private ElectoralAuthorityRepository electoralAuthorityRepository;

    @Override
    public void writeEncryptedBallotBox(final OutputStream stream, final String tenantId, final String electionEventId,
            final String ballotBoxId)
            throws IOException {
        LOG.info("Retrieving cleansed ballot box {} for tenant {} and election event {}.", ballotBoxId, tenantId,
            electionEventId);
        Signature signature = getSignature(tenantId, electionEventId, ballotBoxId);
        writeCleansedBallotBoxItems(new SignatureOutputStream(stream, signature), tenantId, electionEventId,
            ballotBoxId);
        writeSignature(stream, signature);
    }

    @Override
    public Stream<MixingPayload> getMixingPayloads(BallotBoxId ballotBoxId)
            throws PayloadSignatureException, ResourceNotFoundException, CleansedBallotBoxServiceException {
        LOG.info("Requested all votes sets for ballot box {}...", ballotBoxId);

        try {
            LOG.info("Counting votes for ballot box {}...", ballotBoxId);
            // Find out how many vote sets fit the ballot box.
            int voteCount = cleansedBallotBoxRepository.count(ballotBoxId);
            LOG.info("Ballot box {} contains {} votes.", ballotBoxId, voteCount);

            // Get the certificate chain for the election information public
            // key.
            String tenantId = ballotBoxId.getTenantId();
            LOG.info("Finding the validation key certificate chain for ballot box {}...", ballotBoxId);
            X509Certificate[] fullCertificateChain = eiTenantSystemKeys.getSigningCertificateChain(tenantId);
            if (null == fullCertificateChain) {
                throw new PayloadSignatureException(
                    new IllegalArgumentException("A certificate chain was not found for tenant " + tenantId));
            }
            X509Certificate[] certificateChain = new X509Certificate[fullCertificateChain.length - 1];
            System.arraycopy(fullCertificateChain, 0, certificateChain, 0, fullCertificateChain.length - 1);
            LOG.info("Obtained the validation key certificate for tenant {} with {} elements", tenantId,
                certificateChain.length);

            // Get the election information system key to sign the payloads.
            PrivateKey signingKey = eiTenantSystemKeys.getSigningPrivateKey(tenantId);
            LOG.info("Obtained the signing key for tenant {}", tenantId);

            // Get the ballot box's information.
            JsonObject ballotBoxInformation = getBallotBoxInformationJson(ballotBoxId);
            // Get the encryption parameters from the ballot box information.
            JsonObject encryptionParametersJson = ballotBoxInformation.getJsonObject("encryptionParameters");
            ZpSubgroup encryptionParameters = new ZpSubgroup(new BigInteger(encryptionParametersJson.getString("g")),
                new BigInteger(encryptionParametersJson.getString("p")),
                new BigInteger(encryptionParametersJson.getString("q")));
            // Get the electoral authority identifier.
            String electoralAuthorityId = ballotBoxInformation.getString("electoralAuthorityId");

            // Get the vote encryption key, which at this stage is the electoral
            // authority public key.
            ElGamalPublicKey voteEncryptionKey =
                getVoteEncryptionKey(tenantId, ballotBoxId.getElectionEventId(), electoralAuthorityId);

            // Get the layout of the vote sets.
            PartitionCalculator partitionCalculator = new PartitionCalculator(CleansedBallotBox.CHUNK_SIZE);

            // For each of the calculated partitions, get a set of votes, wrap
            // them in a payload object,
            // sign them and make them available as a stream.
            Stream<MixingPayload> mixingPayloadStream = partitionCalculator.getPartitionDetails(voteCount).entrySet()
                .stream().map(entry -> buildSignedMixingPayload(ballotBoxId, entry.getKey(), electoralAuthorityId,
                    signingKey, certificateChain, voteEncryptionKey, encryptionParameters, entry.getValue()));

            LOG.info("{} votes sets from ballot box {} are ready to be streamed.", voteCount, ballotBoxId);

            return mixingPayloadStream;
        } catch (RuntimeException e) {
            // Capture the inner payload signature exception.
            if (e.getCause() instanceof PayloadSignatureException) {
                throw (PayloadSignatureException) e.getCause();
            } else {
                throw e;
            }
        } catch (GeneralCryptoLibException | CleansedBallotBoxRepositoryException | IOException e) {
            LOG.error("Failed retrieving the ballot box {}", ballotBoxId, e);
            throw new CleansedBallotBoxServiceException("Failed retrieving the ballot box " + ballotBoxId, e);
        }
    }

    /**
     * Retrieves a ballot box's information
     *
     * @param ballotBoxId
     *            the ballot box identifier
     * @return a ballot box's information as a JSON object
     * @throws ResourceNotFoundException
     *             if the ballot box is not found
     */
    private JsonObject getBallotBoxInformationJson(BallotBoxId ballotBoxId) throws ResourceNotFoundException {
        return JsonUtils.getJsonObject(ballotBoxInformationRepository.findByTenantIdElectionEventIdBallotBoxId(
            ballotBoxId.getTenantId(), ballotBoxId.getElectionEventId(), ballotBoxId.getId()).getJson());
    }

    /**
     * Constructs a payload for mixing and signs it.
     *
     * @param ballotBoxId
     *            the identifier of the ballot box the votes are retrieved from
     * @param voteSetIndex
     *            the identifier if the vote set within the ballot box
     * @param electoralAuthorityId
     *            the electoral authority identifier
     * @param signingKey
     *            the key used to sign the ballot box
     * @param certificateChain
     *            the full chain certificate of the public key used to verify
     *            the signature
     * @param voteEncryptionKey
     *            the electoral authority key
     * @param encryptionParameters
     *            the Zp group of the votes
     * @param partitionDetails
     *            the details of how the vote sets are split
     */
    private MixingPayload buildSignedMixingPayload(BallotBoxId ballotBoxId, int voteSetIndex,
            String electoralAuthorityId, PrivateKey signingKey, X509Certificate[] certificateChain,
            ElGamalPublicKey voteEncryptionKey, ZpSubgroup encryptionParameters, PartitionDetails partitionDetails) {
        try {
            List<EncryptedVote> votes = cleansedBallotBoxRepository
                .getVoteSet(ballotBoxId, partitionDetails.getOffset(), partitionDetails.getSize())
                .collect(Collectors.toList());
            // Create the payload.
            MixingPayload payload = new MixingPayloadImpl(voteEncryptionKey,
                new VoteSetIdImpl(ballotBoxId, voteSetIndex), votes, electoralAuthorityId, encryptionParameters);
            // Sign the payload.
            payload.setSignature(payloadSigner.sign(payload, signingKey, certificateChain));

            return payload;
        } catch (PayloadSignatureException e) {
        	throw new CryptoLibException("Error trying to get signed mixing payload.", e);
        }
    }

    private Signature getSignature(final String tenantId, final String electionEventId, final String ballotBoxId)
            throws IOException {
        Signature signature = signatureFactory.newSignature();
        try {
            PrivateKey privateKey = privateKeyRepository.findByTenantEEIDObjectIdAlias(tenantId, electionEventId,
                ballotBoxId, KEYSTORE_ALIAS);
            signature.initSign(privateKey);
        } catch (InvalidKeyException | ResourceNotFoundException | CryptographicOperationException e) {
            throw new IOException("Failed to get signature.", e);
        }
        return signature;
    }

    private void writeCleansedBallotBoxItems(final OutputStream stream, final String tenantId,
            final String electionEventId, final String ballotBoxId)
            throws IOException {
        try (CleansedExportedBallotBoxItemWriter writer =
            new CleansedExportedBallotBoxItemWriter(new CloseShieldOutputStream(stream))) {
            int first = 1;
            int last = PAGE_SIZE;
            List<CleansedExportedBallotBoxItem> page =
                cleansedBallotBoxRepository.getEncryptedVotesByTenantIdElectionEventIdBallotBoxId(tenantId,
                    electionEventId, ballotBoxId, first, last);
            while (!page.isEmpty()) {
                for (final CleansedExportedBallotBoxItem item : page) {
                    writer.write(item);
                }
                first += PAGE_SIZE;
                last += PAGE_SIZE;
                page = cleansedBallotBoxRepository.getEncryptedVotesByTenantIdElectionEventIdBallotBoxId(tenantId,
                    electionEventId, ballotBoxId, first, last);
            }
        }
    }

    private void writeSignature(final OutputStream stream, final Signature signature) throws IOException {
        byte[] base64Signature;
        try {
            base64Signature = Base64.getEncoder().encode(signature.sign());
        } catch (SignatureException e) {
            throw new IOException("Failed to sign encrypted ballot box.", e);
        }
        stream.write(LINE_SEPARATOR);
        stream.write(base64Signature);
    }

    @Override
    @TransactionAttribute(TransactionAttributeType.REQUIRED)
    public void storeCleansedVote(Vote vote) throws DuplicateEntryException {
        CleansedBallotBox cleansedBallotBox = new CleansedBallotBox();
        cleansedBallotBox.setTenantId(vote.getTenantId());
        cleansedBallotBox.setElectionEventId(vote.getElectionEventId());
        cleansedBallotBox.setVotingCardId(vote.getVotingCardId());
        cleansedBallotBox.setBallotId(vote.getBallotId());
        cleansedBallotBox.setBallotBoxId(vote.getBallotBoxId());
        String encryptedVote = vote.getEncryptedOptions();

        if (!StringUtils.isEmpty(vote.getEncryptedWriteIns())) {
            encryptedVote += ";" + vote.getEncryptedWriteIns();
        }

        cleansedBallotBox.setEncryptedVote(encryptedVote);
        cleansedBallotBoxAccess.save(cleansedBallotBox);
    }

    /**
     * Store an entry that represents a successful vote. A successful vote entry
     * is composed by the concatenation of the voting card ID, a timestamp and
     * the vote receipt.
     *
     * @param votingCardId
     *            the voting card ID
     * @param receipt
     *            the vote receipt
     * @throws DuplicateEntryException
     *             if the successful vote already exists in the DB
     */
    @Override
    @TransactionAttribute(TransactionAttributeType.REQUIRED)
    public void storeSuccessfulVote(String tenantId, String electionEventId, String ballotBoxId, String votingCardId,
            String receipt)
            throws DuplicateEntryException {
        ZonedDateTime now = ZonedDateTime.now(ZoneOffset.UTC);
        SuccessfulVote successfulVote = new SuccessfulVote();
        successfulVote.setReceipt(receipt);
        successfulVote.setTimestamp(now);
        successfulVote.setTenantId(tenantId);
        successfulVote.setElectionEventId(electionEventId);
        successfulVote.setBallotBoxId(ballotBoxId);
        successfulVote.setVotingCardId(votingCardId);
        successfulVotesAccess.save(successfulVote);
    }

    /**
     * Gets the electoral authority public key.
     *
     * @param tenantId
     *            the identifier of the tenant that ows the electoral authority
     *            key
     * @param electionEventId
     *            the identifier of the election event linked to the electoral
     *            authority
     * @param electoralAuthorityId
     *            the electoral authority identifier
     * @return the electoral authority public key
     * @throws ResourceNotFoundException
     *             if the electoral authority key is not found
     * @throws GeneralCryptoLibException
     *             if there was a problem de-serialising the electoral authority
     *             key
     * @throws IOException
     */
    private ElGamalPublicKey getVoteEncryptionKey(String tenantId, String electionEventId, String electoralAuthorityId)
            throws ResourceNotFoundException, GeneralCryptoLibException, IOException {
        ElectoralAuthorityEntity entity = electoralAuthorityRepository
            .findByTenantIdElectionEventIdElectoralAuthorityId(tenantId, electionEventId, electoralAuthorityId);

        String elGamalPublicKeyJson = new String(
            Base64.getDecoder().decode(new ObjectMapper().readTree(entity.getJson()).get("publicKey").asText()),
            StandardCharsets.UTF_8);

        return ElGamalPublicKey.fromJson(elGamalPublicKeyJson);
    }
}
