tp-ems/core/frontend/src/hooks/useRealtimeWebSocket.ts

import { useEffect, useRef, useState, useCallback } from 'react';
import { getToken } from '../utils/auth';

export interface RealtimeData {
  pv_power: number;
  heatpump_power: number;
  total_load: number;
  grid_power: number;
  active_alarms: number;
  timestamp: string;
}

export interface AlarmEventData {
  id: number;
  title: string;
  severity: string;
  message?: string;
  device_name?: string;
  triggered_at?: string;
}

interface WebSocketMessage {
  type: 'realtime_update' | 'alarm_event' | 'pong';
  data?: RealtimeData | AlarmEventData;
}

interface UseRealtimeWebSocketOptions {
  /** Called when a new alarm event arrives */
  onAlarmEvent?: (alarm: AlarmEventData) => void;
  /** Polling interval in ms when WS is unavailable (default: 15000) */
  fallbackInterval?: number;
  /** Whether the hook is enabled (default: true) */
  enabled?: boolean;
}

interface UseRealtimeWebSocketResult {
  /** Latest realtime data from WebSocket */
  data: RealtimeData | null;
  /** Whether WebSocket is currently connected */
  connected: boolean;
  /** Whether we are using fallback polling */
  usingFallback: boolean;
}

const MAX_RECONNECT_DELAY = 30000;
const INITIAL_RECONNECT_DELAY = 1000;

export default function useRealtimeWebSocket(
  options: UseRealtimeWebSocketOptions = {}
): UseRealtimeWebSocketResult {
  const { onAlarmEvent, fallbackInterval = 15000, enabled = true } = options;
  const [data, setData] = useState<RealtimeData | null>(null);
  const [connected, setConnected] = useState(false);
  const [usingFallback, setUsingFallback] = useState(false);

  const wsRef = useRef<WebSocket | null>(null);
  const reconnectDelayRef = useRef(INITIAL_RECONNECT_DELAY);
  const reconnectTimerRef = useRef<ReturnType<typeof setTimeout> | null>(null);
  const pingTimerRef = useRef<ReturnType<typeof setInterval> | null>(null);
  const fallbackTimerRef = useRef<ReturnType<typeof setInterval> | null>(null);
  const onAlarmEventRef = useRef(onAlarmEvent);
  const mountedRef = useRef(true);

  // Keep callback ref up to date
  useEffect(() => {
    onAlarmEventRef.current = onAlarmEvent;
  }, [onAlarmEvent]);

  const cleanup = useCallback(() => {
    if (reconnectTimerRef.current) {
      clearTimeout(reconnectTimerRef.current);
      reconnectTimerRef.current = null;
    }
    if (pingTimerRef.current) {
      clearInterval(pingTimerRef.current);
      pingTimerRef.current = null;
    }
    if (wsRef.current) {
      wsRef.current.onclose = null;
      wsRef.current.onerror = null;
      wsRef.current.onmessage = null;
      wsRef.current.close();
      wsRef.current = null;
    }
  }, []);

  const startFallbackPolling = useCallback(() => {
    if (fallbackTimerRef.current) return;
    setUsingFallback(true);
    // We don't do actual polling here - the parent component's
    // existing polling handles data fetch. This flag signals the
    // parent to keep its polling active.
  }, []);

  const stopFallbackPolling = useCallback(() => {
    if (fallbackTimerRef.current) {
      clearInterval(fallbackTimerRef.current);
      fallbackTimerRef.current = null;
    }
    setUsingFallback(false);
  }, []);

  const connect = useCallback(() => {
    if (!mountedRef.current || !enabled) return;

    const token = getToken();
    if (!token) {
      startFallbackPolling();
      return;
    }

    cleanup();

    const protocol = window.location.protocol === 'https:' ? 'wss:' : 'ws:';
    const host = window.location.host;
    const wsUrl = `${protocol}//${host}/api/v1/ws/realtime?token=${encodeURIComponent(token)}`;

    try {
      const ws = new WebSocket(wsUrl);
      wsRef.current = ws;

      ws.onopen = () => {
        if (!mountedRef.current) return;
        setConnected(true);
        stopFallbackPolling();
        reconnectDelayRef.current = INITIAL_RECONNECT_DELAY;

        // Ping every 30s to keep alive
        pingTimerRef.current = setInterval(() => {
          if (ws.readyState === WebSocket.OPEN) {
            ws.send('ping');
          }
        }, 30000);
      };

      ws.onmessage = (event) => {
        if (!mountedRef.current) return;
        try {
          const msg: WebSocketMessage = JSON.parse(event.data);
          if (msg.type === 'realtime_update' && msg.data) {
            setData(msg.data as RealtimeData);
          } else if (msg.type === 'alarm_event' && msg.data) {
            onAlarmEventRef.current?.(msg.data as AlarmEventData);
          }
          // pong is just a keepalive ack, ignore
        } catch {
          // ignore parse errors
        }
      };

      ws.onclose = (event) => {
        if (!mountedRef.current) return;
        setConnected(false);
        if (pingTimerRef.current) {
          clearInterval(pingTimerRef.current);
          pingTimerRef.current = null;
        }

        // Don't reconnect if closed intentionally (4001 = auth error)
        if (event.code === 4001) {
          startFallbackPolling();
          return;
        }

        // Reconnect with exponential backoff
        startFallbackPolling();
        const delay = reconnectDelayRef.current;
        reconnectDelayRef.current = Math.min(delay * 2, MAX_RECONNECT_DELAY);
        reconnectTimerRef.current = setTimeout(connect, delay);
      };

      ws.onerror = () => {
        // onclose will fire after this, which handles reconnection
      };
    } catch {
      startFallbackPolling();
      const delay = reconnectDelayRef.current;
      reconnectDelayRef.current = Math.min(delay * 2, MAX_RECONNECT_DELAY);
      reconnectTimerRef.current = setTimeout(connect, delay);
    }
  }, [enabled, cleanup, startFallbackPolling, stopFallbackPolling]);

  useEffect(() => {
    mountedRef.current = true;
    if (enabled) {
      connect();
    }
    return () => {
      mountedRef.current = false;
      cleanup();
      stopFallbackPolling();
    };
  }, [enabled, connect, cleanup, stopFallbackPolling]);

  return { data, connected, usingFallback };
}