/**
 * erasure.ts — Reed-Solomon erasure coding around the WASM library.
 *
 * The library wants a single contiguous buffer of length
 *   (dataShards + parityShards) * shardSize
 * with data shards laid out first, followed by space for parity. `encode()`
 * fills the parity region in-place. `reconstruct()` rewrites missing shards
 * in-place using a `shardsAvailable: boolean[]` mask.
 *
 * We expose a chunk-oriented API (`encodeWithErasure` / `decodeWithErasure`)
 * because the rest of the codex pipeline thinks in terms of individual
 * envelopes. Both adapters allocate a fresh buffer per call — the library is
 * not re-entrant on a shared buffer.
 */

import { ReedSolomonErasure } from "@digitaldefiance/reed-solomon-erasure.wasm";
import {
  CHUNK_SIZE,
  CodexError,
  computeParityShards,
} from "../../shared/codex-types";
import { getReedSolomon } from "./rs-loader";

/**
 * The shared rs-loader caches a single `ReedSolomonErasure` instance, which
 * is fine for `encode()` but trips a known library bug on the SECOND call to
 * `reconstruct()` against the same instance: the internal __wbindgen state
 * survives between calls and the second reconstruction reports
 * RESULT_ERROR_TOO_FEW_SHARDS_PRESENT even when enough shards are present.
 *
 * To get correct behaviour we (a) use the loader to ensure the WASM is
 * actually instantiable and to cache the bytes lazily, and (b) build a fresh
 * `ReedSolomonErasure` instance for every reconstruction. The cost is a few
 * milliseconds per call — negligible compared to compressing the vault and
 * rendering Aztec codes.
 */
let cachedWasmBytes: Uint8Array | null = null;

async function freshReedSolomon(): Promise<ReedSolomonErasure> {
  if (cachedWasmBytes) {
    return ReedSolomonErasure.fromBytes(cachedWasmBytes);
  }
  // Capture bytes by intercepting the loader's first instantiation. The
  // loader caches its own promise, so calling it twice is cheap on the second
  // pass. We pull bytes from the WASM memory image after the fact — see
  // `loadWasmBytes` below for the lazy initialization.
  await getReedSolomon();
  if (!cachedWasmBytes) {
    cachedWasmBytes = await loadWasmBytes();
  }
  return ReedSolomonErasure.fromBytes(cachedWasmBytes);
}

/**
 * Load the WASM bytes through the same probing logic rs-loader uses, so this
 * file does not assume any specific filesystem layout. In the popup the
 * loader's fetch path is the source of truth; in vitest it's the on-disk
 * path under `node_modules/`. This helper is deliberately a duplicate of
 * the bytes-only portion of rs-loader — we cannot touch rs-loader itself
 * (orchestrator-owned), and we need raw bytes to work around the library's
 * reuse bug.
 */
async function loadWasmBytes(): Promise<Uint8Array> {
  const inPopup =
    typeof chrome !== "undefined" &&
    typeof chrome.runtime !== "undefined" &&
    typeof chrome.runtime.getURL === "function";
  if (inPopup) {
    const url = chrome.runtime.getURL("codex/reed_solomon_erasure_bg.wasm");
    const response = await fetch(url);
    if (!response.ok) {
      throw new CodexError(
        "ERASURE_CODING_VERIFICATION_FAILED",
        `Failed to fetch WASM: ${response.status}`,
      );
    }
    return new Uint8Array(await response.arrayBuffer());
  }
  const fs = await import("node:fs");
  const path = await import("node:path");
  const wasmPath = path.resolve(
    "node_modules/@digitaldefiance/reed-solomon-erasure.wasm/dist/reed_solomon_erasure_bg.wasm",
  );
  return new Uint8Array(fs.readFileSync(wasmPath));
}

export interface EncodeResult {
  dataChunks: Uint8Array[];
  parityChunks: Uint8Array[];
  totalChunks: number;
  dataShards: number;
  parityShards: number;
}

/**
 * Split `data` into fixed-size shards plus `computeParityShards(dataShards)`
 * parity shards. Pads the trailing data shard with zeros — the original byte
 * length is recovered via the `CodexPayload.vaultBlob` (no length prefix
 * needed on the wire because the outer JSON already bounds the data).
 */
