/*
 * 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.mixnet.proofs.bg.shuffle;

import java.math.BigInteger;

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

import com.scytl.cryptolib.api.exceptions.CryptoLibException;
import com.scytl.cryptolib.api.exceptions.GeneralCryptoLibException;
import com.scytl.cryptolib.elgamal.cryptoapi.Ciphertext;
import com.scytl.cryptolib.mathematical.groups.impl.Exponent;
import com.scytl.cryptolib.mathematical.groups.impl.ZpSubgroup;
import com.scytl.products.ov.mixnet.commons.beans.proofs.ShuffleProofSecondAnswer;
import com.scytl.products.ov.mixnet.commons.constants.Constants;
import com.scytl.products.ov.mixnet.commons.homomorphic.Cryptosystem;
import com.scytl.products.ov.mixnet.commons.proofs.bg.commitments.CommitmentParams;
import com.scytl.products.ov.mixnet.commons.proofs.bg.commitments.PrivateCommitment;
import com.scytl.products.ov.mixnet.commons.proofs.bg.commitments.PublicCommitment;
import com.scytl.products.ov.mixnet.commons.tools.CiphertextTools;
import com.scytl.products.ov.mixnet.commons.tools.GroupTools;
import com.scytl.products.ov.mixnet.commons.tools.MultiExponentiation;
import com.scytl.products.ov.mixnet.commons.tools.RandomOracleHash;
import com.scytl.products.ov.mixnet.proofs.bg.multiexp.basic.MultiExponentiationBasicProofVerifier;
import com.scytl.products.ov.mixnet.proofs.bg.multiexp.reduction.MultiExponentiationReductionProofVerifier;
import com.scytl.products.ov.mixnet.proofs.bg.product.ProductProofVerifier;

public class ShuffleProofVerifier {

    private static final Logger LOG = LoggerFactory.getLogger("std");

    private final ZpSubgroup group;

    private final Cryptosystem cryptosystem;

    private final CommitmentParams compars;

    private final Ciphertext[][] c;

    private final Ciphertext[][] cPrime;

    private final int m;

    private final int n;

    private final RandomOracleHash rO;

    private final int numiterations;

    private final int mu;

    private final CiphertextTools cipherTextTools;

    private final MultiExponentiation multiExponentiation;

    public ShuffleProofVerifier(final ZpSubgroup group, final Cryptosystem cryptosystem, final CommitmentParams compars,
            final Ciphertext[][] c, final Ciphertext[][] cPrime, final int m, final int n, final int mu,
            final int numiterations, final CiphertextTools cipherTextTools,
            final MultiExponentiation multiExponentiation) {
        this.group = group;
        this.cryptosystem = cryptosystem;
        this.compars = compars;
        this.c = c;
        this.cPrime = cPrime;
        this.m = m;
        this.n = n;
        this.rO = new RandomOracleHash(group.getQ());
        this.mu = mu;
        this.numiterations = numiterations;
        this.cipherTextTools = cipherTextTools;
        this.multiExponentiation = multiExponentiation;
    }

    public ShuffleProofVerifier(final ZpSubgroup group, final Cryptosystem cryptosystem, final CommitmentParams compars,
            final Ciphertext[][] c, final Ciphertext[][] cPrime, final int m, final int n,
            final CiphertextTools cipherTextTools, final MultiExponentiation multiExponentiation) {
        this(group, cryptosystem, compars, c, cPrime, m, n, 0, 0, cipherTextTools, multiExponentiation);
    }

    public boolean verifyProof(final PublicCommitment[] cA, final PublicCommitment[] cB,
            final ShuffleProofSecondAnswer ans) throws GeneralCryptoLibException {

        boolean correct = true;

        // check that the commitments are ok

        if (cA.length != m) {
            correct = false;
            LOG.error("ERROR(Shuffle Argument): cA doesn't have the expected length");
        }

        for (int i = 0; i < cA.length; i++) {

            if (!GroupTools.isGroupElement(cA[i].getElement())) {
                correct = false;
                LOG.error("ERROR(Shuffle Argument): cA[" + i + "] is not a group element");
            }
        }

        rO.addDataToRO(c);
        rO.addDataToRO(cPrime);
        rO.addDataToRO(cA);
        final Exponent challengeX = rO.getHash();

        if (cB.length != m) {
            correct = false;
            LOG.error("ERROR(Shuffle Argument): cA doesn't have the expected length");
        }

        for (int i = 0; i < cB.length; i++) {
            if (!GroupTools.isGroupElement(cB[i].getElement())) {
                correct = false;
                LOG.error("ERROR(Shuffle Argument): cB[" + i + "] is not a group element");
            }
        }

        rO.addDataToRO(cB);
        final Exponent challengeY = rO.getHash();
        rO.addDataToRO("1");
        final Exponent challengeZ = rO.getHash();
        rO.reset();

        // compute new commitments

        final Exponent[] minusZ = new Exponent[n];

        for (int i = 0; i < n; i++) {
            minusZ[i] = challengeZ.negate();
        }

        final PublicCommitment[] cMinusZ = new PublicCommitment[m];

        Exponent zero;
        try {
            zero = new Exponent(group.getQ(), BigInteger.ZERO);
        } catch (GeneralCryptoLibException ex) {
            throw new CryptoLibException(ex);
        }

        for (int i = 0; i < m; i++) {
            cMinusZ[i] = new PrivateCommitment(minusZ, zero, compars,
                multiExponentiation).makePublicCommitment();
        }

        final PublicCommitment[] cD = new PublicCommitment[m];
        for (int i = 0; i < cD.length; i++) {
            cD[i] = cA[i].exponentiate(challengeY).multiply(cB[i]);
        }

        // prepare verifiers for ME, PA

        PublicCommitment[] cPA;
        if (m == 1) {
            cPA = new PublicCommitment[m + 1];
        } else {
            cPA = new PublicCommitment[m];
        }

        for (int i = 0; i < m; i++) {
            cPA[i] = cD[i].multiply(cMinusZ[i]);
        }

        try {
            // handle single dimensional array (m=1)
            if (m == 1) {

                Exponent[] exponents = new Exponent[n];

                for (int i = 0; i < n; i++) {
                    exponents[i] = new Exponent(compars.getGroup().getQ(), Constants.MULTIPLICATIVE_IDENTITY);
                }

                PrivateCommitment privateCommitmentToBeAppended =
                    new PrivateCommitment(exponents, new Exponent(compars.getGroup().getQ(), BigInteger.ZERO),
                        compars, multiExponentiation);

                cPA[m] = privateCommitmentToBeAppended.makePublicCommitment();
            }

        } catch (GeneralCryptoLibException ex) {
            throw new CryptoLibException(ex);
        }

        final Exponent[][] vecX = computeVecX(challengeX);
        final Exponent rhsPA = computeRightHandProductArg(challengeY, challengeZ, vecX);
        final Ciphertext lhsME = computeLeftHandMEArg(vecX);

        final ProductProofVerifier verifPA = new ProductProofVerifier(n, m, compars, cPA, rhsPA, group.getQ(),
            multiExponentiation);

        if (!verifPA.verify(ans.getMsgPA())) {
            correct = false;
            LOG.error("ERROR(Shuffle Argument):  Product Argument didn't verify");
        }

        if (numiterations == 0) {
            final MultiExponentiationBasicProofVerifier verifME = new MultiExponentiationBasicProofVerifier(m, n,
                cryptosystem, compars, cPrime, lhsME, cB, group.getQ(), cipherTextTools, multiExponentiation);

            if (!verifME.verify(ans.getIniMEBasic(), ans.getAnsMEBasic())) {
                correct = false;
                LOG.error("ERROR(Shuffle Argument):  MultiExpo argument didn't verify");
            }
        } else {

            final MultiExponentiationReductionProofVerifier verifME =
                new MultiExponentiationReductionProofVerifier(m, n, cryptosystem, group, compars, cPrime, lhsME,
                    cB, mu, numiterations, cipherTextTools, multiExponentiation);
            if (!verifME.verify(ans.getIniMEReduct(), ans.getAnsMEReduct())) {
                correct = false;
                LOG.error("ERROR(Shuffle Argument):  MultiExpo argument didn't verify");
            }
        }
        return correct;
    }

    public Exponent[][] computeVecX(final Exponent challengeX) {

        try {
            Exponent accum = challengeX;
            final Exponent[][] result = new Exponent[m][n];
            for (int i = 0; i < m; i++) {
                for (int j = 0; j < n; j++) {
                    result[i][j] = accum;
                    accum = accum.multiply(challengeX);
                }
            }
            return result;
        } catch (GeneralCryptoLibException ex) {
            throw new CryptoLibException(ex);
        }
    }

    public Exponent computeRightHandProductArg(final Exponent challengeY, final Exponent challengeZ,
            final Exponent[][] vecX) {

        try {
            Exponent result = new Exponent(group.getQ(), BigInteger.ONE);

            Exponent minusZ = challengeZ.negate();

            for (int i = 0; i < m * n; i++) {

                int value = i + 1;
                final Exponent factor = challengeY.multiply(new Exponent(group.getQ(), BigInteger.valueOf(value)))
                    .add(vecX[i / n][i % n]).add(minusZ);
                result = result.multiply(factor);
            }
            return result;

        } catch (GeneralCryptoLibException ex) {
            throw new CryptoLibException(ex);
        }
    }

    public Ciphertext computeLeftHandMEArg(final Exponent[][] vecX) {
        Ciphertext result = null;
        for (int i = 0; i < m; i++) {
            final Ciphertext aux = cipherTextTools.compVecCiphVecExp(c[i], vecX[i]);

            try {
                result = result == null ? aux : result.multiply(aux);
            } catch (GeneralCryptoLibException ex) {
                throw new CryptoLibException(ex);
            }
        }
        return result;
    }
}
