rfcp/frontend/src/services/websocket.ts

/**
 * Singleton WebSocket service for real-time coverage calculation.
 *
 * Manages a persistent WebSocket connection and provides calculate/cancel
 * methods usable from Zustand stores (not just React components).
 *
 * Usage:
 *   import { wsService } from '@/services/websocket';
 *   wsService.connect();
 *   wsService.calculate(sites, settings, onResult, onError, onProgress);
 */

import { getApiBaseUrl } from '@/lib/desktop.ts';
import type { CoverageResponse } from '@/services/api.ts';

export interface WSProgress {
  phase: string;
  progress: number;
  eta_seconds?: number;
}

export interface WSPartialResults {
  points: Array<Record<string, unknown>>;
  tile: number;
  total_tiles: number;
  progress: number;
}

type ProgressCallback = (progress: WSProgress) => void;
type ResultCallback = (data: CoverageResponse) => void;
type ErrorCallback = (error: string) => void;
type PartialResultsCallback = (data: WSPartialResults) => void;
type ConnectionCallback = (connected: boolean) => void;

interface PendingCalc {
  onProgress?: ProgressCallback;
  onResult: ResultCallback;
  onError: ErrorCallback;
  onPartialResults?: PartialResultsCallback;
}

class WebSocketService {
  private ws: WebSocket | null = null;
  private reconnectTimer: ReturnType<typeof setTimeout> | undefined;
  private pingTimer: ReturnType<typeof setInterval> | undefined;
  private _connected = false;
  private _pendingCalcs = new Map<string, PendingCalc>();
  private _connectionListeners = new Set<ConnectionCallback>();

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

  /** Register a listener for connection state changes. Returns unsubscribe fn. */
  onConnectionChange(cb: ConnectionCallback): () => void {
    this._connectionListeners.add(cb);
    return () => this._connectionListeners.delete(cb);
  }

  private _setConnected(val: boolean): void {
    this._connected = val;
    for (const cb of this._connectionListeners) {
      try { cb(val); } catch { /* ignore */ }
    }
  }

  connect(): void {
    if (this.ws && this.ws.readyState === WebSocket.OPEN) return;
    if (this.ws && this.ws.readyState === WebSocket.CONNECTING) return;

    const apiBase = getApiBaseUrl();
    const url = apiBase.replace(/\/api\/?$/, '').replace(/^http/, 'ws') + '/ws';

    try {
      this.ws = new WebSocket(url);
    } catch {
      this.reconnectTimer = setTimeout(() => this.connect(), 3000);
      return;
    }

    this.ws.onopen = () => {
      this._setConnected(true);
      // Keepalive pings every 30s to prevent connection timeout during long calculations
      if (this.pingTimer) clearInterval(this.pingTimer);
      this.pingTimer = setInterval(() => {
        if (this.ws?.readyState === WebSocket.OPEN) {
          this.ws.send(JSON.stringify({ type: 'ping' }));
        }
      }, 30_000);
    };

    this.ws.onclose = () => {
      this._setConnected(false);
      if (this.pingTimer) { clearInterval(this.pingTimer); this.pingTimer = undefined; }
      // Fail all pending calculations — their callbacks reference the old socket
      this._failPendingCalcs('WebSocket disconnected');
      this.reconnectTimer = setTimeout(() => this.connect(), 2000);
    };

    this.ws.onerror = () => {
      // onclose will handle reconnect
    };

    this.ws.onmessage = (event) => {
      try {
        const msg = JSON.parse(event.data);
        const calcId: string | undefined = msg.calculation_id;
        const pending = calcId ? this._pendingCalcs.get(calcId) : undefined;

        switch (msg.type) {
          case 'progress':
            if (pending?.onProgress) {
              pending.onProgress({
                phase: msg.phase,
                progress: msg.progress,
                eta_seconds: msg.eta_seconds,
              });
            } else {
              console.warn('[WS] progress msg but no pending calc:', calcId, msg.phase, msg.progress);
            }
            break;
          case 'partial_results':
            if (pending?.onPartialResults) {
              pending.onPartialResults({
                points: msg.points,
                tile: msg.tile,
                total_tiles: msg.total_tiles,
                progress: msg.progress,
              });
            }
            break;
          case 'result':
            if (pending) {
              pending.onResult(msg.data);
            } else {
              console.warn('[WS] result msg but no pending calc:', calcId);
            }
            if (calcId) this._pendingCalcs.delete(calcId);
            break;
          case 'error':
            console.error('[WS] error:', msg.message);
            if (pending) {
              pending.onError(msg.message);
            }
            if (calcId) this._pendingCalcs.delete(calcId);
            break;
        }
      } catch {
        // Ignore parse errors
      }
    };
  }

  /** Fail all pending calculations (e.g. on disconnect). */
  private _failPendingCalcs(reason: string): void {
    for (const [calcId, pending] of this._pendingCalcs) {
      try { pending.onError(reason); } catch { /* ignore */ }
      this._pendingCalcs.delete(calcId);
    }
  }

  disconnect(): void {
    if (this.reconnectTimer) clearTimeout(this.reconnectTimer);
    if (this.pingTimer) { clearInterval(this.pingTimer); this.pingTimer = undefined; }
    this._failPendingCalcs('WebSocket disconnected');
    this.ws?.close();
    this.ws = null;
    this._setConnected(false);
  }

  /**
   * Send a coverage calculation request via WebSocket.
   * Returns the calculation ID, or undefined if WS is not connected.
   */
  calculate(
    sites: Array<Record<string, unknown>>,
    settings: Record<string, unknown>,
    onResult: ResultCallback,
    onError: ErrorCallback,
    onProgress?: ProgressCallback,
    onPartialResults?: PartialResultsCallback,
  ): string | undefined {
    if (!this.ws || this.ws.readyState !== WebSocket.OPEN) {
      return undefined;
    }

    const calcId = crypto.randomUUID();
    this._pendingCalcs.set(calcId, { onProgress, onResult, onError, onPartialResults });

    this.ws.send(JSON.stringify({
      type: 'calculate',
      id: calcId,
      sites,
      settings,
    }));

    return calcId;
  }

  /** Cancel a running calculation by ID. */
  cancel(calcId: string): void {
    this.ws?.send(JSON.stringify({ type: 'cancel', id: calcId }));
    this._pendingCalcs.delete(calcId);
  }
}

/** Singleton WebSocket service. */
export const wsService = new WebSocketService();