dashboard/src/hooks/useWebSocket.js

/**
 * useWebSocket — live video and metrics, with automatic HTTP polling fallback.
 *
 * Strategy:
 *   1. Attempt a WebSocket connection.
 *   2. If it fails within WS_PROBE_MS (e.g. HF Spaces blocks wss://), switch
 *      permanently to HTTP polling for this page load.
 *   3. Once in polling mode, video frames come from GET /api/pipeline/frame
 *      and metrics from GET /api/pipeline/metrics.
 */
import { useEffect, useRef, useState, useCallback } from 'react'

const _proto  = window.location.protocol === 'https:' ? 'wss:' : 'ws:'
const BASE_WS = import.meta.env.VITE_API_WS_URL
  || `${_proto}//${window.location.host}/api`

const WS_PROBE_MS      = 4_000   // give WS this long to open before falling back
const VIDEO_POLL_MS    = 200     // ~5 FPS polling
const METRICS_POLL_MS  = 1_000
const WS_RECONNECT_MAX = 30_000


// ── WS reconnect helper ───────────────────────────────────────────────────────

function useExponentialBackoff(connectFn, enabled) {
  const delayRef   = useRef(1_000)
  const timerRef   = useRef(null)
  const socketRef  = useRef(null)
  const enabledRef = useRef(enabled)

  useEffect(() => { enabledRef.current = enabled }, [enabled])

  const connect = useCallback(() => {
    if (!enabledRef.current) return
    const ws = connectFn({
      onClose: () => {
        if (!enabledRef.current) return
        timerRef.current = setTimeout(() => {
          delayRef.current = Math.min(delayRef.current * 2, WS_RECONNECT_MAX)
          connect()
        }, delayRef.current)
      },
      onOpen: () => { delayRef.current = 1_000 },
    })
    socketRef.current = ws
  }, [connectFn])

  useEffect(() => {
    if (enabled) connect()
    return () => {
      enabledRef.current = false
      clearTimeout(timerRef.current)
      socketRef.current?.close()
    }
  }, [enabled, connect])
}


// ── Video stream ──────────────────────────────────────────────────────────────

export function useVideoStream(canvasRef, enabled = true) {
  const [connected, setConnected] = useState(false)
  // null = probing, true = ws, false = polling
  const [mode, setMode] = useState(null)
  const pollRef = useRef(null)
  const probeRef = useRef(null)

  // --- Polling path ---
  const startPolling = useCallback(() => {
    if (pollRef.current) return
    setMode(false)
    const tick = async () => {
      if (!enabled) return
      try {
        const res = await fetch('/api/pipeline/frame')
        if (res.status === 200) {
          const blob = await res.blob()
          const url  = URL.createObjectURL(blob)
          const img  = new Image()
          img.onload = () => {
            const canvas = canvasRef.current
            if (!canvas) return
            canvas.width  = img.width
            canvas.height = img.height
            canvas.getContext('2d').drawImage(img, 0, 0)
            URL.revokeObjectURL(url)
          }
          img.src = url
          setConnected(true)
        } else {
          setConnected(false)
        }
      } catch {
        setConnected(false)
      }
    }
    pollRef.current = setInterval(tick, VIDEO_POLL_MS)
  }, [canvasRef, enabled])

  // --- WS path ---
  const connectFn = useCallback(({ onClose, onOpen }) => {
    const ws = new WebSocket(`${BASE_WS}/ws/video`)
    ws.binaryType = 'blob'

    // Start a probe timer — if WS doesn't open in time, fall back
    probeRef.current = setTimeout(() => {
      if (ws.readyState !== WebSocket.OPEN) {
        ws.close()
        startPolling()
      }
    }, WS_PROBE_MS)

    ws.onopen = () => {
      clearTimeout(probeRef.current)
      setMode(true)
      setConnected(true)
      onOpen()
    }
    ws.onclose = (e) => {
      clearTimeout(probeRef.current)
      setConnected(false)
      // If we never opened successfully, switch to polling instead of retrying
      if (mode !== true) {
        startPolling()
      } else {
        onClose(e.wasClean)
      }
    }
    ws.onerror = () => ws.close()

    ws.onmessage = (e) => {
      const canvas = canvasRef.current
      if (!canvas || !(e.data instanceof Blob)) return
      const url = URL.createObjectURL(e.data)
      const img  = new Image()
      img.onload = () => {
        const ctx = canvas.getContext('2d')
        canvas.width  = img.width
        canvas.height = img.height
        ctx.drawImage(img, 0, 0)
        URL.revokeObjectURL(url)
      }
      img.src = url
    }

    return ws
  }, [canvasRef, mode, startPolling])

  useExponentialBackoff(connectFn, enabled && mode !== false)

  // Start polling directly if WS mode already failed
  useEffect(() => {
    if (enabled && mode === false) startPolling()
    if (!enabled) {
      clearInterval(pollRef.current)
      pollRef.current = null
      setConnected(false)
    }
    return () => {
      if (!enabled) {
        clearInterval(pollRef.current)
        pollRef.current = null
      }
    }
  }, [enabled, mode, startPolling])

  return connected
}


