import { canonicalize, sha256Bytes, utf8 } from "./crypto";
import type { JsonObject } from "./browser-types";

function digestBytes(digest: string): Uint8Array {
  if (!/^sha256:[0-9a-f]{64}$/u.test(digest)) throw new TypeError("invalid SHA-256 digest");
  const hexadecimal = digest.slice(7);
  return Uint8Array.from(
    { length: hexadecimal.length / 2 },
    (_, index) => Number.parseInt(hexadecimal.slice(index * 2, index * 2 + 2), 16),
  );
}

function largestPowerOfTwoBelow(value: number): number {
  let power = 1;
  while (power * 2 < value) power *= 2;
  return power;
}

export async function rfc6962NodeHash(left: string, right: string): Promise<string> {
  const leftBytes = digestBytes(left);
  const rightBytes = digestBytes(right);
  const input = new Uint8Array(1 + leftBytes.length + rightBytes.length);
  input[0] = 1;
  input.set(leftBytes, 1);
  input.set(rightBytes, 1 + leftBytes.length);
  return sha256Bytes(input);
}

export async function witnessLeaf(deed: {
  issuer: string;
  content_hash: string;
  attestation_hash: string;
}): Promise<string> {
  const canonicalDeed = canonicalize(deed as JsonObject);
  const bytes = utf8(canonicalDeed);
  const input = new Uint8Array(bytes.length + 1);
  input[0] = 0;
  input.set(bytes, 1);
  return sha256Bytes(input);
}

export async function merkleRoot(leaves: readonly string[]): Promise<string> {
  if (leaves.length === 0) return sha256Bytes(new Uint8Array());
  if (leaves.length === 1) return leaves[0] as string;
  const split = largestPowerOfTwoBelow(leaves.length);
  return rfc6962NodeHash(
    await merkleRoot(leaves.slice(0, split)),
    await merkleRoot(leaves.slice(split)),
  );
}

export async function inclusionProof(
  leaves: readonly string[],
  leafIndex: number,
): Promise<string[]> {
  if (leafIndex < 0 || leafIndex >= leaves.length) throw new RangeError("leaf index is out of range");
  if (leaves.length === 1) return [];
  const split = largestPowerOfTwoBelow(leaves.length);
  if (leafIndex < split) {
    return [
      ...(await inclusionProof(leaves.slice(0, split), leafIndex)),
      await merkleRoot(leaves.slice(split)),
    ];
  }
  return [
    ...(await inclusionProof(leaves.slice(split), leafIndex - split)),
    await merkleRoot(leaves.slice(0, split)),
  ];
}

export async function verifyInclusionProof(
  leaf: string,
  leafIndex: number,
  treeSize: number,
  proof: readonly string[],
  expectedRoot: string,
): Promise<boolean> {
  if (leafIndex < 0 || leafIndex >= treeSize) return false;
  let fn = leafIndex;
  let sn = treeSize - 1;
  let root = leaf;
  for (const node of proof) {
    if (sn === 0) return false;
    if ((fn & 1) === 1 || fn === sn) {
      root = await rfc6962NodeHash(node, root);
      if ((fn & 1) === 0) {
        while ((fn & 1) === 0 && fn !== 0) {
          fn >>= 1;
          sn >>= 1;
        }
      }
    } else {
      root = await rfc6962NodeHash(root, node);
    }
    fn >>= 1;
    sn >>= 1;
  }
  return sn === 0 && root === expectedRoot;
}

async function subproof(
  oldSize: number,
  leaves: readonly string[],
  complete: boolean,
): Promise<string[]> {
  if (oldSize === leaves.length) return complete ? [] : [await merkleRoot(leaves)];
  const split = largestPowerOfTwoBelow(leaves.length);
  if (oldSize <= split) {
    return [
      ...(await subproof(oldSize, leaves.slice(0, split), complete)),
      await merkleRoot(leaves.slice(split)),
    ];
  }
  return [
    ...(await subproof(oldSize - split, leaves.slice(split), false)),
    await merkleRoot(leaves.slice(0, split)),
  ];
}

export async function consistencyProof(
  leaves: readonly string[],
  oldSize: number,
): Promise<string[]> {
  if (oldSize < 0 || oldSize > leaves.length) throw new RangeError("invalid consistency range");
  if (oldSize === 0 || oldSize === leaves.length) return [];
  return subproof(oldSize, leaves, true);
}

export async function verifyConsistencyProof(
  oldSize: number,
  newSize: number,
  proof: readonly string[],
  oldRoot: string,
  newRoot: string,
): Promise<boolean> {
  if (oldSize > newSize) return false;
  if (oldSize === newSize) return proof.length === 0 && oldRoot === newRoot;
  if (oldSize === 0) return proof.length === 0;
  if (proof.length === 0) return false;
  let node = oldSize - 1;
  let last = newSize - 1;
  while ((node & 1) === 1) {
    node >>= 1;
    last >>= 1;
  }
  let offset = 0;
  let oldHash: string;
  let newHash: string;
  if (node > 0) {
    oldHash = proof[0] as string;
    newHash = oldHash;
    offset = 1;
  } else {
    oldHash = oldRoot;
    newHash = oldRoot;
  }
  for (; offset < proof.length; offset += 1) {
    const current = proof[offset] as string;
    if (last === 0) return false;
    if ((node & 1) === 1 || node === last) {
      oldHash = await rfc6962NodeHash(current, oldHash);
      newHash = await rfc6962NodeHash(current, newHash);
      while ((node & 1) === 0 && node !== 0) {
        node >>= 1;
        last >>= 1;
      }
    } else {
      newHash = await rfc6962NodeHash(newHash, current);
    }
    node >>= 1;
    last >>= 1;
  }
  return last === 0 && oldHash === oldRoot && newHash === newRoot;
}
