/**
 * 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 {
  CHUNK_SIZE,
  CodexError,
  MAX_DATA_SHARDS_PER_STRIPE,
  computeParityShards,
} from "../../shared/codex-types";
import { ReedSolomonErasure } from "./reed-solomon";
import { getReedSolomon } from "./rs-loader";

/**
 * The shared rs-loader caches a single `ReedSolomonErasure` instance. Some
 * builds of the underlying WASM module accumulate __wbindgen state between
 * `reconstruct()` calls and return RESULT_ERROR_TOO_FEW_SHARDS_PRESENT on the
 * second invocation even when enough shards are present. Building a fresh
 * instance per call avoids the failure mode entirely; the cost is a few
 * milliseconds per call — negligible compared to compressing the vault and
 * rendering Aztec codes. We keep the loader's cached instance only as a
 * liveness probe (it confirms the WASM is instantiable on this platform).
 */
let cachedWasmBytes: Uint8Array<ArrayBuffer> | 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 raw WASM bytes for re-instantiation. The popup path resolves
 * against the extension package via `chrome.runtime.getURL`; vitest reads
 * directly from `node_modules`. `freshReedSolomon` needs the raw bytes (not
 * the loader's instance) because each call rebuilds the WASM module to avoid
 * cross-call state accumulation in `reconstruct`.
 */
async function loadWasmBytes(): Promise<Uint8Array<ArrayBuffer>> {
  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());
  }
  return readTestWasmBytes();
}

async function readTestWasmBytes(): Promise<Uint8Array<ArrayBuffer>> {
  if (import.meta.env.MODE !== "test") {
    throw new CodexError(
      "ERASURE_CODING_VERIFICATION_FAILED",
      "Reed-Solomon WASM can only be loaded through chrome.runtime.getURL outside tests",
    );
  }
  const { readNodeReedSolomonWasmBytes } = await import("./node-wasm");
  return readNodeReedSolomonWasmBytes();
}

/**
 * We map Reed-Solomon result codes to human-readable strings here so the
 * surrounding CodexError messages stay specific instead of leaking a bare
 * number.
 */
function resultMessage(code: number): string {
  switch (code) {
    case ReedSolomonErasure.RESULT_OK: return "ok";
    case ReedSolomonErasure.RESULT_ERROR_TOO_FEW_SHARDS: return "too few shards";
    case ReedSolomonErasure.RESULT_ERROR_TOO_MANY_SHARDS: return "too many shards";
    case ReedSolomonErasure.RESULT_ERROR_TOO_FEW_DATA_SHARDS: return "too few data shards";
    case ReedSolomonErasure.RESULT_ERROR_TOO_MANY_DATA_SHARDS: return "too many data shards";
    case ReedSolomonErasure.RESULT_ERROR_TOO_FEW_PARITY_SHARDS: return "too few parity shards";
    case ReedSolomonErasure.RESULT_ERROR_TOO_MANY_PARITY_SHARDS: return "too many parity shards";
    case ReedSolomonErasure.RESULT_ERROR_TOO_FEW_BUFFER_SHARDS: return "too few buffer shards";
    case ReedSolomonErasure.RESULT_ERROR_TOO_MANY_BUFFER_SHARDS: return "too many buffer shards";
    case ReedSolomonErasure.RESULT_ERROR_INCORRECT_SHARD_SIZE: return "incorrect shard size";
    case ReedSolomonErasure.RESULT_ERROR_TOO_FEW_SHARDS_PRESENT: return "too few shards present";
    case ReedSolomonErasure.RESULT_ERROR_EMPTY_SHARD: return "empty shard";
    case ReedSolomonErasure.RESULT_ERROR_INVALID_SHARD_FLAGS: return "invalid shard flags";
    case ReedSolomonErasure.RESULT_ERROR_INVALID_INDEX: return "invalid index";
    default: return `unknown result code ${code}`;
  }
}

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

export interface MultiStripeEncodeResult {
  stripes: EncodeResult[];
  /** Sum of `dataShards` across stripes. */
  totalDataChunks: number;
  /** Sum of `parityShards` across stripes. */
  totalParityChunks: number;
  /** Sum of `totalChunks` across stripes = number of Aztec codes to print. */
  totalChunks: number;
}

/**
 * Plan how many data shards belong to each stripe, balanced as evenly as
 * possible so no single stripe approaches the GF(2⁸) cap. Returns an array of
 * shard counts that sum to `totalDataShards` and where every element is
 * ≤ MAX_DATA_SHARDS_PER_STRIPE.
 */
