package prover;

import com.google.protobuf.ByteString;
import com.google.protobuf.InvalidProtocolBufferException;
import meerkat.crypto.concrete.ECElGamalEncryption;
import meerkat.crypto.mixnet.Mix2ZeroKnowledgeProver;
import meerkat.protobuf.ConcreteCrypto.ElGamalCiphertext;
import meerkat.protobuf.Crypto;
import meerkat.protobuf.Mixing;
import org.bouncycastle.math.ec.ECPoint;
import qilin.primitives.RandomOracle;
import qilin.primitives.concrete.ECGroup;

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

public class Prover implements Mix2ZeroKnowledgeProver {

    ECGroup group;
    RandomOracle randomOracle;
    Random rand;
    ECElGamalEncryption ecElGamalEncryption;
    ECPoint g,h;
    ElGamalProofOrganizer organizer;

    public Prover(Random rand,ECElGamalEncryption encryptor,RandomOracle randomOracle) {

        this.rand = rand;
        this.ecElGamalEncryption = encryptor;
        this.randomOracle = randomOracle;
        this.group = ecElGamalEncryption.getGroup();
        this.g = group.getGenerator();
        this.h = ecElGamalEncryption.getElGamalPK().getPK();
        this.organizer = new ElGamalProofOrganizer(group,g,h);
    }

    public Mixing.ZeroKnowledgeProof prove(Crypto.RerandomizableEncryptedMessage in1,
                                           Crypto.RerandomizableEncryptedMessage in2,
                                           Crypto.RerandomizableEncryptedMessage out1,
                                           Crypto.RerandomizableEncryptedMessage out2,
                                           boolean sw,int i,int j, int layer,
                                           Crypto.EncryptionRandomness r1,
                                           Crypto.EncryptionRandomness r2) throws InvalidProtocolBufferException {
        Mixing.ZeroKnowledgeProof.OrProof first,second,third,fourth;

        ElGamalProofOrganizer.ElGamalProofInput input = organizer.createProofInput(in1,in2,out1,out2,r1,r2,sw);

        first = createOrProofElGamal(input.getOrProofInput(ElGamalProofOrganizer.OrProofOrder.first));
        second = createOrProofElGamal(input.getOrProofInput(ElGamalProofOrganizer.OrProofOrder.second));
        third = createOrProofElGamal(input.getOrProofInput(ElGamalProofOrganizer.OrProofOrder.third));
        fourth = createOrProofElGamal(input.getOrProofInput(ElGamalProofOrganizer.OrProofOrder.fourth));

        Mixing.ZeroKnowledgeProof.Location location = Mixing.ZeroKnowledgeProof.Location.newBuilder()
                .setI(i)
                .setJ(j)
                .setLayer(layer)
                .build();

        Mixing.ZeroKnowledgeProof result = Mixing.ZeroKnowledgeProof.newBuilder()
                .setFirst(first)
                .setSecond(second)
                .setThird(third)
                .setFourth(fourth)
                .setLocation(location)
                .build();
        return  result;
    }


    public BigInteger hash(Mixing.ZeroKnowledgeProof.OrProof.ForRandomOracle input) {
        byte[] arr = input.toByteArray();
        return new BigInteger(this.randomOracle.hash(arr,arr.length));
    }


