MMR.ts

import type { MMRNode, MMRState, Operation } from './types';

export async function sha256(data: Uint8Array): Promise<Uint8Array> {
  // SubtleCrypto wants BufferSource. In TS ≥5.7, Uint8Array is generic over
  // ArrayBufferLike (incl. SharedArrayBuffer), so we narrow with a cast.
  const buf = await crypto.subtle.digest('SHA-256', data as unknown as BufferSource);
  return new Uint8Array(buf);
}

export async function hashLeaf(
  location: number,
  key: string,
  value: string,
): Promise<Uint8Array> {
  const loc = new Uint8Array(8);
  new DataView(loc.buffer).setBigUint64(0, BigInt(location), true);
  const k = new TextEncoder().encode(key);
  const v = new TextEncoder().encode(value);
  const buf = new Uint8Array(8 + k.length + v.length);
  buf.set(loc, 0);
  buf.set(k, 8);
  buf.set(v, 8 + k.length);
  return sha256(buf);
}

export async function hashParent(
  left: Uint8Array,
  right: Uint8Array,
): Promise<Uint8Array> {
  const buf = new Uint8Array(64);
  buf.set(left, 0);
  buf.set(right, 32);
  return sha256(buf);
}

export async function bagPeaks(
  totalNodes: number,
  peakHashes: Uint8Array[],
): Promise<Uint8Array> {
  const total = new Uint8Array(8);
  new DataView(total.buffer).setBigUint64(0, BigInt(totalNodes), true);
  const buf = new Uint8Array(8 + peakHashes.length * 32);
  buf.set(total, 0);
  peakHashes.forEach((h, i) => buf.set(h, 8 + i * 32));
  return sha256(buf);
}

export function emptyState(): MMRState {
  return {
    nodes: [],
    peaks: [],
    totalNodes: 0,
    root: new Uint8Array(32),
    operationLog: [],
    keyIndex: new Map(),
  };
}

/**
 * Append a new operation, advancing the MMR. Returns a brand-new MMRState.
 * Implements the canonical MMR append:
 * 1. Add a leaf node at the next position.
 * 2. While the last two peaks have equal height, merge them into a parent
 *    (which itself becomes a new node at the next position).
 * 3. Recompute the root by bagging all peak hashes with the total node count.
 */
export async function appendOperation(
  state: MMRState,
  op: { type: 'assign' | 'delete'; key: string; value: string },
): Promise<MMRState> {
  const nodes: MMRNode[] = state.nodes.slice();
  const peaks = state.peaks.slice();
  const operationLog = state.operationLog.slice();
  const keyIndex = new Map(state.keyIndex);

  const location = operationLog.length;
  const leafPosition = nodes.length;
  const leafHash = await hashLeaf(location, op.key, op.value);

  const leafNode: MMRNode = {
    position: leafPosition,
    hash: leafHash,
    height: 0,
  };
  nodes.push(leafNode);
  peaks.push(leafPosition);

  // Merge equal-height peaks into parents.
  while (peaks.length >= 2) {
    const right = nodes[peaks[peaks.length - 1]];
    const left = nodes[peaks[peaks.length - 2]];
    if (left.height !== right.height) break;
    const parentHash = await hashParent(left.hash, right.hash);
    const parentPosition = nodes.length;
    const parent: MMRNode = {
      position: parentPosition,
      hash: parentHash,
      height: left.height + 1,
      leftChild: left.position,
      rightChild: right.position,
    };
    nodes.push(parent);
    peaks.pop();
    peaks.pop();
    peaks.push(parentPosition);
  }

  const operation: Operation = {
    location,
    type: op.type,
    key: op.key,
    value: op.value,
    nodePosition: leafPosition,
  };
  operationLog.push(operation);
  keyIndex.set(op.key, location);

  const totalNodes = nodes.length;
  const peakHashes = peaks.map((p) => nodes[p].hash);
  const root =
    peakHashes.length === 0 ? new Uint8Array(32) : await bagPeaks(totalNodes, peakHashes);

  return {
    nodes,
    peaks,
    totalNodes,
    root,
    operationLog,
    keyIndex,
  };
}

/**
 * Build a parent-pointer map from existing internal nodes.
 * Used by the proof generator to walk leaf → peak.
 */
export function parentMap(nodes: MMRNode[]): Map<number, number> {
  const m = new Map<number, number>();
  for (const n of nodes) {
    if (n.leftChild !== undefined) m.set(n.leftChild, n.position);
    if (n.rightChild !== undefined) m.set(n.rightChild, n.position);
  }
  return m;
}

// ----- Snapshot serialization (Map → entries[] → structuredClone → Map) -----

interface SerializedState {
  nodes: MMRNode[];
  peaks: number[];
  totalNodes: number;
  root: Uint8Array;
  operationLog: Operation[];
  keyIndexEntries: Array<[string, number]>;
}

export function cloneState(s: MMRState): MMRState {
  const serialized: SerializedState = {
    nodes: s.nodes.map((n) => ({
      position: n.position,
      hash: new Uint8Array(n.hash),
      height: n.height,
      leftChild: n.leftChild,
      rightChild: n.rightChild,
    })),
    peaks: s.peaks.slice(),
    totalNodes: s.totalNodes,
    root: new Uint8Array(s.root),
    operationLog: s.operationLog.map((o) => ({ ...o })),
    keyIndexEntries: Array.from(s.keyIndex.entries()),
  };
  const cloned: SerializedState = structuredClone(serialized);
  return {
    nodes: cloned.nodes,
    peaks: cloned.peaks,
    totalNodes: cloned.totalNodes,
    root: cloned.root,
    operationLog: cloned.operationLog,
    keyIndex: new Map(cloned.keyIndexEntries),
  };
}

// ----- Hex helpers -----

const HEX = '0123456789abcdef';
export function bytesToHex(b: Uint8Array): string {
  let out = '';
  for (let i = 0; i < b.length; i++) {
    out += HEX[b[i] >> 4] + HEX[b[i] & 0x0f];
  }
  return out;
}

export function hexToBytes(hex: string): Uint8Array | null {
  const clean = hex.toLowerCase().replace(/^0x/, '').trim();
  if (clean.length !== 64) return null;
  if (!/^[0-9a-f]+$/.test(clean)) return null;
  const out = new Uint8Array(32);
  for (let i = 0; i < 32; i++) {
    out[i] = parseInt(clean.substr(i * 2, 2), 16);
  }
  return out;
}

export function shortHex(b: Uint8Array): string {
  const h = bytesToHex(b);
  return `0x${h.slice(0, 8)}…${h.slice(-8)}`;
}

export function fullHex(b: Uint8Array): string {
  return `0x${bytesToHex(b)}`;
}

export function bytesEqual(a: Uint8Array, b: Uint8Array): boolean {
  if (a.length !== b.length) return false;
  for (let i = 0; i < a.length; i++) if (a[i] !== b[i]) return false;
  return true;
}