#include <stdlib.h>
#include <iostream>
#include <fstream>

#include "ecgadget.hpp"
#include "scalarmul.hpp"

using namespace libsnark;
using namespace std;

typedef enum {
    MODE_NONE,
    MODE_PRIV,
    MODE_PUB,
    MODE_CONST
} Mode;

template<typename FieldT>
class verified_encryption_gadget : public gadget<FieldT> {
private:
  const size_t numbits;
  FieldT curve_b, Gx, Gy, Hx, Hy;
  pb_variable<FieldT> r;
  pb_variable<FieldT> xsquared, ysquared;
  pb_variable_array<FieldT> kbits, rbits;
  pb_variable<FieldT> elgx, elgy;
  pb_linear_combination<FieldT> x;
  vector<packing_gadget<FieldT> > packers;
  vector<ec_constant_scalarmul_vec_gadget<FieldT> > constmuls;
  vector<ec_scalarmul_vec_gadget<FieldT> > muls;
  vector<ec_add_gadget<FieldT> > adders;

public:
  const Mode mode;
  const pb_variable<FieldT> C1x, C1y, C2x, C2y, Kx, Ky;
  const pb_variable<FieldT> Px, Py, k, s, y;
  const pb_variable_array<FieldT> Ptable;

  verified_encryption_gadget(protoboard<FieldT> &pb,
              Mode mode,
              const pb_variable<FieldT> &C1x,
              const pb_variable<FieldT> &C1y,
              const pb_variable<FieldT> &C2x,
              const pb_variable<FieldT> &C2y,
              const pb_variable<FieldT> &Kx,
              const pb_variable<FieldT> &Ky,
              const pb_variable<FieldT> &Px,
              const pb_variable<FieldT> &Py,
              const pb_variable_array<FieldT> &Ptable,
              const pb_variable<FieldT> &k,
              const pb_variable<FieldT> &s,
              const pb_variable<FieldT> &y) :
    gadget<FieldT>(pb, "verified_encryption_gadget"),
    // Curve parameters and generators
    numbits(FieldT::num_bits),
    curve_b("7950939520449436327800262930799465135910802758673292356620796789196167463969"),
    Gx(0), Gy("11977228949870389393715360594190192321220966033310912010610740966317727761886"),
    Hx(1), Hy("21803877843449984883423225223478944275188924769286999517937427649571474907279"),
    mode(mode), C1x(C1x), C1y(C1y), C2x(C2x), C2y(C2y),
    Kx(Kx), Ky(Ky), Px(Px), Py(Py), Ptable(Ptable),
    k(k), s(s), y(y)
  {
    r.allocate(pb, "r");
    xsquared.allocate(pb, "xsquared");
    ysquared.allocate(pb, "ysquared");
    kbits.allocate(pb, numbits-8, "kbits");
    rbits.allocate(pb, numbits, "rbits");

    // The unpacking gadgets to turn k and r into bits
    packers.emplace_back(pb, kbits, k);
    packers.emplace_back(pb, rbits, r);

    // The El Gamal first component r*G
    constmuls.emplace_back(pb, C1x, C1y, rbits, Gx, Gy);

    // The El Gamal intermediate value r*P
    elgx.allocate(pb, "elgx");
    elgy.allocate(pb, "elgy");
    if (mode == MODE_CONST) {
        constmuls.emplace_back(pb, elgx, elgy, rbits, Hx, Hy);
    } else {
        muls.emplace_back(pb, elgx, elgy, rbits, Px, Py, Ptable, mode == MODE_PRIV, true);
    }

    // The El Gamal second component r*P + M
    x.assign(pb, k * 256 + s);
    adders.emplace_back(pb, C2x, C2y, elgx, elgy, x, y);

    // The generated public key k*G
    constmuls.emplace_back(pb, Kx, Ky, kbits, Gx, Gy);
  }

  void generate_r1cs_constraints()
  {
    // Prove (256*k+s,y) is on the curve
    this->pb.add_r1cs_constraint(r1cs_constraint<FieldT>(y, y, ysquared));
    this->pb.add_r1cs_constraint(r1cs_constraint<FieldT>(k * 256 + s, k * 256 + s, xsquared));
    this->pb.add_r1cs_constraint(r1cs_constraint<FieldT>(xsquared - 3, k * 256 + s, ysquared - curve_b));

    for (auto&& gadget : packers) {
        gadget.generate_r1cs_constraints(true);
    }

    for (auto&& gadget : constmuls) {
        gadget.generate_r1cs_constraints();
    }

    for (auto&& gadget : muls) {
        gadget.generate_r1cs_constraints();
    }

    for (auto&& gadget : adders) {
        gadget.generate_r1cs_constraints();
    }
  }

  void generate_r1cs_witness()
  {
    this->pb.val(r) = FieldT::random_element();
    this->pb.val(xsquared) = (this->pb.val(k) * 256 + this->pb.val(s)).squared();
    this->pb.val(ysquared) = this->pb.val(y).squared();
    x.evaluate(this->pb);
    for (auto&& gadget : packers) {
        gadget.generate_r1cs_witness_from_packed();
    }

    for (auto&& gadget : constmuls) {
        gadget.generate_r1cs_witness();
    }

    for (auto&& gadget : muls) {
        gadget.generate_r1cs_witness();
    }

    for (auto&& gadget : adders) {
        gadget.generate_r1cs_witness();
    }
  }
};

