import { afterEach, describe, expect, it, vi } from "vitest";
import {
  MAX_PAIRING_FRAME_BYTES,
  MAX_PAIRING_INBOX_BYTES,
  MAX_PAIRING_INBOX_MESSAGES,
  PAIRING_RESOURCE_LIMIT,
  openRelayPairingSocket,
  type PairingSocket,
} from "./pairing-dial";

// PB-26-205: the relay pairing inbox must be bounded before any
// allocate/copy of attacker-controlled frame data.

interface FakeSocket {
  binaryType: string;
  onmessage: ((ev: { data: unknown }) => void) | null;
  onerror: (() => void) | null;
  onclose: ((ev: { code: number }) => void) | null;
  onopen: (() => void) | null;
  sent: unknown[];
  closeCalls: Array<{ code?: number; reason?: string }>;
  send(data: unknown): void;
  close(code?: number, reason?: string): void;
  simulateOpen(): void;
  simulateMessage(data: unknown): void;
  simulateClose(code?: number): void;
}

let fake: FakeSocket;

function createFake(): FakeSocket {
  const sock: FakeSocket = {
    binaryType: "",
    onmessage: null,
    onerror: null,
    onclose: null,
    onopen: null,
    sent: [],
    closeCalls: [],
    send(data: unknown) {
      sock.sent.push(data);
    },
    close(code?: number, reason?: string) {
      sock.closeCalls.push({ code, reason });
      // Like a real browser, firing onclose asynchronously after close().
      queueMicrotask(() => sock.simulateClose(code));
    },
    simulateOpen() {
      sock.onopen?.();
    },
    simulateMessage(data: unknown) {
      sock.onmessage?.({ data });
    },
    simulateClose(code = 1000) {
      sock.onclose?.({ code });
    },
  };
  return sock;
}

async function openSocket(): Promise<PairingSocket> {
  const pending = openRelayPairingSocket("wss://relay.example.test/room");
  fake.simulateOpen();
  return pending;
}

function frame(n: number): ArrayBuffer {
  return new Uint8Array(n).buffer as ArrayBuffer;
}

afterEach(() => {
  vi.unstubAllGlobals();
});

