import { useEffect, useRef } from 'react'
import { useTaskStore } from '../stores/taskStore'
import { useUIStore } from '../stores/uiStore'

export const useWebSocket = (taskId: string | null) => {
  const wsRef = useRef<WebSocket | null>(null)
  const reconnectTimeoutRef = useRef<NodeJS.Timeout | null>(null)
  const reconnectAttempts = useRef(0)

  const { setStatus, setProgress, setMessage, setError, setTracks, addLog } = useTaskStore()
  const { setWsConnected } = useUIStore()

  useEffect(() => {
    if (!taskId) {
      if (wsRef.current) {
        wsRef.current.close()
        wsRef.current = null
      }
      setWsConnected(false)
      return
    }

    const connect = () => {
      const wsUrl = `/ws/${taskId}`
      const ws = new WebSocket(wsUrl)

      ws.onopen = () => {
        console.log(`WebSocket connected for task ${taskId}`)
        setWsConnected(true)
        reconnectAttempts.current = 0
        if (reconnectTimeoutRef.current) {
          clearTimeout(reconnectTimeoutRef.current)
          reconnectTimeoutRef.current = null
        }
      }

      ws.onmessage = (event) => {
        try {
          const data = JSON.parse(event.data)

          if (data.type === 'status') {
            const status = data.data
            setStatus(status.status)
            setProgress(status.progress)
            setMessage(status.message)
            if (status.error) {
              setError(status.error)
            }
            if (status.tracks) {
              setTracks(status.tracks)
            }
          } else if (data.type === 'progress') {
            const progressData = data.data
            setStatus(progressData.status)
            setProgress(progressData.progress)
            setMessage(progressData.message)
            if (progressData.status === 'done') {
              addLog('✅ Split complete!')
            } else if (progressData.status === 'error') {
              setError(progressData.message)
              addLog(`❌ Error: ${progressData.message}`)
            } else {
              addLog(`🔄 ${progressData.message} (${progressData.progress}%)`)
            }
          }
        } catch (error) {
          console.error('Failed to parse WebSocket message:', error)
        }
      }

      ws.onclose = () => {
        console.log(`WebSocket disconnected for task ${taskId}`)
        setWsConnected(false)

        // Try to reconnect if the task is still processing
        // We'll check status via polling if needed
        if (reconnectAttempts.current < 5) {
          reconnectTimeoutRef.current = setTimeout(() => {
            reconnectAttempts.current += 1
            connect()
          }, 3000)
        }
      }

      ws.onerror = (error) => {
        console.error('WebSocket error:', error)
        // The onclose will handle reconnection
      }

      wsRef.current = ws
    }

    connect()

    return () => {
      if (wsRef.current) {
        wsRef.current.close()
        wsRef.current = null
      }
      if (reconnectTimeoutRef.current) {
        clearTimeout(reconnectTimeoutRef.current)
        reconnectTimeoutRef.current = null
      }
      setWsConnected(false)
    }
  }, [taskId, setStatus, setProgress, setMessage, setError, setTracks, addLog, setWsConnected])

  return wsRef.current
}