int main(int argc, char **argv)
{
  Mode mode = MODE_NONE;

  if (argc == 2) {
    if (!strcmp(argv[1], "priv")) {
        mode = MODE_PRIV;
    } else if (!strcmp(argv[1], "pub")) {
        mode = MODE_PUB;
    } else if (!strcmp(argv[1], "const")) {
        mode = MODE_CONST;
    }
  }
  if (mode == MODE_NONE) {
    cerr << "Usage: " << argv[0] << " mode" << endl << endl;
    cerr << "Where mode is one of:" << endl;
    cerr << " priv:  use private Ptable" << endl;
    cerr << " pub:   use public Ptable" << endl;
    cerr << " const: use constant public key (no Ptable)" << endl;
    exit(1);
  }

  // Initialize the curve parameters

  default_r1cs_gg_ppzksnark_pp::init_public_params();

  typedef libff::Fr<default_r1cs_gg_ppzksnark_pp> FieldT;
  
  // Create protoboard

  libff::start_profiling();

  cout << "Keypair" << endl;

  protoboard<FieldT> pb;
  pb_variable<FieldT> C1x, C1y, C2x, C2y, Kx, Ky;
  pb_variable<FieldT> Px, Py;
  pb_variable_array<FieldT> Ptable;
  pb_variable<FieldT> k, s, y, r;

  const size_t numbits = FieldT::num_bits;

  // Allocate variables

  // Public outputs:

  // El Gamal encryption of k under public key P (or H if MODE_CONST)
  // C1 = r*G, C2 = r*P + M  (where M=(256*k+s,y))
  C1x.allocate(pb, "C1x");
  C1y.allocate(pb, "C1y");
  C2x.allocate(pb, "C2x");
  C2y.allocate(pb, "C2y");

  // Public key corresponding to private key k
  // K = k*G
  Kx.allocate(pb, "Kx");
  Ky.allocate(pb, "Ky");

  // Public inputs:

  // The public key P (if not MODE_CONST)
  if (mode != MODE_CONST) {
    Px.allocate(pb, "Px");
    Py.allocate(pb, "Py");

    // The Ptable might be public or private, according to the mode
    Ptable.allocate(pb, 2*numbits, "Ptable");
  }

  // Private inputs:
  // k is a 246-bit random number
  k.allocate(pb, "k");
  // s and y are such that M = (256*k+s,y) is a point on the curve
  s.allocate(pb, "s");
  y.allocate(pb, "y");

  // This sets up the protoboard variables so that the first n of them
  // represent the public input and the rest is private input

  if (mode == MODE_PRIV) {
      pb.set_input_sizes(8);
  } else if (mode == MODE_PUB) {
      pb.set_input_sizes(8+2*numbits);
  } else if (mode == MODE_CONST) {
      pb.set_input_sizes(6);
  }

  // Initialize the gadgets
  verified_encryption_gadget<FieldT> venc(pb, mode, C1x, C1y, C2x, C2y, Kx, Ky, Px, Py, Ptable, k, s, y);
  venc.generate_r1cs_constraints();

  const r1cs_constraint_system<FieldT> constraint_system = pb.get_constraint_system();

  const r1cs_gg_ppzksnark_keypair<default_r1cs_gg_ppzksnark_pp> keypair = r1cs_gg_ppzksnark_generator<default_r1cs_gg_ppzksnark_pp>(constraint_system);

  // Add witness values

  cout << "Prover" << endl;
  
  if (mode != MODE_CONST) {
    // A variable base point P
    pb.val(Px) = FieldT("1095194319010475832867263440470707690447963461907735667341232728633587089702");
    pb.val(Py) = FieldT("9185463202887631101218413269806857706246311016297504828581985913021301344974");
  }
  pb.val(k) = FieldT("31329510635628557928212225120518124937732397714111203844965919301557399521");
  pb.val(s) = FieldT(1);
  pb.val(y) = FieldT("4364798287654239504994818950156019747851405522689486598132350453516910863367");

  venc.generate_r1cs_witness();

  const r1cs_gg_ppzksnark_proof<default_r1cs_gg_ppzksnark_pp> proof = r1cs_gg_ppzksnark_prover<default_r1cs_gg_ppzksnark_pp>(keypair.pk, pb.primary_input(), pb.auxiliary_input());

  cout << "Verifier" << endl;

  bool verified = r1cs_gg_ppzksnark_verifier_strong_IC<default_r1cs_gg_ppzksnark_pp>(keypair.vk, pb.primary_input(), proof);

  cout << "Number of R1CS constraints: " << constraint_system.num_constraints() << endl;
  cout << "Primary (public) input length: " << pb.primary_input().size() << endl;
//  cout << "Primary (public) input: " << pb.primary_input() << endl;
  cout << "Auxiliary (private) input length: " << pb.auxiliary_input().size() << endl;
//  cout << "Auxiliary (private) input: " << pb.auxiliary_input() << endl;
  cout << "Verification status: " << verified << endl;

  ofstream pkfile(string("pk_verifenc_") + argv[1]);
  pkfile << keypair.pk;
  pkfile.close();
  ofstream vkfile(string("vk_verifenc_") + argv[1]);
  vkfile << keypair.vk;
  vkfile.close();
  ofstream pffile(string("proof_verifenc_") + argv[1]);
  pffile << proof;
  pffile.close();

  return 0;
}