// SPDX-License-Identifier: AGPL-3.0-only

//! This is an example program that is used to create the contract calls
//! for our contract. The generated contract calls can be fed into the `drk`
//! wallet utility to create a full transaction.

use std::str::FromStr;
use std::{env, process};

use darkfi::util::encoding::base64;
use darkfi::zk::{empty_witnesses, halo2::Value, Proof, ProvingKey, Witness, ZkCircuit};
use darkfi::zkas::ZkBinary;
use darkfi_sdk::crypto::{ContractId, Keypair, SecretKey};
use darkfi_sdk::{ContractCall, ContractCallImport};
use darkfi_serial::serialize;
use rand::rngs::OsRng;

// This comes from our contract's lib.rs
use membership::{ContractFunction, MembershipParams};

const USAGE: &str = r#"
Commands:
  generate                              - Creates a keypair for the membership proof
  register <contract-id> <secret_key>   - Creates a Register contract call for the given key
  deregister <contract-id> <secret_key> - Creates a Deregister contract call for the given key
"#;

/// We have to include the compiled ZK circuit in order to create ZK proofs
const ZKBIN: &[u8] = include_bytes!("../proof/membership_proof.zk.bin");

fn usage(argv0: &str) {
    eprintln!("Usage: {argv0} <commmand> <args>\n{USAGE}");
    process::exit(1);
}

/// Generate a new membership keypair
fn generate_keypair() {
    let keypair = Keypair::random(&mut OsRng);

    println!("Secret key: {}", keypair.secret);
    println!("Public key: {}", keypair.public);
}

/// Create a Register call
fn register(argv0: &str, cid: ContractId, secret_key: Option<String>) {
    let Some(secret_key) = secret_key else {
        return usage(argv0);
    };

    // Let's parse the compiled ZK circuit and create a proof.
    let zkbin = ZkBinary::decode(ZKBIN, false).unwrap();

    // For the membership, we will use our secret key.
    // We'll decode it from the input.
    let secret_key = SecretKey::from_str(&secret_key).unwrap();
    let keypair = Keypair::new(secret_key);

    // Our circuit derives the public key by multiplying the secret
    // with a generator point on the curve. The `PublicKey` from the
    // `Keypair` object already did this. Manually it would be:
    // public_key = NullifierK.generator() * fp_mod_fv(secret_key.inner());

    // So our public inputs for the circuit are the public key's coordinates.
    // They have to be in the same order as the `constrain_instance` calls
    // in the ZK circuit.
    let (pub_x, pub_y) = keypair.public.xy();
    let public_inputs = vec![pub_x, pub_y];
    let circuit_witnesses = empty_witnesses(&zkbin).unwrap();

    // For our private witness values, we'll use our secret key.
    let prover_witnesses = vec![Witness::Base(Value::known(keypair.secret.inner()))];

    // Now we create the circuit and its ProvingKey so we are able to create
    // the ZK proof. We are using `eprintln` for execution logs so they are
    // printed in `stderr` as we want to print just the encoded transaction
    // in `stdout` for further processing.
    eprintln!("Generating {} circuit ProvingKey...", zkbin.namespace);
    let circuit = ZkCircuit::new(circuit_witnesses, &zkbin);
    let proving_key = ProvingKey::build(zkbin.k, &circuit);

    // Now create the ZK proof
    eprintln!("Generating {} ZK proof...", zkbin.namespace);
    let circuit = ZkCircuit::new(prover_witnesses, &zkbin);
    let proof = Proof::create(&proving_key, &[circuit], &public_inputs, &mut OsRng).unwrap();

    // Create the contract call.
    // We will call the `Register` function, and we'll use our secret key
    // as one of the transaction signatures. This will be recognized by the
    // wallet when it is imported and it will use the key to sign the tx along
    // any other keys it needs from the wallet.
    let params = MembershipParams {
        member: keypair.public,
    };

    // Build the payload
    let mut payload = vec![ContractFunction::Register as u8];
    payload.extend_from_slice(&serialize(&params));

    let call = ContractCall {
        contract_id: cid,
        data: payload,
    };

    // Now we'll pack it all up for the wallet. It is printed on stdout
    // encoded as base64 which can later be imported by the wallet when
    // creating the full transaction.
    let call_import =
        ContractCallImport::new(call, vec![proof.as_ref().to_vec()], vec![keypair.secret]);

    println!("{}", base64::encode(&serialize(&call_import)));
}

/// Create a deregister call
fn deregister(argv0: &str, cid: ContractId, secret_key: Option<String>) {
    // This function is mostly the same as `fn register()` with the only
    // exception being the Contract Function. So for explanations, just
    // reference the above.
    let Some(secret_key) = secret_key else {
        return usage(argv0);
    };

    let zkbin = ZkBinary::decode(ZKBIN, false).unwrap();

    let secret_key = SecretKey::from_str(&secret_key).unwrap();
    let keypair = Keypair::new(secret_key);

    let (pub_x, pub_y) = keypair.public.xy();
    let public_inputs = vec![pub_x, pub_y];
    let circuit_witnesses = empty_witnesses(&zkbin).unwrap();

    let prover_witnesses = vec![Witness::Base(Value::known(keypair.secret.inner()))];

    eprintln!("Generating {} circuit ProvingKey...", zkbin.namespace);
    let circuit = ZkCircuit::new(circuit_witnesses, &zkbin);
    let proving_key = ProvingKey::build(zkbin.k, &circuit);

    eprintln!("Generating {} ZK proof...", zkbin.namespace);
    let circuit = ZkCircuit::new(prover_witnesses, &zkbin);
    let proof = Proof::create(&proving_key, &[circuit], &public_inputs, &mut OsRng).unwrap();

    let params = MembershipParams {
        member: keypair.public,
    };

    let mut payload = vec![ContractFunction::Deregister as u8];
    payload.extend_from_slice(&serialize(&params));

    let call = ContractCall {
        contract_id: cid,
        data: payload,
    };

    let call_import =
        ContractCallImport::new(call, vec![proof.as_ref().to_vec()], vec![keypair.secret]);

    println!("{}", base64::encode(&serialize(&call_import)));
}

fn main() {
    let mut args = env::args();
    let argv0 = args.next().unwrap();

    let Some(command) = args.next() else {
        return usage(&argv0);
    };

    if command.as_str() == "generate" {
        return generate_keypair();
    }

    let Some(contract_id) = args.next() else {
        return usage(&argv0);
    };

    let cid = ContractId::from_str(&contract_id).unwrap();

    match command.as_str() {
        "register" => register(&argv0, cid, args.next()),
        "deregister" => deregister(&argv0, cid, args.next()),
        _ => usage(&argv0),
    }
}