/**
 * generate.ts — full Codex generation pipeline.
 *
 * Runs entirely in the popup. The service-worker hands the popup a `CodexDeps`
 * snapshot via the GENERATE_CODEX message and the popup is responsible for
 * every step that follows: compression, erasure coding, Aztec rasterization,
 * round-trip verification, and PDF rendering.
 *
 * Each step is wrapped in a verification gate. A failure at any stage aborts
 * generation with a typed `CodexError` so the UI can surface a precise reason
 * rather than a generic "something went wrong." This is the
 * physical-disaster-recovery surface — silent corruption here means the
 * user's only escape hatch is poisoned.
 */

import { deflate, inflate, Inflate } from "pako";
import { IMPORT_LIMITS } from "../../shared/import-limits";
import {
  CHUNK_SIZE,
  CODES_PER_PAGE,
  CodexError,
  computeParityShards,
  type ChunkEnvelope,
  type CodexDeps,
  type CodexManifest,
  type CodexPayload,
  type CodexProgress,
} from "../../shared/codex-types";
import { generateVerifiedAztec } from "./aztec";
import {
  decodeWithErasure,
  encodeIntoStripes,
  planStripeShardCounts,
  type EncodeResult,
} from "./erasure";
import {
  base64UrlToBytes,
  bytesToBase64Url,
  buildChunkEnvelope,
  sha256Hex,
} from "./helpers";
import { renderCodexPDF } from "./pdf";
export { pickConfirmationPositions } from "./confirmation";

/**
 * Per-stripe scanned-chunk bundle. Restore code groups envelopes by stripe
 * and passes one of these per stripe to `reconstructPayloadFromStripes`.
 */
export interface StripeChunks {
  /** Within-stripe index → raw chunk bytes. */
  chunks: Map<number, Uint8Array>;
  dataShards: number;
  totalChunks: number;
}

// ── Public surface ──────────────────────────────────────────────────────────

export async function generateCodex(
  deps: CodexDeps,
  onProgress?: CodexProgress,
): Promise<{ pdfBytes: Uint8Array; manifest: CodexManifest }> {
  // STEP 0 — preconditions
  if (!deps.isVaultUnlocked()) {
    throw new CodexError("VAULT_LOCKED", "Vault must be unlocked to generate a Codex");
  }

  // STEP 1 — vault blob
  const { vaultBlob, timestamp, entryCount, vaultFingerprint } = await buildVaultBlob(deps);
  const verifyCode = await deps.getVerifyCode();

  // STEP 2 — canonical payload (no recovery words — those live on paper only)
  const payload: CodexPayload = {
    version: 2,
    verifyCode,
    vaultBlob,
    timestamp,
    entryCount,
    vaultFingerprint,
    generatedAt: Date.now(),
  };

  // STEP 3 — compress + round-trip verify
  onProgress?.("compress");
  const payloadJson = JSON.stringify(payload);
  const payloadBytes = new TextEncoder().encode(payloadJson);
  const compressed = deflate(payloadBytes, { level: 9 });
  assertCompressionRoundTrip(compressed, payloadJson);

  // STEP 4 — erasure encode (multi-stripe) + simulated-loss verify per stripe
  onProgress?.("erasure");
  const stripeResult = await encodeIntoStripes(compressed);
  await verifyErasureCoding(stripeResult.stripes, compressed);
  const stripeCount = stripeResult.stripes.length;
  const totalChunks = stripeResult.totalChunks;

  // STEP 5 — envelopes. One envelope per shard, tagged with its stripe so
  // restore can route it. Within a stripe, data shards come first (indices
  // 0..dataShards-1), then parity (dataShards..totalChunks-1). Across stripes,
  // we emit stripe 0 first, then stripe 1, etc., so the PDF lays them out in
  // a predictable order on the page.
  const envelopes: ChunkEnvelope[] = [];
  for (let sIdx = 0; sIdx < stripeCount; sIdx++) {
    const stripe = stripeResult.stripes[sIdx]!;
    for (let i = 0; i < stripe.dataChunks.length; i++) {
      envelopes.push(
        await buildChunkEnvelope(
          stripe.dataChunks[i]!,
          i,
          stripe.totalChunks,
          false,
          verifyCode,
          sIdx,
          stripeCount,
        ),
      );
    }
    for (let i = 0; i < stripe.parityChunks.length; i++) {
      envelopes.push(
        await buildChunkEnvelope(
          stripe.parityChunks[i]!,
          stripe.dataShards + i,
          stripe.totalChunks,
          true,
          verifyCode,
          sIdx,
          stripeCount,
        ),
      );
    }
  }

  // STEP 6 — verified Aztec PNGs (round-trip per code where DOM is available)
  const aztecPngs: Uint8Array[] = [];
  for (let i = 0; i < envelopes.length; i++) {
    onProgress?.("aztec", i, envelopes.length);
    const img = await generateVerifiedAztec(envelopes[i]!);
    aztecPngs.push(img.pngBytes);
    // Pass a blob URL so the view can render a live gallery of real codes.
    // Only created when OffscreenCanvas is available (same guard as round-trip).
    // Copy into a fresh Uint8Array to satisfy Blob's ArrayBuffer type constraint
    // (img.pngBytes may be backed by a SharedArrayBuffer in some type environments).
    let blobUrl: string | undefined;
    if (typeof URL !== "undefined" && typeof Blob !== "undefined") {
      const owned = new Uint8Array(img.pngBytes);
      blobUrl = URL.createObjectURL(new Blob([owned], { type: "image/png" }));
    }
    onProgress?.("aztec", i + 1, envelopes.length, blobUrl);
  }

  // STEP 7 — full reconstruction simulation (lose the maximum tolerated
  // number of envelopes per stripe, then rebuild end-to-end)
  onProgress?.("verify");
  await simulateAndVerifyReconstruction(envelopes, stripeResult.stripes, payload);

  // STEP 8 — render PDF
  onProgress?.("pdf");
  const manifest: CodexManifest = {
    verifyCode,
    totalChunks,
    dataChunks: stripeResult.totalDataChunks,
    parityChunks: stripeResult.totalParityChunks,
    stripeCount,
    dataPerStripe: stripeResult.stripes.map((s) => s.dataShards),
    parityPerStripe: stripeResult.stripes.map((s) => s.parityShards),
    pageCount: Math.ceil(totalChunks / CODES_PER_PAGE),
    entryCount,
    vaultFingerprint,
    generatedAt: payload.generatedAt,
  };
  const pdfBytes = await renderCodexPDF(aztecPngs, manifest);

  // STEP 9 — sanity check the PDF
  if (!pdfBytes || pdfBytes.length < 1000) {
    throw new CodexError(
      "PDF_GENERATION_FAILED",
      `PDF generation produced an invalid file (${pdfBytes?.length ?? 0} bytes)`,
    );
  }

  return { pdfBytes, manifest };
}

