package meerkat.mixer.proofs.concrete;

import com.google.protobuf.InvalidProtocolBufferException;
import meerkat.crypto.concrete.ECElGamalEncryption;
import meerkat.protobuf.ConcreteCrypto.ElGamalCiphertext;
import meerkat.protobuf.Crypto;
import meerkat.protobuf.Crypto.EncryptionRandomness;
import org.spongycastle.math.ec.ECPoint;
import org.factcenter.qilin.primitives.concrete.ECGroup;

import java.math.BigInteger;

/**
 * used for organizing the input for each ZKP
 *
 * both meerkat.mixProverVerifier.proofs and meerkat.mixProverVerifier.verifier implementation use it
 */
public class Statements {

    final public ECElGamalEncryption encryptor;
    final public ECGroup group;
    final public ECPoint g;
    final public ECPoint h;

    /**
     * @param encryptor encryptor used for encoding/decoding serialized ciphertexts. The group, default generator and
     *                  second base (h) are taken from the encryptor (second base is the public key)
     */
    public Statements(ECElGamalEncryption encryptor){
        this.encryptor = encryptor;
        this.group = encryptor.getGroup();
        this.g = group.getGenerator();
        this.h = encryptor.getElGamalPK().getPK();
    }

    /**
     * can be used by anyone, e.g meerkat.mixer.verifier
     *
     * call to the meerkat.mixer.main overload with trueClauseIndex = false
     */
    public Mix2Statement createStatement(Crypto.RerandomizableEncryptedMessage in1, Crypto.RerandomizableEncryptedMessage in2,
                                         Crypto.RerandomizableEncryptedMessage out1, Crypto.RerandomizableEncryptedMessage out2) throws InvalidProtocolBufferException {

        ElGamalCiphertext in1ElGamal = ECElGamalEncryption.RerandomizableEncryptedMessage2ElGamalCiphertext(in1);
        ElGamalCiphertext in2ElGamal = ECElGamalEncryption.RerandomizableEncryptedMessage2ElGamalCiphertext(in2);
        ElGamalCiphertext out1ElGamal = ECElGamalEncryption.RerandomizableEncryptedMessage2ElGamalCiphertext(out1);
        ElGamalCiphertext out2ElGamal = ECElGamalEncryption.RerandomizableEncryptedMessage2ElGamalCiphertext(out2);
        return new Mix2Statement(in1ElGamal, in2ElGamal, out1ElGamal, out2ElGamal);
    }

    /**

     * convert each encrypted message to ElGamalCiphertext
     *
     * @throws InvalidProtocolBufferException - in case one or more of the encrypted messages isn't
     *                                          ElGamalCiphertext
     */
    protected Mix2StatementWitness createMix2Witness(EncryptionRandomness r1, EncryptionRandomness r2, boolean switched) {
        // if (!switched), dlogW1 is witness for a ~ c, dlogW2 is witness for b ~ d,
        // otherwise, dlogW1 is witness for a ~ d, dlogW2 is witness for b ~ c
        DlogStatementWitness dlogW1 = new DlogStatementWitness(encryptor.extractRandomness(r1));
        DlogStatementWitness dlogW2 = new DlogStatementWitness(encryptor.extractRandomness(r2));

        AndStatementWitness andW = new AndStatementWitness(dlogW1, dlogW2);

        return new Mix2StatementWitness(switched ? 1 : 0, andW);
    }


    /**
     * Mix2Statement to be proved.
     *
     * The actual stored data is a cached representation for use in proofs and verification. This consists of the four substatements that are ORs of the DLOG equality.
     *
     * A Mix2Statement can be constructed only by calling the {@link #createStatement} factory method on an instance of ECElGamalMixStatementGenerator  (all constructors are private)
     *
     */
    public class Mix2Statement {

        public final AndStatement[] clauses = new AndStatement[2];

