/**
 * 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 } from "pako";
import {
  CHUNK_SIZE,
  CODES_PER_PAGE,
  CodexError,
  PARITY_SHARDS,
  type ChunkEnvelope,
  type CodexDeps,
  type CodexManifest,
  type CodexPayload,
  type CodexProgress,
} from "../../shared/codex-types";
import { generateVerifiedAztec } from "./aztec";
import { decodeWithErasure, encodeWithErasure } from "./erasure";
import {
  base64UrlToBytes,
  bytesToBase64Url,
  buildChunkEnvelope,
  sha256Hex,
} from "./helpers";
import { renderCodexPDF } from "./pdf";

// ── 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 + simulated-loss verify
  onProgress?.("erasure");
  const { dataChunks, parityChunks, totalChunks, dataShards } = await encodeWithErasure(compressed);
  await verifyErasureCoding(dataChunks, parityChunks, dataShards, totalChunks, compressed);

  // STEP 5 — envelopes (data first, then parity, preserving global index ordering)
  const envelopes: ChunkEnvelope[] = [];
  for (let i = 0; i < dataChunks.length; i++) {
    envelopes.push(await buildChunkEnvelope(dataChunks[i]!, i, totalChunks, false, verifyCode));
  }
  for (let i = 0; i < parityChunks.length; i++) {
    envelopes.push(
      await buildChunkEnvelope(parityChunks[i]!, dataChunks.length + i, totalChunks, true, verifyCode),
    );
  }

  // 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);
  }
  onProgress?.("aztec", envelopes.length, envelopes.length);

  // STEP 7 — full reconstruction simulation (lose 2 random envelopes, rebuild)
  onProgress?.("verify");
  await simulateAndVerifyReconstruction(envelopes, dataShards, totalChunks, payload);

  // STEP 8 — render PDF
  onProgress?.("pdf");
  const manifest: CodexManifest = {
    verifyCode,
    totalChunks,
    dataChunks: dataShards,
    parityChunks: PARITY_SHARDS,
    pageCount: Math.ceil(totalChunks / CODES_PER_PAGE),
    entryCount,
    vaultFingerprint,
    generatedAt: payload.generatedAt,
  };
  const pdfBytes = await renderCodexPDF(aztecPngs, manifest);

  // STEP 10 — 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 (or raw chunk
 * bytes). Exported because the restore-side agent reuses this exact path to
 * rebuild the payload from scanned codes.
 *
 * Accepts either:
 *   - Map<index, ChunkEnvelope>  — typical scan path
 *   - Map<index, Uint8Array>     — raw bytes, used by the in-flight
 *                                  simulation during generation
 */
export async function reconstructPayloadFromChunks(
  envelopes: Map<number, Uint8Array> | Map<number, ChunkEnvelope>,
  dataShards: number,
  totalChunks: number,
): Promise<CodexPayload> {
  const chunkBytes = new Map<number, Uint8Array>();
  for (const [index, value] of envelopes) {
    if (value instanceof Uint8Array) {
      chunkBytes.set(index, value);
    } else {
      chunkBytes.set(index, base64UrlToBytes(value.d));
    }
  }

  const recovered = await decodeWithErasure(chunkBytes, dataShards, totalChunks);

  // The trailing pad bytes are zeros up to the next chunk boundary. pako's
  // inflate tolerates trailing junk after the deflate stream end marker, so
  // we can hand it the padded buffer directly.
  let inflated: Uint8Array;
  try {
    inflated = inflate(recovered);
  } 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<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 dataShards = Math.max(1, Math.ceil(compressed.length / CHUNK_SIZE));
  const totalChunks = dataShards + PARITY_SHARDS;
  return Math.ceil(totalChunks / CODES_PER_PAGE);
}

/** Pick 3 distinct positions in [1, 12] for the recovery-words confirmation
 *  gate. Sorted ascending so they read naturally in the prompt. */
export function pickConfirmationPositions(): [number, number, number] {
  const positions = Array.from({ length: 12 }, (_, i) => i + 1);
  for (let i = positions.length - 1; i > 0; i--) {
    const j = Math.floor(Math.random() * (i + 1));
    [positions[i], positions[j]] = [positions[j]!, positions[i]!];
  }
  const picked = [positions[0]!, positions[1]!, positions[2]!] as const;
  return [...picked].sort((a, b) => a - b) as [number, number, number];
}

// ── 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));

  // AUD-004: 32 hex chars = 128 bits (was 16 chars / 64 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,
  };
}

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 the encoded shards can be reconstructed from a partial set. We
 * drop two random shards (the maximum loss our PARITY_SHARDS budget tolerates)
 * and check byte-for-byte that the recovered prefix equals the original
 * compressed blob.
 */
async function verifyErasureCoding(
  dataChunks: Uint8Array[],
  parityChunks: Uint8Array[],
  dataShards: number,
  totalChunks: number,
  originalCompressed: Uint8Array,
): Promise<void> {
  const present = new Map<number, Uint8Array>();
  dataChunks.forEach((chunk, i) => present.set(i, chunk));
  parityChunks.forEach((chunk, i) => present.set(dataShards + i, chunk));

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

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

  for (let i = 0; i < originalCompressed.length; i++) {
    if (recovered[i] !== originalCompressed[i]) {
      throw new CodexError(
        "ERASURE_CODING_CONTENT_MISMATCH",
        `Reconstructed data does not match original at byte ${i}`,
      );
    }
  }
}

async function simulateAndVerifyReconstruction(
  envelopes: ChunkEnvelope[],
  dataShards: number,
  totalChunks: number,
  originalPayload: CodexPayload,
): Promise<void> {
  const present = new Map<number, Uint8Array>();
  envelopes.forEach((env, i) => present.set(i, base64UrlToBytes(env.d)));

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

  let reconstructed: CodexPayload;
  try {
    reconstructed = await reconstructPayloadFromChunks(present, dataShards, totalChunks);
  } catch (err) {
    throw new CodexError(
      "RECONSTRUCTION_SIMULATION_FAILED",
      `Codex reconstruction simulation failed. Lost chunks: ${removed.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);
}