1  import { useEffect, useRef } from 'react'
  2  import { KeyboardEvent } from '../ink/events/keyboard-event.js'
  3  // eslint-disable-next-line custom-rules/prefer-use-keybindings -- backward-compat bridge until REPL wires handleKeyDown to <Box onKeyDown>
  4  import { useInput } from '../ink.js'
  5  import {
  6    type AppState,
  7    useAppState,
  8    useSetAppState,
  9  } from '../state/AppState.js'
 10  import {
 11    enterTeammateView,
 12    exitTeammateView,
 13  } from '../state/teammateViewHelpers.js'
 14  import {
 15    getRunningTeammatesSorted,
 16    InProcessTeammateTask,
 17  } from '../tasks/InProcessTeammateTask/InProcessTeammateTask.js'
 18  import {
 19    type InProcessTeammateTaskState,
 20    isInProcessTeammateTask,
 21  } from '../tasks/InProcessTeammateTask/types.js'
 22  import { isBackgroundTask } from '../tasks/types.js'
 23  
 24  // Step teammate selection by delta, wrapping across leader(-1)..teammates(0..n-1)..hide(n).
 25  // First step from a collapsed tree expands it and parks on leader.
 26  function stepTeammateSelection(
 27    delta: 1 | -1,
 28    setAppState: (updater: (prev: AppState) => AppState) => void,
 29  ): void {
 30    setAppState(prev => {
 31      const currentCount = getRunningTeammatesSorted(prev.tasks).length
 32      if (currentCount === 0) return prev
 33  
 34      if (prev.expandedView !== 'teammates') {
 35        return {
 36          ...prev,
 37          expandedView: 'teammates' as const,
 38          viewSelectionMode: 'selecting-agent',
 39          selectedIPAgentIndex: -1,
 40        }
 41      }
 42  
 43      const maxIdx = currentCount // hide row
 44      const cur = prev.selectedIPAgentIndex
 45      const next =
 46        delta === 1
 47          ? cur >= maxIdx
 48            ? -1
 49            : cur + 1
 50          : cur <= -1
 51            ? maxIdx
 52            : cur - 1
 53      return {
 54        ...prev,
 55        selectedIPAgentIndex: next,
 56        viewSelectionMode: 'selecting-agent',
 57      }
 58    })
 59  }
 60  
 61  /**
 62   * Custom hook that handles Shift+Up/Down keyboard navigation for background tasks.
 63   * When teammates (swarm) are present, navigates between leader and teammates.
 64   * When only non-teammate background tasks exist, opens the background tasks dialog.
 65   * Also handles Enter to confirm selection, 'f' to view transcript, and 'k' to kill.
 66   */
 67  export function useBackgroundTaskNavigation(options?: {
 68    onOpenBackgroundTasks?: () => void
 69  }): { handleKeyDown: (e: KeyboardEvent) => void } {
 70    const tasks = useAppState(s => s.tasks)
 71    const viewSelectionMode = useAppState(s => s.viewSelectionMode)
 72    const viewingAgentTaskId = useAppState(s => s.viewingAgentTaskId)
 73    const selectedIPAgentIndex = useAppState(s => s.selectedIPAgentIndex)
 74    const setAppState = useSetAppState()
 75  
 76    // Filter to running teammates and sort alphabetically to match TeammateSpinnerTree display
 77    const teammateTasks = getRunningTeammatesSorted(tasks)
 78    const teammateCount = teammateTasks.length
 79  
 80    // Check for non-teammate background tasks (local_agent, local_bash, etc.)
 81    const hasNonTeammateBackgroundTasks = Object.values(tasks).some(
 82      t => isBackgroundTask(t) && t.type !== 'in_process_teammate',
 83    )
 84  
 85    // Track previous teammate count to detect when teammates are removed
 86    const prevTeammateCountRef = useRef<number>(teammateCount)
 87  
 88    // Clamp selection index if teammates are removed or reset when count becomes 0
 89    useEffect(() => {
 90      const prevCount = prevTeammateCountRef.current
 91      prevTeammateCountRef.current = teammateCount
 92  
 93      setAppState(prev => {
 94        const currentTeammates = getRunningTeammatesSorted(prev.tasks)
 95        const currentCount = currentTeammates.length
 96  
 97        // When teammates are removed (count goes from >0 to 0), reset selection
 98        // Only reset if we previously had teammates (not on initial mount with 0)
 99        // Don't clobber viewSelectionMode if actively viewing a teammate transcript —
100        // the user may be reviewing a completed teammate and needs escape to exit
101        if (
102          currentCount === 0 &&
103          prevCount > 0 &&
104          prev.selectedIPAgentIndex !== -1
105        ) {
106          if (prev.viewSelectionMode === 'viewing-agent') {
107            return {
108              ...prev,
109              selectedIPAgentIndex: -1,
110            }
111          }
112          return {
113            ...prev,
114            selectedIPAgentIndex: -1,
115            viewSelectionMode: 'none',
116          }
117        }
118  
119        // Clamp if index is out of bounds
120        // Max valid index is currentCount (the "hide" row) when spinner tree is shown
121        const maxIndex =
122          prev.expandedView === 'teammates' ? currentCount : currentCount - 1
123        if (currentCount > 0 && prev.selectedIPAgentIndex > maxIndex) {
124          return {
125            ...prev,
126            selectedIPAgentIndex: maxIndex,
127          }
128        }
129  
130        return prev
131      })
132    }, [teammateCount, setAppState])
133  
134    // Get the selected teammate's task info
135    const getSelectedTeammate = (): {
136      taskId: string
137      task: InProcessTeammateTaskState
138    } | null => {
139      if (teammateCount === 0) return null
140      const selectedIndex = selectedIPAgentIndex
141      const task = teammateTasks[selectedIndex]
142      if (!task) return null
143  
144      return { taskId: task.id, task }
145    }
146  
147    const handleKeyDown = (e: KeyboardEvent): void => {
148      // Escape in viewing mode:
149      // - If teammate is running: abort current work only (stops current turn, teammate stays alive)
150      // - If teammate is not running (completed/killed/failed): exit the view back to leader
151      if (e.key === 'escape' && viewSelectionMode === 'viewing-agent') {
152        e.preventDefault()
153        const taskId = viewingAgentTaskId
154        if (taskId) {
155          const task = tasks[taskId]
156          if (isInProcessTeammateTask(task) && task.status === 'running') {
157            // Abort currentWorkAbortController (stops current turn) NOT abortController (kills teammate)
158            task.currentWorkAbortController?.abort()
159            return
160          }
161        }
162        // Teammate is not running or task doesn't exist — exit the view
163        exitTeammateView(setAppState)
164        return
165      }
166  
167      // Escape in selection mode: exit selection without aborting leader
168      if (e.key === 'escape' && viewSelectionMode === 'selecting-agent') {
169        e.preventDefault()
170        setAppState(prev => ({
171          ...prev,
172          viewSelectionMode: 'none',
173          selectedIPAgentIndex: -1,
174        }))
175        return
176      }
177  
178      // Shift+Up/Down for teammate transcript switching (with wrapping)
179      // Index -1 represents the leader, 0+ are teammates
180      // When showSpinnerTree is true, index === teammateCount is the "hide" row
181      if (e.shift && (e.key === 'up' || e.key === 'down')) {
182        e.preventDefault()
183        if (teammateCount > 0) {
184          stepTeammateSelection(e.key === 'down' ? 1 : -1, setAppState)
185        } else if (hasNonTeammateBackgroundTasks) {
186          options?.onOpenBackgroundTasks?.()
187        }
188        return
189      }
190  
191      // 'f' to view selected teammate's transcript (only in selecting mode)
192      if (
193        e.key === 'f' &&
194        viewSelectionMode === 'selecting-agent' &&
195        teammateCount > 0
196      ) {
197        e.preventDefault()
198        const selected = getSelectedTeammate()
199        if (selected) {
200          enterTeammateView(selected.taskId, setAppState)
201        }
202        return
203      }
204  
205      // Enter to confirm selection (only when in selecting mode)
206      if (e.key === 'return' && viewSelectionMode === 'selecting-agent') {
207        e.preventDefault()
208        if (selectedIPAgentIndex === -1) {
209          exitTeammateView(setAppState)
210        } else if (selectedIPAgentIndex >= teammateCount) {
211          // "Hide" row selected - collapse the spinner tree
212          setAppState(prev => ({
213            ...prev,
214            expandedView: 'none' as const,
215            viewSelectionMode: 'none',
216            selectedIPAgentIndex: -1,
217          }))
218        } else {
219          const selected = getSelectedTeammate()
220          if (selected) {
221            enterTeammateView(selected.taskId, setAppState)
222          }
223        }
224        return
225      }
226  
227      // k to kill selected teammate (only in selecting mode)
228      if (
229        e.key === 'k' &&
230        viewSelectionMode === 'selecting-agent' &&
231        selectedIPAgentIndex >= 0
232      ) {
233        e.preventDefault()
234        const selected = getSelectedTeammate()
235        if (selected && selected.task.status === 'running') {
236          void InProcessTeammateTask.kill(selected.taskId, setAppState)
237        }
238        return
239      }
240    }
241  
242    // Backward-compat bridge: REPL.tsx doesn't yet wire handleKeyDown to
243    // <Box onKeyDown>. Subscribe via useInput and adapt InputEvent →
244    // KeyboardEvent until the consumer is migrated (separate PR).
245    // TODO(onKeyDown-migration): remove once REPL passes handleKeyDown.
246    useInput((_input, _key, event) => {
247      handleKeyDown(new KeyboardEvent(event.keypress))
248    })
249  
250    return { handleKeyDown }
251  }