        Mix2Statement(ElGamalCiphertext a, ElGamalCiphertext b,
                              ElGamalCiphertext c, ElGamalCiphertext d){

            ECPoint a1 = encryptor.decodeElement(a.getC1());
            ECPoint a2 = encryptor.decodeElement(a.getC2());
            ECPoint b1 = encryptor.decodeElement(b.getC1());
            ECPoint b2 = encryptor.decodeElement(b.getC2());
            ECPoint c1 = encryptor.decodeElement(c.getC1());
            ECPoint c2 = encryptor.decodeElement(c.getC2());
            ECPoint d1 = encryptor.decodeElement(d.getC1());
            ECPoint d2 = encryptor.decodeElement(d.getC2());

            ECPoint c1_div_a1 = group.add(c1, group.negate(a1));
            ECPoint d1_div_a1 = group.add(d1, group.negate(a1));
            ECPoint c1_div_b1 = group.add(c1, group.negate(b1));
            ECPoint d1_div_b1 = group.add(d1, group.negate(b1));
            ECPoint c2_div_a2 = group.add(c2, group.negate(a2));
            ECPoint d2_div_a2 = group.add(d2, group.negate(a2));
            ECPoint c2_div_b2 = group.add(c2, group.negate(b2));
            ECPoint d2_div_b2 = group.add(d2, group.negate(b2));

            // "Straight assignment" (a ~ c, b ~ d)
            DlogStatement s0_0 = new DlogStatement(c1_div_a1, c2_div_a2);
            DlogStatement s0_1 = new DlogStatement(d1_div_b1, d2_div_b2);
            clauses[0] = new AndStatement(s0_0, s0_1);

            // "Switched assignment" (a ~ d, b ~ c)
            DlogStatement s1_0 = new DlogStatement(d1_div_a1, d2_div_a2);
            DlogStatement s1_1 = new DlogStatement(c1_div_b1, c2_div_b2);
            clauses[1] = new AndStatement(s1_0, s1_1);
        }
    }



    /**
     * Witness information that enables us to generate a ZK proof for the 2-ciphertext mix
     */
    public class Mix2StatementWitness {
        /**
         * Which of the two clauses is the true one
         */
        public final int trueClauseIndex;

        public final AndStatementWitness witness;

        Mix2StatementWitness(int trueClauseIndex, AndStatementWitness witness) {

            this.trueClauseIndex = trueClauseIndex;
            assert (trueClauseIndex >= 0 && trueClauseIndex < 2);
            this.witness = witness;
        }
    }


    /**
     * Statement consisting of a conjunction of two {@link DlogStatement}s
     */
    public class AndStatement {
        public final DlogStatement[] clauses = new DlogStatement[2];

        AndStatement(DlogStatement clause0, DlogStatement clause1) {
            this.clauses[0] = clause0;
            this.clauses[1] = clause1;
        }
    }

    public class AndStatementWitness {
        public final DlogStatementWitness[] witnesses = new DlogStatementWitness[2];

        AndStatementWitness(DlogStatementWitness witness0, DlogStatementWitness witness1) {
            this.witnesses[0] = witness0;
            this.witnesses[1] = witness1;
        }
    }

    /**
     * The parameters for a statement of discrete-log equality
     *
     * log_{g}(a)==log_{h}(b)
     */
    public class DlogStatement {
        public final ECPoint g;
        public final ECPoint a;
        public final ECPoint h;
        public final ECPoint b;


        DlogStatement(ECPoint a, ECPoint b) {
            this.g = Statements.this.g;
            this.a = a;
            this.h = Statements.this.h;
            this.b = b;
        }
    }

    /**
     * Witness for correctness of a {@link DlogStatement}
     */
    public class DlogStatementWitness {
        /**
         * The actual discrete logarithm (i.e., a=g^x, b=h^x)
         */
        BigInteger x;

        public DlogStatementWitness(BigInteger x) {
            this.x = x;
        }
    }
}
