/*
 * 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.channel.cc.choicecodes.verification.service;

import java.math.BigInteger;
import java.nio.charset.StandardCharsets;
import java.security.Security;
import java.util.ArrayList;
import java.util.Base64;
import java.util.List;

import org.bouncycastle.jce.provider.BouncyCastleProvider;
import org.junit.Assert;
import org.junit.BeforeClass;
import org.junit.Test;
import org.junit.runner.RunWith;
import org.mockito.Matchers;
import org.mockito.Mockito;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.test.context.ContextConfiguration;
import org.springframework.test.context.junit4.SpringRunner;

import com.fasterxml.jackson.core.JsonProcessingException;
import com.scytl.cryptolib.api.elgamal.ElGamalServiceAPI;
import com.scytl.cryptolib.api.exceptions.GeneralCryptoLibException;
import com.scytl.cryptolib.api.proofs.ProofsServiceAPI;
import com.scytl.cryptolib.elgamal.bean.ElGamalComputationsValues;
import com.scytl.cryptolib.elgamal.bean.ElGamalEncrypterValues;
import com.scytl.cryptolib.elgamal.bean.ElGamalEncryptionParameters;
import com.scytl.cryptolib.elgamal.bean.ElGamalKeyPair;
import com.scytl.cryptolib.elgamal.bean.WitnessImpl;
import com.scytl.cryptolib.elgamal.cryptoapi.CryptoAPIElGamalEncrypter;
import com.scytl.cryptolib.elgamal.service.ElGamalService;
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.proofs.cryptoapi.ProofProverAPI;
import com.scytl.cryptolib.proofs.cryptoapi.ProofVerifierAPI;
import com.scytl.cryptolib.proofs.proof.Proof;
import com.scytl.cryptolib.proofs.service.ProofsService;
import com.scytl.products.oscore.logging.api.domain.Level;
import com.scytl.products.oscore.logging.api.writer.SecureLoggingWriter;
import com.scytl.products.ov.commons.beans.domain.model.Information.VoterInformation;
import com.scytl.products.ov.commons.beans.domain.model.authentication.AuthenticationToken;
import com.scytl.products.ov.commons.beans.domain.model.vote.Vote;
import com.scytl.products.ov.commons.dto.ChoiceCodesVerificationInput;
import com.scytl.products.ov.commons.ui.ws.rs.ObjectMappers;

@RunWith(SpringRunner.class)
@ContextConfiguration(classes = TestConfiguration.class)
public class ProofValidatorImplTest {

    private static final String VOTING_CARD_ID = "votingCardId";

    private static final String ELECTION_EVENT_ID = "electionEventId";

    private static final String TOKEN_ID = "tokenId";

    private static final String TIMESTAMP = "timestamp";

    private static final String SIGNATURE = "signature";

    private static ElGamalKeyPair compressesedChoiceCodesEncryptionKeyPair;

    private static ElGamalKeyPair elGamalElectoralKeyKeyPair;

    private static ElGamalEncryptionParameters encryptionParameters;

    private static Proof exponentiationProof;

    private static Proof plaintextEqualityProof;

    private static Proof schnorrProof;

    private static String encryptedOptions;

    private static String cipherTextExponentiations;

    private static String encryptedPartialChoiceCodes;

    @Autowired
    private ProofValidator proofValidator;

    @Autowired
    private SecureLoggingWriter secureLoggingWriter;

    @BeforeClass
    public static void createValidVote() throws GeneralCryptoLibException {

        Security.addProvider(new BouncyCastleProvider());

        ProofsServiceAPI proofsService = new ProofsService();

        BigInteger p = new BigInteger("67555108767603707298738874065797424621214999980483524802737421636930165840299");

        BigInteger q = new BigInteger("33777554383801853649369437032898712310607499990241762401368710818465082920149");

        BigInteger g = new BigInteger("3");

        encryptionParameters = new ElGamalEncryptionParameters(p, q, g);

        ProofProverAPI proofProver = proofsService.createProofProverAPI(encryptionParameters.getGroup());
        ProofVerifierAPI proofVerifier = proofsService.createProofVerifierAPI(encryptionParameters.getGroup());

        ElGamalServiceAPI elGamalService = new ElGamalService();
        compressesedChoiceCodesEncryptionKeyPair =
            elGamalService.getElGamalKeyPairGenerator().generateKeys(encryptionParameters, 1);
        elGamalElectoralKeyKeyPair = compressesedChoiceCodesEncryptionKeyPair;

        CryptoAPIElGamalEncrypter encrypter =
            elGamalService.createEncrypter(elGamalElectoralKeyKeyPair.getPublicKeys());
        GroupElementsCompressor<ZpGroupElement> compressor = new GroupElementsCompressor<>();
        List<ZpGroupElement> cleartextOptions = new ArrayList<>();
        cleartextOptions.add(new ZpGroupElement(new BigInteger("3"), (ZpSubgroup) encryptionParameters.getGroup()));
        cleartextOptions.add(new ZpGroupElement(new BigInteger("5"), (ZpSubgroup) encryptionParameters.getGroup()));

        ZpGroupElement compressedElement = compressor.compress(cleartextOptions);
        List<ZpGroupElement> voteCompressedMessage = new ArrayList<>();
        voteCompressedMessage.add(compressedElement);
        ElGamalEncrypterValues voteEncryptGroupElements = encrypter.encryptGroupElements(voteCompressedMessage);
        List<ZpGroupElement> originalZpGroupElements = voteEncryptGroupElements.getComputationValues().getValues();

        List<ZpGroupElement> zpGroupElements = originalZpGroupElements.subList(0, 2);
        List<ZpGroupElement> exponentiatedZpGroupElements =
            exponentiateZpGroupElements(zpGroupElements, elGamalElectoralKeyKeyPair.getPrivateKeys().getKeys().get(0));

        cipherTextExponentiations = convertToStringRepresentation(exponentiatedZpGroupElements);
        encryptedPartialChoiceCodes = convertToStringRepresentation(exponentiatedZpGroupElements);

        createSchnorrProof(proofProver, voteEncryptGroupElements);
        createExponentiationProof(voteEncryptGroupElements, proofProver, proofVerifier, zpGroupElements,
            exponentiatedZpGroupElements);
        createPlainTextEqualityProof(cipherTextExponentiations, proofProver, proofVerifier,
            (ZpSubgroup) encryptionParameters.getGroup(), elGamalElectoralKeyKeyPair, voteEncryptGroupElements);
    }

    private static void createPlainTextEqualityProof(String cipherTextExponentiations, ProofProverAPI proofProver,
            ProofVerifierAPI proofVerifier, ZpSubgroup group_g2_q11, ElGamalKeyPair elGamalKeyPair,
            ElGamalEncrypterValues encryptGroupElements) throws GeneralCryptoLibException {

        PlainTextEqualityProofGenerator generator =
            new PlainTextEqualityProofGenerator(proofProver, proofVerifier, group_g2_q11, elGamalKeyPair);
        plaintextEqualityProof =
            generator.createPlaintextEqualityProof(cipherTextExponentiations, encryptGroupElements);
    }

    private static void createSchnorrProof(ProofProverAPI proofProver, ElGamalEncrypterValues voteEncryptGroupElements)
            throws GeneralCryptoLibException {
        ElGamalComputationsValues voteCiphertext = voteEncryptGroupElements.getComputationValues();
        Exponent voteExponent = voteEncryptGroupElements.getExponent();
        encryptedOptions = convertToStringRepresentation(voteCiphertext.getValues());
        schnorrProof = proofProver.createSchnorrProof(VOTING_CARD_ID, ELECTION_EVENT_ID, voteCiphertext.getGamma(),
            new WitnessImpl(voteExponent));
    }

    private static void createExponentiationProof(ElGamalEncrypterValues voteEncryptGroupElements,
            ProofProverAPI proofProver, ProofVerifierAPI proofVerifier, List<ZpGroupElement> zpGroupElements,
            List<ZpGroupElement> exponentiatedZpGroupElements) throws GeneralCryptoLibException {
        ExponentiationProofGenerator generator =
            new ExponentiationProofGenerator(proofProver, proofVerifier, encryptionParameters);
        exponentiationProof = generator.createExponentiationProof(elGamalElectoralKeyKeyPair, zpGroupElements,
            exponentiatedZpGroupElements);
    }

    @Test
    public void validate_all_proofs_successfully_test() throws GeneralCryptoLibException, JsonProcessingException {
        Mockito.reset(secureLoggingWriter);
        ChoiceCodesVerificationInput choiceCodesVerificationInput =
            createInput(exponentiationProof, plaintextEqualityProof, schnorrProof);
        Assert.assertTrue(
            proofValidator.validate((ZpSubgroup) encryptionParameters.getGroup(), choiceCodesVerificationInput));
        Mockito.verify(secureLoggingWriter, Mockito.times(2)).log(Matchers.eq(Level.INFO), Matchers.any());
    }

    @Test
    public void exponentiation_proof_failure() throws JsonProcessingException, GeneralCryptoLibException {
        Mockito.reset(secureLoggingWriter);
        Proof modified = breakProof(exponentiationProof);
        ChoiceCodesVerificationInput choiceCodesVerificationInput =
            createInput(modified, plaintextEqualityProof, schnorrProof);
        Assert.assertFalse(
            proofValidator.validate((ZpSubgroup) encryptionParameters.getGroup(), choiceCodesVerificationInput));
        Mockito.verify(secureLoggingWriter, Mockito.times(1)).log(Matchers.eq(Level.ERROR), Matchers.any());
    }

    @Test
    public void schnorr_proof_failure_test() throws JsonProcessingException, GeneralCryptoLibException {
        Mockito.reset(secureLoggingWriter);
        Proof modified = breakProof(schnorrProof);
        ChoiceCodesVerificationInput choiceCodesVerificationInput =
            createInput(exponentiationProof, plaintextEqualityProof, modified);
        Assert.assertFalse(
            proofValidator.validate((ZpSubgroup) encryptionParameters.getGroup(), choiceCodesVerificationInput));
        Mockito.verify(secureLoggingWriter, Mockito.times(1)).log(Matchers.eq(Level.ERROR), Matchers.any());
    }

    @Test
    public void plaintext_equality_proof_failure_test() throws JsonProcessingException, GeneralCryptoLibException {
        Mockito.reset(secureLoggingWriter);
        Proof modified = breakProof(plaintextEqualityProof);
        ChoiceCodesVerificationInput choiceCodesVerificationInput =
            createInput(exponentiationProof, modified, schnorrProof);
        Assert.assertFalse(
            proofValidator.validate((ZpSubgroup) encryptionParameters.getGroup(), choiceCodesVerificationInput));
        Mockito.verify(secureLoggingWriter, Mockito.times(1)).log(Matchers.eq(Level.ERROR), Matchers.any());
    }

    private Proof breakProof(Proof original) throws GeneralCryptoLibException {
        List<Exponent> valuesList = original.getValuesList();
        return new Proof(new Exponent(encryptionParameters.getQ(), new BigInteger("1")), valuesList);
    }

    private ChoiceCodesVerificationInput createInput(Proof newExponentiationProof, Proof newPlaintextEqualityProof,
            Proof newSchnorrProof) throws GeneralCryptoLibException, JsonProcessingException {
        Vote vote = new Vote();
        vote.setVotingCardId(VOTING_CARD_ID);
        vote.setElectionEventId(ELECTION_EVENT_ID);
        vote.setEncryptedOptions(encryptedOptions);
        vote.setCipherTextExponentiations(cipherTextExponentiations);
        vote.setEncryptedPartialChoiceCodes(encryptedPartialChoiceCodes);
        vote.setVerificationCardPublicKey(Base64.getEncoder().encodeToString(
            compressesedChoiceCodesEncryptionKeyPair.getPublicKeys().toJson().getBytes(StandardCharsets.UTF_8)));

        vote.setSchnorrProof(schnorrProof.toJson());
        vote.setExponentiationProof(newExponentiationProof.toJson());
        vote.setPlaintextEqualityProof(newPlaintextEqualityProof.toJson());
        vote.setSchnorrProof(newSchnorrProof.toJson());

        VoterInformation voterInformation = new VoterInformation();
        AuthenticationToken authenticationToken =
            new AuthenticationToken(voterInformation, TOKEN_ID, TIMESTAMP, SIGNATURE);
        vote.setAuthenticationToken(ObjectMappers.toJson(authenticationToken));

        ChoiceCodesVerificationInput choiceCodesVerificationInput = new ChoiceCodesVerificationInput();
        choiceCodesVerificationInput
            .setCompressesedChoiceCodesEncryptionPK(compressesedChoiceCodesEncryptionKeyPair.getPublicKeys().toJson());
        choiceCodesVerificationInput.setElGamalElectoralPublicKey(elGamalElectoralKeyKeyPair.getPublicKeys().toJson());
        choiceCodesVerificationInput.setVote(ObjectMappers.toJson(vote));

        return choiceCodesVerificationInput;
    }

    private static String convertToStringRepresentation(List<ZpGroupElement> elements) {
        StringBuilder sb = new StringBuilder();
        for (int i = 0; i < elements.size() && i < 2; i++) {
            sb.append(elements.get(i).getValue().toString());
            if (i < elements.size() - 1)
                sb.append(";");
        }
        return sb.toString();
    }

    private static List<ZpGroupElement> exponentiateZpGroupElements(final List<ZpGroupElement> baseElements,
            final Exponent exponent) throws GeneralCryptoLibException {
        List<ZpGroupElement> exponentiatedElements = new ArrayList<>();
        for (ZpGroupElement baseElement : baseElements) {
            exponentiatedElements.add(baseElement.exponentiate(exponent));
        }
        return exponentiatedElements;
    }
}
