import { beforeEach, describe, expect, it, vi } from "vitest";
import {
  CHUNK_SIZE,
  CodexError,
  computeParityShards,
  STORAGE_KEY_VERIFY_CODE,
  type ChunkEnvelope,
  type CodexDeps,
} from "../../shared/codex-types";
import type { VaultEntry } from "../../shared/types";
import {
  base64UrlToBytes,
  buildChunkEnvelope,
  bytesToBase64Url,
  sha256Hex,
  verifyChunkIntegrity,
} from "./helpers";
import {
  decodeWithErasure,
  encodeIntoStripes,
  encodeWithErasure,
} from "./erasure";
import { buildAztecMatrix, generateVerifiedAztec, verifyMatrixRoundTrip } from "./aztec";
import { generateCodex } from "./generate";

// ── chrome.storage shim ─────────────────────────────────────────────────────

function installChromeStorageMock(): { storage: Record<string, unknown> } {
  const storage: Record<string, unknown> = {};
  const chromeMock = {
    storage: {
      local: {
        get: vi.fn(async (key?: string | string[]) => {
          if (!key) return { ...storage };
          if (typeof key === "string") return { [key]: storage[key] };
          return key.reduce<Record<string, unknown>>((acc, k) => {
            acc[k] = storage[k];
            return acc;
          }, {});
        }),
        set: vi.fn(async (entries: Record<string, unknown>) => {
          Object.assign(storage, entries);
        }),
        remove: vi.fn(async (k: string | string[]) => {
          const keys = Array.isArray(k) ? k : [k];
          for (const key of keys) delete storage[key];
        }),
      },
    },
    // Deliberately omit `runtime.getURL` so rs-loader.ts falls through to the
    // Node WASM-from-disk path. With `chrome.runtime.getURL` defined the
    // loader treats us as an extension popup and tries to `fetch()` the WASM,
    // which fails in vitest's node environment.
    runtime: {},
  };
  // @ts-expect-error — test-only injection of the chrome global
  globalThis.chrome = chromeMock;
  return { storage };
}

// ── Mock CodexDeps factory ──────────────────────────────────────────────────

function makeDeps(overrides: Partial<CodexDeps> = {}): CodexDeps {
  // A realistic-sized encrypted-looking blob, encoded as deterministic-but-
  // high-entropy base64-ish bytes. The pipeline compresses the JSON wrapper,
  // so for the reconstruction simulation to exercise multiple data shards we
  // need post-compression size >= ~2KB. Real encrypted vault data is dense
  // and incompressible — we model that with a pseudo-random byte stream
  // expressed in a base64 alphabet so it survives JSON stringification.
  const alphabet = "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789+/";
  // Linear congruential generator — deterministic per test run, but the
  // output histogram is flat enough that deflate cannot compress it well.
  let seed = 0x1234abcd;
  const slotALen = 8192;
  const chars: string[] = new Array(slotALen);
  for (let i = 0; i < slotALen; i++) {
    seed = (seed * 1103515245 + 12345) >>> 0;
    chars[i] = alphabet[(seed >>> 16) & 0x3f]!;
  }
  const slotA = "ENCRYPTED::" + chars.join("");
  const slotB: string | null = null;
  const entries: VaultEntry[] = [
    { id: "1", title: "GitHub", username: "alice", password: "secret", createdAt: 1, modifiedAt: 2 },
    { id: "2", title: "Email", username: "bob", password: "hunter2", createdAt: 1, modifiedAt: 2 },
    { id: "3", title: "Bank", username: "carol", password: "money", createdAt: 1, modifiedAt: 2 },
  ];

  return {
    getSlotA: async () => slotA,
    getSlotB: async () => slotB,
    getLastModified: async () => "1700000000000",
    getAllEntries: async () => entries,
    getVerifyCode: async () => "PX-TEST-CODE",
    isVaultUnlocked: () => true,
    ...overrides,
  };
}

// ── Tests ───────────────────────────────────────────────────────────────────