export function planStripeShardCounts(totalDataShards: number): number[] {
  if (totalDataShards <= 0) return [];
  const numStripes = Math.max(
    1,
    Math.ceil(totalDataShards / MAX_DATA_SHARDS_PER_STRIPE),
  );
  const base = Math.floor(totalDataShards / numStripes);
  const remainder = totalDataShards % numStripes;
  const counts: number[] = [];
  for (let i = 0; i < numStripes; i++) {
    counts.push(base + (i < remainder ? 1 : 0));
  }
  return counts;
}

/**
 * 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).
 *
 * Single-stripe primitive. The GF(2⁸) field caps `dataShards + parityShards`
 * at 256; anything larger panics inside the WASM module, so call
 * `encodeIntoStripes` instead for payloads that need more total shards.
 */
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;
  if (totalShards > 256) {
    throw new CodexError(
      "ERASURE_CODING_VERIFICATION_FAILED",
      `Single RS stripe cannot exceed 256 shards (got ${totalShards}; ${data.length} bytes / ${CHUNK_SIZE}-byte chunks). Use encodeIntoStripes for large payloads.`,
    );
  }
  const buffer = new Uint8Array(totalShards * CHUNK_SIZE);
  buffer.set(data, 0); // Trailing bytes after `data.length` stay zero — the pad.

  let code: number;
  try {
    code = rs.encode(buffer, dataShards, parityShards);
  } catch (err) {
    // WASM traps (e.g. Rust panics) surface as raw JS errors with messages
    // like "unreachable". Map them to a typed CodexError so the UI can
    // explain what happened instead of leaking the trap string.
    throw new CodexError(
      "ERASURE_CODING_VERIFICATION_FAILED",
      `Reed-Solomon encode trapped (data=${dataShards}, parity=${parityShards}): ${
        err instanceof Error ? err.message : String(err)
      }`,
    );
  }
  if (code !== ReedSolomonErasure.RESULT_OK) {
    throw new CodexError(
      "ERASURE_CODING_VERIFICATION_FAILED",
      `Reed-Solomon encode failed: ${resultMessage(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;
  }

  let code: number;
  try {
    code = rs.reconstruct(buffer, dataShards, parityShards, available);
  } catch (err) {
    throw new CodexError(
      "ERASURE_CODING_VERIFICATION_FAILED",
      `Reed-Solomon reconstruct trapped (data=${dataShards}, parity=${parityShards}, have=${chunks.size}): ${
        err instanceof Error ? err.message : String(err)
      }`,
    );
  }
  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: ${resultMessage(code)}`,
    );
  }

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

/**
 * Multi-stripe encode. Splits `data` into balanced RS stripes (each ≤
 * MAX_DATA_SHARDS_PER_STRIPE data shards) and encodes each independently.
 * The stripes are byte-contiguous: stripe `s` covers exactly
 * `dataShardsForStripe(s) * CHUNK_SIZE` bytes of payload, with the last shard
 * of the last stripe zero-padded when the payload doesn't align. Restore
 * concatenates each stripe's decoded data region and inflates the whole thing
 * — pako tolerates the trailing zeros.
 */
export async function encodeIntoStripes(
  data: Uint8Array,
): Promise<MultiStripeEncodeResult> {
  const totalDataShards = Math.max(1, Math.ceil(data.length / CHUNK_SIZE));
  const stripeShardCounts = planStripeShardCounts(totalDataShards);

  const stripes: EncodeResult[] = [];
  let byteCursor = 0;
  for (const shardCount of stripeShardCounts) {
    const stripeByteCap = shardCount * CHUNK_SIZE;
    const stripeEnd = Math.min(byteCursor + stripeByteCap, data.length);
    const stripeBytes = data.subarray(byteCursor, stripeEnd);
    const result = await encodeWithErasure(stripeBytes);
    // Sanity: encodeWithErasure infers dataShards from byte length; for any
    // stripe but the last, the slice is exactly `shardCount * CHUNK_SIZE`
    // bytes and the inferred count must match the plan. For the last stripe
    // the byte length may be shorter (zero-pad happens inside the encoder)
    // but the shard count still matches because the planner used the same
    // ceil-division.
    if (result.dataShards !== shardCount) {
      throw new CodexError(
        "ERASURE_CODING_VERIFICATION_FAILED",
        `Stripe plan mismatch: planned ${shardCount} data shards, encoder produced ${result.dataShards}`,
      );
    }
    stripes.push(result);
    byteCursor += stripeByteCap;
  }

  const totalDataChunks = stripes.reduce((s, x) => s + x.dataShards, 0);
  const totalParityChunks = stripes.reduce((s, x) => s + x.parityShards, 0);
  const totalChunks = totalDataChunks + totalParityChunks;
  return { stripes, totalDataChunks, totalParityChunks, totalChunks };
}