import pkg from 'hardhat';
import { getSigners, getContractFactory } from '@nomiclabs/hardhat-ethers/src/internal/helpers'
import { expect } from 'chai';
import { BigNumber } from 'ethers';
import { Noir } from '../utils/noir';
import { execSync } from 'child_process';
import { MerkleTree } from '../utils/MerkleTree'
import {
  generateUTXO,
  generateTestTransaction,
  generateTestPublish,
  generateDataToml,
  randomBytesFr,
  generateTreeProof,
  treeConfig,
  dumpToml,
  dumpTomlRecursive
} from '../utils/test_helpers';
import fs from "fs"
// @ts-ignore
import { Fr } from '@aztec/bb.js';

import mainCircuit from '../circuits/main/target/main.json';
import recursiveCircuit from '../circuits/recursion/target/recursion.json';

import { beforeAll, describe } from 'vitest';
import { SignerWithAddress } from '@nomiclabs/hardhat-ethers/src/signers';

let recompile: boolean = true;
let preprocessing: boolean = false;
let signers: SignerWithAddress[], recipients, recipients_new

beforeAll(async () => {
  signers = await getSigners(pkg);
  recipients = new Array(16).fill(`0x` + signers[1].address.slice(2).padStart(64, "0"))
  recipients_new = new Array(16).fill(`0x` + signers[2].address.slice(2).padStart(64, "0"))

  if (recompile) {
    execSync(`cd ./circuits/main/ && nargo check && nargo compile && nargo codegen-verifier`)
    execSync(`cd ./circuits/recursion/ && nargo check && nargo compile && nargo codegen-verifier && cp ./contract/recursion/plonk_vk.sol ../../contracts/plonk_vk.sol`)
    execSync(`npx hardhat compile`)
  }

});

