#!/usr/bin/env node

/*
    Copyright 2018 0KIMS association.

    This file is part of jaz (Zero Knowledge Circuit Compiler).

    jaz is a free software: you can redistribute it and/or modify it
    under the terms of the GNU General Public License as published by
    the Free Software Foundation, either version 3 of the License, or
    (at your option) any later version.

    jaz is distributed in the hope that it will be useful, but WITHOUT
    ANY WARRANTY; without even the implied warranty of MERCHANTABILITY
    or FITNESS FOR A PARTICULAR PURPOSE. See the GNU General Public
    License for more details.

    You should have received a copy of the GNU General Public License
    along with jaz. If not, see <https://www.gnu.org/licenses/>.
*/

/* eslint-disable no-console */

const fs = require("fs");
const path = require("path");

const zkSnark = require("./index.js");
const {stringifyBigInts, unstringifyBigInts} = require("./src/stringifybigint.js");


const version = require("./package").version;

const argv = require("yargs")
    .version(version)
    .usage(`snarkjs <command> <options>

setup command
=============

    snarkjs setup <option>

    Runs a setup for a circuit generating the proving and the verification key.

    -c or --circuit <circuitFile>

        Filename of the compiled circuit file generated by circom.

        Default: circuit.json

    --pk or --provingkey <provingKeyFile>

        Output filename where the proving key will be stored.

        Default: proving_key.json

    --vk or --verificationkey <verificationKeyFile>

        Output Filename where the verification key will be stored.

        Default: verification_key.json

calculate witness command
=========================

    snarkjs calculatewitness <options>

    Calculate the witness of a circuit given an input.

    -c or --circuit <circuitFile>

        Filename of the compiled circuit file generated by circom.

        Default: circuit.json

    -i or --input <inputFile>

        JSON file with the inputs of the circuit.

        Default: input.json

        Example of for a circuit with tow inputs a and b:

            {"a": "22", "b": "33"}

    -w or --witness

        Output filename with the generated witness.

        Default: witness.json

generate a proof command
========================

    snarkjs proof <options>

    -w or --witness

        Input filename used to calculate the proof.

        Default: witness.json

    --pk or --provingkey <provingKeyFile>

        Input filename with the proving key (generated during the setup).

        Default: proving_key.json

    -p or --proof

        Output filenam with the zero knowlage proof.

        Default: proof.json

    --pub or --public <publicFilename>

        Output filename with the value of the public wires/signals.
        This info will be needed to verify the proof.

        Default: public.json

verify command
==============

    snarkjs verify <options>

    The command returns "OK" if the proof is valid
    and  "INVALID" in case it is not a valid proof.

    --vk or --verificationkey <verificationKeyFile>

        Input Filename with the verification key (generated during the setup).

        Default: verification_key.json

    -p or --proof

        Input filenam with the zero knowlage proof you want to verify

        Default: proof.json

    --pub or --public <publicFilename>

        Input filename with the public wires/signals.

        Default: public.json

generate solidity verifier command
==================================

    snarkjs generateverifier <options>

    Generates a solidity smart contract that verifies the zero knowlage proof.

    --vk or --verificationkey <verificationKeyFile>

        Input Filename with the verification key (generated during the setup).

        Default: verification_key.json

    -v or --verifier

        Output file with a solidity smart contract that verifies a zero knowlage proof.

        Default: verifier.sol

generate call parameters
========================

    snarkjs generatecall <options>

    Outputs into the console the raw parameters to be used in 'verifyProof'
    method of the solidity verifier function.

    -p or --proof

        Input filename with the zero knowlage proof you want to use

        Default: proof.json

    --pub or --public <publicFilename>

        Input filename with the public wires/signals.

        Default: public.json

circuit info
============

    snarkjs info <options>

    Print statistics of a circuit

    -c or --circuit <circuitFile>

        Filename of the compiled circuit file generated by circom.

        Default: circuit.json

print constraints
=================

    snarkjs printconstraints <options>

    Print all the constraints of a given circuit

    -c or --circuit <circuitFile>

        Filename of the compiled circuit file generated by circom.

        Default: circuit.json
        `)
    .alias("c", "circuit")
    .alias("pk", "provingkey")
    .alias("vk", "verificationkey")
    .alias("w", "witness")
    .alias("p", "proof")
    .alias("i", "input")
    .alias("pub", "public")
    .alias("v", "verifier")
    .help("h")
    .alias("h", "help")
    .epilogue(`Copyright (C) 2018  0kims association
    This program comes with ABSOLUTELY NO WARRANTY;
    This is free software, and you are welcome to redistribute it
    under certain conditions; see the COPYING file in the official
    repo directory at  https://github.com/iden3/circom `)
    .argv;

