package meerkat.mixer.proofs.generic;

import meerkat.mixer.proofs.Concatenator;
import meerkat.mixer.proofs.SigmaProtocol;

import java.math.BigInteger;

/**
 * Disjunction of 2 Sigma protocol statements
 */
public class SigmaProtocolOr2 {
    public interface ChallengeGenerator {

        /**
         * Generate a random challenge in an appropriate way (e.g., mod group order)
         * @return
         */
        public BigInteger generateChallenge();

        /**
         * Subtract two challenges in an appropriate way (e.g., mod group order)
         * @param c1
         * @param c2
         * @return
         */
        public BigInteger subtractChallenge(BigInteger c1, BigInteger c2);
    }

    static public class Prover <FirstMessageOut, FirstMessageIn, FinalMessageOut, FinalMessageIn>
            implements SigmaProtocol.Prover<FirstMessageOut, FinalMessageOut> {
        final SigmaProtocol.Prover<FirstMessageIn, FinalMessageIn> prover;
        final SigmaProtocol.Simulator<FirstMessageIn, FinalMessageIn> simulator;
        final int proverIdx;

        final Concatenator.Pair<FirstMessageOut, FirstMessageIn, FirstMessageIn> firstMessageConcatenator;
        final Concatenator.Triplet<FinalMessageOut, BigInteger, FinalMessageIn, FinalMessageIn> finalMessageConcatenator;
        final ChallengeGenerator challengeGenerator;

        BigInteger simChallenge;

        /**
         *
         * @param prover
         * @param simulator
         * @param proverIdx Which index should the prover be inserted at (0 means first, 1 means second)
         */
        public Prover(Concatenator.Pair<FirstMessageOut, FirstMessageIn, FirstMessageIn>  firstMessageConcatenator,
                      Concatenator.Triplet<FinalMessageOut, BigInteger, FinalMessageIn, FinalMessageIn> finalMessageConcatenator,
                      ChallengeGenerator challengeGenerator,
                      SigmaProtocol.Prover<FirstMessageIn, FinalMessageIn> prover,
                      SigmaProtocol.Simulator<FirstMessageIn, FinalMessageIn> simulator, int proverIdx) {
            this.firstMessageConcatenator = firstMessageConcatenator;
            this.finalMessageConcatenator = finalMessageConcatenator;
            this.challengeGenerator = challengeGenerator;
            this.prover = prover;
            this.simulator = simulator;
            this.proverIdx = proverIdx;
            assert (proverIdx >= 0 && proverIdx < 2);
        }

        @Override
        public FirstMessageOut getFirstMessage() {
            simChallenge = challengeGenerator.generateChallenge();
            if (proverIdx == 0) {
                return firstMessageConcatenator.concatenate(prover.getFirstMessage(), simulator.getFirstMessage(simChallenge));
            } else {
                return firstMessageConcatenator.concatenate(simulator.getFirstMessage(simChallenge), prover.getFirstMessage());
            }
        }


        @Override
        public FinalMessageOut getFinalMessage(BigInteger challenge) {
            BigInteger realchallenge = challengeGenerator.subtractChallenge(challenge, simChallenge);
            if (proverIdx == 0) {
                return finalMessageConcatenator.concatenate(realchallenge, prover.getFinalMessage(realchallenge), simulator.getFinalMessage());
            } else {
                return finalMessageConcatenator.concatenate(simChallenge, simulator.getFinalMessage(), prover.getFinalMessage(realchallenge));
            }
        }
    }

    static public class Verifier <FirstMessageOut, FirstMessageIn, FinalMessageOut, FinalMessageIn>
            implements SigmaProtocol.Verifier<FirstMessageOut, FinalMessageOut> {

        final Concatenator.Pair<FirstMessageOut, FirstMessageIn, FirstMessageIn> firstMessageConcatenator;
        final Concatenator.Triplet<FinalMessageOut, BigInteger, FinalMessageIn, FinalMessageIn> finalMessageConcatenator;

        final ChallengeGenerator challengeGenerator;

        final SigmaProtocol.Verifier<FirstMessageIn, FinalMessageIn>[] verifiers;

        @SafeVarargs
        public Verifier(Concatenator.Pair<FirstMessageOut, FirstMessageIn, FirstMessageIn>  firstMessageConcatenator,
                        Concatenator.Triplet<FinalMessageOut, BigInteger, FinalMessageIn, FinalMessageIn> finalMessageConcatenator,
                        ChallengeGenerator challengeGenerator,
                        SigmaProtocol.Verifier<FirstMessageIn, FinalMessageIn>... verifiers) {
            this.firstMessageConcatenator = firstMessageConcatenator;
            this.finalMessageConcatenator = finalMessageConcatenator;
            this.challengeGenerator = challengeGenerator;

            this.verifiers = verifiers;
        }

        public boolean verify(FirstMessageOut firstMessage, BigInteger challenge, FinalMessageOut finalMessage) {
            BigInteger firstChallenge = finalMessageConcatenator.getMsg1of3(finalMessage);
            BigInteger secondChallenge = challengeGenerator.subtractChallenge(challenge, firstChallenge);
            return verifiers[0].verify(firstMessageConcatenator.getMsg1(firstMessage),  firstChallenge, finalMessageConcatenator.getMsg2of3(finalMessage)) &
                    verifiers[1].verify(firstMessageConcatenator.getMsg2(firstMessage), secondChallenge, finalMessageConcatenator.getMsg3of3(finalMessage));
        }
    }

    static public class Simulator<FirstMessageOut, FirstMessageIn, FinalMessageOut, FinalMessageIn>
            implements SigmaProtocol.Simulator<FirstMessageOut, FinalMessageOut> {
        final Concatenator.Pair<FirstMessageOut, FirstMessageIn, FirstMessageIn> firstMessageConcatenator;
        final Concatenator.Triplet<FinalMessageOut, BigInteger, FinalMessageIn, FinalMessageIn> finalMessageConcatenator;

        final ChallengeGenerator challengeGenerator;

        final SigmaProtocol.Simulator<FirstMessageIn, FinalMessageIn>[] simulators;


        BigInteger simChallenge0;


        @SafeVarargs
        public Simulator(Concatenator.Pair<FirstMessageOut, FirstMessageIn, FirstMessageIn> firstMessageConcatenator,
                         Concatenator.Triplet<FinalMessageOut, BigInteger, FinalMessageIn, FinalMessageIn> finalMessageConcatenator,
                         ChallengeGenerator challengeGenerator, SigmaProtocol.Simulator<FirstMessageIn, FinalMessageIn>... simulators) {
            this.firstMessageConcatenator = firstMessageConcatenator;
            this.finalMessageConcatenator = finalMessageConcatenator;
            this.challengeGenerator = challengeGenerator;
            this.simulators = simulators;
        }


        @Override
        public FirstMessageOut getFirstMessage(BigInteger challenge) {

            simChallenge0 = challengeGenerator.generateChallenge();
            BigInteger simChallenge1 = challengeGenerator.subtractChallenge(challenge, simChallenge0);

            return firstMessageConcatenator.concatenate(simulators[0].getFirstMessage(simChallenge0), simulators[1].getFirstMessage(simChallenge1));
        }

        @Override
        public FinalMessageOut getFinalMessage() {
            return finalMessageConcatenator.concatenate(simChallenge0, simulators[0].getFinalMessage(), simulators[1].getFinalMessage());
        }
    }

}