/**
 * Reconstruct a CodexPayload from a partial set of envelopes grouped per
 * stripe. Each `StripeChunks` entry carries the chunks scanned for that
 * stripe plus the stripe's (dataShards, totalChunks) shape — restore obtains
 * those from the envelope's `t` field via `deriveShardCounts`. Stripes must
 * be supplied in stripe-index order so the concatenated data matches the
 * original compressed payload byte-for-byte.
 */
export async function reconstructPayloadFromStripes(
  stripes: StripeChunks[],
): Promise<CodexPayload> {
  if (stripes.length === 0) {
    throw new CodexError(
      "INSUFFICIENT_CHUNKS",
      "Cannot reconstruct payload: no stripes provided",
    );
  }

  // Decode each stripe to its `dataShards * shardSize` byte region.
  const stripeData: Uint8Array[] = [];
  for (const stripe of stripes) {
    const recovered = await decodeWithErasure(
      stripe.chunks,
      stripe.dataShards,
      stripe.totalChunks,
    );
    stripeData.push(recovered);
  }

  // Concatenate stripe data in stripe order — the original compressed payload.
  const totalLen = stripeData.reduce((s, b) => s + b.length, 0);
  const combined = new Uint8Array(totalLen);
  let offset = 0;
  for (const part of stripeData) {
    combined.set(part, offset);
    offset += part.length;
  }

  // Trailing zero-pad bytes are tolerated by pako's inflate after the deflate
  // stream end marker, so we can hand it the concatenated buffer directly.
  // MV-4: stream the inflate with a byte counter — one-shot inflate() has no
  // output bound and a crafted Codex payload could otherwise balloon to GiB.
  let inflated: Uint8Array;
  try {
    inflated = inflateBounded(combined, IMPORT_LIMITS.MAX_DECOMPRESSED_BYTES);
  } catch (err) {
    throw new CodexError(
      "RECONSTRUCTION_SIMULATION_FAILED",
      `Decompression failed during reconstruction: ${err instanceof Error ? err.message : String(err)}`,
    );
  }

  const json = new TextDecoder().decode(inflated);
  let payload: CodexPayload;
  try {
    payload = JSON.parse(json) as CodexPayload;
  } catch (err) {
    throw new CodexError(
      "RECONSTRUCTION_SIMULATION_FAILED",
      `Reconstructed payload is not valid JSON: ${err instanceof Error ? err.message : String(err)}`,
    );
  }
  return payload;
}