const circuitName = (argv.circuit) ? argv.circuit : "circuit.json";
const provingKeyName = (argv.provingkey) ? argv.provingkey : "proving_key.json";
const verificationKeyName = (argv.verificationkey) ? argv.verificationkey : "verification_key.json";
const inputName = (argv.input) ? argv.input : "input.json";
const witnessName = (argv.witness) ? argv.witness : "witness.json";
const proofName = (argv.proof) ? argv.proof : "proof.json";
const publicName = (argv.public) ? argv.public : "public.json";
const verifierName = (argv.public) ? argv.public : "verifier.sol";

function p256(n) {
    let nstr = n.toString(16);
    while (nstr.length < 64) nstr = "0"+nstr;
    nstr = `"0x${nstr}"`;
    return nstr;
}

try {
    if (argv._[0].toUpperCase() == "INFO") {
        const cirDef = JSON.parse(fs.readFileSync(circuitName, "utf8"));
        const cir = new zkSnark.Circuit(cirDef);

        console.log(`# Wires: ${cir.nVars}`);
        console.log(`# Constraints: ${cir.nConstraints}`);
        console.log(`# Private Inputs: ${cir.nPrvInputs}`);
        console.log(`# Public Inputs: ${cir.nPubInputs}`);
        console.log(`# Outputs: ${cir.nOutputs}`);

    } else if (argv._[0].toUpperCase() == "PRINTCONSTRAINTS") {
        const cirDef = JSON.parse(fs.readFileSync(circuitName, "utf8"));
        const cir = new zkSnark.Circuit(cirDef);

        cir.printConstraints();

    } else if (argv._[0].toUpperCase() == "SETUP") {
        const cirDef = JSON.parse(fs.readFileSync(circuitName, "utf8"));
        const cir = new zkSnark.Circuit(cirDef);
        const setup = zkSnark.setup(cir);

        fs.writeFileSync(provingKeyName, JSON.stringify(stringifyBigInts(setup.vk_proof), null, 1), "utf-8");
        fs.writeFileSync(verificationKeyName, JSON.stringify(stringifyBigInts(setup.vk_verifier), null, 1), "utf-8");
        process.exit(0);
    } else if (argv._[0].toUpperCase() == "CALCULATEWITNESS") {
        const cirDef = JSON.parse(fs.readFileSync(circuitName, "utf8"));
        const cir = new zkSnark.Circuit(cirDef);
        const input = unstringifyBigInts(JSON.parse(fs.readFileSync(inputName, "utf8")));

        const witness = cir.calculateWitness(input);

        fs.writeFileSync(witnessName, JSON.stringify(stringifyBigInts(witness), null, 1), "utf-8");
        process.exit(0);
    } else if (argv._[0].toUpperCase() == "PROOF") {
        const witness = unstringifyBigInts(JSON.parse(fs.readFileSync(witnessName, "utf8")));
        const provingKey = unstringifyBigInts(JSON.parse(fs.readFileSync(provingKeyName, "utf8")));

        const {proof, publicSignals} = zkSnark.genProof(provingKey, witness);

        fs.writeFileSync(proofName, JSON.stringify(stringifyBigInts(proof), null, 1), "utf-8");
        fs.writeFileSync(publicName, JSON.stringify(stringifyBigInts(publicSignals), null, 1), "utf-8");
        process.exit(0);
    } else if (argv._[0].toUpperCase() == "VERIFY") {
        const public = unstringifyBigInts(JSON.parse(fs.readFileSync(publicName, "utf8")));
        const verificationKey = unstringifyBigInts(JSON.parse(fs.readFileSync(verificationKeyName, "utf8")));
        const proof = unstringifyBigInts(JSON.parse(fs.readFileSync(proofName, "utf8")));
        const isValid = zkSnark.isValid(verificationKey, proof, public);

        if (isValid) {
            console.log("OK");
            process.exit(0);
        } else {
            console.log("INVALID");
            process.exit(1);
        }
    } else if (argv._[0].toUpperCase() == "GENERATEVERIFIER") {

        const verificationKey = unstringifyBigInts(JSON.parse(fs.readFileSync(verificationKeyName, "utf8")));
        let template = fs.readFileSync(path.join( __dirname,  "templates", "verifier.sol"), "utf-8");

        const vka_str = `[${verificationKey.vk_a[0][1].toString()},`+
                         `${verificationKey.vk_a[0][0].toString()}], `+
                        `[${verificationKey.vk_a[1][1].toString()},` +
                         `${verificationKey.vk_a[1][0].toString()}]`;
        template = template.replace("<%vk_a%>", vka_str);

        const vkb_str = `${verificationKey.vk_b[0].toString()},`+
                        `${verificationKey.vk_b[1].toString()}`;
        template = template.replace("<%vk_b%>", vkb_str);

        const vkc_str = `[${verificationKey.vk_c[0][1].toString()},`+
                         `${verificationKey.vk_c[0][0].toString()}], `+
                        `[${verificationKey.vk_c[1][1].toString()},` +
                         `${verificationKey.vk_c[1][0].toString()}]`;
        template = template.replace("<%vk_c%>", vkc_str);

        const vkg_str = `[${verificationKey.vk_g[0][1].toString()},`+
                         `${verificationKey.vk_g[0][0].toString()}], `+
                        `[${verificationKey.vk_g[1][1].toString()},` +
                         `${verificationKey.vk_g[1][0].toString()}]`;
        template = template.replace("<%vk_g%>", vkg_str);

        const vkgb1_str = `${verificationKey.vk_gb_1[0].toString()},`+
                          `${verificationKey.vk_gb_1[1].toString()}`;
        template = template.replace("<%vk_gb1%>", vkgb1_str);

        const vkgb2_str = `[${verificationKey.vk_gb_2[0][1].toString()},`+
                           `${verificationKey.vk_gb_2[0][0].toString()}], `+
                          `[${verificationKey.vk_gb_2[1][1].toString()},` +
                           `${verificationKey.vk_gb_2[1][0].toString()}]`;
        template = template.replace("<%vk_gb2%>", vkgb2_str);

        const vkz_str = `[${verificationKey.vk_z[0][1].toString()},`+
                         `${verificationKey.vk_z[0][0].toString()}], `+
                        `[${verificationKey.vk_z[1][1].toString()},` +
                         `${verificationKey.vk_z[1][0].toString()}]`;
        template = template.replace("<%vk_z%>", vkz_str);

        // The points

        template = template.replace("<%vk_input_length%>", (verificationKey.A.length-1).toString());
        template = template.replace("<%vk_ic_length%>", verificationKey.A.length.toString());
        let vi = "";
        for (let i=0; i<verificationKey.A.length; i++) {
            if (vi != "") vi = vi + "        ";
            vi = vi + `vk.IC[${i}] = Pairing.G1Point(${verificationKey.A[i][0].toString()},`+
                                                    `${verificationKey.A[i][1].toString()});\n`;
        }
        template = template.replace("<%vk_ic_pts%>", vi);


        fs.writeFileSync(verifierName, template, "utf-8");
        process.exit(0);

    } else if (argv._[0].toUpperCase() == "GENERATECALL") {

        const public = unstringifyBigInts(JSON.parse(fs.readFileSync(publicName, "utf8")));
        const proof = unstringifyBigInts(JSON.parse(fs.readFileSync(proofName, "utf8")));

        let inputs = "";
        for (let i=0; i<public.length; i++) {
            if (inputs != "") inputs = inputs + ",";
            inputs = inputs + p256(public[i]);
        }

        const S=`[${p256(proof.pi_a[0])}, ${p256(proof.pi_a[1])}],` +
                `[${p256(proof.pi_ap[0])}, ${p256(proof.pi_ap[1])}],` +
                `[[${p256(proof.pi_b[0][1])}, ${p256(proof.pi_b[0][0])}],[${p256(proof.pi_b[1][1])}, ${p256(proof.pi_b[1][0])}]],` +
                `[${p256(proof.pi_bp[0])}, ${p256(proof.pi_bp[1])}],` +
                `[${p256(proof.pi_c[0])}, ${p256(proof.pi_c[1])}],` +
                `[${p256(proof.pi_cp[0])}, ${p256(proof.pi_cp[1])}],` +
                `[${p256(proof.pi_h[0])}, ${p256(proof.pi_h[1])}],` +
                `[${p256(proof.pi_kp[0])}, ${p256(proof.pi_kp[1])}],` +
                `[${inputs}]` ;

        console.log(S);
        process.exit(0);
    } else {
        throw new Error("Invalid Command");
    }
} catch(err) {
    console.log("ERROR: " + err);
    process.exit(1);
}