import { beforeEach, describe, expect, it, vi } from "vitest";
import {
  CHUNK_SIZE,
  CodexError,
  computeParityShards,
  STORAGE_KEY_VERIFY_CODE,
  type CodexDeps,
} from "../../shared/codex-types";
import type { VaultEntry } from "../../shared/types";
import {
  base64UrlToBytes,
  buildChunkEnvelope,
  bytesToBase64Url,
  sha256Hex,
  verifyChunkIntegrity,
} from "./helpers";
import { decodeWithErasure, encodeWithErasure } from "./erasure";
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(2);
    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);
    expect(computeParityShards(40)).toBe(2);
  });

  it("computeParityShards scales at 5% for medium vaults", () => {
    expect(computeParityShards(56)).toBe(3); // 56 * 0.05 = 2.8 → ceil = 3
    expect(computeParityShards(100)).toBe(5);
    expect(computeParityShards(213)).toBe(11);
  });

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

  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: survives loss of up to parityShards codes", async () => {
    const data = new Uint8Array(200_000);
    for (let i = 0; i < data.length; i++) data[i] = (i * 37) & 0xff;

    const { dataChunks, parityChunks, dataShards, parityShards, totalChunks } =
      await encodeWithErasure(data);

    const allChunks = new Map<number, Uint8Array>(
      [...dataChunks, ...parityChunks].map((c, i) => [i, c]),
    );
    // Drop exactly parityShards codes — the maximum survivable loss.
    for (let i = 0; i < parityShards; i++) allChunks.delete(i);

    const reconstructed = await decodeWithErasure(
      allChunks,
      dataShards,
      totalChunks,
    );
    for (let i = 0; i < data.length; i++) {
      expect(reconstructed[i]).toBe(data[i]);
    }
  }, 60_000);

  it("large vault: fails gracefully when loss exceeds parityShards", async () => {
    const data = new Uint8Array(200_000);
    for (let i = 0; i < data.length; i++) data[i] = (i * 41) & 0xff;
    const { dataChunks, parityChunks, dataShards, parityShards, totalChunks } =
      await encodeWithErasure(data);

    const allChunks = new Map<number, Uint8Array>(
      [...dataChunks, ...parityChunks].map((c, i) => [i, c]),
    );
    for (let i = 0; i < parityShards + 1; i++) allChunks.delete(i);

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