describe('It compiles noir program code, receiving circuit bytes and abi object.', async () => {
  const noirInstances: { main: Noir, recursive: Noir } = {
    main: new Noir(mainCircuit),
    recursive: new Noir(recursiveCircuit)
  }

  const { main: noir } = noirInstances;
  await noir.init();

  let recursiveInputs: string[] = [];
  let trees = {
    utxo_tree: new MerkleTree(treeConfig.utxoDepth, noir.api),
    tx_tree: new MerkleTree(treeConfig.txDepth, noir.api),
    historic_tree: new MerkleTree(treeConfig.stateDepth, noir.api),
    newHistoricRoot: ""
  }

  await trees.utxo_tree.init();
  await trees.tx_tree.init();
  await trees.historic_tree.init();

  let utxoIn: any[] = []
  let utxoOut: any[] = []
  let treeProof: any[] = []
  let data = await generateDataToml("0", "0", trees, noir.api)
  let amountPublic = {
    amountIn: BigInt(0),
    amountOut: new Array(16).fill(BigInt(0))
  }
  let contractInputs: any[]

  let Verifier, verifierContract, TechnicalPreview, technicalPreviewContract, XFTMock, xftMockContract;
  let dataToml

  beforeAll(async () => {
    Verifier = await getContractFactory(pkg, "UltraVerifier");
    verifierContract = await Verifier.deploy();
    let verifierAddr = await verifierContract.deployed();
    console.log(`Verifier deployed to ${verifierAddr.address}`);
    let vk_hash = await verifierContract.getVerificationKeyHash();

    XFTMock = await getContractFactory(pkg, "XFTMock");
    xftMockContract = await XFTMock.deploy(vk_hash);
    let xftMockContractAddr = await xftMockContract.deployed();
    console.log(`Token deployed to ${xftMockContractAddr.address}`);

    TechnicalPreview = await getContractFactory(pkg, "TechnicalPreview");
    technicalPreviewContract = await TechnicalPreview.deploy(verifierContract.address, xftMockContract.address, "0x2f51641a7c20eec5405aedc1309dccfd3841bfd54e87d32957daa0371904fb11");
    await xftMockContract.grantRole(await xftMockContract.MINTER_ROLE(), technicalPreviewContract.address);

    let merkle_root = await technicalPreviewContract.merkleRoot()
    let utxo_leaves = await technicalPreviewContract.getUtxoFromRoot(merkle_root);
    for (let i = 0; i < utxo_leaves.length; i++) {
      await trees.utxo_tree.insert(utxo_leaves[i]);
    }

    let historic_leaves = await technicalPreviewContract.getValidRoots();
    for (let i = 0; i < historic_leaves.length; i++) {
      await trees.historic_tree.insert(historic_leaves[i]);
    }
  })


  it('Should generate valid proof for correct input (deposit)', async () => {

    console.log("** Generating Test Batch **")
    let batchSize0 = 1
    let secret0: Fr[] = []
    for (let s = 0; s < batchSize0; s++) {
      secret0.push(randomBytesFr(32))
    }
    let amountsOutUTXO = new Array(1).fill(BigInt(1e18))
    utxoOut = await generateUTXO(batchSize0, amountsOutUTXO, secret0, noir.api);
    amountPublic.amountIn = BigInt(1e18)
    await generateTestTransaction(utxoIn, utxoOut, trees, treeProof, amountPublic, data, recipients, noir.api, technicalPreviewContract)

    console.log("** Generating Transaction Proof #1 (Deposit) **")
    await generateTestPublish(trees, data, noir.api)
    dumpToml(data)
    const input = [
      data.current_root,
      data.deposit,
      ...data.withdrawals,
      ...data.commitment_out,
      ...data.recipients,
      data.oracle,
      ...data.old_root_proof,
      ...data.nullifier_hashes,
      ...data.secrets,
      ...data.utxo_in,
      ...data.utxo_out,
      ...data.indexes,
      ...data.hash_path
    ];

    const public_input = [
      data.current_root,
      data.deposit,
      ...data.withdrawals,
      ...data.commitment_out,
      ...data.recipients,
      data.oracle,
      ...data.nullifier_hashes
    ]

    const witness = await noir.generateWitness(input);
    const proof = await noir.generateInnerProof(witness);

    expect(proof instanceof Uint8Array).to.be.true;

    const verified = await noir.verifyInnerProof(proof);
    expect(verified).to.be.true;

    const numPublicInputs = public_input.length + 1;
    const { proofAsFields, vkAsFields, vkHash } = await noir.generateInnerProofArtifacts(
      proof,
      numPublicInputs,
    );

    const publicInputs = proofAsFields.slice(0, numPublicInputs)

    expect(vkAsFields).to.be.of.length(114);
    expect(vkHash).to.be.a('string');

    const aggregationObject = Array(16).fill(
      '0x0000000000000000000000000000000000000000000000000000000000000000',
    );
    recursiveInputs = [
      ...vkAsFields.map(e => e.toString()),
      ...proofAsFields,
      ...publicInputs,
      ...aggregationObject,
      vkHash.toString(),
      ...data.tx_in,
      data.old_root,
      data.new_root,
      data.oracle,
    ];

    dataToml = {
      verification_key: vkAsFields.map(e => e.toString()),
      proof: proofAsFields,
      public_inputs: publicInputs,
      input_aggregation_object: aggregationObject,
      key_hash: vkHash.toString(),
      tx_ids: data.tx_in,
      old_root: data.old_root,
      new_root: data.new_root,
      oracle: data.oracle,
    }


    const deposit = publicInputs[1]
    const withdrawals = publicInputs.slice(2, 18)
    const commitment_out = publicInputs.slice(18, 34)
    const recipientPI = publicInputs.slice(34, 50)
    const nullifier_hashes = publicInputs.slice(51, 67)
    const id = publicInputs[67]

    contractInputs = [
      id,
      commitment_out,
      nullifier_hashes,
      recipientPI,
      withdrawals,
      deposit
    ]

    const savedOutputs = {
      data: data,
      dataToml: dataToml,
      recursiveInputs: recursiveInputs,
      contractInputs: contractInputs
    }

    fs.writeFileSync('./outputs.json', JSON.stringify(savedOutputs)); // Caches outputs
    await noir.destroy();
  });

  it('Should verify proof within a proof', async () => {
    const { recursive: noir } = noirInstances;
    await noir.init();

    const savedOutputs = JSON.parse(fs.readFileSync('./outputs.json').toString())
    data = savedOutputs.data
    dataToml = savedOutputs.dataToml
    recursiveInputs = savedOutputs.recursiveInputs
    contractInputs = savedOutputs.contractInputs

    const proofInput = recursiveInputs.slice(114, 275); 
    const witness = await noir.generateWitness(recursiveInputs);
    const proof = await noir.generateOuterProof(witness);
    expect(proof instanceof Uint8Array).to.be.true;

    const numPublicInputs = 36;
    const { proofAsFields } = await noir.generateInnerProofArtifacts(
      proof,
      numPublicInputs,
    );

    const verified = await noir.verifyOuterProof(proof);
    console.log(verified);

    contractInputs.push(proofInput)
    contractInputs.push(proofAsFields.slice(numPublicInputs - 16, numPublicInputs))

    await technicalPreviewContract.enqueue(contractInputs)

    const batch = await technicalPreviewContract.getCurrentBatch();
    fs.writeFileSync('./batch.json', JSON.stringify(batch))

    dumpTomlRecursive(dataToml)
    execSync(`cd ./circuits/recursion/ && nargo prove`)
    const proofString = `0x` + fs.readFileSync('./circuits/recursion/proofs/recursion.proof').toString()


    const batchPublicInputs = [
      dataToml.key_hash,
      dataToml.oracle,
      dataToml.old_root,
      dataToml.new_root,
      batch
    ]

    await technicalPreviewContract.publish(proofString, batchPublicInputs)
    utxoIn = utxoOut
    await noir.destroy();
  });

  it('Should generate valid proof for correct input (withdrawal)', async () => {
    const { main: noir } = noirInstances
    await noir.init();

    trees.utxo_tree = new MerkleTree(treeConfig.utxoDepth, noir.api)
    trees.tx_tree = new MerkleTree(treeConfig.txDepth, noir.api)
    trees.historic_tree = new MerkleTree(treeConfig.stateDepth, noir.api)
    await trees.utxo_tree.init()
    await trees.tx_tree.init()
    await trees.historic_tree.init()

    console.log("Populating Historic Tree")
    const historicRoots = await technicalPreviewContract.getValidRoots()
    for (let r = 0; r < historicRoots.length; r++) {
      await trees.historic_tree.insert(historicRoots[r])
    }

    console.log("** Generating Transaction Proof #2 (Withdraw/Transfer) **")

    let amountsOutUTXO = new Array(1).fill(BigInt(5e17))
    treeProof = []
    for (let i = 0; i < utxoIn.length; i++) {
      let utxoProof = await generateTreeProof(utxoIn[i], noir.api, technicalPreviewContract)
      treeProof.push(utxoProof);
    }

    console.log("Tree proof generated")

    let batchSize0 = 1
    let secret0: Fr[] = []
    for (let s = 0; s < batchSize0; s++) {
      secret0.push(randomBytesFr(32))
    }

    amountPublic.amountIn = BigInt(0);
    amountPublic.amountOut = new Array(5).fill(1e17);

    utxoOut = await generateUTXO(batchSize0, amountsOutUTXO, secret0, noir.api);

    let oldRoot = dataToml.old_root
    let newRoot = dataToml.new_root
    data = await generateDataToml(oldRoot, newRoot, trees, noir.api)

    await generateTestTransaction(utxoIn, utxoOut, trees, treeProof, amountPublic, data, recipients_new, noir.api, technicalPreviewContract)
    console.log("Test transaction generated")

    await generateTestPublish(trees, data, noir.api)
    console.log("Test publish generated")
    dumpToml(data)
    const input = [
      data.current_root,
      data.deposit,
      ...data.withdrawals,
      ...data.commitment_out,
      ...data.recipients,
      data.oracle,
      ...data.old_root_proof,
      ...data.nullifier_hashes,
      ...data.secrets,
      ...data.utxo_in,
      ...data.utxo_out,
      ...data.indexes,
      ...data.hash_path
    ];

    const public_input = [
      data.current_root,
      data.deposit,
      ...data.withdrawals,
      ...data.commitment_out,
      ...data.recipients,
      data.oracle,
      ...data.nullifier_hashes
    ]

    console.log("Input formatted")

    const witness = await noir.generateWitness(input);
    console.log("witness calculated")
    const proof = await noir.generateInnerProof(witness);
    console.log("Proof generated")

    expect(proof instanceof Uint8Array).to.be.true;

    const verified = await noir.verifyInnerProof(proof);
    expect(verified).to.be.true;

    const numPublicInputs = public_input.length + 1; const { proofAsFields, vkAsFields, vkHash } = await noir.generateInnerProofArtifacts(
      proof,
      numPublicInputs,
    );

    const publicInputs = proofAsFields.slice(0, numPublicInputs)

    expect(vkAsFields).to.be.of.length(114);
    expect(vkHash).to.be.a('string');

    const aggregationObject = Array(16).fill(
      '0x0000000000000000000000000000000000000000000000000000000000000000',
    );
    recursiveInputs = [
      ...vkAsFields.map(e => e.toString()),
      ...proofAsFields,
      ...publicInputs,
      ...aggregationObject,
      vkHash.toString(),
      ...data.tx_in,
      data.old_root,
      data.new_root,
      data.oracle,
    ];

    dataToml = {
      verification_key: vkAsFields.map(e => e.toString()),
      proof: proofAsFields,
      public_inputs: publicInputs,
      input_aggregation_object: aggregationObject,
      key_hash: vkHash.toString(),
      tx_ids: data.tx_in,
      old_root: data.old_root,
      new_root: data.new_root,
      oracle: data.oracle,
    }


    const current_root = publicInputs[0]
    const deposit = publicInputs[1]
    const withdrawals = publicInputs.slice(2, 18)
    const commitment_out = publicInputs.slice(18, 34)
    const recipientPI = publicInputs.slice(34, 50)
    const oracle = publicInputs[50]
    const nullifier_hashes = publicInputs.slice(51, 67)
    const id = publicInputs[67]

    contractInputs = [
      id,
      commitment_out,
      nullifier_hashes,
      recipientPI,
      withdrawals,
      deposit
    ]

    const savedOutputs = {
      data: data,
      dataToml: dataToml,
      recursiveInputs: recursiveInputs,
      contractInputs: contractInputs
    }

    fs.writeFileSync('./outputs.json', JSON.stringify(savedOutputs));


    await noir.destroy();
  });

  it('Should verify proof within a proof (withdrawal)', async () => {
    const { recursive: noir } = noirInstances;
    await noir.init();

    const savedOutputs = JSON.parse(fs.readFileSync('./outputs.json').toString())
    data = savedOutputs.data
    dataToml = savedOutputs.dataToml
    recursiveInputs = savedOutputs.recursiveInputs
    contractInputs = savedOutputs.contractInputs

    const proofInput = recursiveInputs.slice(114, 275); 
    const witness = await noir.generateWitness(recursiveInputs);
    const proof = await noir.generateOuterProof(witness);
    expect(proof instanceof Uint8Array).to.be.true;

    const numPublicInputs = 36;
    const { proofAsFields } = await noir.generateInnerProofArtifacts(
      proof,
      numPublicInputs,
    );

    const verified = await noir.verifyOuterProof(proof);
    console.log(verified);

    contractInputs.push(proofInput)
    contractInputs.push(proofAsFields.slice(numPublicInputs - 16, numPublicInputs))

    await technicalPreviewContract.enqueue(contractInputs)

    const batch = await technicalPreviewContract.getCurrentBatch();
    fs.writeFileSync('./batch.json', JSON.stringify(batch))

    dumpTomlRecursive(dataToml)
    execSync(`cd ./circuits/recursion/ && nargo prove`)
    const proofString = `0x` + fs.readFileSync('./circuits/recursion/proofs/recursion.proof').toString()


    const batchPublicInputs = [
      dataToml.key_hash,
      dataToml.oracle,
      dataToml.old_root,
      dataToml.new_root,
      batch
    ]

    await technicalPreviewContract.publish(proofString, batchPublicInputs)
    utxoIn = utxoOut

    const balanceOut = await xftMockContract.balanceOf(signers[2].address)

    expect(balanceOut).to.be.equal(BigNumber.from(BigInt(5e17).toString()))
    await noir.destroy();
  });

});