Menu

'use client'

import {
  createContext,
  use,
  useEffect,
  useId,
  useLayoutEffect,
  useRef,
  type ComponentPropsWithRef,
  type KeyboardEvent,
  type RefObject
} from 'react'
import type { Placement } from '@floating-ui/dom'

import { Slot } from './slot'
import { useComposedRefs } from './sui-refs'
import { useControllableState } from './sui-use-controllable-state'
import { FLOATING_STYLE, useFloating } from './sui-floating'
import { nextIndex, ownItems, useTypeahead } from './sui-keyboard'

type Focus = 'first' | 'last'

interface MenuContextValue {
  open: boolean
  openMenu: (focus: Focus) => void
  closeMenu: (returnFocus: boolean) => void
  focusIntent: RefObject<Focus>
  triggerRef: RefObject<HTMLButtonElement | null>
  triggerId: string
  contentId: string
}

const MenuContext = createContext<MenuContextValue | null>(null)

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

const ITEMS = '[role="menuitem"]'
const CONTENT = '[data-scope="menu"][data-part="content"]'

export interface MenuProps {
  open?: boolean
  defaultOpen?: boolean
  onOpenChange?: (open: boolean) => void
  children?: React.ReactNode
}

/**
 * A menu button (APG Menu Button). The trigger opens a list of actions in
 * the top layer. Arrow keys, Home, End and typeahead move between items;
 * Enter or Space picks one; Esc closes and returns focus to the trigger.
 */
export function Menu({ open, defaultOpen = false, onOpenChange, children }: MenuProps) {
  const [isOpen, setOpen] = useControllableState({
    prop: open,
    defaultProp: defaultOpen,
    onChange: onOpenChange
  })
  const triggerRef = useRef<HTMLButtonElement | null>(null)
  const focusIntent = useRef<Focus>('first')
  const id = useId()

  return (
    <MenuContext
      value={{
        open: isOpen,
        openMenu: focus => {
          focusIntent.current = focus
          setOpen(true)
        },
        closeMenu: returnFocus => {
          setOpen(false)
          if (returnFocus) triggerRef.current?.focus()
        },
        focusIntent,
        triggerRef,
        triggerId: `${id}-trigger`,
        contentId: `${id}-content`
      }}
    >
      {children}
    </MenuContext>
  )
}

export function MenuTrigger({
  asChild,
  onClick,
  onKeyDown,
  ref,
  ...props
}: ComponentPropsWithRef<'button'> & { asChild?: boolean }) {
  const { open, openMenu, closeMenu, triggerRef, triggerId, contentId } =
    useMenu('MenuTrigger')
  const composedRef = useComposedRefs(ref, triggerRef)
  const Comp = asChild ? Slot : 'button'

  return (
    <Comp
      type={asChild ? undefined : 'button'}
      {...props}
      ref={composedRef}
      id={triggerId}
      aria-haspopup="menu"
      aria-expanded={open}
      aria-controls={open ? contentId : undefined}
      data-scope="menu"
      data-part="trigger"
      data-state={open ? 'open' : 'closed'}
      onClick={(event: React.MouseEvent<HTMLButtonElement>) => {
        onClick?.(event)
        if (event.defaultPrevented) return
        if (open) closeMenu(false)
        else openMenu('first')
      }}
      onKeyDown={(event: KeyboardEvent<HTMLButtonElement>) => {
        onKeyDown?.(event)
        if (event.defaultPrevented) return
        if (event.key === 'ArrowDown' || event.key === 'ArrowUp') {
          event.preventDefault()
          openMenu(event.key === 'ArrowDown' ? 'first' : 'last')
        }
      }}
    />
  )
}

export interface MenuContentProps extends ComponentPropsWithRef<'div'> {
  placement?: Placement
}

export function MenuContent(props: MenuContentProps) {
  const { open } = useMenu('MenuContent')
  return open ? <OpenMenuContent {...props} /> : null
}