    private Mixing.ZeroKnowledgeProof.OrProof createOrProofElGamal(ElGamalProofOrganizer.OrProofInput orProofInput) {

        ECPoint g1 = orProofInput.g1;
        ECPoint h1 = orProofInput.h1;
        ECPoint g2 = orProofInput.g2;
        ECPoint h2 = orProofInput.h2;

        ECPoint g1Tag  = orProofInput.g1Tag;
        ECPoint h1Tag  = orProofInput.h1Tag;
        ECPoint g2Tag  = orProofInput.g2Tag;
        ECPoint h2Tag  = orProofInput.h2Tag;

        BigInteger r = new BigInteger(ecElGamalEncryption.generateRandomness(rand).getData().toByteArray()).mod(group.orderUpperBound());
        BigInteger c1,c2,z,zTag;
        ECPoint u,v,uTag,vTag;
        Mixing.ZeroKnowledgeProof.OrProof.ForRandomOracle forRandomOracle;


        switch (orProofInput.flag) {
            case left:
                c2 = new BigInteger(ecElGamalEncryption.generateRandomness(rand).getData().toByteArray()).mod(group.orderUpperBound());
                zTag = new BigInteger(ecElGamalEncryption.generateRandomness(rand).getData().toByteArray()).mod(group.orderUpperBound());
                //step 1
                u = group.multiply(g1, r);
                v = group.multiply(g2, r);
                uTag = group.add(group.multiply(g1Tag, zTag), group.negate(group.multiply(h1Tag, c2)));
                vTag = group.add(group.multiply(g2Tag, zTag), group.negate(group.multiply(h2Tag, c2)));
                //step 2
                // c1 = (hash(input + step1) + group size - c2)% group size
                forRandomOracle =
                        Mixing.ZeroKnowledgeProof.OrProof.ForRandomOracle.newBuilder()
                                .setG1(ByteString.copyFrom(orProofInput.g1Encoded))
                                .setH1(ByteString.copyFrom(orProofInput.h1Encoded))
                                .setG2(ByteString.copyFrom(orProofInput.g2Encoded))
                                .setH2(ByteString.copyFrom(orProofInput.h2Encoded))
                                .setG1Tag(ByteString.copyFrom(orProofInput.g1TagEncoded))
                                .setH1Tag(ByteString.copyFrom(orProofInput.g1TagEncoded))
                                .setG2Tag(ByteString.copyFrom(orProofInput.g1TagEncoded))
                                .setH2Tag(ByteString.copyFrom(orProofInput.g1TagEncoded))
                                .setU(ByteString.copyFrom(group.encode(u)))
                                .setV(ByteString.copyFrom(group.encode(v)))
                                .setUTag(ByteString.copyFrom(group.encode(uTag)))
                                .setVTag(ByteString.copyFrom(group.encode(vTag)))
                                .build();
                c1 = hash(forRandomOracle).add(group.orderUpperBound().subtract(c2)).mod(group.orderUpperBound());
                //step 3
                //z = (r +  c1 * x) % group size;
                z = r.add(c1.multiply(new BigInteger(orProofInput.x.getData().toByteArray()))).mod(group.orderUpperBound());
                break;
            case right:
                c1 = new BigInteger(ecElGamalEncryption.generateRandomness(rand).getData().toByteArray()).mod(group.orderUpperBound());
                z = new BigInteger(ecElGamalEncryption.generateRandomness(rand).getData().toByteArray()).mod(group.orderUpperBound());
                //step 1
                uTag = group.multiply(g1Tag, r);
                vTag = group.multiply(g2Tag, r);
                u = group.add(group.multiply(g1, z), group.negate(group.multiply(h1, c1)));
                v = group.add(group.multiply(g2, z), group.negate(group.multiply(h2, c1)));
                //step 2
                // c1 = (hash(input + step1) + group size - c1)% group size
                forRandomOracle =
                        Mixing.ZeroKnowledgeProof.OrProof.ForRandomOracle.newBuilder()
                                .setG1(ByteString.copyFrom(orProofInput.g1Encoded))
                                .setH1(ByteString.copyFrom(orProofInput.h1Encoded))
                                .setG2(ByteString.copyFrom(orProofInput.g2Encoded))
                                .setH2(ByteString.copyFrom(orProofInput.h2Encoded))
                                .setG1Tag(ByteString.copyFrom(orProofInput.g1TagEncoded))
                                .setH1Tag(ByteString.copyFrom(orProofInput.g1TagEncoded))
                                .setG2Tag(ByteString.copyFrom(orProofInput.g1TagEncoded))
                                .setH2Tag(ByteString.copyFrom(orProofInput.g1TagEncoded))
                                .setU(ByteString.copyFrom(group.encode(u)))
                                .setV(ByteString.copyFrom(group.encode(v)))
                                .setUTag(ByteString.copyFrom(group.encode(uTag)))
                                .setVTag(ByteString.copyFrom(group.encode(vTag)))
                                .build();
                c2 = hash(forRandomOracle).add(group.orderUpperBound().subtract(c1)).mod(group.orderUpperBound());
                //step 3
                //zTag = (r +  c2 * x) % group size;
                zTag = r.add(c2.multiply(new BigInteger(orProofInput.x.getData().toByteArray()))).mod(group.orderUpperBound());
                break;
            default:
                return null;
        }

        //debugging
        //assert (group.multiply(g1, z).equals(group.add(u, group.multiply(h1,c1))));
        //assert (group.multiply(g2, z).equals(group.add(v, group.multiply(h2,c1))));
        //assert (group.multiply(g1Tag, zTag).equals(group.add(uTag, group.multiply(h1Tag,c2))));
        //assert (group.multiply(g2Tag, zTag).equals(group.add(vTag, group.multiply(h2Tag,c2))));


        return Mixing.ZeroKnowledgeProof.OrProof.newBuilder()
                .setG1(forRandomOracle.getG1())
                .setH1(forRandomOracle.getH1())
                .setG2(forRandomOracle.getG2())
                .setH2(forRandomOracle.getH2())
                .setG1Tag(forRandomOracle.getG1())
                .setH1Tag(forRandomOracle.getH1Tag())
                .setG2Tag(forRandomOracle.getG2Tag())
                .setH2Tag(forRandomOracle.getH2Tag())
                .setU(forRandomOracle.getU())
                .setV(forRandomOracle.getV())
                .setUTag(forRandomOracle.getUTag())
                .setVTag(forRandomOracle.getVTag())
                .setC1(ByteString.copyFrom(c1.toByteArray()))
                .setC2(ByteString.copyFrom(c2.toByteArray()))
                .setZ(ByteString.copyFrom(z.toByteArray()))
                .setZTag(ByteString.copyFrom(zTag.toByteArray()))
                .build();
    }
}