WRNexusJS/packages/core/src/realtime.ts

/**
 * Realtime rooms.
 *
 * A file in `app/realtime/` exports `default defineRoom({ onConnect, onMessage,
 * onLeave })` and is served at `ws://host/realtime/<name>`. The framework's
 * client runtime (`/__wrnexus/realtime.js`) handles the browser side, so pages
 * ship NO hand-written WebSocket code.
 *
 * Handlers get a `RoomClient` with everything you need:
 *   client.send(msg)                 → this connection
 *   client.broadcast(msg)            → everyone else in the room
 *   client.room.broadcast(msg)       → everyone (incl. sender)
 *   client.to(id | ids).send(msg)    → specific connection(s)
 *   client.toUser(u | users).send()  → a user / selected users (all their tabs)
 *   client.user = "u1"               → identify a connection for targeting
 *   client.data / client.room.state  → per-connection / shared room state
 *
 * The dynamic route `app/realtime/[room].ts` gives one handler many independent
 * rooms — `/realtime/lobby` and `/realtime/game-7` are separate room instances.
 */

// --- Low-level socket the registry drives (a subset of Bun's ServerWebSocket) ---

export interface RawSocket {
  send(data: string): unknown;
  close(code?: number, reason?: string): void;
}

// --- Legacy raw handler (still supported alongside defineRoom) ---

export interface RealtimeSocket<Data = unknown> {
  readonly data: Data;
  send(data: string | Uint8Array): number;
  subscribe(topic: string): void;
  unsubscribe(topic: string): void;
  publish(topic: string, data: string | Uint8Array): number;
  isSubscribed(topic: string): boolean;
  close(code?: number, reason?: string): void;
}

export interface RealtimeHandler<Data = unknown> {
  open?(ws: RealtimeSocket<Data>): void | Promise<void>;
  message?(ws: RealtimeSocket<Data>, message: string | Uint8Array): void | Promise<void>;
  close?(ws: RealtimeSocket<Data>, code?: number, reason?: string): void | Promise<void>;
  drain?(ws: RealtimeSocket<Data>): void | Promise<void>;
}

// --- Room API ---

export interface Target {
  /** Send a message (objects are JSON-serialized). */
  send(message: unknown): void;
}

export interface Room<TData = Record<string, unknown>> {
  readonly name: string;
  /** Shared, in-memory room state (lives while ≥1 client is connected). */
  readonly state: Record<string, unknown>;
  /** All connected clients. */
  clients(): RoomClient<TData>[];
  /** Number of connected clients. */
  count(): number;
  /** Send to everyone in the room, including the sender. */
  broadcast(message: unknown): void;
  /** Target specific connection id(s). */
  to(id: string | string[]): Target;
  /** Target a user / users by identity (reaches all their connections). */
  toUser(user: string | string[]): Target;
}

export interface RoomClient<TData = Record<string, unknown>> {
  /** Unique per connection (a tab). */
  readonly id: string;
  /** App identity for targeting; assign it in `onConnect`. */
  user: string | undefined;
  /** Query params from the connection URL. */
  readonly query: Record<string, string>;
  /** Per-connection scratch state. */
  readonly data: TData;
  readonly room: Room<TData>;
  /** Send to THIS connection. */
  send(message: unknown): void;
  /** Send to everyone else in the room. */
  broadcast(message: unknown): void;
  /** Target specific connection id(s). */
  to(id: string | string[]): Target;
  /** Target a user / users by identity. */
  toUser(user: string | string[]): Target;
  /** Close this connection. */
  close(code?: number, reason?: string): void;
}

/** Info available when authorizing a connection, before it is accepted. */
export interface RoomAuthInfo {
  /** Authenticated session user id, or `?user=` — undefined when anonymous. */
  user?: string;
  /** Connection URL query params. */
  query: Record<string, string>;
  /** The upgrade request's headers (cookies, etc.). */
  headers: Headers;
}

