import React, { createContext, useContext, useState, useEffect, useRef, useCallback, useMemo } from 'react'
import type {
  BroadcastMessage,
  TaskData,
  PauseState,
  TaskStatus,
} from '@/lib/types'

// ── Stable context (changes at task boundaries, not every progress tick) ──

interface TaskContextType {
  currentTaskId: number | null
  activeRootId: number | null
  isRunning: boolean
  isExclusive: boolean
  isPaused: boolean
  pauseUntil: number | null
  lastTaskCompletedAt: number | null
  lastTaskScheduledAt: number | null
  backendConnected: boolean | null
  stopTask: (taskId: number) => Promise<void>
  pauseTasks: (durationSeconds: number) => Promise<void>
  unpauseTasks: () => Promise<void>
  notifyTaskScheduled: () => void
}

const TaskContext = createContext<TaskContextType | null>(null)

// ── Progress context (updates every ~250ms during active tasks) ──

interface TaskProgressContextType {
  activeTask: TaskData | null
}

const TaskProgressContext = createContext<TaskProgressContextType>({ activeTask: null })

// ── Provider ──

const TERMINAL_STATUSES: TaskStatus[] = ['completed', 'stopped', 'error']

// API endpoints
const WS_PROGRESS_ENDPOINT = '/ws/tasks/progress'
const API_STOP_ENDPOINT = (taskId: number) => `/api/tasks/${taskId}/stop`

