package prover;

import com.google.protobuf.ByteString;
import meerkat.crypto.Encryption;
import meerkat.crypto.mixnet.Mix2ZeroKnowledgeProver;
import meerkat.protobuf.ConcreteCrypto.ElGamalCiphertext;
import meerkat.protobuf.Crypto;
import meerkat.protobuf.Mixing;
import necessary.General;
import necessary.Group;


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

public class Prover implements Mix2ZeroKnowledgeProver {

    Group group;
    General general;
    Random rand;
    Encryption encryptor;

    public Prover(Group group,Random rand,Encryption encryptor,General general) {
        this.group = group;
        this.rand = rand;
        this.encryptor = encryptor;
        this.general = general;
    }

    public Mixing.ZeroKnowledgeProof prove(Crypto.RerandomizableEncryptedMessage in1,
                                           Crypto.RerandomizableEncryptedMessage in2,
                                           Crypto.RerandomizableEncryptedMessage out1,
                                           Crypto.RerandomizableEncryptedMessage out2,
                                           boolean sw,
                                           Crypto.EncryptionRandomness r1,
                                           Crypto.EncryptionRandomness r2) {

        Mixing.ZeroKnowledgeProof.OrProof first,second,third,fourth;
        if (!sw)
        {
            first = createOrProof(in1, out1, in2, out1, r1, true);
            second = createOrProof(in1, out1, in1, out2, r1, true);
            third = createOrProof(in2, out1, in2, out2, r2, false);
            fourth = createOrProof(in1, out2, in2, out2, r2, false);
        }
        else
        {
            first = createOrProof(in1, out1, in2, out1, r2, false);
            second = createOrProof(in1, out1, in1, out2, r1, false);
            third = createOrProof(in2, out1, in2, out2, r2, true);
            fourth = createOrProof(in1, out2, in2, out2, r1, true);
        }
        return Mixing.ZeroKnowledgeProof.newBuilder()
                .setFirst(first)
                .setSecond(second)
                .setThird(third)
                .setFourth(fourth)
                .build();
    }

    private Mixing.ZeroKnowledgeProof.OrProof createOrProof(Crypto.RerandomizableEncryptedMessage e1,
                                                            Crypto.RerandomizableEncryptedMessage e2,
                                                            Crypto.RerandomizableEncryptedMessage e1New,
                                                            Crypto.RerandomizableEncryptedMessage e2New,
                                                            Crypto.EncryptionRandomness x,
                                                            boolean flag){

        ElGamalCiphertext e1ElGamal = general.calcRerandomizable2ElGamal(e1);
        ElGamalCiphertext e2ElGamal = general.calcRerandomizable2ElGamal(e2);
        ElGamalCiphertext e1TagElGamal = general.calcRerandomizable2ElGamal(e1New);
        ElGamalCiphertext e2TagElGamal = general.calcRerandomizable2ElGamal(e2New);


        return createOrProofElGamal(e1ElGamal,e2ElGamal,e1TagElGamal,e2TagElGamal,x,flag);
    }

