lila/apps/web/src/lib/ws-client.ts

import { WsServerMessageSchema } from "@lila/shared";
import type { WsClientMessage, WsServerMessage } from "@lila/shared";

/**
 * Minimal WebSocket client for multiplayer communication.
 *
 * NOTE: Callbacks registered via `on()` are stored by reference.
 * When using in React components, wrap callbacks in `useCallback`
 * to ensure the same reference is passed to both `on()` and `off()`.
 */
export class WsClient {
  private ws: WebSocket | null = null;
  private callbacks = new Map<string, Set<(msg: WsServerMessage) => void>>();

  /**
   * Called when the WebSocket connection closes.
   * Set by WsProvider — do not set directly in components.
   */
  public onError: ((event: Event) => void) | null = null;

  /**
   * Called when the WebSocket connection encounters an error.
   * Set by WsProvider — do not set directly in components.
   */
  public onClose: ((event: CloseEvent) => void) | null = null;

  connect(apiUrl: string): Promise<void> {
    return new Promise((resolve, reject) => {
      if (
        this.ws &&
        (this.ws.readyState === WebSocket.OPEN ||
          this.ws.readyState === WebSocket.CONNECTING)
      ) {
        resolve();
        return;
      }

      let wsUrl: string;
      if (!apiUrl) {
        wsUrl = "/ws";
      } else {
        wsUrl =
          apiUrl
            .replace(/^https:\/\//, "wss://")
            .replace(/^http:\/\//, "ws://") + "/ws";
      }

      this.ws = new WebSocket(wsUrl);

      this.ws.onopen = () => {
        resolve();
      };

      this.ws.onmessage = (event: MessageEvent) => {
        let parsed: unknown;
        try {
          parsed = JSON.parse(event.data as string);
        } catch {
          console.error("WsClient: received invalid JSON", event.data);
          return;
        }

        const result = WsServerMessageSchema.safeParse(parsed);
        if (!result.success) {
          console.error("WsClient: received unknown message shape", parsed);
          return;
        }

        const msg = result.data;
        const handlers = this.callbacks.get(msg.type);
        if (!handlers) return;
        for (const handler of handlers) {
          handler(msg);
        }
      };

      this.ws.onerror = (event: Event) => {
        this.onError?.(event);
        reject(new Error("WebSocket connection failed"));
      };

      this.ws.onclose = (event: CloseEvent) => {
        this.ws = null;
        this.onClose?.(event);
      };
    });
  }

  disconnect(): void {
    if (!this.ws) return;
    this.ws.close();
    this.ws = null;
  }

  isConnected(): boolean {
    return this.ws !== null && this.ws.readyState === WebSocket.OPEN;
  }

  send(message: WsClientMessage): void {
    if (!this.isConnected()) {
      console.warn(
        "WsClient: attempted to send message while disconnected",
        message,
      );
      return;
    }
    this.ws!.send(JSON.stringify(message));
  }

  on<T extends WsServerMessage["type"]>(
    type: T,
    callback: (msg: Extract<WsServerMessage, { type: T }>) => void,
  ): void {
    if (!this.callbacks.has(type)) {
      this.callbacks.set(type, new Set());
    }
    this.callbacks.get(type)!.add(callback as (msg: WsServerMessage) => void);
  }

  off<T extends WsServerMessage["type"]>(
    type: T,
    callback: (msg: Extract<WsServerMessage, { type: T }>) => void,
  ): void {
    const handlers = this.callbacks.get(type);
    if (!handlers) return;
    handlers.delete(callback as (msg: WsServerMessage) => void);
    if (handlers.size === 0) {
      this.callbacks.delete(type);
    }
  }

  clearCallbacks(): void {
    this.callbacks.clear();
  }
}