audio_splitter/web/frontend/src/hooks/useWebSocket.ts

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

export const useWebSocket = (taskId: string | null) => {
  const wsRef = useRef<WebSocket | null>(null)
  const reconnectTimeoutRef = useRef<ReturnType<typeof setTimeout> | null>(null)
  const reconnectAttempts = useRef(0)
  const pollingRef = useRef<ReturnType<typeof setTimeout> | null>(null)

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

  const pollStatus = async () => {
    if (!taskId) return
  
    try {
      const response = await getStatus(taskId)
      console.log('[Polling] Status response:', response)
  
      // Update all state fields together
      setStatus(response.status)
      setProgress(response.progress)
      setMessage(response.message)
  
      // Explicitly set tracks if present
      if (response.tracks && response.tracks.length > 0) {
        console.log('[Polling] Setting tracks:', response.tracks)
        setTracks(response.tracks)
      }
  
      if (response.status === 'done') {
        console.log('[Polling] Split complete, tracks set:', response.tracks)
        setIsProcessing(false)
        // Ensure tracks are set one more time (safety)
        if (response.tracks && response.tracks.length > 0) {
          setTracks(response.tracks)
        }
        return
      }
  
      if (response.status === 'error') {
        setError(response.error || 'Split failed')
        setIsProcessing(false)
        return
      }
  
      // Still processing – poll again
      if (pollingRef.current) {
        clearTimeout(pollingRef.current)
      }
      pollingRef.current = setTimeout(pollStatus, 2000)
    } catch (error) {
      console.error('[Polling] Error:', error)
      if (pollingRef.current) {
        clearTimeout(pollingRef.current)
      }
      pollingRef.current = setTimeout(pollStatus, 3000)
    }
  }


  useEffect(() => {
    if (!taskId) {
      if (wsRef.current) {
        wsRef.current.close()
        wsRef.current = null
      }
      setWsConnected(false)
      // Clear polling
      if (pollingRef.current) {
        clearTimeout(pollingRef.current)
        pollingRef.current = null
      }
      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
        }
        // Start polling when connection is established
        // This ensures we get the final status even if WebSocket fails
        setTimeout(() => {
          pollStatus()
        }, 1000)
      }

      ws.onmessage = (event) => {
        try {
          const data = JSON.parse(event.data)
          console.log('[WebSocket] Message:', data)

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

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

        // If task is not done and we have a taskId, start polling
        // We check the status store to see if it's already done
        if (taskId && status !== 'done' && status !== 'error') {
          console.log('[WebSocket] Disconnected while processing, starting polling...')
          setTimeout(pollStatus, 1000)
        }
      }

      ws.onerror = (error) => {
        console.error('[WebSocket] Error:', error)
        // 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
      }
      if (pollingRef.current) {
        clearTimeout(pollingRef.current)
        pollingRef.current = null
      }
      setWsConnected(false)
    }
  }, [taskId, setStatus, setProgress, setMessage, setError, setTracks, addLog, setWsConnected, setIsProcessing, status])

  return wsRef.current
}