/**
 * RFC 6962-compatible Merkle Hash Tree (MTH) over an append-only list
 * of pre-hashed leaves.
 *
 * Recurrence (RFC 6962 §2.1):
 *
 *   MTH({})         = SHA-256()
 *   MTH({d(0)})     = leaf_hash(d(0))
 *   MTH(D[n])       = node_hash( MTH(D[0:k]), MTH(D[k:n]) )
 *   where k = largest power of 2 < n
 *
 * `MerkleLog` stores **leaf hashes** (already prefixed with 0x00) and
 * recomputes the tree on demand. Storage is O(N); audit-path / consistency
 * computation is O(log N) per leaf. This is acceptable for prekey-server
 * scale (≤ ~5M leaves over a decade for 100k addresses).
 */

import { emptyRootHash, leafHash, nodeHash } from './hashes.js';
import { constantTimeEqual } from './util.js';

/** Largest power of two strictly less than n (for n ≥ 2). */
function largestPow2LessThan(n: number): number {
  if (n < 2) throw new Error('largestPow2LessThan requires n >= 2');
  let k = 1;
  while (k < n) k <<= 1;
  return k >>> 1;
}

/**
 * Compute the Merkle Tree Hash (MTH) over a slice of pre-hashed leaves.
 * Used internally by audit-path / consistency-proof builders.
 */
export function mth(leaves: Uint8Array[], lo: number, hi: number): Uint8Array {
  const n = hi - lo;
  if (n === 0) return emptyRootHash();
  if (n === 1) return leaves[lo]!;
  const k = largestPow2LessThan(n);
  return nodeHash(mth(leaves, lo, lo + k), mth(leaves, lo + k, hi));
}

/**
 * Build the audit path for the leaf at index `m` in a tree of size `n`.
 *
 * RFC 6962 §2.1.1:
 *   PATH(m, D[n]) = PATH(m, D[0:k]) : MTH(D[k:n])         if m < k
 *   PATH(m, D[n]) = PATH(m-k, D[k:n]) : MTH(D[0:k])       if m >= k
 *
 * The returned array is ordered from leaf-sibling outward to root-sibling.
 * Each entry is the *hash of the sibling subtree*; the verifier reconstructs
 * the root using `audit_path_hash` below.
 */
export function auditPath(leaves: Uint8Array[], m: number, n: number): Uint8Array[] {
  if (n <= 0) throw new Error('auditPath requires n > 0');
  if (m < 0 || m >= n) throw new Error(`m out of range: ${m} of ${n}`);
  return auditPathInner(leaves, m, 0, n);
}

function auditPathInner(
  leaves: Uint8Array[],
  m: number,
  lo: number,
  hi: number,
): Uint8Array[] {
  const n = hi - lo;
  if (n === 1) return [];
  const k = largestPow2LessThan(n);
  if (m < k) {
    return [...auditPathInner(leaves, m, lo, lo + k), mth(leaves, lo + k, hi)];
  }
  return [...auditPathInner(leaves, m - k, lo + k, hi), mth(leaves, lo, lo + k)];
}

/**
 * Reconstruct the Merkle root from a leaf, its index, the tree size, and
 * an audit path. RFC 6962 §2.1.1.
 *
 * Returns the recomputed root; the caller compares (constant-time) against
 * the STH root to verify inclusion.
 */
export function recomputeRootFromAuditPath(
  leaf: Uint8Array,
  m: number,
  n: number,
  path: Uint8Array[],
): Uint8Array {
  if (n <= 0) throw new Error('recomputeRoot requires n > 0');
  if (m < 0 || m >= n) throw new Error(`m out of range: ${m} of ${n}`);
  return recomputeRootInner(leaf, m, 0, n, path, 0).root;
}

function recomputeRootInner(
  leaf: Uint8Array,
  m: number,
  lo: number,
  hi: number,
  path: Uint8Array[],
  pathIdx: number,
): { root: Uint8Array; pathIdx: number } {
  const n = hi - lo;
  if (n === 1) return { root: leaf, pathIdx };
  const k = largestPow2LessThan(n);
  if (m < k) {
    const left = recomputeRootInner(leaf, m, lo, lo + k, path, pathIdx);
    const sibling = path[left.pathIdx];
    if (!sibling) throw new Error('audit path too short');
    return { root: nodeHash(left.root, sibling), pathIdx: left.pathIdx + 1 };
  }
  const right = recomputeRootInner(leaf, m - k, lo + k, hi, path, pathIdx);
  const sibling = path[right.pathIdx];
  if (!sibling) throw new Error('audit path too short');
  return { root: nodeHash(sibling, right.root), pathIdx: right.pathIdx + 1 };
}

/**
 * Build a consistency proof between tree sizes m (older) and n (newer).
 * RFC 6962 §2.1.2.
 */
