KSP-MissionControl/apps/live-map/src/hooks/useLiveState.ts

/**
 * useLiveState — subscribe to the kerbal-rt API over WebSocket and
 * expose the latest UniverseSnapshot plus connection status.
 *
 * Auto-reconnects with exponential backoff on drop. Falls back to
 * polling /api/v1/state every 5s if WebSocket fails to connect after
 * 3 retries (e.g. a proxy doesn't support upgrade).
 */
import { useEffect, useRef, useState } from 'react';
import type { LiveMessage, UniverseSnapshot } from '@kerbal-rt/shared-types';

export type ConnectionState = 'connecting' | 'open' | 'closed' | 'fallback';

export interface LiveState {
  snapshot: UniverseSnapshot | null;
  status: ConnectionState;
  lastUpdate: string | null;
  messageCount: number;
  error: string | null;
}

const RECONNECT_DELAYS_MS = [1000, 2000, 5000, 10_000, 30_000];

export function useLiveState(apiUrl: string): LiveState {
  const [snapshot, setSnapshot] = useState<UniverseSnapshot | null>(null);
  const [status, setStatus] = useState<ConnectionState>('connecting');
  const [lastUpdate, setLastUpdate] = useState<string | null>(null);
  const [messageCount, setMessageCount] = useState(0);
  const [error, setError] = useState<string | null>(null);
  const reconnectAttempt = useRef(0);
  const stopped = useRef(false);

  useEffect(() => {
    stopped.current = false;
    let pollTimer: number | null = null;
    let reconnectTimer: number | null = null;
    let ws: WebSocket | null = null;

    const startPollingFallback = () => {
      if (pollTimer) return;
      setStatus('fallback');
      const tick = async () => {
        try {
          const res = await fetch(`${apiUrl}/api/v1/state`);
          if (res.ok) {
            const body = await res.json();
            if (!body.error && body.data) {
              setSnapshot(body.data as UniverseSnapshot);
              setLastUpdate(new Date().toISOString());
              setMessageCount((n) => n + 1);
            }
          }
        } catch (e) {
          setError(String((e as Error).message ?? e));
        }
      };
      void tick();
      pollTimer = window.setInterval(tick, 5000);
    };

    const scheduleReconnect = () => {
      if (stopped.current) return;
      if (reconnectAttempt.current >= RECONNECT_DELAYS_MS.length) {
        // Gave up on WS, fall back to polling.
        startPollingFallback();
        return;
      }
      const delay = RECONNECT_DELAYS_MS[reconnectAttempt.current] ?? 30_000;
      reconnectAttempt.current += 1;
      setStatus('closed');
      reconnectTimer = window.setTimeout(connect, delay);
    };

    const connect = () => {
      if (stopped.current) return;
      setStatus('connecting');
      setError(null);
      const wsUrl = apiUrl.replace(/^http/, 'ws') + '/api/v1/live';
      try {
        ws = new WebSocket(wsUrl);
      } catch (e) {
        setError(String((e as Error).message ?? e));
        scheduleReconnect();
        return;
      }
      ws.onopen = () => {
        setStatus('open');
        reconnectAttempt.current = 0; // success — reset backoff
      };
      ws.onmessage = (event) => {
        setMessageCount((n) => n + 1);
        setLastUpdate(new Date().toISOString());
        try {
          const msg = JSON.parse(event.data) as LiveMessage;
          if (msg.type === 'snapshot') {
            setSnapshot(msg.snapshot);
          }
        } catch (e) {
          setError(`bad message: ${String((e as Error).message ?? e)}`);
        }
      };
      ws.onerror = () => {
        // close will fire too
      };
      ws.onclose = () => {
        scheduleReconnect();
      };
    };

    connect();
    return () => {
      stopped.current = true;
      if (reconnectTimer) window.clearTimeout(reconnectTimer);
      if (pollTimer) window.clearInterval(pollTimer);
      if (ws) {
        ws.onclose = null;
        ws.close();
      }
    };
  }, [apiUrl]);

  return { snapshot, status, lastUpdate, messageCount, error };
}