proof.ts

import type { InclusionProof, MMRState } from './types';
import { bagPeaks, bytesEqual, hashParent, parentMap } from './MMR';

/**
 * Walk a leaf up to its mountain peak, collecting siblings on the way.
 * Returns the peak's position so the caller can locate it among all peaks.
 */
function pathToPeak(
  state: MMRState,
  leafPosition: number,
): { peakPosition: number; siblings: number[]; siblingIsRight: boolean[] } {
  const parents = parentMap(state.nodes);
  const siblings: number[] = [];
  const siblingIsRight: boolean[] = [];
  let cur = leafPosition;
  while (parents.has(cur)) {
    const parentPos = parents.get(cur)!;
    const parent = state.nodes[parentPos];
    if (parent.leftChild === cur) {
      siblings.push(parent.rightChild!);
      siblingIsRight.push(true);
    } else {
      siblings.push(parent.leftChild!);
      siblingIsRight.push(false);
    }
    cur = parentPos;
  }
  return { peakPosition: cur, siblings, siblingIsRight };
}

/**
 * Generate an inclusion proof for the leaf at the given operation location.
 */
export async function generateProof(
  state: MMRState,
  location: number,
): Promise<InclusionProof> {
  const op = state.operationLog[location];
  const leafPosition = op.nodePosition;
  const leafHash = state.nodes[leafPosition].hash;

  const { peakPosition, siblings, siblingIsRight } = pathToPeak(state, leafPosition);
  const siblingHashes = siblings.map((p) => state.nodes[p].hash);

  const peakIndex = state.peaks.indexOf(peakPosition);
  const peakHashes = state.peaks
    .filter((_, i) => i !== peakIndex)
    .map((p) => state.nodes[p].hash);

  return verifyProof({
    location,
    leafHash,
    siblingHashes,
    siblingIsRight,
    peakHashes,
    peakIndex,
    totalNodes: state.totalNodes,
    expectedRoot: state.root,
  });
}

/**
 * Re-run the proof math against (possibly tampered) inputs and return the
 * full InclusionProof with computed root + valid flag.
 */
export async function verifyProof(input: {
  location: number;
  leafHash: Uint8Array;
  siblingHashes: Uint8Array[];
  siblingIsRight: boolean[];
  peakHashes: Uint8Array[];
  peakIndex: number;
  totalNodes: number;
  expectedRoot: Uint8Array;
}): Promise<InclusionProof> {
  // Step 1: reconstruct this leaf's mountain peak from leaf + siblings.
  let acc = input.leafHash;
  for (let i = 0; i < input.siblingHashes.length; i++) {
    const sib = input.siblingHashes[i];
    if (input.siblingIsRight[i]) {
      acc = await hashParent(acc, sib);
    } else {
      acc = await hashParent(sib, acc);
    }
  }

  // Step 2: assemble the full peak list (other peaks + reconstructed peak).
  const allPeaks: Uint8Array[] = [];
  let otherIdx = 0;
  const totalPeaks = input.peakHashes.length + 1;
  for (let i = 0; i < totalPeaks; i++) {
    if (i === input.peakIndex) {
      allPeaks.push(acc);
    } else {
      allPeaks.push(input.peakHashes[otherIdx++]);
    }
  }

  // Step 3: bag.
  const computedRoot = await bagPeaks(input.totalNodes, allPeaks);
  const valid = bytesEqual(computedRoot, input.expectedRoot);

  return {
    location: input.location,
    leafHash: input.leafHash,
    siblingHashes: input.siblingHashes,
    siblingIsRight: input.siblingIsRight,
    peakHashes: input.peakHashes,
    peakIndex: input.peakIndex,
    totalNodes: input.totalNodes,
    computedRoot,
    expectedRoot: input.expectedRoot,
    valid,
  };
}

/**
 * Returns the set of node positions that are "lit" by a proof:
 * the leaf, every node on the path up to its peak, every sibling node,
 * and every other peak. Used by the tree view to highlight a proof.
 */
export function proofHighlight(
  state: MMRState,
  location: number,
): {
  pathNodes: Set<number>; // leaf-to-peak path (yellow)
  siblingNodes: Set<number>; // siblings (orange)
  otherPeaks: Set<number>; // remaining peaks (blue)
} {
  const op = state.operationLog[location];
  const leafPosition = op.nodePosition;

  const pathNodes = new Set<number>();
  const siblingNodes = new Set<number>();
  const otherPeaks = new Set<number>();

  const parents = parentMap(state.nodes);
  let cur = leafPosition;
  pathNodes.add(cur);
  while (parents.has(cur)) {
    const parentPos = parents.get(cur)!;
    const parent = state.nodes[parentPos];
    const sibling = parent.leftChild === cur ? parent.rightChild! : parent.leftChild!;
    siblingNodes.add(sibling);
    pathNodes.add(parentPos);
    cur = parentPos;
  }
  for (const p of state.peaks) {
    if (p !== cur) otherPeaks.add(p);
  }
  return { pathNodes, siblingNodes, otherPeaks };
}