package meerkat.mixer.proofs;

import com.google.protobuf.Message;
import meerkat.mixer.proofs.generic.SigmaFiatShamir;
import org.factcenter.qilin.util.Pair;
import org.junit.Before;
import org.junit.Test;

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

import static org.junit.Assert.assertTrue;

/**
 * Generic test for Sigma Protocol
 */
abstract public class SigmaProtocolTest<M1 extends Message, M2> {
    public final int NUM_REPEAT = 10;

    abstract protected void generateRandomTrueStatement();

    abstract protected void generateRandomFalseStatement();

    abstract protected SigmaProtocol.Prover<M1, M2> getNewProver();
    abstract protected SigmaProtocol.Verifier<M1, M2> getNewVerifier();
    abstract protected SigmaProtocol.Simulator<M1, M2> getNewSimulator();

    public class NIZKConcat implements Concatenator.Pair<Pair<M1,M2>, M1, M2> {
        @Override
        public Pair<M1, M2> concatenate(M1 msg1, M2 msg2) { return new Pair<M1, M2>(msg1, msg2);  }

        @Override
        public M1 getMsg1(Pair<M1,M2> msg) { return msg.a; }

        @Override
        public M2 getMsg2(Pair<M1,M2> msg) { return msg.b; }
    }

    final protected NIZKConcat nizkConcat = new NIZKConcat();

    abstract protected BigInteger getChallengeModulus();

    SigmaProtocol.Prover<M1, M2> prover;
    SigmaProtocol.Verifier<M1, M2> verifier;
    SigmaProtocol.Simulator<M1, M2> simulator;
    BigInteger challengeModulus;
    int challengeBits;

    protected Random rand = new Random(1);;

    @Before
    public void setup() {
        challengeModulus = getChallengeModulus();
        challengeBits = challengeModulus.bitLength() + 1;
    }

    @Test
    public void testProverCompleteness() {
        for (int i = 0; i < NUM_REPEAT; ++i) {
            generateRandomTrueStatement();

            prover = getNewProver();
            verifier = getNewVerifier();

            M1 msg1 = prover.getFirstMessage();
            BigInteger c = new BigInteger(challengeBits, rand).mod(challengeModulus);
            M2 msg2 = prover.getFinalMessage(c);

            assertTrue(String.format("Verification error in iteration %d", i), verifier.verify(msg1, c, msg2));
        }
    }


    @Test
    public void testSimulatorCompleteness() {
        for (int i = 0; i < NUM_REPEAT; ++i) {
            generateRandomFalseStatement();

            simulator = getNewSimulator();
            verifier = getNewVerifier();

            BigInteger c = new BigInteger(challengeBits, rand).mod(challengeModulus);

            M1 msg1 = simulator.getFirstMessage(c);
            M2 msg2 = simulator.getFinalMessage();

            assertTrue(String.format("Simulator Verification error in iteration %d", i), verifier.verify(msg1, c, msg2));
        }
    }

    @Test
    public void testNIZKCompleteness() {
        for (int i = 0; i < NUM_REPEAT; ++i) {
            generateRandomTrueStatement();

            SigmaFiatShamir<Pair<M1,M2>, M1, M2> fiatShamir = new SigmaFiatShamir<Pair<M1,M2>, M1, M2>(nizkConcat);

            prover = getNewProver();
            Pair<M1,M2> nizk = fiatShamir.generateNizk(prover);

            verifier = getNewVerifier();
            boolean nizkOk = fiatShamir.verifyNizk(nizk, verifier);

            assertTrue(String.format("NIZK Verification error in iteration %d", i), nizkOk);
        }
    }


}