export interface RealtimeSecurityOptions {
  /** Maximum inbound or outbound serialized message size. Defaults to 64 KiB. */
  maxMessageBytes?: number;
  /** Maximum messages accepted per connection per rolling second. Defaults to 30. */
  maxMessagesPerSecond?: number;
  /** Maximum live connections in one room. Defaults to 1,000. */
  maxConnectionsPerRoom?: number;
  /** Maximum connections for one authenticated user in a room. Defaults to 10. */
  maxConnectionsPerUser?: number;
  /** Reject anonymous connections before onConnect. */
  requireUser?: boolean;
  /** Maximum nested JSON depth. Defaults to 32. */
  maxJsonDepth?: number;
  /** Optional message schema/authorization predicate. */
  validateMessage?(message: unknown, client: RoomClient): boolean | Promise<boolean>;
  /** Called when a connection is rejected or closed for a policy violation. */
  onViolation?(reason: string, client?: RoomClient): void;
}

export interface RoomHandlers<TData = Record<string, unknown>, TMessage = any> {
  /** Per-room abuse and payload controls. */
  security?: RealtimeSecurityOptions;
  /**
   * Gate the connection BEFORE it is accepted. Return false to reject the
   * upgrade with 403 (e.g. `authorize: (info) => !!info.user` to require auth).
   */
  authorize?(info: RoomAuthInfo): boolean | Promise<boolean>;
  /** A client connected (a new tab joined the room). */
  onConnect?(client: RoomClient<TData>): void | Promise<void>;
  /** A message arrived (JSON is parsed; non-JSON arrives as a string). */
  onMessage?(client: RoomClient<TData>, message: TMessage): void | Promise<void>;
  /** A client disconnected. */
  onLeave?(client: RoomClient<TData>): void | Promise<void>;
}

export interface RoomDefinition<TData = Record<string, unknown>, TMessage = any> {
  readonly __wrnexusRoom: true;
  readonly handlers: RoomHandlers<TData, TMessage>;
}

/** Define a realtime room. Export the result as the `default` of a realtime file. */
export function defineRoom<TData = Record<string, unknown>, TMessage = any>(
  handlers: RoomHandlers<TData, TMessage>,
): RoomDefinition<TData, TMessage> {
  return { __wrnexusRoom: true, handlers };
}

export function isRoomDefinition(value: unknown): value is RoomDefinition {
  return (
    !!value &&
    typeof value === "object" &&
    (value as { __wrnexusRoom?: unknown }).__wrnexusRoom === true
  );
}

// --- Registry (server-side connection manager) ---

interface Conn {
  id: string;
  messageWindowStartedAt: number;
  messageCount: number;
  user?: string;
  data: Record<string, unknown>;
  query: Record<string, string>;
  socket: RawSocket;
  roomName: string;
  client: RoomClient;
}

interface RoomImpl {
  name: string;
  state: Record<string, unknown>;
  def: RoomDefinition;
  conns: Map<string, Conn>;
  users: Map<string, Set<string>>; // user identity → connection ids
}

export interface RealtimeConnectMeta {
  room: string;
  def: RoomDefinition;
  query?: Record<string, string>;
  user?: string;
}

/** One cross-instance message: a room broadcast, or a targeted user send. */
export interface RealtimeEnvelope {
  room: string;
  /** If set, deliver only to these user identities; otherwise the whole room. */
  users?: string[];
  message: unknown;
}

/**
 * A pub/sub bridge for horizontal scaling. Wire the registry to a shared bus
 * (Redis pub/sub, NATS, …): local broadcasts/`toUser` sends are published to
 * peers, and messages received from peers are delivered via `registry.deliver`.
 * Connection-targeted sends (`send`, `to(id)`) stay local (ids are per-process).
 */
export interface RealtimeBridge {
  publish(envelope: RealtimeEnvelope): void;
}