/**
 * Predict how many pages a Codex will occupy without doing the work to
 * generate one. Used by the popup's ESTIMATE message to drive the "this will
 * be N pages" UI hint before the user commits to printing. Speculative
 * compression is fine here — we never persist the result.
 */
export async function estimateCodexPages(
  slotA: string,
  slotB: string | null,
): Promise<{ pages: number; parityShards: number }> {
  const combined = JSON.stringify({ a: slotA, b: slotB ?? null });
  const vaultBlob = bytesToBase64Url(new TextEncoder().encode(combined));
  // Stand-in payload — the exact verify code / timestamp / fingerprint don't
  // affect the compressed size enough to matter for a page count estimate.
  const stub: CodexPayload = {
    version: 2,
    verifyCode: "PX-XXXX-XXXX",
    vaultBlob,
    timestamp: 0,
    entryCount: 0,
    vaultFingerprint: "0000000000000000",
    generatedAt: 0,
  };
  const compressed = deflate(new TextEncoder().encode(JSON.stringify(stub)), { level: 9 });
  const totalDataShards = Math.max(1, Math.ceil(compressed.length / CHUNK_SIZE));
  // Stripe planning matches encodeIntoStripes so the estimate stays honest
  // for vaults large enough to be split.
  const stripeShardCounts = planStripeShardCounts(totalDataShards);
  let totalChunks = 0;
  let totalParity = 0;
  for (const dataShards of stripeShardCounts) {
    const parity = computeParityShards(dataShards);
    totalChunks += dataShards + parity;
    totalParity += parity;
  }
  const pages = Math.ceil(totalChunks / CODES_PER_PAGE);
  return { pages, parityShards: totalParity };
}

// ── Internal helpers ────────────────────────────────────────────────────────

async function buildVaultBlob(deps: CodexDeps): Promise<{
  vaultBlob: string;
  timestamp: number;
  entryCount: number;
  vaultFingerprint: string;
}> {
  const slotA = await deps.getSlotA();
  if (!slotA) {
    throw new CodexError("VAULT_EMPTY", "No vault data found");
  }
  const slotB = await deps.getSlotB();
  const lastModified = await deps.getLastModified();
  const entries = await deps.getAllEntries();

  const combined = JSON.stringify({ a: slotA, b: slotB ?? null });
  const vaultBlob = bytesToBase64Url(new TextEncoder().encode(combined));

  // 32 hex chars = 128 bits.
  const fingerprint = (await sha256Hex(new TextEncoder().encode(vaultBlob))).slice(0, 32);

  const timestamp = lastModified ? Number.parseInt(lastModified, 10) : Date.now();
  const liveEntryCount = entries.filter((e) => !e.deletedAt).length;

  return {
    vaultBlob,
    timestamp: Number.isFinite(timestamp) ? timestamp : Date.now(),
    entryCount: liveEntryCount,
    vaultFingerprint: fingerprint,
  };
}

/** Inflate with a hard output ceiling — aborts mid-stream past maxBytes.
 *  Exported for MV-4 bound tests. */
export function inflateBounded(input: Uint8Array, maxBytes: number): Uint8Array {
  const chunks: Uint8Array[] = [];
  let total = 0;
  const inf = new Inflate();
  inf.onData = (chunk: Uint8Array) => {
    total += chunk.length;
    if (total > maxBytes) {
      throw new Error(`IMPORT_TOO_LARGE: inflated output exceeds ${maxBytes} bytes`);
    }
    chunks.push(chunk);
  };
  inf.push(input, true);
  if (inf.err) throw new Error(inf.msg || `inflate error ${inf.err}`);
  const out = new Uint8Array(total);
  let off = 0;
  for (const c of chunks) { out.set(c, off); off += c.length; }
  return out;
}

function assertCompressionRoundTrip(compressed: Uint8Array, expectedJson: string): void {
  const inflated = inflate(compressed);
  const restored = new TextDecoder().decode(inflated);
  if (restored !== expectedJson) {
    throw new CodexError(
      "COMPRESSION_ROUNDTRIP_FAILED",
      "Compression verification failed: inflate(deflate(payload)) did not match the original",
    );
  }
}

/**
 * Confirm each stripe's encoded shards can be reconstructed after losing the
 * maximum number of chunks the stripe's parity budget tolerates. Checks the
 * recovered bytes against the original compressed payload slice for that
 * stripe — byte-for-byte equality across the populated bytes; trailing
 * zero-pad in the last stripe is ignored.
 */
