import { afterEach, describe, expect, test } from "vitest";

import {
  addListenerTyped,
  addRequestListenerTyped,
  parserFromSchema,
  sendRequestTyped,
} from "../src/communicator.js";
import {
  create,
  fromBinary,
  PacketBaseSchema,
  TestDataSchema,
  toBinary,
  type TestData,
} from "../src/packets/message_common_pb.js";
import { BrowserWsClient, connectBrowserWs } from "../src/browser_ws_client.js";

type Listener = (event: { data?: unknown; code?: number; reason?: string }) => void;

class FakeWebSocket {
  static instances: FakeWebSocket[] = [];

  static readonly CONNECTING = 0 as const;
  static readonly OPEN = 1 as const;
  static readonly CLOSING = 2 as const;
  static readonly CLOSED = 3 as const;

  readonly CONNECTING = FakeWebSocket.CONNECTING;
  readonly OPEN = FakeWebSocket.OPEN;
  readonly CLOSING = FakeWebSocket.CLOSING;
  readonly CLOSED = FakeWebSocket.CLOSED;

  url: string;
  binaryType = "blob";
  readyState: number = FakeWebSocket.CONNECTING;
  sentFrames: Uint8Array[] = [];
  sendShouldThrow: Error | null = null;
  closedWith: { code: number; reason: string } | null = null;

  private readonly listeners = new Map<string, Set<Listener>>();

  constructor(url: string) {
    this.url = url;
    FakeWebSocket.instances.push(this);
  }

  addEventListener(type: string, fn: Listener): void {
    if (!this.listeners.has(type)) {
      this.listeners.set(type, new Set());
    }
    this.listeners.get(type)!.add(fn);
  }
  removeEventListener(type: string, fn: Listener): void {
    this.listeners.get(type)?.delete(fn);
  }

  send(data: ArrayBuffer | Uint8Array): void {
    if (this.sendShouldThrow !== null) {
      throw this.sendShouldThrow;
    }
    const buf =
      data instanceof Uint8Array ? new Uint8Array(data) : new Uint8Array(data as ArrayBuffer);
    this.sentFrames.push(buf);
  }

  close(code = 1000, reason = ""): void {
    if (this.readyState === FakeWebSocket.CLOSED) {
      return;
    }
    this.closedWith = { code, reason };
    this.readyState = FakeWebSocket.CLOSED;
    this.dispatch("close", { code, reason });
  }

  fireOpen(): void {
    this.readyState = FakeWebSocket.OPEN;
    this.dispatch("open", {});
  }
  fireMessage(data: unknown): void {
    this.dispatch("message", { data });
  }
  fireError(): void {
    this.dispatch("error", {});
  }

  private dispatch(type: string, event: { data?: unknown; code?: number; reason?: string }): void {
    const set = this.listeners.get(type);
    if (set === undefined) {
      return;
    }
    for (const fn of [...set]) {
      fn(event);
    }
  }
}

function parserMap() {
  return new Map([[TestDataSchema.typeName, parserFromSchema(TestDataSchema)]]);
}

function decodeFrame(buf: Uint8Array): { typeName: string; nonce: number; responseNonce: number; payload: Uint8Array } {
  const base = fromBinary(PacketBaseSchema, buf);
  return {
    typeName: base.typeName,
    nonce: base.nonce,
    responseNonce: base.responseNonce,
    payload: base.data,
  };
}

function encodeTestData(msg: TestData, nonce: number, responseNonce = 0): Uint8Array {
  return toBinary(
    PacketBaseSchema,
    create(PacketBaseSchema, {
      typeName: TestDataSchema.typeName,
      nonce,
      responseNonce,
      data: toBinary(TestDataSchema, msg),
    }),
  );
}

