// SPDX-FileCopyrightText: 2025 Weibo, Inc.
//
// SPDX-License-Identifier: Apache-2.0

'use client'

/**
 * Socket.IO Context Provider
 *
 * Manages Socket.IO connection at the application level.
 * Provides connection state and socket instance to child components.
 * Auto-connects when user is authenticated.
 */

import React, {
  createContext,
  useContext,
  useEffect,
  useState,
  useCallback,
  useRef,
  ReactNode,
} from 'react'
import { io, Socket } from 'socket.io-client'
import { getToken, removeToken } from '@/apis/user'
import {
  ClientEvents,
  ServerEvents,
  ClientSkillEvents,
  ChatStartPayload,
  ChatChunkPayload,
  ChatDonePayload,
  ChatErrorPayload,
  ChatCancelledPayload,
  ChatMessagePayload,
  ChatBlockCreatedPayload,
  ChatBlockUpdatedPayload,
  ChatSendPayload,
  ChatSendAck,
  ChatGuidePayload,
  ChatGuideAck,
  ChatGuidanceQueuedPayload,
  ChatGuidanceAppliedPayload,
  ChatGuidanceExpiredPayload,
  TaskCreatedPayload,
  TaskStatusPayload,
  TaskInvitedPayload,
  TaskAppUpdatePayload,
  SkillRequestPayload,
  SkillResponsePayload,
  CorrectionStartPayload,
  CorrectionProgressPayload,
  CorrectionChunkPayload,
  CorrectionDonePayload,
  CorrectionErrorPayload,
  BackgroundExecutionUpdatePayload,
  AuthErrorPayload,
} from '@/types/socket'

import { fetchRuntimeConfig, getSocketUrl } from '@/lib/runtime-config'
import { paths } from '@/config/paths'
import { POST_LOGIN_REDIRECT_KEY } from '@/features/login/constants'
import { taskStateManager } from '@/features/tasks/state'

const SOCKETIO_PATH = '/socket.io'

/** Callback type for reconnect event */
export type ReconnectCallback = () => void

interface SocketContextType {
  /** Socket.IO instance */
  socket: Socket | null
  /** Whether connected to server */
  isConnected: boolean
  /** Connection error if any */
  connectionError: Error | null
  /** Reconnect attempt count */
  reconnectAttempts: number
  /** Connect to Socket.IO server */
  connect: (token: string) => void
  /** Disconnect from server */
  disconnect: () => void
  /**
   * Join a task room.
   * @param taskId - Task ID to join
   * @param options - Join options
   * @param options.forceRefresh - If true, always emit task:join to get streaming status
   * @param options.afterMessageId - If provided, only return messages after this ID (for incremental sync)
   */
  joinTask: (
    taskId: number,
    options?: {
      forceRefresh?: boolean
      afterMessageId?: number
    }
  ) => Promise<{
    streaming?: {
      subtask_id: number
      offset: number
      cached_content: string
      started_at?: string
      last_activity_at?: string
    }
    /** Subtasks data for immediate message sync (same format as task detail API) */
    subtasks?: Array<Record<string, unknown>>
    error?: string
  }>
  /** Leave a task room */
  leaveTask: (taskId: number) => void
  /** Send a chat message via WebSocket */
  sendChatMessage: (payload: ChatSendPayload) => Promise<ChatSendAck>
  /** Send Chat Shell guidance via WebSocket */
  sendChatGuidance: (payload: ChatGuidePayload) => Promise<ChatGuideAck>
  /** Cancel a chat stream via WebSocket */
  cancelChatStream: (
    subtaskId: number,
    partialContent?: string,
    shellType?: string
  ) => Promise<{ success: boolean; error?: string }>
  /** Close a device task session via WebSocket */
  closeTaskSession: (taskId: number) => Promise<{ success: boolean; error?: string }>
  /** Retry a failed message via WebSocket */
  retryMessage: (
    taskId: number,
    subtaskId: number,
    modelId?: string,
    modelType?: string,
    forceOverride?: boolean
  ) => Promise<{ success: boolean; error?: string }>
  /** Register chat event handlers */
  registerChatHandlers: (handlers: ChatEventHandlers) => () => void
  /** Register task event handlers */
  registerTaskHandlers: (handlers: TaskEventHandlers) => () => void
  /** Register skill event handlers */
  registerSkillHandlers: (handlers: SkillEventHandlers) => () => void
  /** Send skill response back to server */
  sendSkillResponse: (payload: SkillResponsePayload) => void
  /** Register correction event handlers */
  registerCorrectionHandlers: (handlers: CorrectionEventHandlers) => () => void
  /** Register background execution event handlers */
  registerBackgroundExecutionHandlers: (handlers: BackgroundExecutionEventHandlers) => () => void
  /** Register a callback to be called when WebSocket reconnects */
  onReconnect: (callback: ReconnectCallback) => () => void
}