describe("helpers", () => {
  beforeEach(() => {
    installChromeStorageMock();
  });

  it("base64url round-trips arbitrary binary", () => {
    const original = new Uint8Array([0, 1, 2, 250, 251, 252, 253, 254, 255, 0x3f, 0x2b, 0x2f]);
    const encoded = bytesToBase64Url(original);
    expect(encoded).not.toMatch(/[+/=]/);
    const decoded = base64UrlToBytes(encoded);
    expect(Array.from(decoded)).toEqual(Array.from(original));
  });

  it("sha256Hex returns uppercase 64-char hex for known input", async () => {
    const empty = new Uint8Array();
    const hex = await sha256Hex(empty);
    expect(hex).toBe("E3B0C44298FC1C149AFBF4C8996FB92427AE41E4649B934CA495991B7852B855");
  });

  it("buildChunkEnvelope + verifyChunkIntegrity agree", async () => {
    const chunk = new Uint8Array(CHUNK_SIZE).fill(0x42);
    const env = await buildChunkEnvelope(chunk, 3, 7, false, "PX-AAAA-BBBB");
    expect(env.v).toBe(3);
    expect(env.i).toBe(3);
    expect(env.t).toBe(7);
    expect(env.p).toBe(false);
    expect(env.c).toBe("PX-AAAA-BBBB");
    expect(env.h).toHaveLength(32); // AUD-004: 128-bit hash prefix
    expect(await verifyChunkIntegrity(env)).toBe(true);
  });

  it("verifyChunkIntegrity rejects a corrupted envelope", async () => {
    const chunk = new Uint8Array(CHUNK_SIZE).fill(0x42);
    const env = await buildChunkEnvelope(chunk, 0, 4, false, "PX-AAAA-BBBB");
    // Corrupt the data payload after-the-fact — the recorded hash no longer matches.
    const tampered = { ...env, d: bytesToBase64Url(new Uint8Array(CHUNK_SIZE).fill(0x43)) };
    expect(await verifyChunkIntegrity(tampered)).toBe(false);
  });
});

describe("erasure", () => {
  beforeEach(() => {
    installChromeStorageMock();
  });

  it("round-trips with the full chunk set", async () => {
    // Three full data shards plus tail
    const original = new Uint8Array(CHUNK_SIZE * 3 + 17);
    for (let i = 0; i < original.length; i++) original[i] = (i * 31) & 0xff;

    const { dataChunks, parityChunks, dataShards, parityShards, totalChunks } =
      await encodeWithErasure(original);
    expect(dataShards).toBe(4);
    expect(parityShards).toBe(computeParityShards(dataShards));
    expect(parityChunks.length).toBe(parityShards);
    expect(totalChunks).toBe(dataShards + parityShards);

    const present = new Map<number, Uint8Array>();
    dataChunks.forEach((c, i) => present.set(i, c));
    parityChunks.forEach((c, i) => present.set(dataShards + i, c));

    const recovered = await decodeWithErasure(present, dataShards, totalChunks);
    for (let i = 0; i < original.length; i++) {
      expect(recovered[i]).toBe(original[i]);
    }
  });

  it("reconstructs after losing two random data chunks", async () => {
    const original = new Uint8Array(CHUNK_SIZE * 4);
    for (let i = 0; i < original.length; i++) original[i] = (i * 7 + 5) & 0xff;

    const { dataChunks, parityChunks, dataShards, totalChunks } = await encodeWithErasure(original);
    const present = new Map<number, Uint8Array>();
    dataChunks.forEach((c, i) => present.set(i, c));
    parityChunks.forEach((c, i) => present.set(dataShards + i, c));

    // Drop two distinct data shards
    present.delete(0);
    present.delete(2);

    const recovered = await decodeWithErasure(present, dataShards, totalChunks);
    for (let i = 0; i < original.length; i++) {
      expect(recovered[i]).toBe(original[i]);
    }
  });

  it("throws INSUFFICIENT_CHUNKS when too many are missing", async () => {
    const original = new Uint8Array(CHUNK_SIZE * 4);
    const { dataChunks, parityChunks, dataShards, totalChunks } = await encodeWithErasure(original);
    const present = new Map<number, Uint8Array>();
    dataChunks.forEach((c, i) => present.set(i, c));
    parityChunks.forEach((c, i) => present.set(dataShards + i, c));

    // Drop three shards — exceeds parity budget
    present.delete(0);
    present.delete(1);
    present.delete(2);

    await expect(decodeWithErasure(present, dataShards, totalChunks)).rejects.toMatchObject({
      name: "CodexError",
      code: "INSUFFICIENT_CHUNKS",
    });
  });
});