function OpenMenuContent({
  placement = 'bottom-start',
  onKeyDown,
  style,
  ref,
  ...props
}: MenuContentProps) {
  const { closeMenu, focusIntent, triggerRef, triggerId, contentId } =
    useMenu('MenuContent')
  const contentRef = useRef<HTMLDivElement | null>(null)
  const composedRef = useComposedRefs(ref, contentRef)
  const typeahead = useTypeahead()

  useFloating(triggerRef, contentRef, true, placement)

  useLayoutEffect(() => {
    const content = contentRef.current
    if (!content) return
    const items = ownItems(content, ITEMS, CONTENT)
    const target = focusIntent.current === 'last' ? items.at(-1) : items[0]
    ;(target ?? content).focus()
  }, [focusIntent])

  /* A press anywhere outside the menu and its trigger closes it. */
  useEffect(() => {
    const onPointerDown = (event: PointerEvent) => {
      const target = event.target as Node
      if (contentRef.current?.contains(target)) return
      if (triggerRef.current?.contains(target)) return
      closeMenu(false)
    }
    document.addEventListener('pointerdown', onPointerDown)
    return () => document.removeEventListener('pointerdown', onPointerDown)
  }, [closeMenu, triggerRef])

  const handleKeyDown = (event: KeyboardEvent<HTMLDivElement>) => {
    onKeyDown?.(event)
    if (event.defaultPrevented) return
    if (event.key === 'Escape') {
      event.preventDefault()
      closeMenu(true)
      return
    }
    if (event.key === 'Tab') {
      closeMenu(false)
      return
    }
    const items = ownItems(event.currentTarget, ITEMS, CONTENT)
    const index = items.indexOf(document.activeElement as HTMLElement)
    let next = nextIndex(event.key, index, items.length, 'vertical')
    if (next === null) {
      next = typeahead(
        event.key,
        items.map(item => item.textContent ?? ''),
        index
      )
    }
    if (next === null) return
    event.preventDefault()
    items[next].focus()
  }

  return (
    <div
      {...props}
      ref={composedRef}
      id={contentId}
      role="menu"
      aria-labelledby={triggerId}
      aria-orientation="vertical"
      tabIndex={-1}
      popover="manual"
      style={{ ...FLOATING_STYLE, ...style }}
      data-scope="menu"
      data-part="content"
      data-state="open"
      onKeyDown={handleKeyDown}
    />
  )
}

export interface MenuItemProps
  extends Omit<ComponentPropsWithRef<'div'>, 'onSelect'> {
  onSelect?: () => void
  disabled?: boolean
}

export function MenuItem({
  onSelect,
  disabled = false,
  onClick,
  onKeyDown,
  onPointerMove,
  ...props
}: MenuItemProps) {
  const { closeMenu } = useMenu('MenuItem')

  const select = () => {
    if (disabled) return
    onSelect?.()
    closeMenu(true)
  }

  return (
    <div
      {...props}
      role="menuitem"
      tabIndex={-1}
      aria-disabled={disabled || undefined}
      data-scope="menu"
      data-part="item"
      data-disabled={disabled ? '' : undefined}
      onClick={event => {
        onClick?.(event)
        if (!event.defaultPrevented) select()
      }}
      onKeyDown={event => {
        onKeyDown?.(event)
        if (event.defaultPrevented) return
        if (event.key === 'Enter' || event.key === ' ') {
          event.preventDefault()
          select()
        }
      }}
      /* Pointer and keyboard share one highlight: the focused item. */
      onPointerMove={event => {
        onPointerMove?.(event)
        if (!disabled && document.activeElement !== event.currentTarget) {
          event.currentTarget.focus()
        }
      }}
    />
  )
}

export function MenuSeparator(props: ComponentPropsWithRef<'div'>) {
  return (
    <div
      {...props}
      role="separator"
      aria-orientation="horizontal"
      data-scope="menu"
      data-part="separator"
    />
  )
}

import {
  Children,
  cloneElement,
  isValidElement,
  type CSSProperties,
  type ReactElement,
  type ReactNode,
  type Ref
} from 'react'

import { composeRefs } from './sui-refs'

type AnyProps = Record<string, unknown>

/**
 * Renders its only child with the slot's props merged in, so any component
 * can take over the element (`asChild`). Event handlers run child first,
 * then the slot's, unless the child's handler called preventDefault.
 * Class names join, styles merge, refs compose. No hooks, so it works in
 * server components too.
 */