/** Chat event handlers for streaming */
export interface ChatEventHandlers {
  onChatStart?: (data: ChatStartPayload) => void
  onChatChunk?: (data: ChatChunkPayload) => void
  onChatDone?: (data: ChatDonePayload) => void
  onChatError?: (data: ChatErrorPayload) => void
  onChatCancelled?: (data: ChatCancelledPayload) => void
  /** Handler for chat:message event (other users' messages in group chat) */
  onChatMessage?: (data: ChatMessagePayload) => void
  /** Handler for chat:block_created event (new block added) */
  onBlockCreated?: (data: ChatBlockCreatedPayload) => void
  /** Handler for chat:block_updated event (block content/status updated) */
  onBlockUpdated?: (data: ChatBlockUpdatedPayload) => void
  /** Handler for chat:guidance_queued event */
  onGuidanceQueued?: (data: ChatGuidanceQueuedPayload) => void
  /** Handler for chat:guidance_applied event */
  onGuidanceApplied?: (data: ChatGuidanceAppliedPayload) => void
  /** Handler for chat:guidance_expired event */
  onGuidanceExpired?: (data: ChatGuidanceExpiredPayload) => void
}

/** Task event handlers for task list updates */
export interface TaskEventHandlers {
  onTaskCreated?: (data: TaskCreatedPayload) => void
  onTaskInvited?: (data: TaskInvitedPayload) => void
  onTaskStatus?: (data: TaskStatusPayload) => void
  /** Handler for task:app_update event (app preview data updated, sent to task room) */
  onTaskAppUpdate?: (data: TaskAppUpdatePayload) => void
}

/** Skill event handlers for generic skill requests */
export interface SkillEventHandlers {
  /** Handler for skill:request event (server requests frontend to perform a skill action) */
  onSkillRequest?: (data: SkillRequestPayload) => void
}

/** Correction event handlers for cross-validation progress */
export interface CorrectionEventHandlers {
  onCorrectionStart?: (data: CorrectionStartPayload) => void
  onCorrectionProgress?: (data: CorrectionProgressPayload) => void
  onCorrectionChunk?: (data: CorrectionChunkPayload) => void
  onCorrectionDone?: (data: CorrectionDonePayload) => void
  onCorrectionError?: (data: CorrectionErrorPayload) => void
}

/** Background execution event handlers for subscription execution updates */
export interface BackgroundExecutionEventHandlers {
  onBackgroundExecutionUpdate?: (data: BackgroundExecutionUpdatePayload) => void
}

const SocketContext = createContext<SocketContextType | undefined>(undefined)

