usda-vision/vision-system-remote/src/hooks/useWebSocket.ts

import { useEffect, useRef, useState, useCallback } from 'react'

// WebSocket message types from the API
export interface WebSocketEvent {
  type: 'event'
  event_type: string
  source: string
  data: any
  timestamp: string
}

type EventHandler = (event: WebSocketEvent) => void

export function useWebSocket(url: string, options?: { enabled?: boolean }) {
  const [isConnected, setIsConnected] = useState(false)
  const [reconnectAttempts, setReconnectAttempts] = useState(0)
  const wsRef = useRef<WebSocket | null>(null)
  const reconnectTimeoutRef = useRef<NodeJS.Timeout | null>(null)
  const handlersRef = useRef<Map<string, Set<EventHandler>>>(new Map())
  const enabled = options?.enabled !== false

  const connect = useCallback(() => {
    if (!enabled || wsRef.current?.readyState === WebSocket.OPEN) {
      return
    }

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

      ws.onopen = () => {
        setIsConnected(true)
        setReconnectAttempts(0)
      }

      ws.onmessage = (event) => {
        try {
          const message: WebSocketEvent = JSON.parse(event.data)
          
          // Call all handlers for this event type
          const handlers = handlersRef.current.get(message.event_type)
          if (handlers) {
            handlers.forEach(handler => handler(message))
          }

          // Also call handlers for 'all' type
          const allHandlers = handlersRef.current.get('all')
          if (allHandlers) {
            allHandlers.forEach(handler => handler(message))
          }
        } catch (err) {
          console.error('Failed to parse WebSocket message:', err)
        }
      }

      ws.onerror = (error) => {
        console.error('WebSocket error:', error)
      }

      ws.onclose = () => {
        setIsConnected(false)
        
        // Reconnect with exponential backoff (max 10 attempts)
        if (enabled && reconnectAttempts < 10) {
          const delay = Math.min(1000 * Math.pow(2, reconnectAttempts), 30000)
          reconnectTimeoutRef.current = setTimeout(() => {
            setReconnectAttempts(prev => prev + 1)
            connect()
          }, delay)
        }
      }
    } catch (err) {
      console.error('Failed to create WebSocket connection:', err)
      setIsConnected(false)
    }
  }, [url, enabled, reconnectAttempts])

  useEffect(() => {
    if (enabled) {
      connect()
    }

    return () => {
      if (reconnectTimeoutRef.current) {
        clearTimeout(reconnectTimeoutRef.current)
      }
      if (wsRef.current) {
        wsRef.current.close()
        wsRef.current = null
      }
      setIsConnected(false)
    }
  }, [connect, enabled])

  const subscribe = useCallback((eventType: string | 'all', handler: EventHandler) => {
    if (!handlersRef.current.has(eventType)) {
      handlersRef.current.set(eventType, new Set())
    }
    handlersRef.current.get(eventType)!.add(handler)

    return () => {
      const handlers = handlersRef.current.get(eventType)
      if (handlers) {
        handlers.delete(handler)
      }
    }
  }, [])

  const send = useCallback((message: any) => {
    if (wsRef.current?.readyState === WebSocket.OPEN) {
      wsRef.current.send(JSON.stringify(message))
    } else {
      console.warn('WebSocket is not connected')
    }
  }, [])

  return {
    isConnected,
    subscribe,
    send,
    reconnect: connect,
  }
}