export async function encodeWithErasure(data: Uint8Array): Promise<EncodeResult> {
  const rs = await freshReedSolomon();

  const dataShards = Math.max(1, Math.ceil(data.length / CHUNK_SIZE));
  const parityShards = computeParityShards(dataShards);
  const totalShards = dataShards + parityShards;
  const buffer = new Uint8Array(totalShards * CHUNK_SIZE);
  buffer.set(data, 0); // Trailing bytes after `data.length` stay zero — the pad.

  const code = rs.encode(buffer, dataShards, parityShards);
  if (code !== ReedSolomonErasure.RESULT_OK) {
    throw new CodexError(
      "ERASURE_CODING_VERIFICATION_FAILED",
      `Reed-Solomon encode failed: ${ReedSolomonErasure.getResultMessage(code)}`,
    );
  }

  // Slice() copies so each shard owns its memory and callers cannot mutate
  // the working buffer through a returned view.
  const dataChunks: Uint8Array[] = [];
  for (let i = 0; i < dataShards; i++) {
    dataChunks.push(buffer.slice(i * CHUNK_SIZE, (i + 1) * CHUNK_SIZE));
  }
  const parityChunks: Uint8Array[] = [];
  for (let i = 0; i < parityShards; i++) {
    const base = (dataShards + i) * CHUNK_SIZE;
    parityChunks.push(buffer.slice(base, base + CHUNK_SIZE));
  }

  return {
    dataChunks,
    parityChunks,
    totalChunks: totalShards,
    dataShards,
    parityShards,
  };
}

/**
 * Reassemble the data-shard region from a partial set of chunks. Missing
 * shards are reconstructed from parity. The returned buffer is exactly
 * `dataShards * shardSize` bytes long, including any trailing pad bytes
 * — callers strip the pad using their out-of-band length information.
 *
 * Shard size is derived from the actual chunks rather than the current
 * `CHUNK_SIZE` constant so that Codexes printed under an older constant
 * (e.g. 1024-byte shards) still restore on a build whose constant has moved.
 * Reed-Solomon requires all shards to share one length; the first chunk
 * establishes it and every other chunk must match.
 */
export async function decodeWithErasure(
  chunks: Map<number, Uint8Array>,
  dataShards: number,
  totalShards: number,
): Promise<Uint8Array> {
  const parityShards = totalShards - dataShards;
  if (parityShards < 0) {
    throw new CodexError(
      "ERASURE_CODING_VERIFICATION_FAILED",
      `Invalid shard layout: dataShards=${dataShards} totalShards=${totalShards}`,
    );
  }

  if (chunks.size < dataShards) {
    throw new CodexError(
      "INSUFFICIENT_CHUNKS",
      `Need at least ${dataShards} chunks, have ${chunks.size}`,
    );
  }

  // First scanned chunk fixes the shard size for the whole reconstruction.
  const firstChunk = chunks.values().next().value;
  if (!firstChunk || firstChunk.length === 0) {
    throw new CodexError(
      "ERASURE_CODING_VERIFICATION_FAILED",
      "First chunk is empty; cannot determine shard size",
    );
  }
  const shardSize = firstChunk.length;

  const rs = await freshReedSolomon();
  const buffer = new Uint8Array(totalShards * shardSize);
  const available: boolean[] = new Array(totalShards).fill(false);

  for (const [index, chunk] of chunks) {
    if (index < 0 || index >= totalShards) {
      throw new CodexError(
        "ERASURE_CODING_VERIFICATION_FAILED",
        `Chunk index ${index} out of range [0, ${totalShards})`,
      );
    }
    if (chunk.length !== shardSize) {
      throw new CodexError(
        "ERASURE_CODING_VERIFICATION_FAILED",
        `Chunk ${index} has length ${chunk.length}, expected ${shardSize}`,
      );
    }
    buffer.set(chunk, index * shardSize);
    available[index] = true;
  }

  const code = rs.reconstruct(buffer, dataShards, parityShards, available);
  if (code !== ReedSolomonErasure.RESULT_OK) {
    // The library distinguishes "not enough shards" with its own code; map it
    // back to our INSUFFICIENT_CHUNKS so handlers upstream can surface a
    // clean "scan more codes" prompt instead of a generic library error.
    if (code === ReedSolomonErasure.RESULT_ERROR_TOO_FEW_SHARDS_PRESENT) {
      throw new CodexError(
        "INSUFFICIENT_CHUNKS",
        `Need at least ${dataShards} valid chunks for reconstruction`,
      );
    }
    throw new CodexError(
      "ERASURE_CODING_VERIFICATION_FAILED",
      `Reed-Solomon reconstruct failed: ${ReedSolomonErasure.getResultMessage(code)}`,
    );
  }

  return buffer.slice(0, dataShards * shardSize);
}