describe("generateCodex", () => {
  beforeEach(() => {
    installChromeStorageMock();
  });

  it("produces a non-trivial PDF and a coherent manifest", async () => {
    const deps = makeDeps();
    const { pdfBytes, manifest } = await generateCodex(deps);

    expect(pdfBytes.byteLength).toBeGreaterThan(1000);
    // PDF magic header
    expect(pdfBytes[0]).toBe(0x25); // '%'
    expect(pdfBytes[1]).toBe(0x50); // 'P'
    expect(pdfBytes[2]).toBe(0x44); // 'D'
    expect(pdfBytes[3]).toBe(0x46); // 'F'

    expect(manifest.verifyCode).toBe("PX-TEST-CODE");
    expect(manifest.entryCount).toBe(3);
    expect(manifest.parityChunks).toBe(computeParityShards(manifest.dataChunks));
    expect(manifest.totalChunks).toBe(manifest.dataChunks + manifest.parityChunks);
    expect(manifest.pageCount).toBeGreaterThanOrEqual(1);
    expect(manifest.vaultFingerprint).toMatch(/^[0-9A-F]{32}$/); // AUD-004
  }, 60_000);

  it("rejects generation when the vault is locked", async () => {
    const deps = makeDeps({ isVaultUnlocked: () => false });
    await expect(generateCodex(deps)).rejects.toMatchObject({
      name: "CodexError",
      code: "VAULT_LOCKED",
    });
  });

  it("reports progress at named milestones", async () => {
    const deps = makeDeps();
    const steps: string[] = [];
    await generateCodex(deps, (step) => {
      if (!steps.includes(step)) steps.push(step);
    });
    expect(steps).toEqual(["compress", "erasure", "aztec", "verify", "pdf"]);
  }, 60_000);

  it("persists the storage_key_verify_code constant correctly", () => {
    // Smoke check the constant import — catches a rename / re-export regression
    // before tests that depend on chrome.storage start failing in confusing ways.
    expect(STORAGE_KEY_VERIFY_CODE).toBe("peach_codex_verify_code");
    expect(CodexError).toBeTypeOf("function");
  });
});

describe("parity scaling", () => {
  beforeEach(() => {
    installChromeStorageMock();
  });

  it("computeParityShards returns minimum 2 for very small vaults", () => {
    expect(computeParityShards(1)).toBe(2);
    expect(computeParityShards(10)).toBe(2);
    expect(computeParityShards(13)).toBe(2);
  });

  it("computeParityShards scales at 15% for medium vaults", () => {
    expect(computeParityShards(56)).toBe(9); // 56 * 0.15 = 8.4 → ceil = 9
    expect(computeParityShards(100)).toBe(15);
    expect(computeParityShards(213)).toBe(32);
  });

  it("computeParityShards caps at 35 for large vaults", () => {
    expect(computeParityShards(600)).toBe(35);
    expect(computeParityShards(660)).toBe(35);
    expect(computeParityShards(1000)).toBe(35);
  });

  it("encodeWithErasure returns scaled parityShards", async () => {
    const data = new Uint8Array(50_000);
    for (let i = 0; i < data.length; i++) data[i] = (i * 37) & 0xff;
    const { dataShards, parityShards } = await encodeWithErasure(data);
    expect(parityShards).toBe(computeParityShards(dataShards));
  });

  it("large vault: each stripe survives loss of up to its parityShards codes", async () => {
    // 200 KB / 800-byte chunks = 250 data shards → split into 2 stripes of
    // 125 each (well under the GF(2⁸) 256-shard cap). Each stripe gets its
    // own parity budget; we drop exactly parityShards from every stripe.
    const data = new Uint8Array(200_000);
    for (let i = 0; i < data.length; i++) data[i] = (i * 37) & 0xff;

    const { stripes } = await encodeIntoStripes(data);
    expect(stripes.length).toBeGreaterThan(1);

    let byteCursor = 0;
    for (const stripe of stripes) {
      const present = new Map<number, Uint8Array>();
      stripe.dataChunks.forEach((c, i) => present.set(i, c));
      stripe.parityChunks.forEach((c, i) =>
        present.set(stripe.dataShards + i, c),
      );
      // Drop exactly parityShards codes from this stripe — maximum survivable.
      for (let i = 0; i < stripe.parityShards; i++) present.delete(i);

      const recovered = await decodeWithErasure(
        present,
        stripe.dataShards,
        stripe.totalChunks,
      );
      const stripeRealBytes = Math.min(
        stripe.dataShards * CHUNK_SIZE,
        Math.max(0, data.length - byteCursor),
      );
      for (let i = 0; i < stripeRealBytes; i++) {
        expect(recovered[i]).toBe(data[byteCursor + i]);
      }
      byteCursor += stripe.dataShards * CHUNK_SIZE;
    }
  }, 60_000);

  it("large vault: fails gracefully when stripe loss exceeds parityShards", async () => {
    const data = new Uint8Array(200_000);
    for (let i = 0; i < data.length; i++) data[i] = (i * 41) & 0xff;
    const { stripes } = await encodeIntoStripes(data);
    expect(stripes.length).toBeGreaterThan(1);
    const stripe = stripes[0]!;
    const present = new Map<number, Uint8Array>();
    stripe.dataChunks.forEach((c, i) => present.set(i, c));
    stripe.parityChunks.forEach((c, i) =>
      present.set(stripe.dataShards + i, c),
    );
    for (let i = 0; i < stripe.parityShards + 1; i++) present.delete(i);

    await expect(
      decodeWithErasure(present, stripe.dataShards, stripe.totalChunks),
    ).rejects.toThrow(CodexError);
  }, 60_000);
});