export function consistencyProof(
  leaves: Uint8Array[],
  m: number,
  n: number,
): Uint8Array[] {
  if (m < 0 || n < m) throw new Error(`invalid m,n: ${m},${n}`);
  if (m === 0 || m === n) return [];
  return subProof(leaves, m, 0, n, true);
}

function subProof(
  leaves: Uint8Array[],
  m: number,
  lo: number,
  hi: number,
  isOriginalRoot: boolean,
): Uint8Array[] {
  const n = hi - lo;
  if (m === n) {
    return isOriginalRoot ? [] : [mth(leaves, lo, hi)];
  }
  const k = largestPow2LessThan(n);
  if (m <= k) {
    return [...subProof(leaves, m, lo, lo + k, isOriginalRoot), mth(leaves, lo + k, hi)];
  }
  return [...subProof(leaves, m - k, lo + k, hi, false), mth(leaves, lo, lo + k)];
}

/**
 * Verify a consistency proof. Given:
 *  - oldRoot = MTH(D[0:m])
 *  - newRoot = MTH(D[0:n])  with n >= m
 *  - proof   = consistencyProof(leaves, m, n)
 *
 * Returns true if the proof is valid (i.e. D[0:n] really is an extension
 * of D[0:m]). RFC 6962 §2.1.2.
 */
export function verifyConsistencyProof(
  m: number,
  n: number,
  oldRoot: Uint8Array,
  newRoot: Uint8Array,
  proof: Uint8Array[],
): boolean {
  if (m < 0 || n < m) return false;
  if (m === 0) return true; // any newRoot is consistent with empty old tree
  if (m === n) return proof.length === 0 && constantTimeEqual(oldRoot, newRoot);

  // RFC 6962 verification recurrence
  let path = proof;
  if (isPowerOfTwo(m)) {
    path = [oldRoot, ...path];
  }

  let fn = m - 1;
  let sn = n - 1;
  while ((fn & 1) === 1) {
    fn >>= 1;
    sn >>= 1;
  }

  if (path.length === 0) return false;
  let fr = path[0]!;
  let sr = path[0]!;
  let i = 1;

  while (sn > 0) {
    if ((fn & 1) === 1 || fn === sn) {
      if (i >= path.length) return false;
      const c = path[i++]!;
      fr = nodeHash(c, fr);
      sr = nodeHash(c, sr);
      while ((fn & 1) === 0 && fn !== 0) {
        fn >>= 1;
        sn >>= 1;
      }
    } else {
      if (i >= path.length) return false;
      sr = nodeHash(sr, path[i++]!);
    }
    fn >>= 1;
    sn >>= 1;
  }

  if (i !== path.length) return false;
  return constantTimeEqual(fr, oldRoot) && constantTimeEqual(sr, newRoot);
}

function isPowerOfTwo(n: number): boolean {
  return n > 0 && (n & (n - 1)) === 0;
}

/**
 * Append-only Merkle log used server-side. Holds leaf hashes in memory
 * and recomputes paths on demand.
 *
 * For production deployments the caller wraps this with a persistent
 * `KTLogStore` (see `store.ts`) — this class is the algorithmic core.
 */
export class MerkleLog {
  private readonly leaves: Uint8Array[] = [];

  /** Number of leaves currently in the tree. */
  get size(): number {
    return this.leaves.length;
  }

  /** Append a *raw leaf data* (will be domain-separated and hashed). */
  appendData(data: Uint8Array): { index: number; leafHash: Uint8Array } {
    const lh = leafHash(data);
    const index = this.leaves.length;
    this.leaves.push(lh);
    return { index, leafHash: lh };
  }

  /** Append a leaf that has *already been hashed* (rebuild path). */
  appendLeafHash(lh: Uint8Array): number {
    const index = this.leaves.length;
    this.leaves.push(lh);
    return index;
  }

  /** Current root hash (MTH of all leaves). */
  rootHash(): Uint8Array {
    return mth(this.leaves, 0, this.leaves.length);
  }

  /** Audit path for leaf at `index`. */
  auditPath(index: number): Uint8Array[] {
    return auditPath(this.leaves, index, this.leaves.length);
  }

  /** Consistency proof from `oldSize` to current size. */
  consistencyProof(oldSize: number): Uint8Array[] {
    return consistencyProof(this.leaves, oldSize, this.leaves.length);
  }

  /** Snapshot the leaf hash at `index` (read-only). */
  leafHashAt(index: number): Uint8Array {
    const lh = this.leaves[index];
    if (!lh) throw new Error(`no leaf at index ${index}`);
    return lh;
  }

  /** Defensive copy of all leaves (used by persistent stores on hot-load). */
  exportLeaves(): Uint8Array[] {
    return this.leaves.map((l) => new Uint8Array(l));
  }

  /** Hot-load from persisted leaf hashes. */
  static fromLeaves(leaves: Uint8Array[]): MerkleLog {
    const log = new MerkleLog();
    for (const l of leaves) log.appendLeafHash(l);
    return log;
  }
}