export interface RealtimeRegistryOptions extends RealtimeSecurityOptions {
  now?: () => number;
}

export interface RealtimeRegistry {
  open(socket: RawSocket, meta: RealtimeConnectMeta): void | Promise<void>;
  message(socket: RawSocket, raw: string | Uint8Array): void | Promise<void>;
  close(socket: RawSocket): void | Promise<void>;
  /** Attach a cross-instance bridge (call once at startup). */
  setBridge(bridge: RealtimeBridge): void;
  /** Deliver an envelope received from a peer to LOCAL connections only. */
  deliver(envelope: RealtimeEnvelope): void;
  /** Number of live connections (across all rooms) — for tests/metrics. */
  size(): number;
}

const DANGEROUS_REALTIME_KEYS = new Set(["__proto__", "prototype", "constructor"]);

function assertRealtimePayload(value: unknown, maxDepth: number, depth = 0): void {
  if (depth > maxDepth) throw new Error("Realtime payload nesting limit exceeded.");
  if (!value || typeof value !== "object") return;
  for (const [key, child] of Object.entries(value as Record<string, unknown>)) {
    if (DANGEROUS_REALTIME_KEYS.has(key)) throw new Error(`Dangerous realtime key '${key}'.`);
    assertRealtimePayload(child, maxDepth, depth + 1);
  }
}

function byteLength(value: string | Uint8Array): number {
  return typeof value === "string" ? new TextEncoder().encode(value).byteLength : value.byteLength;
}

function serialize(message: unknown): string {
  return typeof message === "string" ? message : JSON.stringify(message);
}

function mergedRealtimeSecurity(
  globalOptions: RealtimeRegistryOptions,
  room: RoomDefinition,
): Required<
  Pick<
    RealtimeSecurityOptions,
    | "maxMessageBytes"
    | "maxMessagesPerSecond"
    | "maxConnectionsPerRoom"
    | "maxConnectionsPerUser"
    | "requireUser"
    | "maxJsonDepth"
  >
> &
  RealtimeSecurityOptions {
  return {
    maxMessageBytes: 64 * 1024,
    maxMessagesPerSecond: 30,
    maxConnectionsPerRoom: 1_000,
    maxConnectionsPerUser: 10,
    requireUser: false,
    maxJsonDepth: 32,
    ...globalOptions,
    ...room.handlers.security,
  };
}

