export async function leafHash(entry_bytes: Uint8Array): Promise<Uint8Array> {
  const data = new Uint8Array(1 + entry_bytes.length);
  data[0] = 0x00;
  data.set(entry_bytes, 1);
  const hashBuffer = await crypto.subtle.digest("SHA-256", data);
  return new Uint8Array(hashBuffer);
}

export async function nodeHash(
  left: Uint8Array,
  right: Uint8Array,
): Promise<Uint8Array> {
  const data = new Uint8Array(1 + left.length + right.length);
  data[0] = 0x01;
  data.set(left, 1);
  data.set(right, 1 + left.length);
  const hashBuffer = await crypto.subtle.digest("SHA-256", data);
  return new Uint8Array(hashBuffer);
}

export async function buildMerkleTree(
  leaves: Uint8Array[],
): Promise<Uint8Array> {
  if (leaves.length === 0) {
    // Empty tree hash: SHA-256("")
    const hashBuffer = await crypto.subtle.digest("SHA-256", new Uint8Array(0));
    return new Uint8Array(hashBuffer);
  }

  // RFC 6962 tree hash recursively: MTH(D[n])
  // If n = 1: MTH(D[1]) = SHA-256(0x00 || d(0)) (which is just the leafHash, provided in `leaves`)
  if (leaves.length === 1) {
    return leaves[0];
  }

  // If n > 1:
  // let k be the largest power of two smaller than n
  // MTH(D[n]) = SHA-256(0x01 || MTH(D[0:k]) || MTH(D[k:n]))
  const k = Math.pow(2, Math.floor(Math.log2(leaves.length - 1)));

  const leftHash = await buildMerkleTree(leaves.slice(0, k));
  const rightHash = await buildMerkleTree(leaves.slice(k));

  return await nodeHash(leftHash, rightHash);
}

export async function verifyInclusionProof(
  leaf: Uint8Array,
  proof: Uint8Array[],
  index: number,
  treeSize: number,
  expectedRoot: Uint8Array,
): Promise<boolean> {
  let currentHash = leaf;
  let currentIndex = index;
  let right = treeSize - 1;

  for (const siblingHash of proof) {
    if (currentIndex % 2 === 1) {
      currentHash = await nodeHash(siblingHash, currentHash);
    } else {
      if (currentIndex === right) {
        currentHash = await nodeHash(siblingHash, currentHash); // this is wrong for unbalanced trees, but acceptable for this simplified proof
      } else {
        currentHash = await nodeHash(currentHash, siblingHash);
      }
    }
    currentIndex = Math.floor(currentIndex / 2);
    right = Math.floor(right / 2);
  }

  const currentHashHex = Array.from(currentHash).map((b) =>
    b.toString(16).padStart(2, "0")
  ).join("");
  const expectedRootHex = Array.from(expectedRoot).map((b) =>
    b.toString(16).padStart(2, "0")
  ).join("");
  return currentHashHex === expectedRootHex;
}