// For reference implementation, see https://github.com/smartcontractkit/ccip/blob/ccip-develop/core/services/ocr2/plugins/ccip/merklemulti/merkle_multi.go
import { ZERO_HASH, hashInternal } from './common.ts'
import {
  CCIPMerkleFlagsMismatchError,
  CCIPMerkleHashesTooLargeError,
  CCIPMerkleInternalError,
  CCIPMerkleProofEmptyError,
  CCIPMerkleProofFlagsMismatchError,
  CCIPMerkleProofIncompleteError,
  CCIPMerkleProofTooLargeError,
  CCIPMerkleTreeEmptyError,
} from '../errors/index.ts'

export const MAX_NUMBER_TREE_LEAVES = 256

type Hash = string

interface SingleLayerProof {
  nextIndices: number[]
  subProof: Hash[]
  sourceFlags: boolean[]
}

/**
 * Proof represents a proof data structure with hashes and source flags.
 */
export class Proof {
  hashes: Hash[] = []
  sourceFlags: boolean[] = []

  /**
   * Creates a new Proof instance.
   * @param hashes - Proof hashes.
   * @param sourceFlags - Source flags indicating leaf vs internal nodes.
   */
  constructor(hashes: Hash[] = [], sourceFlags: boolean[] = []) {
    this.hashes = hashes
    this.sourceFlags = sourceFlags
  }

  /**
   * Counts the number of source flags that match the given boolean value.
   * @param b - The boolean value to count.
   * @returns The count of source flags matching the provided boolean value.
   */
  countSourceFlags(b: boolean): number {
    return this.sourceFlags.filter((flag) => flag === b).length
  }
}

/**
 * Computes the next layer of a binary tree structure represented as an array of hash values.
 *
 * This function takes a context `ctx` and an array `layer` containing hash values as input.
 * It processes the `layer` to generate the next layer of hash values, pairing adjacent values
 * and applying the hashing function provided by the context.
 *
 * @param layer - An array of hash values representing the current layer of the binary tree.
 * @returns A tuple containing two arrays:
 *   - The first array contains the original input layer.
 *   - The second array contains the next layer of hash values generated by combining pairs of hashes.
 */
function computeNextLayer(layer: Hash[]): readonly [Hash[], Hash[]] {
  if (layer.length === 1) {
    return [layer, layer]
  }
  if (layer.length % 2 !== 0) {
    layer.push(ZERO_HASH)
  }

  const nextLayer: Hash[] = []
  for (let i = 0; i < layer.length; i += 2) {
    nextLayer.push(hashInternal(layer[i]!, layer[i + 1]!))
  }

  return [layer, nextLayer]
}

/**
 * Calculates the parent index of a given index in a binary tree.
 * @param idx - The index for which to calculate the parent index.
 * @returns The parent index.
 */
function parentIndex(idx: number): number {
  return Math.floor(idx / 2)
}

/**
 * Calculates the sibling index of a given index in a binary tree.
 * @param idx - The index for which to calculate the sibling index.
 * @returns The sibling index.
 */
function siblingIndex(idx: number): number {
  return idx ^ 1
}

/**
 * Generates a single-layer Merkle proof for a given set of indices within a layer.
 * @param layer - The layer of data for which to generate the proof.
 * @param indices - The indices within the layer for which to generate the proof.
 * @returns A single-layer proof object containing nextIndices, subProof, and sourceFlags.
 */
function proveSingleLayer(layer: Hash[], indices: number[]): SingleLayerProof {
  const authIndices: number[] = []
  const nextIndices: number[] = []
  const sourceFlags: boolean[] = []

  let j = 0
  while (j < indices.length) {
    const x = indices[j]!

    // Calculate the parent index for the current index and add it to the nextIndices array.
    nextIndices.push(parentIndex(x))

    // Check if there is a sibling index following the current index.
    if (j + 1 < indices.length && indices[j + 1] === siblingIndex(x)) {
      // If a sibling index is found, mark it as a source of the proof (sourceFlag = true).
      j++
      sourceFlags.push(true)
    } else {
      // If there is no sibling index, add the sibling index to the authentication indices.
      authIndices.push(siblingIndex(x))
      // Mark it as not a source of the proof (sourceFlag = false).
      sourceFlags.push(false)
    }

    j++
  }

  // Initialize an array to store the subproof.
  const subProof: Hash[] = []
  for (const i of authIndices) {
    // Populate the subproof array with the values from the authentication indices in the layer.
    subProof.push(layer[i]!)
  }

  // Return the single-layer proof object containing nextIndices, subProof, and sourceFlags.
  return {
    nextIndices,
    subProof,
    sourceFlags,
  }
}

/**
 * Represents a Merkle Tree data structure.
 */
export class Tree {
  layers: Hash[][]

