import { createORPCClient } from "@orpc/client";
import { ClientRetryPlugin } from "@orpc/client/plugins";
import { RPCLink } from "@orpc/client/websocket";
import { os } from "@orpc/server";
import { RPCHandler } from "@orpc/server/bun-ws";
import { WebSocket } from "partysocket";

const unhandled: string[] = [];
process.on("unhandledRejection", (e: any) => unhandled.push(`${e?.name}: ${e?.message}`));

let hang = false;
const handler = new RPCHandler({
  subscribe: os.handler(async function* () {
    while (true) {
      await Bun.sleep(300);
      yield "tick";
    }
  }),
  updateToken: os.handler(async () => {
    if (hang) await new Promise(() => {});
  }),
});
const sockets = new Set<any>();
const server = Bun.serve({
  port: 0,
  fetch: (req, srv) => (srv.upgrade(req) ? undefined : new Response(null, { status: 400 })),
  websocket: {
    open: (ws) => void sockets.add(ws),
    message: (ws, msg) => void handler.message(ws, msg, { context: {} }),
    close: (ws) => {
      sockets.delete(ws);
      handler.close(ws);
    },
  },
});

function connect() {
  let attempts = 0;
  const websocket = new WebSocket(`ws://localhost:${server.port}`);
  const client: any = createORPCClient(
    new RPCLink({
      websocket,
      clientInterceptors: [async (o) => (attempts++, o.next())],
      plugins: [new ClientRetryPlugin()],
    }),
  );
  return { websocket, client, attempts: () => attempts };
}

// 1a. Refresh with the socket closed, as the plugin does it.
{
  const { websocket, client } = connect();
  await client.updateToken({ token: "a" });
  websocket.close();
  setTimeout(async () => void (await client.updateToken({ token: "b" })), 0);
  await Bun.sleep(200);
  console.log("1a", unhandled.splice(0)); // [ "Error: Cannot send message, WebSocket is not open." ]
}

// 1b. The connection drops while the update is in flight.
{
  const { client } = connect();
  await client.updateToken({ token: "a" });
  hang = true;
  setTimeout(async () => void (await client.updateToken({ token: "b" })), 0);
  await Bun.sleep(100);
  for (const ws of sockets) ws.terminate();
  await Bun.sleep(300);
  hang = false;
  console.log("1b", unhandled.splice(0)); // [ "AbortError: [AsyncIdQueue] Queue[2] was closed or aborted while waiting for pulling." ]
}

// 2. Subscribe with infinite retry (as the plugin does), then close the socket (closeConnection()).
{
  const { websocket, client, attempts } = connect();
  const iterator = await client.subscribe(undefined, { context: { retry: Number.POSITIVE_INFINITY } });
  const consumer = (async () => {
    for await (const _ of iterator);
    return "settled";
  })();
  await Bun.sleep(1000);
  const before = attempts();
  websocket.close();
  await Bun.sleep(9000);
  console.log("2", attempts() - before, "new attempts in 9 s;", await Promise.race([consumer, Bun.sleep(20).then(() => "pending")]));
  // 4 new attempts in 9 s; pending
}

server.stop(true);
process.exit(0);