// ── Metrics stream ────────────────────────────────────────────────────────────

const METRICS_HISTORY = 60

const DEFAULT_METRICS = {
  latest: null,
  history: [],
  alerts: [],
}

export function useMetricsStream(enabled = true) {
  const [connected, setConnected]   = useState(false)
  const [metrics, setMetrics]       = useState(DEFAULT_METRICS)
  const [mode, setMode]             = useState(null)  // null/true/false
  const pollRef  = useRef(null)
  const probeRef = useRef(null)

  const applyMsg = useCallback((msg) => {
    setMetrics(prev => {
      const history   = [...prev.history, msg].slice(-METRICS_HISTORY)
      const newAlerts = msg.alerts?.length
        ? [...prev.alerts, ...msg.alerts].slice(-50)
        : prev.alerts
      return { latest: msg, history, alerts: newAlerts }
    })
  }, [])

  // --- Polling path ---
  const startPolling = useCallback(() => {
    if (pollRef.current) return
    setMode(false)
    const tick = async () => {
      if (!enabled) return
      try {
        const res = await fetch('/api/pipeline/metrics')
        if (res.status === 200) {
          const msg = await res.json()
          applyMsg(msg)
          setConnected(true)
        } else {
          setConnected(false)
        }
      } catch {
        setConnected(false)
      }
    }
    pollRef.current = setInterval(tick, METRICS_POLL_MS)
  }, [enabled, applyMsg])

  // --- WS path ---
  const connectFn = useCallback(({ onClose, onOpen }) => {
    const ws = new WebSocket(`${BASE_WS}/ws/metrics`)

    probeRef.current = setTimeout(() => {
      if (ws.readyState !== WebSocket.OPEN) {
        ws.close()
        startPolling()
      }
    }, WS_PROBE_MS)

    ws.onopen = () => {
      clearTimeout(probeRef.current)
      setMode(true)
      setConnected(true)
      onOpen()
    }
    ws.onclose = (e) => {
      clearTimeout(probeRef.current)
      setConnected(false)
      if (mode !== true) {
        startPolling()
      } else {
        onClose(e.wasClean)
      }
    }
    ws.onerror = () => ws.close()

    ws.onmessage = (e) => {
      try { applyMsg(JSON.parse(e.data)) } catch { /* ignore malformed */ }
    }

    return ws
  }, [mode, startPolling, applyMsg])

  useExponentialBackoff(connectFn, enabled && mode !== false)

  useEffect(() => {
    if (enabled && mode === false) startPolling()
    if (!enabled) {
      clearInterval(pollRef.current)
      pollRef.current = null
      setConnected(false)
    }
    return () => {
      if (!enabled) {
        clearInterval(pollRef.current)
        pollRef.current = null
      }
    }
  }, [enabled, mode, startPolling])

  return { connected, metrics }
}