describe("PB-26-205 pairing relay inbox limits", () => {
  it("accepts a frame of exactly MAX_PAIRING_FRAME_BYTES", async () => {
    vi.stubGlobal("WebSocket", vi.fn().mockImplementation(function () {
      fake = createFake();
      return fake;
    }));
    const sock = await openSocket();
    fake.simulateMessage(frame(MAX_PAIRING_FRAME_BYTES));
    const got = await sock.receive();
    expect(got.byteLength).toBe(MAX_PAIRING_FRAME_BYTES);
    expect(fake.closeCalls).toEqual([]);
  });

  it("rejects a max+1 frame with a 1009 close before copying", async () => {
    vi.stubGlobal("WebSocket", vi.fn().mockImplementation(function () {
      fake = createFake();
      return fake;
    }));
    const sock = await openSocket();
    const pending = sock.receive();
    const rejection = expect(pending).rejects.toThrow(PAIRING_RESOURCE_LIMIT);
    fake.simulateMessage(frame(MAX_PAIRING_FRAME_BYTES + 1));
    await rejection;
    expect(fake.closeCalls[0]?.code).toBe(1009);
    // Queued references are released: nothing remains deliverable.
    await expect(sock.receive()).rejects.toThrow(PAIRING_RESOURCE_LIMIT);
  });

  it("rejects non-ArrayBuffer frames without copying", async () => {
    vi.stubGlobal("WebSocket", vi.fn().mockImplementation(function () {
      fake = createFake();
      return fake;
    }));
    const sock = await openSocket();
    const pending = sock.receive();
    const rejection = expect(pending).rejects.toThrow(PAIRING_RESOURCE_LIMIT);
    fake.simulateMessage("not-binary");
    await rejection;
    expect(fake.closeCalls[0]?.code).toBe(1009);
  });

  it("closes when individually valid frames exceed the inbox byte budget", async () => {
    vi.stubGlobal("WebSocket", vi.fn().mockImplementation(function () {
      fake = createFake();
      return fake;
    }));
    const sock = await openSocket();
    const frames = Math.floor(MAX_PAIRING_INBOX_BYTES / MAX_PAIRING_FRAME_BYTES);
    for (let i = 0; i < frames; i++) fake.simulateMessage(frame(MAX_PAIRING_FRAME_BYTES));
    // One more byte over budget trips the limit even though the frame is valid.
    fake.simulateMessage(frame(1));
    expect(fake.closeCalls[0]?.code).toBe(1009);
    await expect(sock.receive()).rejects.toThrow(PAIRING_RESOURCE_LIMIT);
  });

  it("enforces the message-count limit with tiny frames", async () => {
    vi.stubGlobal("WebSocket", vi.fn().mockImplementation(function () {
      fake = createFake();
      return fake;
    }));
    const sock = await openSocket();
    for (let i = 0; i < MAX_PAIRING_INBOX_MESSAGES; i++) fake.simulateMessage(frame(1));
    expect(fake.closeCalls).toEqual([]);
    fake.simulateMessage(frame(1));
    expect(fake.closeCalls[0]?.code).toBe(1009);
    await expect(sock.receive()).rejects.toThrow(PAIRING_RESOURCE_LIMIT);
  });

  it("dequeue relieves the budget so later valid frames are accepted", async () => {
    vi.stubGlobal("WebSocket", vi.fn().mockImplementation(function () {
      fake = createFake();
      return fake;
    }));
    const sock = await openSocket();
    for (let i = 0; i < MAX_PAIRING_INBOX_MESSAGES; i++) fake.simulateMessage(frame(1));
    const first = await sock.receive();
    expect(first.byteLength).toBe(1);
    // Count and byte budget both dropped by one: this frame fits again.
    fake.simulateMessage(frame(1));
    expect(fake.closeCalls).toEqual([]);
    const rest: number[] = [];
    for (let i = 0; i < MAX_PAIRING_INBOX_MESSAGES; i++) rest.push((await sock.receive()).byteLength);
    expect(rest).toHaveLength(MAX_PAIRING_INBOX_MESSAGES);
  });

  it("gives every pending reader the same terminal error and ignores later events", async () => {
    vi.stubGlobal("WebSocket", vi.fn().mockImplementation(function () {
      fake = createFake();
      return fake;
    }));
    const sock = await openSocket();
    const r1 = sock.receive();
    const r2 = sock.receive();
    fake.simulateMessage(frame(MAX_PAIRING_FRAME_BYTES + 1));
    await expect(r1).rejects.toThrow(PAIRING_RESOURCE_LIMIT);
    await expect(r2).rejects.toThrow(PAIRING_RESOURCE_LIMIT);
    // Later frames and the resulting onclose are ignored: still one close.
    fake.simulateMessage(frame(1));
    fake.simulateClose(1009);
    await new Promise((r) => setTimeout(r, 0));
    expect(fake.closeCalls).toHaveLength(1);
    await expect(sock.receive()).rejects.toThrow(PAIRING_RESOURCE_LIMIT);
    expect(sock.closedBeforeFrame()).toBe(false);
  });

  it("refuses oversized outbound frames before sending", async () => {
    vi.stubGlobal("WebSocket", vi.fn().mockImplementation(function () {
      fake = createFake();
      return fake;
    }));
    const sock = await openSocket();
    expect(() => sock.send(new Uint8Array(frame(MAX_PAIRING_FRAME_BYTES + 1)))).toThrow(
      PAIRING_RESOURCE_LIMIT,
    );
    expect(fake.sent).toHaveLength(0);
    // The socket stays usable for honest frames.
    sock.send(new Uint8Array(frame(16)));
    expect(fake.sent).toHaveLength(1);
  });

  it("keeps normal pairing and wrong-code behavior unchanged", async () => {
    vi.stubGlobal("WebSocket", vi.fn().mockImplementation(function () {
      fake = createFake();
      return fake;
    }));
    const sock = await openSocket();
    fake.simulateMessage(frame(64));
    fake.simulateMessage(frame(128));
    expect((await sock.receive()).byteLength).toBe(64);
    expect((await sock.receive()).byteLength).toBe(128);
    // Wrong/expired code: silent close with nothing queued.
    const pending = sock.receive();
    const rejection = expect(pending).rejects.toThrow("socket closed");
    fake.simulateClose(1000);
    await rejection;
    expect(sock.closedBeforeFrame()).toBe(true);
  });
});