Listbox

'use client'

import {
  createContext,
  use,
  useId,
  type ComponentPropsWithRef,
  type KeyboardEvent
} from 'react'

import { useControllableState } from './sui-use-controllable-state'
import { idSafe, nextIndex, ownItems, useTypeahead } from './sui-keyboard'

export interface ListboxKeyAction {
  key: string
  value: string
  label: string
}

interface ListboxContextValue {
  baseId: string
  value: string
  select: (value: string) => void
}

const ListboxContext = createContext<ListboxContextValue | null>(null)

const useListbox = (part: string) => {
  const context = use(ListboxContext)
  if (!context) throw new Error(`<${part}> must be used inside <Listbox>`)
  return context
}

const optionId = (baseId: string, value: string) =>
  `${baseId}-option-${idSafe(value)}`

const OPTIONS = '[role="option"]'
const ROOT = '[data-scope="listbox"][data-part="root"]'

export interface ListboxProps
  extends Omit<ComponentPropsWithRef<'div'>, 'defaultValue' | 'onChange'> {
  value?: string
  defaultValue?: string
  onValueChange?: (value: string) => void
  /** Called after a key moves the selection, with the key and the new option. */
  onKeyAction?: (action: ListboxKeyAction) => void
}

/**
 * A single-select listbox (APG Listbox). Focus stays on the list and
 * aria-activedescendant points at the selected option, which follows the
 * arrow keys, Home, End and typeahead. Name it with aria-label or
 * aria-labelledby.
 */
export function Listbox({
  value,
  defaultValue = '',
  onValueChange,
  onKeyAction,
  onKeyDown,
  children,
  ...props
}: ListboxProps) {
  const baseId = useId()
  const [current, setCurrent] = useControllableState({
    prop: value,
    defaultProp: defaultValue,
    onChange: onValueChange
  })
  const typeahead = useTypeahead()

  const handleKeyDown = (event: KeyboardEvent<HTMLDivElement>) => {
    onKeyDown?.(event)
    if (event.defaultPrevented) return
    const options = ownItems(event.currentTarget, OPTIONS, ROOT)
    const index = options.findIndex(option => option.dataset.value === current)
    let next = nextIndex(event.key, index, options.length, 'vertical')
    if (next === null) {
      next = typeahead(
        event.key,
        options.map(option => option.textContent ?? ''),
        index
      )
    }
    if (next === null) return
    event.preventDefault()
    const option = options[next]
    const nextValue = option.dataset.value ?? ''
    setCurrent(nextValue)
    option.scrollIntoView?.({ block: 'nearest' })
    onKeyAction?.({
      key: event.key,
      value: nextValue,
      label: option.textContent?.trim() ?? ''
    })
  }

  return (
    <ListboxContext value={{ baseId, value: current, select: setCurrent }}>
      <div
        tabIndex={0}
        {...props}
        role="listbox"
        aria-activedescendant={current ? optionId(baseId, current) : undefined}
        data-scope="listbox"
        data-part="root"
        onKeyDown={handleKeyDown}
      >
        {children}
      </div>
    </ListboxContext>
  )
}

export interface ListboxOptionProps extends ComponentPropsWithRef<'div'> {
  value: string
  disabled?: boolean
}

export function ListboxOption({
  value,
  disabled = false,
  onClick,
  ...props
}: ListboxOptionProps) {
  const { baseId, value: current, select } = useListbox('ListboxOption')
  const selected = current === value

  return (
    <div
      {...props}
      id={optionId(baseId, value)}
      role="option"
      aria-selected={selected}
      aria-disabled={disabled || undefined}
      data-scope="listbox"
      data-part="option"
      data-state={selected ? 'selected' : 'unselected'}
      data-disabled={disabled ? '' : undefined}
      data-value={value}
      onClick={event => {
        onClick?.(event)
        if (!event.defaultPrevented && !disabled) select(value)
      }}
    />
  )
}

import { useCallback, useState } from 'react'

/**
 * Controlled when `prop` is defined, uncontrolled otherwise. `onChange`
 * fires for every change either way.
 */
export function useControllableState<T>({
  prop,
  defaultProp,
  onChange
}: {
  prop: T | undefined
  defaultProp: T
  onChange?: (value: T) => void
}) {
  const [uncontrolled, setUncontrolled] = useState(defaultProp)
  const controlled = prop !== undefined
  const value = controlled ? prop : uncontrolled

  const setValue = useCallback(
    (next: T) => {
      if (!controlled) setUncontrolled(next)
      if (!Object.is(next, value)) onChange?.(next)
    },
    [controlled, onChange, value]
  )

  return [value, setValue] as const
}

import { useRef } from 'react'

export type Orientation = 'horizontal' | 'vertical'

/**
 * The next index for a list navigated with arrow keys, Home and End, or
 * null when the key is not a navigation key. Lists wrap, as the APG
 * patterns used here recommend. `both` accepts all four arrows (radio).
 */
export function nextIndex(
  key: string,
  current: number,
  count: number,
  orientation: Orientation | 'both'
): number | null {
  if (count === 0) return null
  const forward =
    (orientation !== 'horizontal' && key === 'ArrowDown') ||
    (orientation !== 'vertical' && key === 'ArrowRight')
  const backward =
    (orientation !== 'horizontal' && key === 'ArrowUp') ||
    (orientation !== 'vertical' && key === 'ArrowLeft')

  if (key === 'Home') return 0
  if (key === 'End') return count - 1
  if (forward) return current < 0 ? 0 : (current + 1) % count
  if (backward) return current < 0 ? count - 1 : (current - 1 + count) % count
  return null
}

/**
 * Items that belong to this container, skipping disabled ones and any that
 * sit inside a nested instance of the same component.
 */
export function ownItems(
  container: HTMLElement,
  itemSelector: string,
  containerSelector: string
): HTMLElement[] {
  return Array.from(container.querySelectorAll<HTMLElement>(itemSelector)).filter(
    el =>
      el.closest(containerSelector) === container &&
      !el.hasAttribute('disabled') &&
      el.getAttribute('aria-disabled') !== 'true'
  )
}

/**
 * Typeahead: printable keys typed within half a second build a search
 * string, and the first matching label after the current item wins.
 */
export function useTypeahead() {
  const buffer = useRef('')
  const timer = useRef<ReturnType<typeof setTimeout> | undefined>(undefined)

  return (key: string, labels: string[], current: number): number | null => {
    if (key.length !== 1 || key === ' ') return null
    clearTimeout(timer.current)
    timer.current = setTimeout(() => (buffer.current = ''), 500)
    buffer.current += key.toLowerCase()

    const search = buffer.current
    const repeated = search.split('').every(char => char === search[0])
    const query = repeated ? search[0] : search
    const start = repeated ? current + 1 : current
    for (let i = 0; i < labels.length; i++) {
      const index = (start + i + labels.length) % labels.length
      if (labels[index].trim().toLowerCase().startsWith(query)) return index
    }
    return null
  }
}

/** Makes a value safe to use inside an id attribute. */
export const idSafe = (value: string) => value.replace(/[^\w-]/g, '_')