// crypto types
import * as Types 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 { MerkleTree } from 'fixed-merkle-tree'
import { genWitnessAndProve, toSolidityInput } from 'websnark/src/utils'

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

// Parse some files
import { Files } 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<Buffer> {
    return await Files.loadRaw('circuits/tornado.json')
  }

  /**
   * @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 }
    return await buildGroth16(defaultParams)
  }
}

export namespace Primitives {
  export function calcPedersenHash(
    pedersenHashData: Types.InputFor.PedersenHash
  ): Types.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): Types.ZKDepositData {
    const _hexNote = hexNote.split('_')[1] ?? hexNote
    const buffer = Buffer.from(_hexNote, 'hex')
    return createDeposit({
      // @ts-expect-error
      nullifier: bigInt.leBuff2int(buffer.subarray(0, 31)),
      // @ts-expect-error
      secret: bigInt.leBuff2int(buffer.subarray(32, 62))
    })
  }

  export function createDeposit(input?: Types.InputFor.CreateDeposit): Types.ZKDepositData {
    if (!input?.nullifier || !input?.secret)
      input = {
        nullifier: NumberUtils.randomBigInteger(31),
        secret: NumberUtils.randomBigInteger(31)
      }
    // @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: Types.InputFor.BuildMerkleTree): MerkleTree {
    return new MerkleTree(inputs.height, inputs.leaves)
  }

  export async function calcDepositProofs(
    inputs: Array<Types.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]

      // Compute Merkle Proof
      const { pathElements, pathIndex } = input.public.tree.path(input.public.leafIndex)

      proofs.push([])

      const 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: pathIndex
        },
        circuit,
        provingKey
      )

      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 { Types }