/** Create the registry that maps sockets ↔ rooms and drives room handlers. */
export function createRealtimeRegistry(options: RealtimeRegistryOptions = {}): RealtimeRegistry {
  const rooms = new Map<string, RoomImpl>();
  const bySocket = new Map<RawSocket, Conn>();
  let bridge: RealtimeBridge | null = null;
  let applyingRemote = false; // true while delivering a peer envelope (no re-publish)
  const now = options.now ?? Date.now;

  const publish = (envelope: RealtimeEnvelope): void => {
    if (bridge && !applyingRemote) bridge.publish(envelope);
  };

  const send = (conn: Conn | undefined, payload: string): void => {
    if (!conn) return;
    const room = rooms.get(conn.roomName);
    const policy = room
      ? mergedRealtimeSecurity(options, room.def)
      : ({ maxMessageBytes: options.maxMessageBytes ?? 64 * 1024 } as ReturnType<
          typeof mergedRealtimeSecurity
        >);
    if (byteLength(payload) > policy.maxMessageBytes) {
      policy.onViolation?.("outbound-message-too-large", conn.client);
      conn.socket.close(1009, "Message too large");
      return;
    }
    try {
      conn.socket.send(payload);
    } catch {
      /* socket already gone */
    }
  };

  const reindexUser = (room: RoomImpl, conn: Conn, next: string | undefined): void => {
    if (conn.user === next) return;
    if (conn.user) {
      const set = room.users.get(conn.user);
      if (set) {
        set.delete(conn.id);
        if (!set.size) room.users.delete(conn.user);
      }
    }
    conn.user = next;
    if (next) {
      let set = room.users.get(next);
      if (!set) room.users.set(next, (set = new Set()));
      set.add(conn.id);
    }
  };

  const idsForUsers = (room: RoomImpl, user: string | string[]): string[] => {
    const out: string[] = [];
    for (const u of Array.isArray(user) ? user : [user]) {
      const set = room.users.get(u);
      if (set) out.push(...set);
    }
    return out;
  };

  const makeRoomApi = (room: RoomImpl): Room => ({
    name: room.name,
    state: room.state,
    clients: () => Array.from(room.conns.values(), (c) => c.client),
    count: () => room.conns.size,
    broadcast: (message) => {
      const payload = serialize(message);
      for (const c of room.conns.values()) send(c, payload);
      publish({ room: room.name, message });
    },
    to: (id) => ({
      send: (message) => {
        // Connection-targeted: local only (ids are per-process).
        const payload = serialize(message);
        for (const cid of Array.isArray(id) ? id : [id]) send(room.conns.get(cid), payload);
      },
    }),
    toUser: (user) => ({
      send: (message) => {
        const payload = serialize(message);
        for (const cid of idsForUsers(room, user)) send(room.conns.get(cid), payload);
        publish({ room: room.name, users: Array.isArray(user) ? user : [user], message });
      },
    }),
  });

  const makeClientApi = (room: RoomImpl, conn: Conn): RoomClient => {
    const roomApi = makeRoomApi(room);
    return {
      id: conn.id,
      get user() {
        return conn.user;
      },
      set user(value: string | undefined) {
        reindexUser(room, conn, value);
      },
      query: conn.query,
      data: conn.data,
      room: roomApi,
      send: (message) => send(conn, serialize(message)),
      broadcast: (message) => {
        const payload = serialize(message);
        for (const c of room.conns.values()) if (c.id !== conn.id) send(c, payload);
        // Peers deliver to all their conns (all "others" relative to this one).
        publish({ room: room.name, message });
      },
      to: roomApi.to,
      toUser: roomApi.toUser,
      close: (code, reason) => conn.socket.close(code, reason),
    };
  };

  return {
    async open(socket, meta) {
      const policy = mergedRealtimeSecurity(options, meta.def);
      if (policy.requireUser && !meta.user) {
        policy.onViolation?.("authentication-required");
        socket.close(1008, "Authentication required");
        return;
      }
      let room = rooms.get(meta.room);
      if (!room) {
        room = { name: meta.room, state: {}, def: meta.def, conns: new Map(), users: new Map() };
        rooms.set(meta.room, room);
      }
      if (room.conns.size >= policy.maxConnectionsPerRoom) {
        policy.onViolation?.("room-connection-limit");
        socket.close(1013, "Room is at capacity");
        return;
      }
      if (meta.user && (room.users.get(meta.user)?.size ?? 0) >= policy.maxConnectionsPerUser) {
        policy.onViolation?.("user-connection-limit");
        socket.close(1008, "Too many connections");
        return;
      }
      const conn: Conn = {
        id: randomId(),
        messageWindowStartedAt: now(),
        messageCount: 0,
        data: {},
        query: meta.query ?? {},
        socket,
        roomName: meta.room,
        client: null as unknown as RoomClient,
      };
      conn.client = makeClientApi(room, conn);
      room.conns.set(conn.id, conn);
      bySocket.set(socket, conn);
      if (meta.user) reindexUser(room, conn, meta.user);
      await room.def.handlers.onConnect?.(conn.client);
    },

    async message(socket, raw) {
      const conn = bySocket.get(socket);
      if (!conn) return;
      const room = rooms.get(conn.roomName);
      if (!room) return;
      const policy = mergedRealtimeSecurity(options, room.def);
      if (byteLength(raw) > policy.maxMessageBytes) {
        policy.onViolation?.("inbound-message-too-large", conn.client);
        socket.close(1009, "Message too large");
        return;
      }
      const timestamp = now();
      if (timestamp - conn.messageWindowStartedAt >= 1_000) {
        conn.messageWindowStartedAt = timestamp;
        conn.messageCount = 0;
      }
      conn.messageCount += 1;
      if (conn.messageCount > policy.maxMessagesPerSecond) {
        policy.onViolation?.("message-rate-limit", conn.client);
        socket.close(1008, "Message rate exceeded");
        return;
      }
      const text = typeof raw === "string" ? raw : new TextDecoder().decode(raw);
      let message: unknown;
      try {
        message = JSON.parse(text);
      } catch {
        message = text;
      }
      try {
        assertRealtimePayload(message, policy.maxJsonDepth);
      } catch {
        policy.onViolation?.("invalid-message-shape", conn.client);
        socket.close(1008, "Invalid message");
        return;
      }
      if (policy.validateMessage && !(await policy.validateMessage(message, conn.client))) {
        policy.onViolation?.("message-validation-failed", conn.client);
        socket.close(1008, "Message rejected");
        return;
      }
      await room.def.handlers.onMessage?.(conn.client, message);
    },

    async close(socket) {
      const conn = bySocket.get(socket);
      if (!conn) return;
      bySocket.delete(socket);
      const room = rooms.get(conn.roomName);
      if (!room) return;
      try {
        await room.def.handlers.onLeave?.(conn.client);
      } finally {
        room.conns.delete(conn.id);
        reindexUser(room, conn, undefined);
        if (room.conns.size === 0) rooms.delete(room.name);
      }
    },

    setBridge(b) {
      bridge = b;
    },

    deliver(envelope) {
      const room = rooms.get(envelope.room);
      if (!room) return;
      applyingRemote = true; // suppress re-publishing what we received
      try {
        const payload = serialize(envelope.message);
        if (envelope.users) {
          for (const cid of idsForUsers(room, envelope.users)) send(room.conns.get(cid), payload);
        } else {
          for (const c of room.conns.values()) send(c, payload);
        }
      } finally {
        applyingRemote = false;
      }
    },

    size: () => bySocket.size,
  };
}

