// crypto types
import { InputFor, OutputOf, ZKDepositData } from 'types/sdk/crypto'

// External crypto
import circomlib from 'circomlib'
import { bigInt } from 'snarkjs'
import { Groth16 } from 'websnark/src/groth16'
import { buildGroth16 } from 'websnark'
import MerkleTreeDefault, { MerkleTree } from 'fixed-merkle-tree'
import { genWitnessAndProve, toSolidityInput } from 'websnark/src/utils'

// Some utils to work with hex numbers
import { ErrorUtils, HexUtils, NumberUtils } from 'lib/utils'

// Parse some files
import { Files, Json } from 'lib/data'

// ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ SETUP ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
/**
 * Several objects have to be set up, like the groth16 prover. All related types are
 * (will be) contained within this namespace.
 */
export namespace Setup {
  export async function getProvingKey(): Promise<ArrayBufferLike> {
    return (await Files.loadRaw('circuits/tornadoProvingKey.bin')).buffer
  }

  export async function getTornadoCircuit(): Promise<any> {
    return Json.load('circuits/tornado.json')
  }

  let cachedGroth16Prover: Groth16 | null = null

  /**
   * @note The following is a comment from tornado-cli: `groth16 initialises a lot of Promises that will never be resolved, that's why we need to use process.exit to terminate the CLI`. They literally didn't check the code to see that these are just worker threads and that `groth16` has a `terminate()` function to remove them. 🤦
   */
  export async function getGroth16(): Promise<Groth16> {
    const defaultParams = { wasmInitialMemory: 5000 }
    if (!cachedGroth16Prover) cachedGroth16Prover = await buildGroth16(defaultParams)
    return cachedGroth16Prover
  }
}

export namespace Primitives {
  export function calcPedersenHash(pedersenHashData: InputFor.PedersenHash): OutputOf.PedersenHash {
    return circomlib.babyJub.unpackPoint(circomlib.pedersenHash.hash(pedersenHashData.msg))[0]
  }

  export function createNote(msg: Buffer): string {
    return HexUtils.bufferToHex(msg, 62)
  }

  export function parseNote(hexNote: string): ZKDepositData {
    const _hexNote = hexNote.split('_')[1] ?? hexNote
    const buffer = Buffer.from(_hexNote.slice(2), 'hex')
    return createDeposit({
      // @ts-expect-error
      nullifier: bigInt.leBuff2int(buffer.subarray(0, 31)),
      // @ts-expect-error
      secret: bigInt.leBuff2int(buffer.subarray(31, 62))
    })
  }

  export function createDeposit(
    input: InputFor.CreateDeposit = {
      nullifier: NumberUtils.randomBigInteger(31),
      secret: NumberUtils.randomBigInteger(31)
    }
  ): ZKDepositData {
    // @ts-expect-error
    let preimage = Buffer.concat([input.nullifier.leInt2Buff(31), input.secret.leInt2Buff(31)])
    let commitment = calcPedersenHash({ msg: preimage })
    let commitmentHex = HexUtils.bigIntToHex(commitment)
    // @ts-expect-error
    let nullifierHash = calcPedersenHash({ msg: input.nullifier.leInt2Buff(31) })
    let nullifierHex = HexUtils.bigIntToHex(nullifierHash)
    return {
      nullifier: input.nullifier!,
      secret: input.secret!,
      preimage: preimage,
      commitment: commitment,
      hexCommitment: commitmentHex,
      nullifierHash: nullifierHash,
      hexNullifierHash: nullifierHex
    }
  }

  export function buildMerkleTree(inputs: InputFor.BuildMerkleTree): MerkleTree {
    // @ts-expect-error
    return new MerkleTreeDefault(inputs.height, inputs.leaves)
  }

  export async function calcDepositProofs(inputs: Array<InputFor.ZKProof>): Promise<Array<Array<string>>> {
    const proofs: string[][] = []
    const groth16 = await Setup.getGroth16()
    const circuit = await Setup.getTornadoCircuit()
    const provingKey = await Setup.getProvingKey()

    for (let i = 0, len = inputs.length; i < len; i++) {
      const input = inputs[i]

      // TODO: remove try and const again after fixing
      let proofData

      // Compute Merkle Proof
      // The ts return is noted as `pathIndex` but using this we get an undefined because it is really `pathIndices`???
      // TODO: Bug that needs to be fixed (above)
      // @ts-expect-error
      const { pathElements, pathIndices } = input.public.tree.path(input.public.leafIndex)

      proofs.push([])

      try {
        proofData = await genWitnessAndProve(
          groth16,
          {
            // Public inputs
            root: input.public.root,
            // @ts-expect-error
            nullifierHash: bigInt(input.public.hexNullifierHash),
            // @ts-expect-error
            recipient: bigInt(input.public.recipientAddress),
            // @ts-expect-error
            relayer: bigInt(input.public.relayerAddress),
            // @ts-expect-error
            fee: bigInt(input.public.fee),
            //
            refund: input.public.refund,

            // Private inputs
            nullifier: input.private.nullifier,
            secret: input.private.secret,
            pathElements: pathElements,
            pathIndices: pathIndices
          },
          circuit,
          provingKey
        )
      } catch (err) {
        groth16.terminate()
        throw ErrorUtils.ensureError(err)
      }

      proofs[i].push(toSolidityInput(proofData).proof)

      proofs[i].push(
        input.public.root,
        input.public.hexNullifierHash,
        '0x' + HexUtils.prepareAddress(input.public.recipientAddress, 20),
        '0x' + HexUtils.prepareAddress(input.public.relayerAddress, 20),
        HexUtils.bigIntToHex(input.public.fee),
        HexUtils.bigIntToHex(input.public.refund)
      )
    }

    // Done. 🤷‍♀️
    groth16.terminate()

    return proofs
  }
}

// TODO: implement and decide whether to add in declarations an ambient namespace and merge it here

// export function buildMerkleTree(deposit: Crypto.TornadoDeposit): Crypto.MerkleTree {}

// export function calcMerkleProof(tree: Crypto.MerkleTree): Crypto.MerkleProof {}

// export function calcDepositProof(merkleProof: Crypto.InputFor.DepositProof): Crypto.OutputOf.DepositProof {}

// Namespace exports
export { InputFor, OutputOf, ZKDepositData }