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

export function useWebSocket(onMessage) {
  const [status, setStatus] = useState('disconnected')
  const wsRef = useRef(null)
  const reconnectTimeoutRef = useRef(null)

  const connect = useCallback(() => {
    const protocol = window.location.protocol === 'https:' ? 'wss:' : 'ws:'
    // Always use root WebSocket path (works with proxy)
    const wsUrl = `${protocol}//${window.location.host}/ws/stream`

    console.log('[WS] Connecting to:', wsUrl)
    const ws = new WebSocket(wsUrl)
    ws.binaryType = 'arraybuffer'

    ws.onopen = () => {
      console.log('[WS] Connected')
      setStatus('connected')
    }

    ws.onclose = () => {
      console.log('[WS] Disconnected')
      setStatus('disconnected')
      reconnectTimeoutRef.current = setTimeout(connect, 3000)
    }

    ws.onerror = (e) => {
      console.error('[WS] Error:', e)
      setStatus('error')
    }

    ws.onmessage = async (event) => {
      if (event.data instanceof ArrayBuffer) {
        onMessage?.({ type: 'binary', data: event.data })
      } else {
        try {
          const data = JSON.parse(event.data)
          onMessage?.(data)
        } catch (e) {
          console.error('[WS] Parse error:', e)
        }
      }
    }

    wsRef.current = ws

    // Keep alive
    const pingInterval = setInterval(() => {
      if (ws.readyState === WebSocket.OPEN) {
        ws.send(JSON.stringify({ type: 'ping' }))
      }
    }, 10000)

    return () => {
      clearInterval(pingInterval)
      ws.close()
    }
  }, [onMessage])

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

  const send = useCallback((data) => {
    if (wsRef.current?.readyState === WebSocket.OPEN) {
      if (typeof data === 'string') {
        wsRef.current.send(data)
      } else {
        wsRef.current.send(JSON.stringify(data))
      }
    }
  }, [])

  return { status, send }
}