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

export interface WebSocketMessage {
  type: 'order_update' | 'position_update' | 'price_update' | 'alert_triggered' | 'strategy_signal' | 'system_event'
  [key: string]: any
}

export function useWebSocket(url: string) {
  const [isConnected, setIsConnected] = useState(false)
  const [lastMessage, setLastMessage] = useState<WebSocketMessage | null>(null)
  const [messageHistory, setMessageHistory] = useState<WebSocketMessage[]>([])
  const wsRef = useRef<WebSocket | null>(null)
  const reconnectTimeoutRef = useRef<NodeJS.Timeout>()
  const messageHandlersRef = useRef<Map<string, (message: WebSocketMessage) => void>>(new Map())
  const isConnectingRef = useRef(false)

  useEffect(() => {
    let isMounted = true

    const connect = () => {
      // Prevent duplicate connections
      if (isConnectingRef.current || wsRef.current?.readyState === WebSocket.OPEN || wsRef.current?.readyState === WebSocket.CONNECTING) {
        return
      }

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

        ws.onopen = () => {
          if (isMounted) {
            isConnectingRef.current = false
            setIsConnected(true)
            console.log('WebSocket connected')
          }
        }

        ws.onmessage = (event) => {
          if (!isMounted) return
          try {
            const message = JSON.parse(event.data) as WebSocketMessage
            setLastMessage(message)
            setMessageHistory((prev) => [...prev.slice(-99), message]) // Keep last 100 messages

            // Call registered handlers for this message type
            const handlers = messageHandlersRef.current.get(message.type)
            if (handlers) {
              handlers(message)
            }
          } catch (error) {
            console.error('Failed to parse WebSocket message:', error)
          }
        }

        ws.onerror = (error) => {
          isConnectingRef.current = false
          // Only log error if we're still mounted (avoid noise from StrictMode cleanup)
          if (isMounted) {
            console.error('WebSocket error:', error)
          }
        }

        ws.onclose = () => {
          isConnectingRef.current = false
          if (isMounted) {
            setIsConnected(false)
            console.log('WebSocket disconnected')
            // Reconnect after 3 seconds
            reconnectTimeoutRef.current = setTimeout(() => {
              if (isMounted) connect()
            }, 3000)
          }
        }
      } catch (error) {
        isConnectingRef.current = false
        console.error('Failed to create WebSocket:', error)
      }
    }

    connect()

    return () => {
      isMounted = false
      isConnectingRef.current = false
      if (reconnectTimeoutRef.current) {
        clearTimeout(reconnectTimeoutRef.current)
      }
      if (wsRef.current) {
        wsRef.current.close()
        wsRef.current = null
      }
    }
  }, [url])

  const sendMessage = useCallback((message: any) => {
    if (wsRef.current && wsRef.current.readyState === WebSocket.OPEN) {
      wsRef.current.send(JSON.stringify(message))
    }
  }, [])

  const subscribe = useCallback((messageType: string, handler: (message: WebSocketMessage) => void) => {
    messageHandlersRef.current.set(messageType, handler)
    return () => {
      messageHandlersRef.current.delete(messageType)
    }
  }, [])

  return {
    isConnected,
    lastMessage,
    messageHistory,
    sendMessage,
    subscribe,
  }
}