// ── Aztec round-trip verification ─────────────────────────────────────────────

describe("aztec round-trip", () => {
  beforeEach(() => {
    installChromeStorageMock();
  });

  it("buildAztecMatrix + verifyMatrixRoundTrip: 200 varied payloads all pass", () => {
    // Exercise the core encode→decode path for a wide range of envelope-like
    // JSON strings. This is fast (synchronous ZXing in-memory, no canvas, no
    // bwip-js) and covers the payload patterns that previously caused
    // AZTEC_CONTENT_MISMATCH with HybridBinarizer.
    for (let i = 0; i < 200; i++) {
      // Simulate an envelope JSON payload. Use CHUNK_SIZE-sized chunk data with
      // a per-chunk bit pattern so each payload is unique.
      const data = new Uint8Array(CHUNK_SIZE);
      for (let j = 0; j < CHUNK_SIZE; j++) data[j] = ((i * 37 + j * 7) ^ (i << 3)) & 0xff;

      // Build a short envelope-like JSON string resembling the real envelope format.
      // Base64url encode the chunk data to keep it printable ASCII (matching the
      // real pipeline which uses bytesToBase64Url).
      const b64 = btoa(String.fromCharCode(...data.slice(0, 60))).replace(/\+/g, "-").replace(/\//g, "_").replace(/=/g, "");
      const payload = JSON.stringify({
        v: 3, i, t: 200, p: i >= 188, c: "PX-TEST-CODE",
        d: b64,
        h: "abcdef1234567890abcdef1234567890",
        s: 0, st: 1,
      });

      const matrix = buildAztecMatrix(payload);
      const ok = verifyMatrixRoundTrip(matrix, payload);
      if (!ok) {
        throw new Error(`verifyMatrixRoundTrip failed for chunk ${i} (payload len=${payload.length})`);
      }
    }
  });

  it("generateVerifiedAztec: 20 real envelopes survive round-trip", async () => {
    // End-to-end test of generateVerifiedAztec using real chunk envelopes built
    // by the same helpers used in the live pipeline. Any AZTEC_ROUNDTRIP_FAILED
    // or AZTEC_CONTENT_MISMATCH throw will fail the test immediately.
    const envelopes: ChunkEnvelope[] = [];
    for (let i = 0; i < 20; i++) {
      const data = new Uint8Array(CHUNK_SIZE);
      for (let j = 0; j < CHUNK_SIZE; j++) data[j] = ((i * 53 + j * 11) ^ (i * 19)) & 0xff;
      envelopes.push(await buildChunkEnvelope(data, i, 20, i >= 17, "PX-RTRP-CODE", 0, 1));
    }

    for (const env of envelopes) {
      const img = await generateVerifiedAztec(env);
      expect(img.pngBytes.length).toBeGreaterThan(0);
    }
  }, 120_000);

  it("generateCodex: large vault (100 KB) completes without AZTEC errors", async () => {
    // A realistic-sized vault that exercises multi-stripe encoding and produces
    // O(100) Aztec codes. All codes must pass in-memory round-trip verification.
    installChromeStorageMock();

    // 80 KB of pseudo-random incompressible base64 data → ~100 data shards
    const alphabet = "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789+/";
    let seed = 0xdeadbeef;
    const chars: string[] = new Array(80_000);
    for (let i = 0; i < 80_000; i++) {
      seed = (seed * 1664525 + 1013904223) >>> 0;
      chars[i] = alphabet[(seed >>> 16) & 0x3f]!;
    }
    const slotA = "ENCRYPTED::" + chars.join("");

    const deps: CodexDeps = {
      getSlotA: async () => slotA,
      getSlotB: async () => null,
      getLastModified: async () => "1700000000000",
      getAllEntries: async () => [
        { id: "1", title: "GitHub", username: "alice", password: "secret", createdAt: 1, modifiedAt: 2 },
      ],
      getVerifyCode: async () => "PX-LARG-TEST",
      isVaultUnlocked: () => true,
    };

    const steps: string[] = [];
    const { pdfBytes, manifest } = await generateCodex(deps, (step) => {
      if (!steps.includes(step)) steps.push(step);
    });

    expect(pdfBytes[0]).toBe(0x25); // %PDF
    expect(manifest.totalChunks).toBeGreaterThan(50);
    expect(steps).toEqual(["compress", "erasure", "aztec", "verify", "pdf"]);
  }, 300_000); // 5 min — bwip-js PNG for ~100 chunks dominates in node
});