1  import {
  2    useCallback,
  3    useEffect,
  4    useMemo,
  5    useReducer,
  6    useRef,
  7    useState,
  8  } from 'react'
  9  import { isDeepStrictEqual } from 'util'
 10  import OptionMap from './option-map.js'
 11  import type { OptionWithDescription } from './select.js'
 12  
 13  type State<T> = {
 14    /**
 15     * Map where key is option's value and value is option's index.
 16     */
 17    optionMap: OptionMap<T>
 18  
 19    /**
 20     * Number of visible options.
 21     */
 22    visibleOptionCount: number
 23  
 24    /**
 25     * Value of the currently focused option.
 26     */
 27    focusedValue: T | undefined
 28  
 29    /**
 30     * Index of the first visible option.
 31     */
 32    visibleFromIndex: number
 33  
 34    /**
 35     * Index of the last visible option.
 36     */
 37    visibleToIndex: number
 38  }
 39  
 40  type Action<T> =
 41    | FocusNextOptionAction
 42    | FocusPreviousOptionAction
 43    | FocusNextPageAction
 44    | FocusPreviousPageAction
 45    | SetFocusAction<T>
 46    | ResetAction<T>
 47  
 48  type SetFocusAction<T> = {
 49    type: 'set-focus'
 50    value: T
 51  }
 52  
 53  type FocusNextOptionAction = {
 54    type: 'focus-next-option'
 55  }
 56  
 57  type FocusPreviousOptionAction = {
 58    type: 'focus-previous-option'
 59  }
 60  
 61  type FocusNextPageAction = {
 62    type: 'focus-next-page'
 63  }
 64  
 65  type FocusPreviousPageAction = {
 66    type: 'focus-previous-page'
 67  }
 68  
 69  type ResetAction<T> = {
 70    type: 'reset'
 71    state: State<T>
 72  }
 73  
 74  const reducer = <T>(state: State<T>, action: Action<T>): State<T> => {
 75    switch (action.type) {
 76      case 'focus-next-option': {
 77        if (state.focusedValue === undefined) {
 78          return state
 79        }
 80  
 81        const item = state.optionMap.get(state.focusedValue)
 82  
 83        if (!item) {
 84          return state
 85        }
 86  
 87        // Wrap to first item if at the end
 88        const next = item.next || state.optionMap.first
 89  
 90        if (!next) {
 91          return state
 92        }
 93  
 94        // When wrapping to first, reset viewport to start
 95        if (!item.next && next === state.optionMap.first) {
 96          return {
 97            ...state,
 98            focusedValue: next.value,
 99            visibleFromIndex: 0,
100            visibleToIndex: state.visibleOptionCount,
101          }
102        }
103  
104        const needsToScroll = next.index >= state.visibleToIndex
105  
106        if (!needsToScroll) {
107          return {
108            ...state,
109            focusedValue: next.value,
110          }
111        }
112  
113        const nextVisibleToIndex = Math.min(
114          state.optionMap.size,
115          state.visibleToIndex + 1,
116        )
117  
118        const nextVisibleFromIndex = nextVisibleToIndex - state.visibleOptionCount
119  
120        return {
121          ...state,
122          focusedValue: next.value,
123          visibleFromIndex: nextVisibleFromIndex,
124          visibleToIndex: nextVisibleToIndex,
125        }
126      }
127  
128      case 'focus-previous-option': {
129        if (state.focusedValue === undefined) {
130          return state
131        }
132  
133        const item = state.optionMap.get(state.focusedValue)
134  
135        if (!item) {
136          return state
137        }
138  
139        // Wrap to last item if at the beginning
140        const previous = item.previous || state.optionMap.last
141  
142        if (!previous) {
143          return state
144        }
145  
146        // When wrapping to last, reset viewport to end
147        if (!item.previous && previous === state.optionMap.last) {
148          const nextVisibleToIndex = state.optionMap.size
149          const nextVisibleFromIndex = Math.max(
150            0,
151            nextVisibleToIndex - state.visibleOptionCount,
152          )
153          return {
154            ...state,
155            focusedValue: previous.value,
156            visibleFromIndex: nextVisibleFromIndex,
157            visibleToIndex: nextVisibleToIndex,
158          }
159        }
160  
161        const needsToScroll = previous.index <= state.visibleFromIndex
162  
163        if (!needsToScroll) {
164          return {
165            ...state,
166            focusedValue: previous.value,
167          }
168        }
169  
170        const nextVisibleFromIndex = Math.max(0, state.visibleFromIndex - 1)
171  
172        const nextVisibleToIndex = nextVisibleFromIndex + state.visibleOptionCount
173  
174        return {
175          ...state,
176          focusedValue: previous.value,
177          visibleFromIndex: nextVisibleFromIndex,
178          visibleToIndex: nextVisibleToIndex,
179        }
180      }
181  
182      case 'focus-next-page': {
183        if (state.focusedValue === undefined) {
184          return state
185        }
186  
187        const item = state.optionMap.get(state.focusedValue)
188  
189        if (!item) {
190          return state
191        }
192  
193        // Move by a full page (visibleOptionCount items)
194        const targetIndex = Math.min(
195          state.optionMap.size - 1,
196          item.index + state.visibleOptionCount,
197        )
198  
199        // Find the item at the target index
200        let targetItem = state.optionMap.first
201        while (targetItem && targetItem.index < targetIndex) {
202          if (targetItem.next) {
203            targetItem = targetItem.next
204          } else {
205            break
206          }
207        }
208  
209        if (!targetItem) {
210          return state
211        }
212  
213        // Update the visible range to include the new focused item
214        const nextVisibleToIndex = Math.min(
215          state.optionMap.size,
216          targetItem.index + 1,
217        )
218        const nextVisibleFromIndex = Math.max(
219          0,
220          nextVisibleToIndex - state.visibleOptionCount,
221        )
222  
223        return {
224          ...state,
225          focusedValue: targetItem.value,
226          visibleFromIndex: nextVisibleFromIndex,
227          visibleToIndex: nextVisibleToIndex,
228        }
229      }
230  
231      case 'focus-previous-page': {
232        if (state.focusedValue === undefined) {
233          return state
234        }
235  
236        const item = state.optionMap.get(state.focusedValue)
237  
238        if (!item) {
239          return state
240        }
241  
242        // Move by a full page (visibleOptionCount items)
243        const targetIndex = Math.max(0, item.index - state.visibleOptionCount)
244  
245        // Find the item at the target index
246        let targetItem = state.optionMap.first
247        while (targetItem && targetItem.index < targetIndex) {
248          if (targetItem.next) {
249            targetItem = targetItem.next
250          } else {
251            break
252          }
253        }
254  
255        if (!targetItem) {
256          return state
257        }
258  
259        // Update the visible range to include the new focused item
260        const nextVisibleFromIndex = Math.max(0, targetItem.index)
261        const nextVisibleToIndex = Math.min(
262          state.optionMap.size,
263          nextVisibleFromIndex + state.visibleOptionCount,
264        )
265  
266        return {
267          ...state,
268          focusedValue: targetItem.value,
269          visibleFromIndex: nextVisibleFromIndex,
270          visibleToIndex: nextVisibleToIndex,
271        }
272      }
273  
274      case 'reset': {
275        return action.state
276      }
277  
278      case 'set-focus': {
279        // Early return if already focused on this value
280        if (state.focusedValue === action.value) {
281          return state
282        }
283  
284        const item = state.optionMap.get(action.value)
285        if (!item) {
286          return state
287        }
288  
289        // Check if the item is already in view
290        if (
291          item.index >= state.visibleFromIndex &&
292          item.index < state.visibleToIndex
293        ) {
294          // Already visible, just update focus
295          return {
296            ...state,
297            focusedValue: action.value,
298          }
299        }
300  
301        // Need to scroll to make the item visible
302        // Scroll as little as possible - put item at edge of viewport
303        let nextVisibleFromIndex: number
304        let nextVisibleToIndex: number
305  
306        if (item.index < state.visibleFromIndex) {
307          // Item is above viewport - scroll up to put it at the top
308          nextVisibleFromIndex = item.index
309          nextVisibleToIndex = Math.min(
310            state.optionMap.size,
311            nextVisibleFromIndex + state.visibleOptionCount,
312          )
313        } else {
314          // Item is below viewport - scroll down to put it at the bottom
315          nextVisibleToIndex = Math.min(state.optionMap.size, item.index + 1)
316          nextVisibleFromIndex = Math.max(
317            0,
318            nextVisibleToIndex - state.visibleOptionCount,
319          )
320        }
321  
322        return {
323          ...state,
324          focusedValue: action.value,
325          visibleFromIndex: nextVisibleFromIndex,
326          visibleToIndex: nextVisibleToIndex,
327        }
328      }
329    }
330  }
331  
332  export type UseSelectNavigationProps<T> = {
333    /**
334     * Number of items to display.
335     *
336     * @default 5
337     */
338    visibleOptionCount?: number
339  
340    /**
341     * Options.
342     */
343    options: OptionWithDescription<T>[]
344  
345    /**
346     * Initially focused option's value.
347     */
348    initialFocusValue?: T
349  
350    /**
351     * Callback for focusing an option.
352     */
353    onFocus?: (value: T) => void
354  
355    /**
356     * Value to focus
357     */
358    focusValue?: T
359  }
360  
361  export type SelectNavigation<T> = {
362    /**
363     * Value of the currently focused option.
364     */
365    focusedValue: T | undefined
366  
367    /**
368     * 1-based index of the focused option in the full list.
369     * Returns 0 if no option is focused.
370     */
371    focusedIndex: number
372  
373    /**
374     * Index of the first visible option.
375     */
376    visibleFromIndex: number
377  
378    /**
379     * Index of the last visible option.
380     */
381    visibleToIndex: number
382  
383    /**
384     * All options.
385     */
386    options: OptionWithDescription<T>[]
387  
388    /**
389     * Visible options.
390     */
391    visibleOptions: Array<OptionWithDescription<T> & { index: number }>
392  
393    /**
394     * Whether the focused option is an input type.
395     */
396    isInInput: boolean
397  
398    /**
399     * Focus next option and scroll the list down, if needed.
400     */
401    focusNextOption: () => void
402  
403    /**
404     * Focus previous option and scroll the list up, if needed.
405     */
406    focusPreviousOption: () => void
407  
408    /**
409     * Focus next page and scroll the list down by a page.
410     */
411    focusNextPage: () => void
412  
413    /**
414     * Focus previous page and scroll the list up by a page.
415     */
416    focusPreviousPage: () => void
417  
418    /**
419     * Focus a specific option by value.
420     */
421    focusOption: (value: T | undefined) => void
422  }
423  
424  const createDefaultState = <T>({
425    visibleOptionCount: customVisibleOptionCount,
426    options,
427    initialFocusValue,
428    currentViewport,
429  }: Pick<UseSelectNavigationProps<T>, 'visibleOptionCount' | 'options'> & {
430    initialFocusValue?: T
431    currentViewport?: { visibleFromIndex: number; visibleToIndex: number }
432  }): State<T> => {
433    const visibleOptionCount =
434      typeof customVisibleOptionCount === 'number'
435        ? Math.min(customVisibleOptionCount, options.length)
436        : options.length
437  
438    const optionMap = new OptionMap<T>(options)
439    const focusedItem =
440      initialFocusValue !== undefined && optionMap.get(initialFocusValue)
441    const focusedValue = focusedItem ? initialFocusValue : optionMap.first?.value
442  
443    let visibleFromIndex = 0
444    let visibleToIndex = visibleOptionCount
445  
446    // When there's a valid focused item, adjust viewport to show it
447    if (focusedItem) {
448      const focusedIndex = focusedItem.index
449  
450      if (currentViewport) {
451        // If focused item is already in the current viewport range, try to preserve it
452        if (
453          focusedIndex >= currentViewport.visibleFromIndex &&
454          focusedIndex < currentViewport.visibleToIndex
455        ) {
456          // Keep the same viewport if it's valid
457          visibleFromIndex = currentViewport.visibleFromIndex
458          visibleToIndex = Math.min(
459            optionMap.size,
460            currentViewport.visibleToIndex,
461          )
462        } else {
463          // Need to adjust viewport to show focused item
464          // Use minimal scrolling - put item at edge of viewport
465          if (focusedIndex < currentViewport.visibleFromIndex) {
466            // Item is above current viewport - scroll up to put it at the top
467            visibleFromIndex = focusedIndex
468            visibleToIndex = Math.min(
469              optionMap.size,
470              visibleFromIndex + visibleOptionCount,
471            )
472          } else {
473            // Item is below current viewport - scroll down to put it at the bottom
474            visibleToIndex = Math.min(optionMap.size, focusedIndex + 1)
475            visibleFromIndex = Math.max(0, visibleToIndex - visibleOptionCount)
476          }
477        }
478      } else if (focusedIndex >= visibleOptionCount) {
479        // No current viewport but focused item is outside default viewport
480        // Scroll to show the focused item at the bottom of the viewport
481        visibleToIndex = Math.min(optionMap.size, focusedIndex + 1)
482        visibleFromIndex = Math.max(0, visibleToIndex - visibleOptionCount)
483      }
484  
485      // Ensure viewport bounds are valid
486      visibleFromIndex = Math.max(
487        0,
488        Math.min(visibleFromIndex, optionMap.size - 1),
489      )
490      visibleToIndex = Math.min(
491        optionMap.size,
492        Math.max(visibleOptionCount, visibleToIndex),
493      )
494    }
495  
496    return {
497      optionMap,
498      visibleOptionCount,
499      focusedValue,
500      visibleFromIndex,
501      visibleToIndex,
502    }
503  }
504  
505  export function useSelectNavigation<T>({
506    visibleOptionCount = 5,
507    options,
508    initialFocusValue,
509    onFocus,
510    focusValue,
511  }: UseSelectNavigationProps<T>): SelectNavigation<T> {
512    const [state, dispatch] = useReducer(
513      reducer<T>,
514      {
515        visibleOptionCount,
516        options,
517        initialFocusValue: focusValue || initialFocusValue,
518      } as Parameters<typeof createDefaultState<T>>[0],
519      createDefaultState<T>,
520    )
521  
522    // Store onFocus in a ref to avoid re-running useEffect when callback changes
523    const onFocusRef = useRef(onFocus)
524    onFocusRef.current = onFocus
525  
526    const [lastOptions, setLastOptions] = useState(options)
527  
528    if (options !== lastOptions && !isDeepStrictEqual(options, lastOptions)) {
529      dispatch({
530        type: 'reset',
531        state: createDefaultState({
532          visibleOptionCount,
533          options,
534          initialFocusValue:
535            focusValue ?? state.focusedValue ?? initialFocusValue,
536          currentViewport: {
537            visibleFromIndex: state.visibleFromIndex,
538            visibleToIndex: state.visibleToIndex,
539          },
540        }),
541      })
542  
543      setLastOptions(options)
544    }
545  
546    const focusNextOption = useCallback(() => {
547      dispatch({
548        type: 'focus-next-option',
549      })
550    }, [])
551  
552    const focusPreviousOption = useCallback(() => {
553      dispatch({
554        type: 'focus-previous-option',
555      })
556    }, [])
557  
558    const focusNextPage = useCallback(() => {
559      dispatch({
560        type: 'focus-next-page',
561      })
562    }, [])
563  
564    const focusPreviousPage = useCallback(() => {
565      dispatch({
566        type: 'focus-previous-page',
567      })
568    }, [])
569  
570    const focusOption = useCallback((value: T | undefined) => {
571      if (value !== undefined) {
572        dispatch({
573          type: 'set-focus',
574          value,
575        })
576      }
577    }, [])
578  
579    const visibleOptions = useMemo(() => {
580      return options
581        .map((option, index) => ({
582          ...option,
583          index,
584        }))
585        .slice(state.visibleFromIndex, state.visibleToIndex)
586    }, [options, state.visibleFromIndex, state.visibleToIndex])
587  
588    // Validate that focusedValue exists in current options.
589    // This handles the case where options change during render but the reset
590    // action hasn't been processed yet - without this, the cursor would disappear
591    // because focusedValue points to an option that no longer exists.
592    const validatedFocusedValue = useMemo(() => {
593      if (state.focusedValue === undefined) {
594        return undefined
595      }
596      const exists = options.some(opt => opt.value === state.focusedValue)
597      if (exists) {
598        return state.focusedValue
599      }
600      // Fall back to first option if focused value doesn't exist
601      return options[0]?.value
602    }, [state.focusedValue, options])
603  
604    const isInInput = useMemo(() => {
605      const focusedOption = options.find(
606        opt => opt.value === validatedFocusedValue,
607      )
608      return focusedOption?.type === 'input'
609    }, [validatedFocusedValue, options])
610  
611    // Call onFocus with the validated value (what's actually displayed),
612    // not the internal state value which may be stale if options changed.
613    // Use ref to avoid re-running when callback reference changes.
614    useEffect(() => {
615      if (validatedFocusedValue !== undefined) {
616        onFocusRef.current?.(validatedFocusedValue)
617      }
618    }, [validatedFocusedValue])
619  
620    // Allow parent to programmatically set focus via focusValue prop
621    useEffect(() => {
622      if (focusValue !== undefined) {
623        dispatch({
624          type: 'set-focus',
625          value: focusValue,
626        })
627      }
628    }, [focusValue])
629  
630    // Compute 1-based focused index for scroll position display
631    const focusedIndex = useMemo(() => {
632      if (validatedFocusedValue === undefined) {
633        return 0
634      }
635      const index = options.findIndex(opt => opt.value === validatedFocusedValue)
636      return index >= 0 ? index + 1 : 0
637    }, [validatedFocusedValue, options])
638  
639    return {
640      focusedValue: validatedFocusedValue,
641      focusedIndex,
642      visibleFromIndex: state.visibleFromIndex,
643      visibleToIndex: state.visibleToIndex,
644      visibleOptions,
645      isInInput: isInInput ?? false,
646      focusNextOption,
647      focusPreviousOption,
648      focusNextPage,
649      focusPreviousPage,
650      focusOption,
651      options,
652    }
653  }