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

// Some utils to work with hex numbers
import { ErrorUtils, HexUtils, NumberUtils } from '@tornado/sdk-utils'

// Parse some files
import { Files, Json } from '@tornado/sdk-data'

// ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ DECLARATIONS ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~

// Interfaces for the cryptographic primitives.

export type bigInt = typeof bigInt

// ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ OUTPUTS ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
/**
 * These are types which define and scope (!) the outputs of the cryptograhic
 * functions. They have been scoped such that there isn't much naming confusion
 * between what a function takes and not. Common aliases are used there where deemed
 * appropriate.
 */
export namespace OutputOf {
  export type PedersenHash = bigInt

  export interface CreateDeposit {
    nullifier: bigInt
    secret: bigInt
    preimage: Buffer
    commitment: PedersenHash
    hexCommitment: string
    nullifierHash: PedersenHash
    hexNullifierHash: string
  }

  export interface MerkleProof {
    root: any
    path: {
      indices: any
      elements: any[]
    }
  }

  export interface Groth16Proof {
    pi_a: string
    pi_b: string
    pi_c: string
  }
}

// ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ OUTPUT ALIASES ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
/**
 * There may be also multiple output aliases for one output, to make the code more readable,
 * in a certain sense of speaking. The delimiter below just makes this comment not pass on to
 * `MerkleTree`.
 */
type __OutputAliasDelimiter = null

export type MerkleProof = OutputOf.MerkleProof
export type DepositProof = OutputOf.Groth16Proof

// ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ INPUTS ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
/**
 * These are types which define and scope (!) the inputs of the cryptograhic
 * functions. The same as above applies, but note that some outputs are inputs
 * for other functions. For example, MerkleTree is an output of the function
 * which builds the Tornado Cash merkle tree, and this is passed as input,
 * to the merkle proof builder.
 */
export namespace InputFor {
  export interface PedersenHash {
    msg: Buffer
  }

  export interface CreateDeposit {
    nullifier?: bigInt
    secret?: bigInt
  }

  export interface BuildMerkleTree {
    height: number
    leaves: Array<Element>
  }

  export interface DepositProof {
    public: {
      root: Element
      tree: MerkleTree
      leafIndex: number
      hexNullifierHash: string
      recipientAddress: string
      relayerAddress: string
      fee: bigInt
      refund: bigInt
    }
    private: {
      nullifier: bigInt
      secret: bigInt
    }
  }

  interface PublicGroth16 {
    root: any
    nullifierHash: PedersenHash
    recipient: bigInt
    relayer: bigInt
    refund: bigInt
    fee: bigInt
  }

  interface PrivateGroth16 {
    nullifier: bigInt
    secret: bigInt
    pathIndices: number[]
    pathElements: string[]
  }

  export type Groth16 = PublicGroth16 & PrivateGroth16
}

// ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ INPUT ALIASES ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
/**
 * There may be also multiple input aliases for one input, to make the code more readable,
 * in a certain sense of speaking. The delimiter below just makes this comment not pass on to
 * `TornadoDeposit`.
 */
type __InputAliasDelimiter = null

export type DepositInfo = OutputOf.CreateDeposit

// ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ 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 function getZeroElement(): string {
    return '21663839004416932945382355908790599225266501822907911457504978515578255421292'
  }

  export function getDefaultHash(): HashFunction<Element> {
    // @ts-ignore
    return (left, right) => mimcsponge.multiHash([bigInt(left), bigInt(right)]).toString()
  }
  export function getProvingKey(): ArrayBufferLike {
    return Files.loadRawSync('circuits/tornadoProvingKey.bin').buffer
  }

  export function getTornadoCircuit(): any {
    return Json.loadSync('circuits/tornado.json')
  }

  let cachedGroth16Prover: Groth16 | null = null

  export async function getGroth16(): Promise<Groth16> {
    const defaultParams = { wasmInitialMemory: 5000 }
    if (!cachedGroth16Prover) cachedGroth16Prover = await buildGroth16(defaultParams)
    return cachedGroth16Prover
  }

  export function terminateGroth16(): void {
    cachedGroth16Prover!.terminate()
    cachedGroth16Prover = null
  }
}

// ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ Primitives ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
/**
 * These are the important cryptographic primitives Tornado Cash uses.
 */
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): DepositInfo {
    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)
    }
  ): DepositInfo {
    // @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 {
    return new MerkleTree(inputs.height, inputs.leaves, {
      hashFunction: Setup.getDefaultHash(),
      zeroElement: Setup.getZeroElement()
    })
  }

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

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

      let proofData

      // Compute Merkle Proof
      // The ts return is noted as `pathIndex` but using this we get an undefined because it is really `pathIndices`???
      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) {
        Setup.terminateGroth16()
        throw ErrorUtils.ensureError(err)
      }

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

      proofs[i].push(
        String(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.
    Setup.terminateGroth16()

    return proofs
  }
}