/**
 * A minimal pub/sub bus (structurally satisfied by `@wrnexus/pubsub`). Used to
 * bridge realtime broadcasts across processes without a hard dependency.
 */
export interface RealtimeBus {
  publish(topic: string, message: unknown): void | Promise<void>;
  subscribe(topic: string, handler: (message: unknown, topic: string) => void): () => void;
}

/**
 * Bridge a realtime registry across processes/instances via a pub/sub bus (use
 * the Redis driver so it crosses machines). After this, `client.room.broadcast`
 * and `client.toUser(...)` reach connected clients on **every** app process/
 * instance subscribed to the same bus — the foundation for realtime that works
 * with multiple running apps behind the gateway. Connection-targeted sends
 * (`send`, `to(id)`) stay local. Returns an unsubscribe function.
 *
 *   import { createRealtimeRegistry, bridgeRealtime } from "@wrnexus/core";
 *   import { createPubSub } from "@wrnexus/pubsub";
 *   import { redisDriver } from "@wrnexus/pubsub/redis";
 *   bridgeRealtime(registry, createPubSub(redisDriver(process.env.REDIS_URL)));
 */
export function bridgeRealtime(
  registry: RealtimeRegistry,
  bus: RealtimeBus,
  topic = "wrnexus:realtime",
): () => void {
  registry.setBridge({ publish: (envelope) => void bus.publish(topic, envelope) });
  return bus.subscribe(topic, (message) => registry.deliver(message as RealtimeEnvelope));
}

function randomId(): string {
  const bytes = new Uint8Array(12);
  crypto.getRandomValues(bytes);
  let out = "";
  for (const b of bytes) out += b.toString(16).padStart(2, "0");
  return out;
}