package meerkat.mixer.proofs.concrete;

import meerkat.crypto.concrete.ECElGamalEncryption;
import meerkat.crypto.concrete.Util;
import meerkat.mixer.proofs.SigmaProtocol;
import meerkat.protobuf.ConcreteCrypto;
import meerkat.protobuf.Mixing;
import org.spongycastle.math.ec.ECPoint;
import org.factcenter.qilin.primitives.concrete.ECGroup;

import java.math.BigInteger;
import java.util.Random;

/**
 * Schnorr's ZK proof of Dlog equivlance
 */
public class SchnorrDlogEquivalence {

    public static class Verifier implements SigmaProtocol.Verifier<Mixing.Mix2Proof.DlogProof.FirstMessage, Mixing.Mix2Proof.DlogProof.FinalMessage> {
        ECElGamalEncryption encryptor;
        ECGroup group;
        final Statements.DlogStatement statement;

        public Verifier(ECElGamalEncryption encryptor, Statements.DlogStatement statement) {
            this.encryptor = encryptor;
            group = encryptor.getGroup();
            this.statement = statement;
        }

        @Override
        public boolean verify(Mixing.Mix2Proof.DlogProof.FirstMessage firstMessage, BigInteger challenge,
                              Mixing.Mix2Proof.DlogProof.FinalMessage finalMessage) {
            ConcreteCrypto.GroupElement grEncoded = firstMessage.getGr();
            ECPoint gr = encryptor.decodeElement(grEncoded);

            ConcreteCrypto.GroupElement hrEncoded = firstMessage.getHr();
            ECPoint hr = encryptor.decodeElement(hrEncoded);

            BigInteger xcr = Util.decodeBigInteger(finalMessage.getXcr());

            boolean gGood = group.add(gr, group.multiply(statement.a, challenge)).equals(group.multiply(statement.g, xcr));
            boolean hGood = group.add(hr, group.multiply(statement.b, challenge)).equals(group.multiply(statement.h, xcr));
            return gGood && hGood;
        }

    }

    public static class Simulator implements SigmaProtocol.Simulator<Mixing.Mix2Proof.DlogProof.FirstMessage, Mixing.Mix2Proof.DlogProof.FinalMessage> {
        ECElGamalEncryption encryptor;
        ECGroup group;
        Random rand;

        Statements.DlogStatement statement;
        BigInteger response = null;

        public Simulator(ECElGamalEncryption encryptor, Random rand, Statements.DlogStatement statement) {
            this.encryptor = encryptor;
            group = encryptor.getGroup();
            this.rand = rand;
            this.statement = statement;
        }

        @Override
        public Mixing.Mix2Proof.DlogProof.FirstMessage getFirstMessage(BigInteger challenge) {
            response = encryptor.generateRandomExponent(rand);

            ECPoint u = group.multiply(statement.g, response).subtract(group.multiply(statement.a, challenge));
            ECPoint v = group.multiply(statement.h, response).subtract(group.multiply(statement.b, challenge));
            return Mixing.Mix2Proof.DlogProof.FirstMessage.newBuilder()
                    .setGr(encryptor.encodeElement(u))
                    .setHr(encryptor.encodeElement(v))
                    .build();
        }

        @Override
        public Mixing.Mix2Proof.DlogProof.FinalMessage getFinalMessage() {

            return Mixing.Mix2Proof.DlogProof.FinalMessage.newBuilder()
                    .setXcr(Util.encodeBigInteger(response))
                    .build();
        }
    }

    public static class Prover implements SigmaProtocol.Prover<Mixing.Mix2Proof.DlogProof.FirstMessage, Mixing.Mix2Proof.DlogProof.FinalMessage> {
        ECElGamalEncryption encryptor;
        ECGroup group;
        Random rand;
        Statements.DlogStatement statement;
        Statements.DlogStatementWitness witness;

        BigInteger r = null;

        public Prover(ECElGamalEncryption encryptor, Random rand,
                      Statements.DlogStatement statement, Statements.DlogStatementWitness witness) {
            this.encryptor = encryptor;
            group = encryptor.getGroup();
            this.rand = rand;
            this.statement = statement;
            this.witness = witness;
        }

        @Override
        public Mixing.Mix2Proof.DlogProof.FirstMessage getFirstMessage() {
            r = encryptor.generateRandomExponent(rand);
            ECPoint gr = group.multiply(statement.g, r);
            ECPoint hr = group.multiply(statement.h, r);

            Mixing.Mix2Proof.DlogProof.FirstMessage firstMessage = Mixing.Mix2Proof.DlogProof.FirstMessage.newBuilder()
                    .setGr(encryptor.encodeElement(gr))
                    .setHr(encryptor.encodeElement(hr))
                    .build();

            return firstMessage;
        }

        @Override
        public Mixing.Mix2Proof.DlogProof.FinalMessage getFinalMessage(BigInteger challenge) {
            return Mixing.Mix2Proof.DlogProof.FinalMessage.newBuilder()
                    .setXcr(Util.encodeBigInteger(challenge.multiply(witness.x).add(r).mod(group.orderUpperBound())))
                    .build();
        }

    }
}