describe("BrowserWsClient", () => {
  const cleanup: Array<() => Promise<void>> = [];

  afterEach(async () => {
    FakeWebSocket.instances.length = 0;
    while (cleanup.length > 0) {
      const fn = cleanup.pop();
      if (fn) {
        await fn();
      }
    }
  });

  test("connectBrowserWs resolves after open and sends binary frames", async () => {
    const pending = connectBrowserWs("ws://test/", 0, 0, parserMap(), {
      WebSocketCtor: FakeWebSocket as unknown as typeof WebSocket,
    });
    expect(FakeWebSocket.instances).toHaveLength(1);
    const ws = FakeWebSocket.instances[0]!;
    expect(ws.binaryType).toBe("arraybuffer");
    ws.fireOpen();
    const client = await pending;
    cleanup.push(async () => client.close());
    expect(client).toBeInstanceOf(BrowserWsClient);

    await client.communicator.send(create(TestDataSchema, { index: 11, message: "hi" }));
    expect(ws.sentFrames).toHaveLength(1);
    const decoded = decodeFrame(ws.sentFrames[0]!);
    expect(decoded.typeName).toBe(TestDataSchema.typeName);
  });

  test("routes ArrayBuffer message to communicator listener", async () => {
    const pending = connectBrowserWs("ws://test/", 0, 0, parserMap(), {
      WebSocketCtor: FakeWebSocket as unknown as typeof WebSocket,
    });
    const ws = FakeWebSocket.instances[0]!;
    ws.fireOpen();
    const client = await pending;
    cleanup.push(async () => client.close());

    const received = new Promise<TestData>((resolve) => {
      addListenerTyped(client.communicator, TestDataSchema, resolve);
    });

    const frame = encodeTestData(create(TestDataSchema, { index: 7, message: "push" }), 1);
    ws.fireMessage(frame.buffer.slice(frame.byteOffset, frame.byteOffset + frame.byteLength));

    await expect(received).resolves.toMatchObject({ index: 7, message: "push" });
  });

  test("routes Blob message", async () => {
    const pending = connectBrowserWs("ws://test/", 0, 0, parserMap(), {
      WebSocketCtor: FakeWebSocket as unknown as typeof WebSocket,
    });
    const ws = FakeWebSocket.instances[0]!;
    ws.fireOpen();
    const client = await pending;
    cleanup.push(async () => client.close());

    const received = new Promise<TestData>((resolve) => {
      addListenerTyped(client.communicator, TestDataSchema, resolve);
    });

    const frame = encodeTestData(create(TestDataSchema, { index: 8, message: "blob" }), 1);
    ws.fireMessage(new Blob([new Uint8Array(frame).slice().buffer]));

    await expect(received).resolves.toMatchObject({ index: 8, message: "blob" });
  });

  test("sendRequest resolves on response frame", async () => {
    const pending = connectBrowserWs("ws://test/", 0, 0, parserMap(), {
      WebSocketCtor: FakeWebSocket as unknown as typeof WebSocket,
    });
    const ws = FakeWebSocket.instances[0]!;
    ws.fireOpen();
    const client = await pending;
    cleanup.push(async () => client.close());

    const promise = sendRequestTyped(
      client.communicator,
      create(TestDataSchema, { index: 21, message: "req" }),
      TestDataSchema,
      500,
    );

    // wait a microtask for queued send
    await new Promise<void>((r) => setTimeout(r, 0));
    expect(ws.sentFrames).toHaveLength(1);
    const sent = decodeFrame(ws.sentFrames[0]!);
    const responseFrame = encodeTestData(
      create(TestDataSchema, { index: 42, message: "echo: req" }),
      99,
      sent.nonce,
    );
    ws.fireMessage(
      responseFrame.buffer.slice(responseFrame.byteOffset, responseFrame.byteOffset + responseFrame.byteLength),
    );

    await expect(promise).resolves.toMatchObject({ index: 42, message: "echo: req" });
  });

  test("disconnects on string frame", async () => {
    const pending = connectBrowserWs("ws://test/", 0, 0, parserMap(), {
      WebSocketCtor: FakeWebSocket as unknown as typeof WebSocket,
    });
    const ws = FakeWebSocket.instances[0]!;
    ws.fireOpen();
    const client = await pending;
    cleanup.push(async () => client.close());

    expect(client.communicator.isAlive()).toBe(true);
    ws.fireMessage("not binary");
    await new Promise<void>((r) => setTimeout(r, 0));
    expect(client.communicator.isAlive()).toBe(false);
  });

  test("server registers request handlers via communicator", async () => {
    const pending = connectBrowserWs("ws://test/", 0, 0, parserMap(), {
      WebSocketCtor: FakeWebSocket as unknown as typeof WebSocket,
    });
    const ws = FakeWebSocket.instances[0]!;
    ws.fireOpen();
    const client = await pending;
    cleanup.push(async () => client.close());

    addRequestListenerTyped(client.communicator, TestDataSchema, (req) =>
      create(TestDataSchema, { index: req.index * 2, message: `echo: ${req.message}` }),
    );

    const reqFrame = encodeTestData(create(TestDataSchema, { index: 3, message: "ping" }), 77);
    ws.fireMessage(reqFrame.buffer.slice(reqFrame.byteOffset, reqFrame.byteOffset + reqFrame.byteLength));

    await new Promise<void>((r) => setTimeout(r, 5));
    expect(ws.sentFrames.length).toBeGreaterThanOrEqual(1);
    const responses = ws.sentFrames.map((f) => decodeFrame(f));
    const match = responses.find((r) => r.responseNonce === 77);
    expect(match).toBeDefined();
  });

  test("disconnects immediately when staging backlog exceeds 64", async () => {
    const pending = connectBrowserWs("ws://test/", 0, 0, parserMap(), {
      WebSocketCtor: FakeWebSocket as unknown as typeof WebSocket,
    });
    const ws = FakeWebSocket.instances[0]!;
    ws.fireOpen();
    const client = await pending;
    cleanup.push(async () => client.close());

    let resolveHandler: (() => void) | null = null;
    const handlerPromise = new Promise<void>((r) => {
      resolveHandler = r;
    });
    addRequestListenerTyped(client.communicator, TestDataSchema, async () => {
      await handlerPromise;
      return create(TestDataSchema, { index: 100, message: "done" });
    });

    const frame1 = encodeTestData(create(TestDataSchema, { index: 1, message: "block" }), 1);
    ws.fireMessage(frame1.buffer.slice(frame1.byteOffset, frame1.byteOffset + frame1.byteLength));

    for (let i = 2; i <= 64; i++) {
      const frame = encodeTestData(create(TestDataSchema, { index: i, message: "staging" }), i);
      ws.fireMessage(frame.buffer.slice(frame.byteOffset, frame.byteOffset + frame.byteLength));
    }

    const overflowFrame = encodeTestData(create(TestDataSchema, { index: 65, message: "overflow" }), 65);
    ws.fireMessage(overflowFrame.buffer.slice(overflowFrame.byteOffset, overflowFrame.byteOffset + overflowFrame.byteLength));

    await new Promise<void>((r) => setTimeout(r, 10));

    expect(client.communicator.isAlive()).toBe(false);
    expect(ws.closedWith).not.toBeNull();
    expect(ws.closedWith?.code).toBe(1008);
    expect(ws.closedWith?.reason).toBe("inbound staging buffer overflow");

    resolveHandler!();
  });
});