async function verifyErasureCoding(
  stripes: EncodeResult[],
  originalCompressed: Uint8Array,
): Promise<void> {
  let byteCursor = 0;
  for (let sIdx = 0; sIdx < stripes.length; sIdx++) {
    const stripe = stripes[sIdx]!;
    const present = new Map<number, Uint8Array>();
    stripe.dataChunks.forEach((chunk, i) => present.set(i, chunk));
    stripe.parityChunks.forEach((chunk, i) =>
      present.set(stripe.dataShards + i, chunk),
    );

    const removed = pickRandomIndices(
      stripe.totalChunks,
      Math.min(stripe.parityShards, stripe.totalChunks - 1),
    );
    for (const idx of removed) present.delete(idx);

    let recovered: Uint8Array;
    try {
      recovered = await decodeWithErasure(
        present,
        stripe.dataShards,
        stripe.totalChunks,
      );
    } catch (err) {
      if (err instanceof CodexError) {
        throw new CodexError(
          "ERASURE_CODING_VERIFICATION_FAILED",
          `Erasure coding simulation failed for stripe ${sIdx} when losing chunks ${removed.join(", ")}: ${err.message}`,
        );
      }
      throw err;
    }

    // Number of real payload bytes in this stripe (rest is zero pad).
    const stripeRealBytes = Math.min(
      stripe.dataShards * CHUNK_SIZE,
      Math.max(0, originalCompressed.length - byteCursor),
    );
    for (let i = 0; i < stripeRealBytes; i++) {
      if (recovered[i] !== originalCompressed[byteCursor + i]) {
        throw new CodexError(
          "ERASURE_CODING_CONTENT_MISMATCH",
          `Reconstructed stripe ${sIdx} differs from original at byte ${byteCursor + i}`,
        );
      }
    }
    byteCursor += stripe.dataShards * CHUNK_SIZE;
  }
}

async function simulateAndVerifyReconstruction(
  envelopes: ChunkEnvelope[],
  stripes: EncodeResult[],
  originalPayload: CodexPayload,
): Promise<void> {
  // Group envelopes by stripe so we can drop up to `parityShards` per stripe
  // and still reconstruct.
  const perStripe: StripeChunks[] = stripes.map((s) => ({
    chunks: new Map<number, Uint8Array>(),
    dataShards: s.dataShards,
    totalChunks: s.totalChunks,
  }));
  for (const env of envelopes) {
    const stripe = env.v === 3 ? (env.s ?? 0) : 0;
    perStripe[stripe]!.chunks.set(env.i, base64UrlToBytes(env.d));
  }

  const removedSummary: string[] = [];
  for (let sIdx = 0; sIdx < stripes.length; sIdx++) {
    const stripe = stripes[sIdx]!;
    const removed = pickRandomIndices(
      stripe.totalChunks,
      Math.min(stripe.parityShards, stripe.totalChunks - 1),
    );
    for (const idx of removed) perStripe[sIdx]!.chunks.delete(idx);
    removedSummary.push(`stripe ${sIdx}: [${removed.join(", ")}]`);
  }

  let reconstructed: CodexPayload;
  try {
    reconstructed = await reconstructPayloadFromStripes(perStripe);
  } catch (err) {
    throw new CodexError(
      "RECONSTRUCTION_SIMULATION_FAILED",
      `Codex reconstruction simulation failed. Lost chunks: ${removedSummary.join("; ")}. ${err instanceof Error ? err.message : String(err)}`,
    );
  }

  if (reconstructed.verifyCode !== originalPayload.verifyCode) {
    throw new CodexError(
      "RECONSTRUCTION_VERIFY_CODE_MISMATCH",
      "Reconstructed Codex verify code does not match original",
    );
  }
  if (reconstructed.vaultFingerprint !== originalPayload.vaultFingerprint) {
    throw new CodexError(
      "RECONSTRUCTION_FINGERPRINT_MISMATCH",
      "Reconstructed vault fingerprint does not match original",
    );
  }
  if (reconstructed.entryCount !== originalPayload.entryCount) {
    throw new CodexError(
      "RECONSTRUCTION_COUNT_MISMATCH",
      `Reconstructed entry count (${reconstructed.entryCount}) does not match original (${originalPayload.entryCount})`,
    );
  }
}

/**
 * Pick `count` distinct indices uniformly at random from [0, total). Used by
 * the verification simulations to model "any two chunks could be lost." Not
 * security-sensitive — we just need diverse coverage across regenerations.
 */
function pickRandomIndices(total: number, count: number): number[] {
  const indices = Array.from({ length: total }, (_, i) => i);
  for (let i = indices.length - 1; i > 0; i--) {
    const j = Math.floor(Math.random() * (i + 1));
    const tmp = indices[i]!;
    indices[i] = indices[j]!;
    indices[j] = tmp;
  }
  return indices.slice(0, count);
}