foxhunt/web-dashboard/src/lib/websocket.ts

import type { ServerMessage } from './types';

type MessageHandler = (msg: ServerMessage) => void;

/** WebSocket connection manager with auto-reconnect and topic subscriptions */
export class WsManager {
  private ws: WebSocket | null = null;
  private url: string;
  private handlers: Set<MessageHandler> = new Set();
  private reconnectTimer: ReturnType<typeof setTimeout> | null = null;
  private backoff = 1000;
  private maxBackoff = 30000;
  private maxAttempts = 50;
  private attempts = 0;
  private subscribedTopics: Set<string> = new Set();
  private _connected = false;

  constructor(url?: string) {
    const envWsUrl = import.meta.env.VITE_WS_URL as string | undefined;
    if (url) {
      this.url = url;
    } else if (envWsUrl) {
      this.url = envWsUrl;
    } else {
      const protocol = window.location.protocol === 'https:' ? 'wss:' : 'ws:';
      this.url = `${protocol}//${window.location.host}/api/ws`;
    }
  }

  get connected(): boolean {
    return this._connected;
  }

  /** Connect to the WebSocket server */
  connect(): void {
    const token = localStorage.getItem('foxhunt_token');
    if (!token) return;

    const wsUrl = `${this.url}?token=${encodeURIComponent(token)}`;
    this.ws = new WebSocket(wsUrl);

    this.ws.onopen = () => {
      this._connected = true;
      this.backoff = 1000;
      this.attempts = 0;
      // Re-subscribe to previously subscribed topics
      if (this.subscribedTopics.size > 0) {
        this.send({ type: 'subscribe', topics: Array.from(this.subscribedTopics) });
      }
    };

    this.ws.onmessage = (event) => {
      try {
        const msg = JSON.parse(event.data) as ServerMessage;
        this.handlers.forEach((h) => h(msg));
      } catch {
        // Ignore malformed messages
      }
    };

    this.ws.onclose = () => {
      this._connected = false;
      this.scheduleReconnect();
    };

    this.ws.onerror = () => {
      this.ws?.close();
    };
  }

  /** Disconnect and stop reconnecting */
  disconnect(): void {
    if (this.reconnectTimer) {
      clearTimeout(this.reconnectTimer);
      this.reconnectTimer = null;
    }
    this.ws?.close();
    this.ws = null;
    this._connected = false;
  }

  /** Subscribe to data topics */
  subscribe(topics: string[]): void {
    topics.forEach((t) => this.subscribedTopics.add(t));
    if (this._connected) {
      this.send({ type: 'subscribe', topics });
    }
  }

  /** Unsubscribe from data topics */
  unsubscribe(topics: string[]): void {
    topics.forEach((t) => this.subscribedTopics.delete(t));
    if (this._connected) {
      this.send({ type: 'unsubscribe', topics });
    }
  }

  /** Register a message handler */
  onMessage(handler: MessageHandler): () => void {
    this.handlers.add(handler);
    return () => this.handlers.delete(handler);
  }

  private send(msg: unknown): void {
    if (this.ws?.readyState === WebSocket.OPEN) {
      this.ws.send(JSON.stringify(msg));
    }
  }

  private scheduleReconnect(): void {
    if (this.reconnectTimer) return;
    this.attempts += 1;
    if (this.attempts > this.maxAttempts) {
      // Stop reconnecting after too many failures
      return;
    }
    const delay = this.backoff;
    this.backoff = Math.min(this.backoff * 2, this.maxBackoff);
    this.reconnectTimer = setTimeout(() => {
      this.reconnectTimer = null;
      this.connect();
    }, delay);
  }
}

/** Singleton WebSocket manager instance */
export const wsManager = new WsManager();