export function TaskProvider({ children }: { children: React.ReactNode }) {
  // Progress state (changes every broadcast)
  const [activeTask, setActiveTask] = useState<TaskData | null>(null)

  // Stable state (changes only at task/pause boundaries)
  const [currentTaskId, setCurrentTaskId] = useState<number | null>(null)
  const [activeRootId, setActiveRootId] = useState<number | null>(null)
  const [isRunning, setIsRunning] = useState(false)
  const [isExclusive, setIsExclusive] = useState(false)
  const [pauseState, setPauseState] = useState<PauseState | null>(null)
  const [lastTaskCompletedAt, setLastTaskCompletedAt] = useState<number | null>(null)
  const [lastTaskScheduledAt, setLastTaskScheduledAt] = useState<number | null>(null)
  const [backendConnected, setBackendConnected] = useState<boolean | null>(null)

  const isPaused = pauseState?.paused ?? false
  const pauseUntil = pauseState?.pauseUntil ?? null

  const wsRef = useRef<WebSocket | null>(null)
  const reconnectTimerRef = useRef<number | null>(null)
  const wsIdRef = useRef(0)
  const lastProcessedState = useRef<{ task_id: number | null; status: TaskStatus | null }>({ task_id: null, status: null })

  // Update stable fields only when they actually change
  const updateStableFields = useCallback((task: TaskData | null) => {
    const newTaskId = task?.task_id ?? null
    const newRootId = task?.active_root_id ?? null
    const newIsRunning = task !== null
    const newIsExclusive = task?.is_exclusive ?? false

    setCurrentTaskId(prev => prev === newTaskId ? prev : newTaskId)
    setActiveRootId(prev => prev === newRootId ? prev : newRootId)
    setIsRunning(prev => prev === newIsRunning ? prev : newIsRunning)
    setIsExclusive(prev => prev === newIsExclusive ? prev : newIsExclusive)
  }, [])

  const handleBroadcastMessage = useCallback((message: BroadcastMessage) => {
    if (message.type === 'paused') {
      setPauseState({ paused: true, pauseUntil: message.pause_until })
      setActiveTask(null)
      updateStableFields(null)
      return
    }

    if (message.type === 'no_active_task') {
      // Signal task completion if we were tracking an active task
      if (lastProcessedState.current.task_id !== null) {
        setLastTaskCompletedAt(Date.now())
        lastProcessedState.current = { task_id: null, status: null }
      }
      setActiveTask(null)
      updateStableFields(null)
      setPauseState(null)
      return
    }

    const task = message.task
    const isTerminal = TERMINAL_STATUSES.includes(task.status)

    // Deduplicate terminal states
    if (isTerminal) {
      const last = lastProcessedState.current
      if (last.task_id === task.task_id && last.status === task.status) {
        return
      }
      lastProcessedState.current = { task_id: task.task_id, status: task.status }
    }

    // Signal task completion when task changes or reaches terminal state
    // (moved out of setActiveTask updater — updaters must be pure)
    const prevTaskId = lastProcessedState.current.task_id
    if (!isTerminal) {
      lastProcessedState.current = { task_id: task.task_id, status: task.status }
    }

    if (prevTaskId !== null && prevTaskId !== task.task_id) {
      setLastTaskCompletedAt(Date.now())
    }

    if (isTerminal) {
      setLastTaskCompletedAt(Date.now())
    }

    const taskData: TaskData = {
      task_id: task.task_id,
      task_type: task.task_type,
      active_root_id: task.active_root_id,
      is_exclusive: task.is_exclusive,
      is_stoppable: task.is_stoppable,
      is_pausable: task.is_pausable,
      action: task.action,
      target: task.target,
      status: task.status,
      error_message: task.error_message,
      breadcrumbs: task.breadcrumbs ?? [],
      phase: task.phase,
      progress_bar: task.progress_bar,
      thread_states: task.thread_states ?? [],
    }

    setActiveTask(taskData)
    // For terminal statuses, immediately clear stable fields so that the
    // idle message (which arrives shortly after) won't trigger a second
    // stableValue change.  This batches all stableValue updates into one
    // render and prevents consumers from flashing/redrawing twice.
    updateStableFields(isTerminal ? null : taskData)
  }, [updateStableFields])

  const connectWebSocket = useCallback(() => {
    if (wsRef.current?.readyState === WebSocket.OPEN) return
    if (wsRef.current) wsRef.current.close()

    const myId = ++wsIdRef.current

    const protocol = window.location.protocol === 'https:' ? 'wss:' : 'ws:'
    const ws = new WebSocket(`${protocol}//${window.location.host}${WS_PROGRESS_ENDPOINT}`)
    wsRef.current = ws

    ws.onopen = () => {
      if (myId !== wsIdRef.current) return
      setBackendConnected(true)
      if (reconnectTimerRef.current) {
        clearTimeout(reconnectTimerRef.current)
        reconnectTimerRef.current = null
      }
    }

    ws.onmessage = (event) => {
      if (myId !== wsIdRef.current) return
      try {
        handleBroadcastMessage(JSON.parse(event.data))
      } catch (error) {
        console.error('Failed to parse WebSocket message:', error)
      }
    }

    ws.onerror = (error) => console.error('WebSocket error:', error)

    ws.onclose = () => {
      if (myId !== wsIdRef.current) return
      setBackendConnected(false)
      if (!reconnectTimerRef.current) {
        reconnectTimerRef.current = window.setTimeout(async () => {
          reconnectTimerRef.current = null
          // Check if the server is migrating before attempting WebSocket reconnect.
          // During migration, the maintenance middleware intercepts WebSocket upgrades,
          // so we navigate to let the browser load the server's maintenance page instead.
          try {
            const resp = await fetch('/health')
            if (resp.ok && (await resp.text()) === 'migrating') {
              window.location.href = '/'
              return
            }
          } catch {
            // Server unreachable — fall through to WebSocket reconnect
          }
          connectWebSocket()
        }, 2000)
      }
    }
  }, [handleBroadcastMessage])

  const stopTask = useCallback(async (taskId: number) => {
    const response = await fetch(API_STOP_ENDPOINT(taskId), { method: 'POST' })
    if (!response.ok) throw new Error('Failed to stop task')
  }, [])

  const pauseTasks = useCallback(async (durationSeconds: number) => {
    const response = await fetch('/api/pause', {
      method: 'POST',
      headers: { 'Content-Type': 'application/json' },
      body: JSON.stringify({ duration_seconds: durationSeconds })
    })
    if (!response.ok) throw new Error('Failed to pause')
  }, [])

  const unpauseTasks = useCallback(async () => {
    const response = await fetch('/api/pause', { method: 'DELETE' })
    if (!response.ok) {
      const text = await response.text()
      throw new Error(text || 'Failed to unpause')
    }
  }, [])

  const notifyTaskScheduled = useCallback(() => {
    setLastTaskScheduledAt(Date.now())
  }, [])

  useEffect(() => {
    connectWebSocket()
    return () => {
      // eslint-disable-next-line react-hooks/exhaustive-deps
      wsIdRef.current++  // Invalidate any in-flight socket callbacks
      wsRef.current?.close()
      if (reconnectTimerRef.current) clearTimeout(reconnectTimerRef.current)
    }
  }, [connectWebSocket])

  const stableValue = useMemo<TaskContextType>(() => ({
    currentTaskId,
    activeRootId,
    isRunning,
    isExclusive,
    isPaused,
    pauseUntil,
    lastTaskCompletedAt,
    lastTaskScheduledAt,
    backendConnected,
    stopTask,
    pauseTasks,
    unpauseTasks,
    notifyTaskScheduled,
  }), [
    currentTaskId, activeRootId, isRunning, isExclusive,
    isPaused, pauseUntil, lastTaskCompletedAt, lastTaskScheduledAt,
    backendConnected, stopTask, pauseTasks, unpauseTasks, notifyTaskScheduled,
  ])

  const progressValue = useMemo<TaskProgressContextType>(() => ({
    activeTask,
  }), [activeTask])

  return (
    <TaskContext.Provider value={stableValue}>
      <TaskProgressContext.Provider value={progressValue}>
        {children}
      </TaskProgressContext.Provider>
    </TaskContext.Provider>
  )
}

export function useTaskContext() {
  const context = useContext(TaskContext)
  if (!context) {
    throw new Error('useTaskContext must be used within TaskProvider')
  }
  return context
}

export function useTaskProgress() {
  return useContext(TaskProgressContext)
}
