planet/frontend/src/hooks/useWebSocket.ts

import { useEffect, useRef, useState, useCallback } from 'react'
import { useAuthStore } from '../stores/auth'

const DEFAULT_WS_URL = (() => {
  if (typeof window === 'undefined') {
    return 'ws://localhost:8000/ws'
  }
  const protocol = window.location.protocol === 'https:' ? 'wss:' : 'ws:'
  return `${protocol}//${window.location.host}/ws`
})()

const WS_URL = (import.meta as any).env?.VITE_WS_URL || DEFAULT_WS_URL

function buildWebSocketCandidates(): string[] {
  if ((import.meta as any).env?.VITE_WS_URL) {
    return [WS_URL]
  }

  if (typeof window === 'undefined') {
    return ['ws://localhost:8000/ws']
  }

  const protocol = window.location.protocol === 'https:' ? 'wss:' : 'ws:'
  const sameOrigin = `${protocol}//${window.location.host}/ws`
  const candidates = [sameOrigin]

  if (window.location.hostname === 'localhost' || window.location.hostname === '127.0.0.1') {
    const directBackend = `${protocol}//${window.location.hostname}:8000/ws`
    if (!candidates.includes(directBackend)) {
      candidates.push(directBackend)
    }
  }

  return candidates
}

interface WebSocketMessage {
  type: string
  channel?: string
  timestamp?: string
  data?: Record<string, unknown>
  payload?: Record<string, unknown>
}

interface UseWebSocketOptions {
  autoConnect?: boolean
  autoSubscribe?: string[]
  heartbeatInterval?: number
  onMessage?: (message: WebSocketMessage) => void
  onConnect?: () => void
  onDisconnect?: () => void
  onError?: (error: Error) => void
}

interface UseWebSocketReturn {
  connected: boolean
  connecting: boolean
  status: 'connecting' | 'connected' | 'disconnected'
  lastMessage: WebSocketMessage | null
  sendMessage: (message: Record<string, unknown>) => void
  subscribe: (channels: string[]) => void
  connect: () => void
  disconnect: () => void
}