    private Mixing.ZeroKnowledgeProof.OrProof createOrProofElGamal(ElGamalCiphertext e1,
                                                                   ElGamalCiphertext e2,
                                                                   ElGamalCiphertext e1New,
                                                                   ElGamalCiphertext e2New,
                                                                   Crypto.EncryptionRandomness x,
                                                                   boolean flag) {
        ByteString g1 = general.getG();
        ByteString h1 = group.div(e1New.getC1(),e1.getC1());
        ByteString g2 = general.getH();
        ByteString h2 = group.div(e1New.getC2(),e1.getC2());

        ByteString g1Tag = general.getG();
        ByteString h1Tag = group.div(e2New.getC1(),e2.getC1());
        ByteString g2Tag = general.getH();
        ByteString h2Tag = group.div(e2New.getC2(),e2.getC2());

        BigInteger r = new BigInteger(encryptor.generateRandomness(rand).getData().toByteArray()).mod(group.groupSize());
        BigInteger c1,c2,z,zTag;
        ByteString u,v,uTag,vTag;
        if (flag)
        {
            c2 = new BigInteger(encryptor.generateRandomness(rand).getData().toByteArray()).mod(group.groupSize());
            zTag = new BigInteger(encryptor.generateRandomness(rand).getData().toByteArray()).mod(group.groupSize());
            //step 1
            u = group.pow(g1, r);
            v = group.pow(g2, r);
            uTag = group.div(group.pow(g1Tag, zTag), group.pow(h1Tag, c2));
            vTag = group.div(group.pow(g2Tag, zTag), group.pow(h2Tag, c2));
            //step 2
            // c1 = (hash(input + step1) + group size - c2)% group size
            Mixing.ZeroKnowledgeProof.OrProof.ForRandomOracle forRandomOracle =
                    Mixing.ZeroKnowledgeProof.OrProof.ForRandomOracle.newBuilder()
                            .setG1(g1)
                            .setH1(h1)
                            .setG2(g2)
                            .setH2(h2)
                            .setG1Tag(g1Tag)
                            .setH1Tag(h1Tag)
                            .setG2Tag(g2Tag)
                            .setH2Tag(h2Tag)
                            .setU(u)
                            .setV(v)
                            .setUTag(uTag)
                            .setVTag(vTag)
                            .build();
            c1 = general.hash(forRandomOracle).add(group.groupSize().subtract(c2)).mod(group.groupSize());
            //step 3
            //z = (r +  c1 * x) % group size;
            z = r.add(c1.multiply(new BigInteger(x.getData().toByteArray()))).mod(group.groupSize());
        }
        else
        {
            c1 = new BigInteger(encryptor.generateRandomness(rand).getData().toByteArray()).mod(group.groupSize());
            z = new BigInteger(encryptor.generateRandomness(rand).getData().toByteArray()).mod(group.groupSize());
            //step 1
            uTag = group.pow(g1Tag, r);
            vTag = group.pow(g2Tag, r);
            u = group.div(group.pow(g1, z), group.pow(h1, c1));
            v = group.div(group.pow(g2, z), group.pow(h2, c1));
            //step 2
            // c1 = (hash(input + step1) + group size - c1)% group size
            Mixing.ZeroKnowledgeProof.OrProof.ForRandomOracle forRandomOracle =
                    Mixing.ZeroKnowledgeProof.OrProof.ForRandomOracle.newBuilder()
                            .setG1(g1)
                            .setH1(h1)
                            .setG2(g2)
                            .setH2(h2)
                            .setG1Tag(g1Tag)
                            .setH1Tag(h1Tag)
                            .setG2Tag(g2Tag)
                            .setH2Tag(h2Tag)
                            .setU(u)
                            .setV(v)
                            .setUTag(uTag)
                            .setVTag(vTag)
                            .build();
            c2 = general.hash(forRandomOracle).add(group.groupSize().subtract(c1)).mod(group.groupSize());
            //step 3
            //zTag = (r +  c2 * x) % group size;
            zTag = r.add(c2.multiply(new BigInteger(x.getData().toByteArray()))).mod(group.groupSize());
        }
        return Mixing.ZeroKnowledgeProof.OrProof.newBuilder()
                .setG1(g1)
                .setH1(h1)
                .setG2(g2)
                .setH2(h2)
                .setG1Tag(g1Tag)
                .setH1Tag(h1Tag)
                .setG2Tag(g2Tag)
                .setH2Tag(h2Tag)
                .setU(u)
                .setV(v)
                .setUTag(uTag)
                .setVTag(vTag)
                .setC1(ByteString.copyFrom(c1.toByteArray()))
                .setC2(ByteString.copyFrom(c2.toByteArray()))
                .setZ(ByteString.copyFrom(z.toByteArray()))
                .setZTag(ByteString.copyFrom(zTag.toByteArray()))
                .build();
    }


}