zmb-webui/frontend/lib/useWebSocket.ts

"use client"

import { useEffect, useRef, useCallback } from "react"

export type WsMessage = {
  type: "pool_status" | "scrub_progress" | "snapshot_created" | "alert"
  data: unknown
}

type Handler = (msg: WsMessage) => void

export function useWebSocket(onMessage: Handler) {
  const wsRef = useRef<WebSocket | null>(null)
  const retryRef = useRef<ReturnType<typeof setTimeout> | null>(null)
  const delayRef = useRef(2000)
  const handlerRef = useRef(onMessage)
  handlerRef.current = onMessage

  const connect = useCallback(() => {
    if (typeof window === "undefined") return

    const token = localStorage.getItem("access_token")
    if (!token) return

    // Derive WS URL from current page origin
    const proto = window.location.protocol === "https:" ? "wss" : "ws"
    const wsUrl = `${proto}://${window.location.host}/ws`

    const ws = new WebSocket(wsUrl)
    wsRef.current = ws

    ws.onopen = () => {
      delayRef.current = 2000 // Reset backoff on success
    }

    ws.onmessage = (event) => {
      try {
        const msg: WsMessage = JSON.parse(event.data)
        handlerRef.current(msg)
      } catch {
        // ignore malformed messages
      }
    }

    ws.onclose = () => {
      // Reconnect with exponential backoff (max 30s)
      retryRef.current = setTimeout(() => {
        delayRef.current = Math.min(delayRef.current * 2, 30000)
        connect()
      }, delayRef.current)
    }

    ws.onerror = () => {
      ws.close()
    }
  }, [])

  useEffect(() => {
    connect()
    return () => {
      if (retryRef.current) clearTimeout(retryRef.current)
      wsRef.current?.close()
    }
  }, [connect])
}