export function Slot({
  children,
  ...slotProps
}: AnyProps & { children?: ReactNode; ref?: Ref<unknown> }) {
  const child = Children.only(children)
  if (!isValidElement(child)) return null
  const element = child as ReactElement<AnyProps>
  const childProps = element.props

  const merged: AnyProps = { ...slotProps, ...childProps }

  for (const key of Object.keys(slotProps)) {
    const slotValue = slotProps[key]
    const childValue = childProps[key]
    if (/^on[A-Z]/.test(key) && typeof slotValue === 'function') {
      merged[key] =
        typeof childValue === 'function'
          ? (...args: unknown[]) => {
              childValue(...args)
              const event = args[0] as { defaultPrevented?: boolean } | undefined
              if (!event?.defaultPrevented) slotValue(...args)
            }
          : slotValue
    } else if (key === 'className') {
      merged.className = [slotValue, childValue].filter(Boolean).join(' ')
    } else if (key === 'style') {
      merged.style = {
        ...(slotValue as CSSProperties),
        ...(childValue as CSSProperties)
      }
    }
  }

  /* Only touch the ref when there is one: a server component rendering
     <Slot><Link /></Slot> must not hand a function ref to a client child. */
  const slotRef = slotProps.ref as Ref<unknown> | undefined
  const childRef = childProps.ref as Ref<unknown> | undefined
  if (slotRef && childRef) merged.ref = composeRefs(slotRef, childRef)
  else if (slotRef || childRef) merged.ref = slotRef ?? childRef
  else delete merged.ref

  return cloneElement(element, merged)
}

import { useCallback, type Ref, type RefCallback } from 'react'

const setRef = <T,>(ref: Ref<T> | undefined, value: T | null) => {
  if (typeof ref === 'function') ref(value)
  else if (ref) ref.current = value
}

/** One callback ref that feeds several refs. */
export const composeRefs =
  <T,>(...refs: (Ref<T> | undefined)[]): RefCallback<T> =>
  node => {
    for (const ref of refs) setRef(ref, node)
  }

/** A user's ref and a component's own ref, as one stable callback ref. */
export const useComposedRefs = <T,>(
  outer: Ref<T> | undefined,
  inner: Ref<T> | undefined
) => useCallback((node: T | null) => composeRefs(outer, inner)(node), [outer, inner])

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 { useLayoutEffect, type CSSProperties, type RefObject } from 'react'
import {
  autoUpdate,
  computePosition,
  flip,
  offset,
  shift,
  type Placement
} from '@floating-ui/dom'

/**
 * Inline styles every floating part needs: fixed positioning, and no UA
 * popover centring. Coordinates are written imperatively below.
 */
export const FLOATING_STYLE: CSSProperties = {
  position: 'fixed',
  inset: 'auto',
  margin: 0,
  top: 0,
  left: 0
}

/**
 * While open, lifts the floating element into the top layer with the
 * Popover API (so no overflow or z-index can clip it) and keeps it placed
 * next to its anchor. Browsers without the API still get fixed placement.
 */
export function useFloating(
  anchor: RefObject<HTMLElement | null>,
  floating: RefObject<HTMLElement | null>,
  open: boolean,
  placement: Placement
) {
  useLayoutEffect(() => {
    const reference = anchor.current
    const element = floating.current
    if (!open || !reference || !element) return

    try {
      element.showPopover?.()
    } catch {
      /* already open, or no popover attribute */
    }

    const update = () =>
      computePosition(reference, element, {
        placement,
        strategy: 'fixed',
        middleware: [offset(6), flip(), shift({ padding: 8 })]
      }).then(({ x, y, placement: final }) => {
        element.style.left = `${x}px`
        element.style.top = `${y}px`
        element.dataset.side = final.split('-')[0]
      })

    const stop = autoUpdate(reference, element, update)
    return () => {
      stop()
      try {
        element.hidePopover?.()
      } catch {
        /* already hidden */
      }
    }
  }, [anchor, floating, open, placement])
}

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, '_')