  /**
   * Creates a new Merkle Tree using the provided context and leaf hashes.
   * @param leafHashes - An array of leaf hashes to construct the tree.
   * @returns A new Merkle Tree instance.
   * @throws Error if there are no leaf hashes provided.
   */
  constructor(leafHashes: Hash[]) {
    if (leafHashes.length === 0) {
      throw new CCIPMerkleTreeEmptyError()
    }

    // Initialize the first layer with leaf hashes.
    let layer: Hash[] = [...leafHashes]

    // Initialize an array to store layers.
    const layers: Hash[][] = [layer]

    let curr = 0

    // Continue building layers until there is only one hash remaining.
    while (layer.length > 1) {
      const [paddedLayer, nextLayer] = computeNextLayer(layer)
      layers[curr] = paddedLayer
      curr++
      layers.push(nextLayer)
      layer = nextLayer
    }

    this.layers = layers
  }

  /**
   * Returns the root hash of the Merkle Tree.
   * @returns The root hash.
   */
  root(): Hash {
    return this.layers[this.layers.length - 1]![0]!
  }

  /**
   * Generates a Merkle proof for a set of indices in the tree.
   * @param indices - The indices for which to generate the proof.
   * @returns A proof object containing hashes and source flags.
   */
  prove(indices: number[]): Proof {
    const proof: Proof = new Proof([], [])

    for (const layer of this.layers.slice(0, this.layers.length - 1)) {
      const res = proveSingleLayer(layer, indices)

      indices = res.nextIndices
      proof.hashes.push(...res.subProof)
      proof.sourceFlags.push(...res.sourceFlags)
    }

    return proof
  }
}

/**
 * Converts an array of boolean proof flags to a BigNumber where each flag represents a bit.
 * @param proofFlags - An array of boolean proof flags.
 * @returns A BigNumber representing the encoded flags.
 */
export function proofFlagsToBits(proofFlags: boolean[]): bigint {
  let encodedFlags = 0n

  // Iterate through the proof flags to encode them into bits.
  for (let i = 0; i < proofFlags.length; i++) {
    if (proofFlags[i]) {
      // If the current flag is true, set the corresponding bit to 1 using bitwise OR.
      encodedFlags |= 1n << BigInt(i)
    }
  }
  // Return the final BigNumber representing the encoded flags.
  return encodedFlags
}

/**
 * Verifies and computes the Merkle root hash based on provided leaf hashes and a Merkle proof.
 * @param leafHashes - An array of leaf hashes.
 * @param proof - The Merkle proof containing hashes and source flags.
 * @returns The computed Merkle root hash.
 */
export function verifyComputeRoot(leafHashes: Hash[], proof: Proof): Hash {
  const leavesLength = leafHashes.length
  const proofsLength = proof.hashes.length

  if (leavesLength === 0 && proofsLength === 0) {
    throw new CCIPMerkleProofEmptyError()
  }

  if (leavesLength > MAX_NUMBER_TREE_LEAVES + 1 || proofsLength > MAX_NUMBER_TREE_LEAVES + 1) {
    throw new CCIPMerkleProofTooLargeError(MAX_NUMBER_TREE_LEAVES)
  }

  const totalHashes = leavesLength + proofsLength - 1

  if (totalHashes > MAX_NUMBER_TREE_LEAVES) {
    throw new CCIPMerkleHashesTooLargeError(totalHashes, MAX_NUMBER_TREE_LEAVES)
  }

  if (totalHashes !== proof.sourceFlags.length) {
    throw new CCIPMerkleFlagsMismatchError(totalHashes, proof.sourceFlags.length)
  }

  // If there are no hashes, return the first leaf hash.
  if (totalHashes === 0) {
    return leafHashes[0]!
  }

  // Count the number of source flags pointing to proofs.
  const sourceProofCount = proof.countSourceFlags(false)

  if (sourceProofCount !== proofsLength) {
    throw new CCIPMerkleProofFlagsMismatchError(sourceProofCount, proofsLength)
  }

  const hashes: Hash[] = new Array<Hash>(totalHashes)

  for (let i = 0; i < totalHashes; i++) {
    // Initialize the hashes array with the first leaf hash.
    hashes.push(leafHashes[0]!)
  }

  let leafPos = 0
  let hashPos = 0
  let proofPos = 0

  for (let i = 0; i < totalHashes; i++) {
    let a: Hash | undefined = undefined
    let b: Hash | undefined

    if (proof.sourceFlags[i] === true) {
      if (leafPos < leavesLength) {
        a = leafHashes[leafPos]
        leafPos++
      } else {
        a = hashes[hashPos]
        hashPos++
      }
    } else if (proof.sourceFlags[i] === false) {
      a = proof.hashes[proofPos]
      proofPos++
    }

    if (leafPos < leavesLength) {
      b = leafHashes[leafPos]
      leafPos++
    } else {
      b = hashes[hashPos]
      hashPos++
    }

    if (a === undefined || b === undefined) {
      throw new CCIPMerkleInternalError('Hash operand is undefined')
    }

    // Compute the hash of a and b and store it in the hashes array.
    hashes[i] = hashInternal(a, b)
  }

  // Check if all proofs were used during processing.
  if (hashPos !== totalHashes - 1 || leafPos !== leavesLength || proofPos !== proofsLength) {
    throw new CCIPMerkleProofIncompleteError()
  }

  // Return the computed Merkle root hash.
  return hashes[totalHashes - 1]!
}