import { useEffect, useRef, useCallback, useState } from 'react'
import { io, Socket } from 'socket.io-client'
import toast from 'react-hot-toast'
import { useSystemStore } from '../stores/systemStore'
import type { WebSocketMessage, ProcessStatus, HardwareStatus, SafetyStatus } from '../types'

interface UseWebSocketReturn {
  socket: Socket | null
  isConnected: boolean
  connect: () => void
  disconnect: () => void
  emit: (event: string, data?: any) => void
}

export const useWebSocket = (): UseWebSocketReturn => {
  const socketRef = useRef<Socket | null>(null)
  const [isConnected, setIsConnected] = useState(false)
  const reconnectTimeoutRef = useRef<NodeJS.Timeout>()
  const reconnectAttempts = useRef(0)
  const maxReconnectAttempts = 10

  // Store actions
  const { updateSystemStatus, updateConnectionStatus, setError } = useSystemStore()

  // Get WebSocket URL from environment
  const getSocketUrl = useCallback(() => {
    const baseUrl = import.meta.env.VITE_API_BASE_URL || 'http://localhost:8000'
    return baseUrl.replace(/^http/, 'ws')
  }, [])

  // Handle reconnection with exponential backoff
  const scheduleReconnect = useCallback(() => {
    if (reconnectAttempts.current >= maxReconnectAttempts) {
      console.error('Max reconnection attempts reached')
      setError('Lost connection to server. Please refresh the page.')
      return
    }

    const delay = Math.min(1000 * Math.pow(2, reconnectAttempts.current), 30000) // Max 30 seconds
    reconnectAttempts.current += 1

    console.log(`Scheduling reconnection attempt ${reconnectAttempts.current} in ${delay}ms`)

    reconnectTimeoutRef.current = setTimeout(() => {
      connect()
    }, delay)
  }, [setError])

  // Connect to WebSocket
  const connect = useCallback(() => {
    if (socketRef.current?.connected) {
      console.log('WebSocket already connected')
      return
    }

    try {
      const socketUrl = getSocketUrl()
      console.log('Connecting to WebSocket:', socketUrl)

      const socket = io(socketUrl, {
        transports: ['websocket', 'polling'],
        timeout: 20000,
        forceNew: true,
        reconnection: false, // We handle reconnection manually
      })

      // Connection events
      socket.on('connect', () => {
        console.log('WebSocket connected')
        setIsConnected(true)
        updateConnectionStatus(true)
        reconnectAttempts.current = 0 // Reset attempts on successful connection
        
        // Clear any existing reconnection timeout
        if (reconnectTimeoutRef.current) {
          clearTimeout(reconnectTimeoutRef.current)
          reconnectTimeoutRef.current = undefined
        }

        toast.success('Connected to tempering system', {
          duration: 2000,
          position: 'top-right',
        })

        // Subscribe to real-time updates
        socket.emit('subscribe', { topics: ['process', 'hardware', 'safety', 'system'] })
      })

      socket.on('disconnect', (reason) => {
        console.log('WebSocket disconnected:', reason)
        setIsConnected(false)
        updateConnectionStatus(false)

        if (reason === 'io server disconnect') {
          // Server initiated disconnect, don't reconnect automatically
          toast.error('Disconnected from server')
        } else {
          // Client-side disconnect or network issue, attempt to reconnect
          scheduleReconnect()
        }
      })

      socket.on('connect_error', (error) => {
        console.error('WebSocket connection error:', error)
        setIsConnected(false)
        updateConnectionStatus(false)
        scheduleReconnect()
      })

      // Message handlers
      socket.on('message', (message: WebSocketMessage) => {
        handleWebSocketMessage(message)
      })

      // Specific event handlers for better performance
      socket.on('process_update', (data: ProcessStatus) => {
        updateSystemStatus({
          system_status: 'healthy',
          process_status: data.status as any,
          hardware_status: 'connected',
          safety_status: 'safe',
          active_alarms: data.error_count + data.warning_count,
          uptime_seconds: 0,
          last_updated: new Date().toISOString(),
        })
      })

      socket.on('hardware_update', (data: HardwareStatus) => {
        // Handle hardware status updates
        console.log('Hardware update:', data)
      })

      socket.on('safety_alert', (data: SafetyStatus) => {
        // Handle safety alerts
        if (data.overall_status === 'alarm') {
          toast.error('Safety alarm triggered!', {
            duration: 0, // Keep visible until dismissed
            style: {
              background: '#ef4444',
              color: 'white',
              fontWeight: 'bold',
              fontSize: '18px',
            },
          })
        } else if (data.overall_status === 'warning') {
          toast.warning('Safety warning', {
            duration: 5000,
          })
        }
      })

      socket.on('system_status', (data: any) => {
        updateSystemStatus(data)
      })

      socketRef.current = socket
    } catch (error) {
      console.error('Failed to create WebSocket connection:', error)
      scheduleReconnect()
    }
  }, [getSocketUrl, updateConnectionStatus, updateSystemStatus, scheduleReconnect])

  // Disconnect from WebSocket
  const disconnect = useCallback(() => {
    if (reconnectTimeoutRef.current) {
      clearTimeout(reconnectTimeoutRef.current)
      reconnectTimeoutRef.current = undefined
    }

    if (socketRef.current) {
      console.log('Disconnecting WebSocket')
      socketRef.current.disconnect()
      socketRef.current = null
    }

    setIsConnected(false)
    updateConnectionStatus(false)
  }, [updateConnectionStatus])

  // Emit message to server
  const emit = useCallback((event: string, data?: any) => {
    if (socketRef.current?.connected) {
      socketRef.current.emit(event, data)
    } else {
      console.warn('Cannot emit event: WebSocket not connected')
    }
  }, [])

  // Handle incoming WebSocket messages
  const handleWebSocketMessage = useCallback((message: WebSocketMessage) => {
    console.log('WebSocket message:', message)

    switch (message.type) {
      case 'process_update':
        // Handle process status updates
        break

      case 'hardware_update':
        // Handle hardware status updates
        break

      case 'safety_alert':
        // Handle safety alerts
        if (message.data.severity === 'critical') {
          toast.error(`Critical Safety Alert: ${message.data.message}`, {
            duration: 0, // Keep visible
          })
        }
        break

      case 'system_status':
        updateSystemStatus(message.data)
        break

      default:
        console.log('Unknown message type:', message.type)
    }
  }, [updateSystemStatus])

  // Cleanup on unmount
  useEffect(() => {
    return () => {
      disconnect()
    }
  }, [disconnect])

  return {
    socket: socketRef.current,
    isConnected,
    connect,
    disconnect,
    emit,
  }
}
