mt.ak

use aiken/builtin
use aiken/bytearray
use aiken/hash.{Hash, Sha2_256}
use aiken/list
use aiken/string

// Merkle Tree in Aiken (ported from: https://github.com/input-output-hk/hydra/blob/master/plutus-merkle-tree/src/Plutus/MerkleTree.hs)

pub type MerkleTree<alg> {
  // represents no value (null object pattern)
  MerkleEmpty

  MerkleLeaf { hash: Hash<alg, ByteArray> }

  MerkleNode {
    hash: Hash<alg, ByteArray>,
    left: MerkleTree<alg>,
    right: MerkleTree<alg>,
  }
}

fn node_hash(t: MerkleTree<alg>) -> Hash<alg, ByteArray> {
  when t is {
    MerkleEmpty -> #""
    MerkleLeaf { hash } -> hash
    MerkleNode { hash, .. } -> hash
  }
}

pub fn size(t: MerkleTree<alg>) -> Int {
  when t is {
    MerkleEmpty -> 0
    MerkleLeaf{..} -> 1
    MerkleNode { left, right, .. } -> size(left) + size(right)
  }
}

// shallow check if node is is equal to another node
fn node_equals(t1: MerkleTree<alg>, t2: MerkleTree<alg>) -> Bool {
  let together = (t1, t2)

  when together is {
    (MerkleLeaf { hash: hash1 }, MerkleLeaf { hash: hash2 }) -> hash1 == hash2
    (MerkleNode { hash: hash1, .. }, MerkleNode { hash: hash2, .. }) ->
      hash1 == hash2
    _ -> False
  }
}

fn get_proof(root: MerkleTree<alg>, item: ByteArray) -> Option<List<a>> {
  todo
}

fn combine_hash(
  h1: Hash<alg, ByteArray>,
  h2: Hash<alg, ByteArray>,
  hash_fn: fn(ByteArray) -> Hash<alg, String>,
) -> Hash<alg, ByteArray> {
  hash_fn(bytearray.concat(h1, h2))
}

fn do_from(
  items: List<Hash<alg, ByteArray>>,
  len: Int,
  hash_fn: fn(ByteArray) -> Hash<alg, ByteArray>,
) -> MerkleTree<alg> {
  when items is {
    // empty list
    [] -> MerkleEmpty
    // single element list
    [item] -> MerkleLeaf { hash: item }
    // all elements from the list
    all -> {
      let cutoff: Int = builtin.divide_integer(len, 2)
      let l = list.take(all, cutoff)
      let r = list.drop(all, cutoff)
      let l_node = do_from(l, cutoff, hash_fn)
      let r_node = do_from(r, len - cutoff, hash_fn)
      MerkleNode {
        hash: combine_hash(node_hash(l_node), node_hash(r_node), hash_fn),
        left: l_node,
        right: r_node,
      }
    }
  }
}

pub fn from(
  items: List<a>,
  serializer_fn: fn(a) -> ByteArray,
  hash_fn: fn(ByteArray) -> Hash<alg, ByteArray>,
) -> MerkleTree<alg> {
  let serialized_items: List<Hash<ByteArray, alg>> =
    list.map(items, serializer_fn)

  let len: Int = list.length(serialized_items)

  do_from(serialized_items, len, hash_fn)
}

test first_test() {
  let items: List<String> = ["a", "b", "c"]

  let tree_root: MerkleTree<Sha2_256> =
    from(
      items: items,
      serializer_fn: fn(item) { string.to_bytearray(item) },
      hash_fn: fn(item) { hash.sha2_256(item) },
    )

  size(tree_root) == 3
}