#Remix 3

// app/live/public/live.ts (trimmed: the example builds more on this client)
// Each client entry hydrates on its own, but they all import this module, so they share one
// client (and one socket). It imports `react-websocket/core`, which has no React dependency.
// The server imports it too, to render the client entries, but nothing connects there: the
// components only hold the connection once they're in the browser (see bindings.ts).
import { createWebSocket, createWebSocketState } from 'react-websocket/core';
import { routes } from '../../routes.ts';
import { serverMessage, type ClientMessage, type ServerMessage } from './protocol.ts';

// Called before every connection attempt, including reconnects, so each one gets a fresh token.
async function url() {
  const response = await fetch(routes.token.href());
  if (!response.ok) throw new Error(`Token request failed: HTTP ${response.status}`);
  const { token } = (await response.json()) as { token: string };
  const protocol = location.protocol === 'https:' ? 'wss:' : 'ws:';
  return `${protocol}//${location.host}/ws?token=${encodeURIComponent(token)}`;
}

export const live = createWebSocket<ServerMessage, ClientMessage>(url, {
  json: true,
  // Every message is validated; one that doesn't match is reported to `error` listeners.
  schema: serverMessage,
  // The demo server answers 'heartbeat' with 'heartbeat'. Three quiet seconds after a ping
  // count as a dead connection.
  heartbeat: { interval: 4000, timeout: 3000, message: 'heartbeat', pong: 'heartbeat' },
  // 4003 means the server refused this client: retrying won't help.
  reconnect: { maxDelay: 8000, shouldReconnect: close => close.code !== 4003 },
  queue: { limit: 50 },
  logger: console.debug
});

/**
 * The name the server gave this connection; a new connection gets a new name. Kept here rather
 * than in a component: client entries hydrate in no fixed order, so the `welcome` message may
 * arrive before the one that shows it is listening.
 */
export const me = createWebSocketState(live, {
  initialState: null as string | null,
  reduce: (state, message) => (message.type === 'welcome' ? message.name : state)
});

// app/live/public/bindings.ts (trimmed)
import type { Handle } from 'remix/component';
import type { EventMap, Status, WebSocketClient } from 'react-websocket/core';

type AnyHandle = Handle<any, any>;
type Client = WebSocketClient<any, any>;

/** Run `start` in the browser after the first render; its cleanups run when the component goes. */
function whileMounted(handle: AnyHandle, start: () => (() => void)[]) {
  handle.queueTask(() => {
    for (const stop of start()) handle.signal.addEventListener('abort', stop, { once: true });
  });
}

/**
 * Re-render when `read()` returns something new (by `Object.is`), however often the store
 * changes, and hold `client`'s connection. Like `useSyncExternalStore` with a selector.
 */
function select<T>(
  handle: AnyHandle,
  client: Client,
  subscribe: (listener: () => void) => () => void,
  read: () => T,
  initial: T
): () => T {
  let value = initial;
  const check = () => {
    const next = read();
    if (Object.is(next, value)) return;
    value = next;
    void handle.update();
  };
  whileMounted(handle, () => {
    const stops = [subscribe(check), client.acquire()];
    check(); // it may have changed since the server rendered `initial`
    return stops;
  });
  return () => value;
}

/** `useWebSocketStatus(client)`: re-renders on status changes only. */
export function watchStatus(handle: AnyHandle, client: Client) {
  return select<Status>(
    handle,
    client,
    client.subscribe,
    () => client.getSnapshot().status,
    'idle'
  );
}

/** `useWebSocketEvent(client, type, listener)`: listens without re-rendering. */
export function listen<In, K extends keyof EventMap<In>>(
  handle: AnyHandle,
  client: WebSocketClient<In, any>,
  type: K,
  listener: EventMap<In>[K]
) {
  whileMounted(handle, () => [client.on(type, listener), client.acquire()]);
}

// app/assets.ts (trimmed)
import { createAssetServer } from 'remix/assets';

// Compiles the browser modules on demand. The allow lists are the security boundary: only
// `public/` directories under app/ and the listed packages can reach the browser, so
// app/live/token.ts (which holds the token secret) can't.
export const assets = createAssetServer({
  basePath: '/assets',
  rootDir: process.cwd(),
  allowFiles: ['app/routes.ts', 'app/**/public/**'],
  // `react-websocket/core` and the schema it validates with. React isn't listed: nothing in the
  // browser imports it.
  allowPackages: ['remix', 'react-websocket', 'zod']
});

// server.ts
// One process serves the Remix app and the demo's WebSocket endpoint on the same port. The
// router turns Requests into Responses, and a WebSocket upgrade isn't one, so upgrades are
// handled on the Node server underneath.
import * as http from 'node:http';
import { createRequestListener } from 'remix/node-fetch-server';
import { router } from './app/router.ts';
import { handleUpgrade } from './demo-server.ts';

const port = Number(process.env.PORT ?? 44100);

const server = http.createServer(createRequestListener(router.fetch));
server.on('upgrade', (request, socket, head) => {
  if (!handleUpgrade(request, socket, head)) socket.destroy();
});
server.listen(port, () => {
  console.log(
    `> react-websocket Remix demo on http://localhost:${port} (${process.env.NODE_ENV ?? 'development'})`
  );
});