export function SocketProvider({ children }: { children: ReactNode }) {
  const [socket, setSocket] = useState<Socket | null>(null)
  const [isConnected, setIsConnected] = useState(false)
  const [connectionError, setConnectionError] = useState<Error | null>(null)
  const [reconnectAttempts, setReconnectAttempts] = useState(0)

  // Track current joined tasks
  const joinedTasksRef = useRef<Set<number>>(new Set())
  // Use ref for socket to avoid dependency issues in connect callback
  const socketRef = useRef<Socket | null>(null)
  // Store reconnect callbacks - single source of truth for reconnection events
  const reconnectCallbacksRef = useRef<Set<ReconnectCallback>>(new Set())

  /**
   * Internal function to create socket connection
   */
  const createSocketConnection = useCallback((token: string, socketUrl: string) => {
    // Create new socket connection
    // Transport strategy:
    // 1. Try WebSocket first (preferred for load-balanced environments without sticky sessions)
    // 2. If WebSocket fails (e.g., load balancer doesn't support it), fall back to polling
    // Note: Polling requires sticky sessions in load-balanced environments
    const newSocket = io(socketUrl + '/chat', {
      path: SOCKETIO_PATH,
      auth: { token },
      autoConnect: true,
      reconnection: true,
      reconnectionAttempts: Infinity,
      reconnectionDelay: 1000,
      reconnectionDelayMax: 5000,
      // Try websocket first, then fall back to polling if websocket fails
      // This handles cases where load balancer doesn't support WebSocket upgrade
      transports: ['websocket', 'polling'],
      // Increase timeout for mobile networks which may have higher latency
      timeout: 20000,
      // Force new connection to avoid stale connections on mobile
      forceNew: false,
      // Disable automatic upgrade from polling to websocket
      // This prevents "Invalid transport" errors when switching transports
      upgrade: true,
    })

    // Store in ref immediately
    socketRef.current = newSocket

    // Connection event handlers
    newSocket.on('connect', () => {
      setIsConnected(true)
      setConnectionError(null)
      setReconnectAttempts(0)
    })

    newSocket.on('disconnect', (_: string) => {
      setIsConnected(false)
      // Don't clear joinedTasksRef here - we need it for rejoining after reconnect
    })

    newSocket.io.on('reconnect_attempt', (attempt: number) => {
      setReconnectAttempts(attempt)
    })

    newSocket.io.on('reconnect', (_attempt: number) => {
      setIsConnected(true)
      setConnectionError(null)
      setReconnectAttempts(0)

      // Trigger recovery for all active tasks via TaskStateManager
      // This handles message state recovery after WebSocket reconnection
      if (taskStateManager.isInitialized()) {
        taskStateManager.recoverAll().catch(err => {
          console.error('[Socket.IO] Error recovering tasks after reconnect:', err)
        })
      }

      // Notify all registered reconnect callbacks
      // This is the single source of truth for reconnection events
      reconnectCallbacksRef.current.forEach(callback => {
        try {
          callback()
        } catch (err) {
          console.error('[Socket.IO] Error in reconnect callback:', err)
        }
      })
    })

    newSocket.io.on('reconnect_error', (error: Error) => {
      console.error('[Socket.IO] Reconnect error:', error)
      setConnectionError(error)
    })

    // Handle authentication errors (token expired during session)
    newSocket.on(ServerEvents.AUTH_ERROR, (_: AuthErrorPayload) => {
      // Remove token and redirect to login
      removeToken()
      newSocket.disconnect()

      const loginPath = paths.auth.login.getHref()
      if (typeof window !== 'undefined' && window.location.pathname !== loginPath) {
        // Save current path for redirect after login
        const currentPath = window.location.pathname + window.location.search
        sessionStorage.setItem(POST_LOGIN_REDIRECT_KEY, currentPath)
        window.location.href = loginPath
      }
    })

    // Handle connect_error for initial connection auth failures
    newSocket.on('connect_error', (error: Error) => {
      console.error('[Socket.IO] Connection error:', error)
      setConnectionError(error)
      setIsConnected(false)

      // Check if error message indicates auth failure
      // Use specific auth-related error patterns to avoid false positives
      const errorMsg = error.message?.toLowerCase() || ''
      const isAuthError =
        errorMsg.includes('expired') ||
        errorMsg.includes('unauthorized') ||
        errorMsg.includes('jwt') ||
        errorMsg.includes('authentication')

      if (isAuthError) {
        removeToken()

        const loginPath = paths.auth.login.getHref()
        if (typeof window !== 'undefined' && window.location.pathname !== loginPath) {
          // Save current path for redirect after login (consistent with AUTH_ERROR handler)
          const currentPath = window.location.pathname + window.location.search
          sessionStorage.setItem(POST_LOGIN_REDIRECT_KEY, currentPath)
          window.location.href = loginPath
        }
      }
    })

    setSocket(newSocket)
  }, []) // No dependencies - use refs instead

  /**
   * Connect to Socket.IO server
   * Fetches runtime config first to allow runtime URL changes
   */
  const connect = useCallback(
    (token: string) => {
      // Check if already connected using ref
      if (socketRef.current?.connected) {
        return
      }

      // Disconnect existing socket if any
      if (socketRef.current) {
        socketRef.current.disconnect()
        socketRef.current = null
      }

      // Fetch runtime config then connect
      // This allows RUNTIME_SOCKET_DIRECT_URL to be changed without rebuilding
      fetchRuntimeConfig().then(config => {
        const socketUrl = config.socketDirectUrl || getSocketUrl()
        createSocketConnection(token, socketUrl)
      })
    },
    [createSocketConnection]
  )

  /**
   * Disconnect from server
   */
  const disconnect = useCallback(() => {
    if (socket) {
      socket.disconnect()
      setSocket(null)
      setIsConnected(false)
      joinedTasksRef.current.clear()
    }
  }, [socket])

  /**
   * Join a task room
   * Prevents duplicate joins by checking if already joined
   * @param taskId - The task ID to join
   * @param options - Join options
   * @param options.forceRefresh - If true, always emit task:join to get streaming status
   * @param options.afterMessageId - If provided, only return messages after this ID (for incremental sync)
   */
  const joinTask = useCallback(
    async (
      taskId: number,
      options?: {
        forceRefresh?: boolean
        afterMessageId?: number
      }
    ): Promise<{
      streaming?: {
        subtask_id: number
        offset: number
        cached_content: string
        started_at?: string
        last_activity_at?: string
      }
      subtasks?: Array<Record<string, unknown>>
      error?: string
    }> => {
      const { forceRefresh = false, afterMessageId } = options || {}

      // Use socketRef for reliable access (socket state may be stale)
      const currentSocket = socketRef.current
      if (!currentSocket?.connected) {
        return { error: 'Not connected' }
      }

      // Check if already joined this task room to prevent duplicate joins.
      // State-machine recovery passes forceRefresh/afterMessageId, so it never uses
      // this dedupe path for refresh or reconnect recovery.
      const alreadyJoined = joinedTasksRef.current.has(taskId)
      const shouldSkip = alreadyJoined && !forceRefresh && afterMessageId === undefined

      if (shouldSkip) {
        return {}
      }

      // Add to set IMMEDIATELY to prevent concurrent duplicate joins
      // This is crucial because the socket.emit is async and multiple calls
      // could pass the above check before any callback completes
      joinedTasksRef.current.add(taskId)

      // Build payload with optional after_message_id for incremental sync
      const payload: { task_id: number; after_message_id?: number } = { task_id: taskId }
      if (afterMessageId !== undefined) {
        payload.after_message_id = afterMessageId
      }

      return new Promise(resolve => {
        currentSocket.emit(
          'task:join',
          payload,
          (response: {
            streaming?: {
              subtask_id: number
              offset: number
              cached_content: string
              started_at?: string
              last_activity_at?: string
            }
            subtasks?: Array<Record<string, unknown>>
            error?: string
          }) => {
            const subtasksCount = Array.isArray(response.subtasks) ? response.subtasks.length : null
            const firstSubtask = Array.isArray(response.subtasks) ? response.subtasks[0] : undefined
            const lastSubtask =
              Array.isArray(response.subtasks) && response.subtasks.length > 0
                ? response.subtasks[response.subtasks.length - 1]
                : undefined

            if (response.streaming) {
              console.info('[StreamingJoinDebug] task:join ack', {
                taskId,
                forceRefresh,
                afterMessageId,
                subtasksCount,
                firstMessageId: firstSubtask?.message_id,
                lastMessageId: lastSubtask?.message_id,
                subtaskId: response.streaming.subtask_id,
                startedAt: response.streaming.started_at,
                lastActivityAt: response.streaming.last_activity_at,
                cachedContentLength: response.streaming.cached_content?.length || 0,
              })
            } else {
              console.info('[StreamingJoinDebug] task:join ack (no streaming)', {
                taskId,
                forceRefresh,
                afterMessageId,
                subtasksCount,
                firstMessageId: firstSubtask?.message_id,
                lastMessageId: lastSubtask?.message_id,
                hasError: Boolean(response.error),
              })
            }

            // If there was an error, remove from the set so it can be retried
            if (response.error) {
              joinedTasksRef.current.delete(taskId)
            }
            resolve(response)
          }
        )
      })
    },
    [] // No dependencies - use socketRef for stable reference
  )

  /**
   * Leave a task room
   */
  const leaveTask = useCallback((taskId: number) => {
    // Use socketRef for reliable access (socket state may be stale)
    const currentSocket = socketRef.current
    if (currentSocket?.connected) {
      currentSocket.emit('task:leave', { task_id: taskId })
      joinedTasksRef.current.delete(taskId)
    }
  }, []) // No dependencies - use socketRef for stable reference

  /**
   * Send a chat message via WebSocket
   */
  const sendChatMessage = useCallback(
    async (payload: ChatSendPayload): Promise<ChatSendAck> => {
      // Use socketRef for reliable access (socket state may be stale)
      const currentSocket = socketRef.current

      if (!currentSocket?.connected) {
        console.error('[Socket.IO] sendChatMessage failed: not connected', {
          hasSocket: !!currentSocket,
          isConnected: currentSocket?.connected,
        })
        return { error: 'Not connected to server' }
      }

      return new Promise(resolve => {
        currentSocket.emit(ClientEvents.CHAT_SEND, payload, (response: ChatSendAck) => {
          resolve(response)
        })
      })
    },
    [] // No dependencies - use socketRef
  )

  /**
   * Send Chat Shell guidance via WebSocket
   */
  const sendChatGuidance = useCallback(async (payload: ChatGuidePayload): Promise<ChatGuideAck> => {
    const currentSocket = socketRef.current

    if (!currentSocket?.connected) {
      console.error('[Socket.IO] sendChatGuidance failed: not connected', {
        hasSocket: !!currentSocket,
        isConnected: currentSocket?.connected,
      })
      return { success: false, error: 'Not connected to server' }
    }

    return new Promise(resolve => {
      currentSocket.emit(ClientEvents.CHAT_GUIDE, payload, (response: ChatGuideAck) => {
        resolve(response)
      })
    })
  }, [])

  /**
   * Cancel a chat stream via WebSocket
   */
  const cancelChatStream = useCallback(
    async (
      subtaskId: number,
      partialContent?: string,
      shellType?: string
    ): Promise<{ success: boolean; error?: string }> => {
      if (!socket?.connected) {
        console.error('[Socket.IO] cancelChatStream failed - not connected')
        return { success: false, error: 'Not connected to server' }
      }

      return new Promise(resolve => {
        socket.emit(
          'chat:cancel',
          {
            subtask_id: subtaskId,
            partial_content: partialContent,
            shell_type: shellType,
          },
          (response: { success?: boolean; error?: string }) => {
            resolve({ success: response.success ?? true, error: response.error })
          }
        )
      })
    },
    [socket]
  )

  /**
   * Close a device task session via WebSocket
   */
  const closeTaskSession = useCallback(
    async (taskId: number): Promise<{ success: boolean; error?: string }> => {
      if (!socket?.connected) {
        console.error('[Socket.IO] closeTaskSession failed - not connected')
        return { success: false, error: 'Not connected to server' }
      }

      return new Promise(resolve => {
        socket.emit(
          'task:close-session',
          { task_id: taskId },
          (response: { success?: boolean; error?: string }) => {
            resolve({ success: response.success ?? true, error: response.error })
          }
        )
      })
    },
    [socket]
  )

  /**
   * Retry a failed message via WebSocket
   */
  const retryMessage = useCallback(
    async (
      taskId: number,
      subtaskId: number,
      modelId?: string,
      modelType?: string,
      forceOverride: boolean = false
    ): Promise<{ success: boolean; error?: string }> => {
      if (!socket?.connected) {
        console.error('[Socket.IO] retryMessage failed - not connected')
        return { success: false, error: 'Not connected to server' }
      }

      const payload = {
        task_id: taskId,
        subtask_id: subtaskId,
        force_override_bot_model: modelId,
        force_override_bot_model_type: modelType,
        use_model_override: forceOverride,
      }

      return new Promise(resolve => {
        socket.emit(
          'chat:retry',
          payload,
          (response: { success?: boolean; error?: string } | undefined) => {
            // Handle undefined response (backend error or no acknowledgment)
            if (!response) {
              console.error('[Socket.IO] chat:retry received undefined response')
              resolve({
                success: false,
                error: 'No response from server',
              })
              return
            }

            resolve({
              success: response.success ?? false,
              error: response.error,
            })
          }
        )
      })
    },
    [socket]
  )

  /**
   * Register chat event handlers
   * Returns a cleanup function to unregister handlers
   */
  const registerChatHandlers = useCallback(
    (handlers: ChatEventHandlers): (() => void) => {
      if (!socket) {
        return () => {}
      }

      const {
        onChatStart,
        onChatChunk,
        onChatDone,
        onChatError,
        onChatCancelled,
        onChatMessage,
        onBlockCreated,
        onBlockUpdated,
        onGuidanceQueued,
        onGuidanceApplied,
        onGuidanceExpired,
      } = handlers

      if (onChatStart) socket.on(ServerEvents.CHAT_START, onChatStart)
      if (onChatChunk) socket.on(ServerEvents.CHAT_CHUNK, onChatChunk)
      if (onChatDone) socket.on(ServerEvents.CHAT_DONE, onChatDone)
      if (onChatError) socket.on(ServerEvents.CHAT_ERROR, onChatError)
      if (onChatCancelled) socket.on(ServerEvents.CHAT_CANCELLED, onChatCancelled)
      if (onChatMessage) socket.on(ServerEvents.CHAT_MESSAGE, onChatMessage)
      if (onBlockCreated) socket.on(ServerEvents.CHAT_BLOCK_CREATED, onBlockCreated)
      if (onBlockUpdated) socket.on(ServerEvents.CHAT_BLOCK_UPDATED, onBlockUpdated)
      if (onGuidanceQueued) socket.on(ServerEvents.CHAT_GUIDANCE_QUEUED, onGuidanceQueued)
      if (onGuidanceApplied) socket.on(ServerEvents.CHAT_GUIDANCE_APPLIED, onGuidanceApplied)
      if (onGuidanceExpired) socket.on(ServerEvents.CHAT_GUIDANCE_EXPIRED, onGuidanceExpired)

      // Return cleanup function
      return () => {
        if (onChatStart) socket.off(ServerEvents.CHAT_START, onChatStart)
        if (onChatChunk) socket.off(ServerEvents.CHAT_CHUNK, onChatChunk)
        if (onChatDone) socket.off(ServerEvents.CHAT_DONE, onChatDone)
        if (onChatError) socket.off(ServerEvents.CHAT_ERROR, onChatError)
        if (onChatCancelled) socket.off(ServerEvents.CHAT_CANCELLED, onChatCancelled)
        if (onChatMessage) socket.off(ServerEvents.CHAT_MESSAGE, onChatMessage)
        if (onBlockCreated) socket.off(ServerEvents.CHAT_BLOCK_CREATED, onBlockCreated)
        if (onBlockUpdated) socket.off(ServerEvents.CHAT_BLOCK_UPDATED, onBlockUpdated)
        if (onGuidanceQueued) socket.off(ServerEvents.CHAT_GUIDANCE_QUEUED, onGuidanceQueued)
        if (onGuidanceApplied) socket.off(ServerEvents.CHAT_GUIDANCE_APPLIED, onGuidanceApplied)
        if (onGuidanceExpired) socket.off(ServerEvents.CHAT_GUIDANCE_EXPIRED, onGuidanceExpired)
      }
    },
    [socket]
  )

  /**
   * Register task event handlers for task list updates
   * Returns a cleanup function to unregister handlers
   */
  const registerTaskHandlers = useCallback(
    (handlers: TaskEventHandlers): (() => void) => {
      if (!socket) {
        return () => {}
      }

      const { onTaskCreated, onTaskInvited, onTaskStatus, onTaskAppUpdate } = handlers

      if (onTaskCreated) socket.on(ServerEvents.TASK_CREATED, onTaskCreated)
      if (onTaskInvited) socket.on(ServerEvents.TASK_INVITED, onTaskInvited)
      if (onTaskStatus) socket.on(ServerEvents.TASK_STATUS, onTaskStatus)
      if (onTaskAppUpdate) socket.on(ServerEvents.TASK_APP_UPDATE, onTaskAppUpdate)

      // Return cleanup function
      return () => {
        if (onTaskCreated) socket.off(ServerEvents.TASK_CREATED, onTaskCreated)
        if (onTaskInvited) socket.off(ServerEvents.TASK_INVITED, onTaskInvited)
        if (onTaskStatus) socket.off(ServerEvents.TASK_STATUS, onTaskStatus)
        if (onTaskAppUpdate) socket.off(ServerEvents.TASK_APP_UPDATE, onTaskAppUpdate)
      }
    },
    [socket]
  )

  /**
   * Register correction event handlers for cross-validation progress
   * Returns a cleanup function to unregister handlers
   */
  const registerCorrectionHandlers = useCallback(
    (handlers: CorrectionEventHandlers): (() => void) => {
      if (!socket) {
        return () => {}
      }

      const {
        onCorrectionStart,
        onCorrectionProgress,
        onCorrectionChunk,
        onCorrectionDone,
        onCorrectionError,
      } = handlers

      if (onCorrectionStart) socket.on(ServerEvents.CORRECTION_START, onCorrectionStart)
      if (onCorrectionProgress) socket.on(ServerEvents.CORRECTION_PROGRESS, onCorrectionProgress)
      if (onCorrectionChunk) socket.on(ServerEvents.CORRECTION_CHUNK, onCorrectionChunk)
      if (onCorrectionDone) socket.on(ServerEvents.CORRECTION_DONE, onCorrectionDone)
      if (onCorrectionError) socket.on(ServerEvents.CORRECTION_ERROR, onCorrectionError)

      // Return cleanup function
      return () => {
        if (onCorrectionStart) socket.off(ServerEvents.CORRECTION_START, onCorrectionStart)
        if (onCorrectionProgress) socket.off(ServerEvents.CORRECTION_PROGRESS, onCorrectionProgress)
        if (onCorrectionChunk) socket.off(ServerEvents.CORRECTION_CHUNK, onCorrectionChunk)
        if (onCorrectionDone) socket.off(ServerEvents.CORRECTION_DONE, onCorrectionDone)
        if (onCorrectionError) socket.off(ServerEvents.CORRECTION_ERROR, onCorrectionError)
      }
    },
    [socket]
  )

  /**
   * Register skill event handlers for generic skill requests
   * Returns a cleanup function to unregister handlers
   */
  const registerSkillHandlers = useCallback(
    (handlers: SkillEventHandlers): (() => void) => {
      if (!socket) {
        return () => {}
      }

      const { onSkillRequest } = handlers

      if (onSkillRequest) socket.on(ServerEvents.SKILL_REQUEST, onSkillRequest)

      // Return cleanup function
      return () => {
        if (onSkillRequest) socket.off(ServerEvents.SKILL_REQUEST, onSkillRequest)
      }
    },
    [socket]
  )

  /**
   * Send skill response back to server
   */
  const sendSkillResponse = useCallback(
    (payload: SkillResponsePayload): void => {
      const currentSocket = socketRef.current

      if (!currentSocket?.connected) {
        console.error('[Socket.IO] sendSkillResponse failed: not connected')
        return
      }

      currentSocket.emit(ClientSkillEvents.SKILL_RESPONSE, payload)
    },
    [] // No dependencies - use socketRef
  )
  /**
   * Register background execution event handlers for subscription execution updates
   * Returns a cleanup function to unregister handlers
   */
  const registerBackgroundExecutionHandlers = useCallback(
    (handlers: BackgroundExecutionEventHandlers): (() => void) => {
      if (!socket) {
        return () => {}
      }

      const { onBackgroundExecutionUpdate } = handlers

      if (onBackgroundExecutionUpdate)
        socket.on(ServerEvents.BACKGROUND_EXECUTION_UPDATE, onBackgroundExecutionUpdate)

      // Return cleanup function
      return () => {
        if (onBackgroundExecutionUpdate)
          socket.off(ServerEvents.BACKGROUND_EXECUTION_UPDATE, onBackgroundExecutionUpdate)
      }
    },
    [socket]
  )

  /**
   * Register a callback to be called when WebSocket reconnects
   * This is the single source of truth for reconnection events in the app.
   * Returns a cleanup function to unregister the callback.
   */
  const onReconnect = useCallback((callback: ReconnectCallback): (() => void) => {
    reconnectCallbacksRef.current.add(callback)
    return () => {
      reconnectCallbacksRef.current.delete(callback)
    }
  }, [])

  // Auto-connect when component mounts if token is available
  useEffect(() => {
    // Only run on client side
    if (typeof window === 'undefined') {
      return
    }

    // Check if already connected
    if (socketRef.current?.connected) {
      return
    }

    const token = getToken()
    if (token) {
      connect(token)
    } else {
      console.error('[Socket.IO] No token found, skipping auto-connect')
    }
  }, [connect])

  // Listen for token changes (login/logout) - works across tabs
  // Also poll for token changes in current tab since storage event doesn't fire for same-tab changes
  useEffect(() => {
    // Handle cross-tab storage changes
    const handleStorageChange = (e: StorageEvent) => {
      if (e.key === 'auth_token') {
        if (e.newValue) {
          // Token was set (login from another tab)
          connect(e.newValue)
        } else {
          // Token was removed (logout from another tab)
          disconnect()
        }
      }
    }

    window.addEventListener('storage', handleStorageChange)
    return () => {
      window.removeEventListener('storage', handleStorageChange)
    }
  }, [connect, disconnect])

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

  return (
    <SocketContext.Provider
      value={{
        socket,
        isConnected,
        connectionError,
        reconnectAttempts,
        connect,
        disconnect,
        joinTask,
        leaveTask,
        sendChatMessage,
        sendChatGuidance,
        cancelChatStream,
        closeTaskSession,
        retryMessage,
        registerChatHandlers,
        registerTaskHandlers,
        registerSkillHandlers,
        sendSkillResponse,
        registerCorrectionHandlers,
        registerBackgroundExecutionHandlers,
        onReconnect,
      }}
    >
      {children}
    </SocketContext.Provider>
  )
}

/**
 * Hook to use socket context
 */
export function useSocket(): SocketContextType {
  const context = useContext(SocketContext)
  if (!context) {
    throw new Error('useSocket must be used within a SocketProvider')
  }
  return context
}

/**
 * Hook to auto-connect socket when token is available
 * @deprecated Socket now auto-connects in SocketProvider
 */
export function useSocketAutoConnect(token: string | null) {
  const { connect, disconnect: _disconnect, isConnected } = useSocket()

  useEffect(() => {
    if (token && !isConnected) {
      connect(token)
    }
    return () => {
      // Don't disconnect on cleanup - let the provider manage lifecycle
    }
  }, [token, connect, isConnected])
}