export function useWebSocket(options: UseWebSocketOptions = {}): UseWebSocketReturn {
  const {
    autoConnect = true,
    autoSubscribe = [],
    heartbeatInterval = 25000,
    onMessage,
    onConnect,
    onDisconnect,
    onError,
  } = options

  const { token } = useAuthStore()

  const wsRef = useRef<WebSocket | null>(null)
  const [connected, setConnected] = useState(false)
  const [connecting, setConnecting] = useState(false)
  const [lastMessage, setLastMessage] = useState<WebSocketMessage | null>(null)
  const reconnectTimeoutRef = useRef<ReturnType<typeof setTimeout> | null>(null)
  const heartbeatTimerRef = useRef<ReturnType<typeof setInterval> | null>(null)
  const activeWsUrlRef = useRef<string | null>(null)
  const intentionalCloseRef = useRef(false)
  const pendingCloseSocketRef = useRef<WebSocket | null>(null)
  const autoSubscribeRef = useRef(autoSubscribe)
  const onMessageRef = useRef(onMessage)
  const onConnectRef = useRef(onConnect)
  const onDisconnectRef = useRef(onDisconnect)
  const onErrorRef = useRef(onError)

  useEffect(() => {
    autoSubscribeRef.current = autoSubscribe
    onMessageRef.current = onMessage
    onConnectRef.current = onConnect
    onDisconnectRef.current = onDisconnect
    onErrorRef.current = onError
  }, [autoSubscribe, onMessage, onConnect, onDisconnect, onError])

  const connect = useCallback(() => {
    if (!token) {
      setConnected(false)
      setConnecting(false)
      return
    }

    intentionalCloseRef.current = false
    setConnected(false)
    setConnecting(true)
    const candidates = buildWebSocketCandidates()
    let candidateIndex = 0
    let opened = false

    const tryConnect = () => {
      const baseUrl = candidates[candidateIndex]
      const wsUrl = `${baseUrl}?token=${encodeURIComponent(token)}`
      activeWsUrlRef.current = baseUrl
      const ws = new WebSocket(wsUrl)

      ws.onopen = () => {
        if (intentionalCloseRef.current || pendingCloseSocketRef.current === ws) {
          pendingCloseSocketRef.current = null
          ws.close()
          return
        }
        opened = true
        setConnected(true)
        setConnecting(false)
        if (autoSubscribeRef.current.length > 0) {
          ws.send(JSON.stringify({ type: 'subscribe', data: { channels: autoSubscribeRef.current } }))
        }
        if (heartbeatTimerRef.current) {
          clearInterval(heartbeatTimerRef.current)
        }
        heartbeatTimerRef.current = setInterval(() => {
          if (ws.readyState === WebSocket.OPEN) {
            ws.send(JSON.stringify({ type: 'heartbeat' }))
          }
        }, heartbeatInterval)
        onConnectRef.current?.()
      }

      ws.onmessage = (event) => {
        try {
          const message: WebSocketMessage = JSON.parse(event.data)
          if (message.type === 'heartbeat' && message.data?.action === 'ping' && ws.readyState === WebSocket.OPEN) {
            ws.send(JSON.stringify({ type: 'heartbeat' }))
          }
          setLastMessage(message)
          onMessageRef.current?.(message)
        } catch {
          console.error('Failed to parse WebSocket message')
        }
      }

      ws.onclose = () => {
        if (wsRef.current === ws) {
          wsRef.current = null
        }
        if (pendingCloseSocketRef.current === ws) {
          pendingCloseSocketRef.current = null
        }
        setConnected(false)
        if (heartbeatTimerRef.current) {
          clearInterval(heartbeatTimerRef.current)
          heartbeatTimerRef.current = null
        }

        if (!opened && candidateIndex < candidates.length - 1) {
          candidateIndex += 1
          setConnecting(true)
          tryConnect()
          return
        }

        setConnecting(false)

        if (intentionalCloseRef.current) {
          return
        }

        onDisconnectRef.current?.()

        if (autoConnect && token) {
          reconnectTimeoutRef.current = setTimeout(() => {
            connect()
          }, 3000)
        }
      }

      ws.onerror = (error) => {
        setConnected(false)
        if (opened || candidateIndex >= candidates.length - 1) {
          setConnecting(false)
        }
        if (intentionalCloseRef.current || ws.readyState === WebSocket.CLOSING || ws.readyState === WebSocket.CLOSED) {
          return
        }
        if (candidateIndex >= candidates.length - 1) {
          console.warn('[WebSocket] Connection error', { url: baseUrl, error })
          onErrorRef.current?.(new Error('WebSocket error'))
        }
      }

      wsRef.current = ws
    }

    try {
      tryConnect()
    } catch (error) {
      setConnected(false)
      setConnecting(false)
      console.warn('[WebSocket] Failed to initialize connection', { url: activeWsUrlRef.current, error })
      if (autoConnect && token) {
        reconnectTimeoutRef.current = setTimeout(() => {
          connect()
        }, 3000)
      }
    }
  }, [token, autoConnect, heartbeatInterval])

  const disconnect = useCallback(() => {
    intentionalCloseRef.current = true
    if (reconnectTimeoutRef.current) {
      clearTimeout(reconnectTimeoutRef.current)
      reconnectTimeoutRef.current = null
    }
    if (heartbeatTimerRef.current) {
      clearInterval(heartbeatTimerRef.current)
      heartbeatTimerRef.current = null
    }
    const socket = wsRef.current
    if (socket) {
      if (socket.readyState === WebSocket.CONNECTING) {
        pendingCloseSocketRef.current = socket
      } else if (socket.readyState === WebSocket.OPEN) {
        socket.close()
      }
    }
    setConnected(false)
    setConnecting(false)
  }, [])

  const sendMessage = useCallback((message: Record<string, unknown>) => {
    wsRef.current?.send(JSON.stringify(message))
  }, [])

  const subscribe = useCallback((channels: string[]) => {
    wsRef.current?.send(JSON.stringify({
      type: 'subscribe',
      data: { channels }
    }))
  }, [])

  useEffect(() => {
    if (autoConnect && token) {
      connect()
    }
    return () => {
      disconnect()
    }
  }, [autoConnect, token, connect, disconnect])

  return {
    connected,
    connecting,
    status: connected ? 'connected' : connecting ? 'connecting' : 'disconnected',
    lastMessage,
    sendMessage,
    subscribe,
    connect,
    disconnect,
  }
}