claude-code/packages/@ant/ink/src/keybindings/KeybindingContext.tsx

import React, {
  createContext,
  type RefObject,
  useContext,
  useLayoutEffect,
  useMemo,
} from 'react'
import type { Key } from '../core/events/input-event.js'
import {
  type ChordResolveResult,
  getBindingDisplayText,
  resolveKeyWithChordState,
} from './resolver.js'
import type {
  KeybindingContextName,
  ParsedBinding,
  ParsedKeystroke,
} from './types.js'

/** Handler registration for action callbacks */
type HandlerRegistration = {
  action: string
  context: KeybindingContextName
  handler: () => void
}

type KeybindingContextValue = {
  /** Resolve a key input to an action name (with chord support) */
  resolve: (
    input: string,
    key: Key,
    activeContexts: KeybindingContextName[],
  ) => ChordResolveResult

  /** Update the pending chord state */
  setPendingChord: (pending: ParsedKeystroke[] | null) => void

  /** Get display text for an action (e.g., "ctrl+t") */
  getDisplayText: (
    action: string,
    context: KeybindingContextName,
  ) => string | undefined

  /** All parsed bindings (for help display) */
  bindings: ParsedBinding[]

  /** Current pending chord keystrokes (null if not in a chord) */
  pendingChord: ParsedKeystroke[] | null

  /** Currently active keybinding contexts (for priority resolution) */
  activeContexts: Set<KeybindingContextName>

  /** Register a context as active (call on mount) */
  registerActiveContext: (context: KeybindingContextName) => void

  /** Unregister a context (call on unmount) */
  unregisterActiveContext: (context: KeybindingContextName) => void

  /** Register a handler for an action (used by useKeybinding) */
  registerHandler: (registration: HandlerRegistration) => () => void

  /** Invoke all handlers for an action (used by ChordInterceptor) */
  invokeAction: (action: string) => boolean
}

const KeybindingContext = createContext<KeybindingContextValue | null>(null)

type ProviderProps = {
  bindings: ParsedBinding[]
  /** Ref for immediate access to pending chord (avoids React state delay) */
  pendingChordRef: RefObject<ParsedKeystroke[] | null>
  /** State value for re-renders (UI updates) */
  pendingChord: ParsedKeystroke[] | null
  setPendingChord: (pending: ParsedKeystroke[] | null) => void
  activeContexts: Set<KeybindingContextName>
  registerActiveContext: (context: KeybindingContextName) => void
  unregisterActiveContext: (context: KeybindingContextName) => void
  /** Ref to handler registry (used by ChordInterceptor) */
  handlerRegistryRef: RefObject<Map<string, Set<HandlerRegistration>>>
  children: React.ReactNode
}

export function KeybindingProvider({
  bindings,
  pendingChordRef,
  pendingChord,
  setPendingChord,
  activeContexts,
  registerActiveContext,
  unregisterActiveContext,
  handlerRegistryRef,
  children,
}: ProviderProps): React.ReactNode {
  const value = useMemo<KeybindingContextValue>(() => {
    const getDisplay = (action: string, context: KeybindingContextName) =>
      getBindingDisplayText(action, context, bindings)

    // Register a handler for an action
    const registerHandler = (registration: HandlerRegistration) => {
      const registry = handlerRegistryRef.current
      if (!registry) return () => {}

      if (!registry.has(registration.action)) {
        registry.set(registration.action, new Set())
      }
      registry.get(registration.action)!.add(registration)

      // Return unregister function
      return () => {
        const handlers = registry.get(registration.action)
        if (handlers) {
          handlers.delete(registration)
          if (handlers.size === 0) {
            registry.delete(registration.action)
          }
        }
      }
    }

    // Invoke all handlers for an action
    const invokeAction = (action: string): boolean => {
      const registry = handlerRegistryRef.current
      if (!registry) return false

      const handlers = registry.get(action)
      if (!handlers || handlers.size === 0) return false

      // Find handlers whose context is active
      for (const registration of handlers) {
        if (activeContexts.has(registration.context)) {
          registration.handler()
          return true
        }
      }
      return false
    }

    return {
      // Use ref for immediate access to pending chord, avoiding React state delay
      // This is critical for chord sequences where the second key might be pressed
      // before React re-renders with the updated pendingChord state
      resolve: (input, key, contexts) =>
        resolveKeyWithChordState(
          input,
          key,
          contexts,
          bindings,
          pendingChordRef.current,
        ),
      setPendingChord,
      getDisplayText: getDisplay,
      bindings,
      pendingChord,
      activeContexts,
      registerActiveContext,
      unregisterActiveContext,
      registerHandler,
      invokeAction,
    }
  }, [
    bindings,
    pendingChordRef,
    pendingChord,
    setPendingChord,
    activeContexts,
    registerActiveContext,
    unregisterActiveContext,
    handlerRegistryRef,
  ])

  return (
    <KeybindingContext.Provider value={value}>
      {children}
    </KeybindingContext.Provider>
  )
}

export function useKeybindingContext(): KeybindingContextValue {
  const ctx = useContext(KeybindingContext)
  if (!ctx) {
    throw new Error(
      'useKeybindingContext must be used within KeybindingProvider',
    )
  }
  return ctx
}

/**
 * Optional hook that returns undefined outside of KeybindingProvider.
 * Useful for components that may render before provider is available.
 */
export function useOptionalKeybindingContext(): KeybindingContextValue | null {
  return useContext(KeybindingContext)
}

/**
 * Hook to register a keybinding context as active while the component is mounted.
 *
 * When a context is registered, its keybindings take precedence over Global bindings.
 * This allows context-specific bindings (like ThemePicker's ctrl+t) to override
 * global bindings (like the todo toggle) when the context is active.
 *
 * @example
 * ```tsx
 * function ThemePicker() {
 *   useRegisterKeybindingContext('ThemePicker')
 *   // Now ThemePicker's ctrl+t binding takes precedence over Global
 * }
 * ```
 */
export function useRegisterKeybindingContext(
  context: KeybindingContextName,
  isActive: boolean = true,
): void {
  const keybindingContext = useOptionalKeybindingContext()

  useLayoutEffect(() => {
    if (!keybindingContext || !isActive) return

    keybindingContext.registerActiveContext(context)
    return () => {
      keybindingContext.unregisterActiveContext(context)
    }
  }, [context, keybindingContext, isActive])
}