Dropdown menu for React

'use client';

import { cloneElement, isValidElement, useCallback, useEffect, useId, useRef, useState } from 'react';
import { cn } from '@/lib/kit';
import { Icon, type IconName } from '@/components/ui/icon';
import { Floating, floatingPop, originFor, panelSurface, useDismiss, useMounted, useOpen, usePresence, type Align, type Side } from '@/components/ui/popover';

/* Items arrive one after another, a few pixels from the trigger side, when a menu opens */
const CSS = '.kitfx-menu-item{animation:kitfx-menu-in .32s var(--ease-kit) both;animation-delay:calc(var(--i,0) * 24ms + 30ms)}@keyframes kitfx-menu-in{from{opacity:0;transform:translateY(-4px)}}@media (prefers-reduced-motion:reduce){.kitfx-menu-item{animation:none}}';

export type MenuItem =
  | { type?: 'item'; label: string; icon?: IconName; description?: string; shortcut?: string; href?: string; disabled?: boolean; destructive?: boolean; onSelect?: () => void }
  | { type: 'checkbox'; label: string; icon?: IconName; checked: boolean; onCheckedChange?: (checked: boolean) => void; disabled?: boolean }
  | { type: 'radio'; label: string; value: string; group: string; checked: boolean; onSelect?: (value: string) => void; disabled?: boolean }
  | { type: 'label'; label: string }
  | { type: 'separator' };

type Actionable = Exclude<MenuItem, { type: 'label' } | { type: 'separator' }>;
interface MenuListProps {
  items: MenuItem[];
  /** Something above the items, e.g. the signed-in user */
  header?: React.ReactNode;
  onClose: (focusTrigger: boolean) => void;
  /** Focus the first (or last) item on mount */
  autoFocus?: 'first' | 'last' | 'none' | 'skip';
  labelledBy?: string;
  id?: string;
  className?: string;
  style?: React.CSSProperties;
  /** Set while the menu animates out */
  closing?: boolean;
}

/** The menu panel with full keyboard support: arrows, Home, End, typeahead, Enter, Space, Escape, Tab. */
export function MenuList({ items, header, onClose, autoFocus = 'first', labelledBy, id, className, style, closing }: MenuListProps) {
  const ref = useRef<HTMLDivElement>(null);
  const typed = useRef({ text: '', at: 0 });
  const nodes = () => Array.from(ref.current?.querySelectorAll<HTMLElement>('[data-menu-item]:not([aria-disabled="true"])') ?? []);

  useEffect(() => {
    // After the event that opened the menu, so its own focus handling cannot win
    if (autoFocus === 'skip') return;
    const id = requestAnimationFrame(() => {
      if (autoFocus === 'none') return ref.current?.focus();
      const list = nodes();
      (autoFocus === 'last' ? list[list.length - 1] : list[0])?.focus();
    });
    return () => cancelAnimationFrame(id);
  }, [autoFocus]);

  const move = (delta: number | 'first' | 'last') => {
    const list = nodes();
    if (!list.length) return;
    const i = list.indexOf(document.activeElement as HTMLElement);
    const next = delta === 'first' ? 0 : delta === 'last' ? list.length - 1 : (i + delta + list.length) % list.length;
    list[next].focus();
  };

  const onKeyDown = (e: React.KeyboardEvent) => {
    if (e.key === 'ArrowDown') { e.preventDefault(); move(1); }
    else if (e.key === 'ArrowUp') { e.preventDefault(); move(-1); }
    else if (e.key === 'Home') { e.preventDefault(); move('first'); }
    else if (e.key === 'End') { e.preventDefault(); move('last'); }
    else if (e.key === 'Tab') onClose(false);
    else if (e.key.length === 1 && !e.ctrlKey && !e.metaKey && !e.altKey && e.key !== ' ') {
      const now = Date.now();
      typed.current = { text: (now - typed.current.at > 600 ? '' : typed.current.text) + e.key.toLowerCase(), at: now };
      const hit = nodes().find((n) => n.textContent?.trim().toLowerCase().startsWith(typed.current.text));
      hit?.focus();
    }
  };

  const select = (item: Actionable) => {
    if (item.disabled) return;
    if (item.type === 'checkbox') {
      item.onCheckedChange?.(!item.checked);
      return; // stays open, so several boxes can be toggled
    }
    if (item.type === 'radio') item.onSelect?.(item.value);
    else item.onSelect?.();
    onClose(true);
  };

  return (
    <div
      ref={ref}
      id={id}
      role="menu"
      aria-labelledby={labelledBy}
      tabIndex={-1}
      onKeyDown={onKeyDown}
      data-closing={closing || undefined}
      inert={closing || undefined}
      className={cn('min-w-56 max-w-[calc(100vw-16px)] p-1.5 outline-none', panelSurface, className)}
      style={style}
    >
      <style href="kitfx-menu" precedence="default">
        {CSS}
      </style>
      {header && <div className="kitfx-menu-item -mx-1.5 -mt-1.5 mb-1.5 border-b border-border px-4 py-3">{header}</div>}
      {items.map((item, i) => {
        const order = { '--i': i + (header ? 1 : 0) } as React.CSSProperties;
        if (item.type === 'separator') return <div key={i} role="separator" className="kitfx-menu-item -mx-1.5 my-1.5 h-px bg-border" style={order} />;
        if (item.type === 'label')
          return (
            <div key={i} role="presentation" style={order} className="kitfx-menu-item px-2.5 pb-1 pt-2 text-[11px] font-medium uppercase tracking-[0.08em] text-subtle-foreground">
              {item.label}
            </div>
          );
        const role = item.type === 'checkbox' ? 'menuitemcheckbox' : item.type === 'radio' ? 'menuitemradio' : 'menuitem';
        const checked = item.type === 'checkbox' || item.type === 'radio' ? item.checked : undefined;
        const destructive = item.type === undefined || item.type === 'item' ? item.destructive : false;
        const icon = item.type === 'radio' ? undefined : item.icon;
        const common = {
          role,
          tabIndex: -1,
          'data-menu-item': '',
          'aria-disabled': item.disabled || undefined,
          'aria-checked': checked,
          style: order,
          onClick: () => select(item),
          onKeyDown: (e: React.KeyboardEvent) => {
            if (e.key === 'Enter' || e.key === ' ') {
              e.preventDefault();
              select(item);
              if (e.key === 'Enter' && (item.type === undefined || item.type === 'item') && item.href) window.location.assign(item.href);
            }
          },
          onPointerMove: (e: React.PointerEvent<HTMLElement>) => e.pointerType === 'mouse' && !item.disabled && e.currentTarget.focus(),
          className: cn(
            'kitfx-menu-item group/item relative flex min-h-9 w-full cursor-pointer select-none items-center gap-2.5 rounded-md px-2.5 py-1.5 text-left text-sm text-foreground outline-none transition-colors max-sm:min-h-11',
            'focus:bg-accent aria-disabled:pointer-events-none aria-disabled:opacity-45',
            destructive && 'text-destructive focus:bg-[color-mix(in_oklab,var(--destructive)_10%,transparent)]',
          ),
        };
        const content = (
          <>
            {(item.type === 'checkbox' || item.type === 'radio') && (
              <span className="grid size-4 shrink-0 place-items-center text-foreground" aria-hidden>
                <span className={cn('grid place-items-center transition-[scale,opacity] duration-200 ease-[cubic-bezier(0.34,1.56,0.64,1)] motion-reduce:transition-none', item.checked ? 'scale-100 opacity-100' : 'scale-50 opacity-0')}>
                  {item.type === 'checkbox' ? <Icon name="check" size={15} strokeWidth={2} /> : <span className="size-[7px] rounded-full bg-current" />}
                </span>
              </span>
            )}
            {icon && <Icon name={icon} size={16} className={cn('shrink-0 transition-[translate,color] duration-200 ease-[var(--ease-kit)] group-focus/item:translate-x-0.5 motion-reduce:transition-none', destructive ? 'text-destructive' : 'text-muted-foreground group-focus/item:text-foreground')} />}
            <span className="min-w-0 flex-1">
              <span className="block truncate">{item.label}</span>
              {(item.type === undefined || item.type === 'item') && item.description && <span className="block truncate text-xs text-muted-foreground">{item.description}</span>}
            </span>
            {(item.type === undefined || item.type === 'item') && item.shortcut && (
              <span className="ms-4 shrink-0 text-xs tracking-wide text-subtle-foreground">{item.shortcut}</span>
            )}
          </>
        );
        if ((item.type === undefined || item.type === 'item') && item.href)
          return (
            <a key={i} href={item.href} {...common}>
              {content}
            </a>
          );
        return (
          <div key={i} {...common}>
            {content}
          </div>
        );
      })}
    </div>
  );
}

export interface DropdownMenuProps {
  /** The button that opens the menu */
  trigger: React.ReactElement;
  items: MenuItem[];
  header?: React.ReactNode;
  side?: Side;
  align?: Align;
  open?: boolean;
  defaultOpen?: boolean;
  onOpenChange?: (open: boolean) => void;
  className?: string;
}

/**
 * A menu of actions behind a button.
 *   <DropdownMenu trigger={<Button variant="outline">Options</Button>} items={[{ label: 'Edit', icon: 'edit', shortcut: '⌘E' }, { type: 'separator' }, …]} />
 */
export function DropdownMenu({ trigger, items, header, side = 'bottom', align = 'start', open: openProp, defaultOpen, onOpenChange, className }: DropdownMenuProps) {
  const [open, setOpen] = useOpen(openProp, defaultOpen, onOpenChange);
  const mounted = useMounted();
  const anchorRef = useRef<HTMLSpanElement>(null);
  const panelRef = useRef<HTMLDivElement>(null);
  const [focus, setFocus] = useState<'first' | 'last' | 'none' | 'skip'>(defaultOpen ? 'skip' : 'none');
  const id = useId();
  const triggerId = `${id}-trigger`;
  const { present, closing } = usePresence(open, 150);

  const focusTrigger = () => anchorRef.current?.querySelector<HTMLElement>('button,a,[tabindex]')?.focus();
  const close = useCallback(
    (reason?: 'escape' | 'outside') => {
      setOpen(false);
      if (reason === 'escape') focusTrigger();
    },
    [setOpen],
  );
  const refs = useRef([anchorRef, panelRef]).current;
  useDismiss(open, close, refs);

  const triggerEl = isValidElement<Record<string, unknown>>(trigger)
    ? cloneElement(trigger, {
        id: triggerId,
        'aria-haspopup': 'menu',
        'aria-expanded': open,
        'aria-controls': open ? id : undefined,
        onClick: (e: React.MouseEvent) => {
          (trigger.props.onClick as ((e: React.MouseEvent) => void) | undefined)?.(e);
          setFocus(e.detail === 0 ? 'first' : 'none');
          setOpen(!open);
        },
        onKeyDown: (e: React.KeyboardEvent) => {
          if (e.key === 'ArrowDown' || e.key === 'ArrowUp') {
            e.preventDefault();
            setFocus(e.key === 'ArrowDown' ? 'first' : 'last');
            setOpen(true);
          }
        },
      })
    : trigger;

  return (
    <>
      <span ref={anchorRef} className="inline-flex">
        {triggerEl}
      </span>
      {mounted && present && (
        <Floating anchor={anchorRef} side={side} align={align} offset={6} ref={panelRef}>
          {(pos) => (
            <MenuList
              closing={closing}
              id={id}
              labelledBy={triggerId}
              items={items}
              header={header}
              autoFocus={focus}
              onClose={(f) => {
                setOpen(false);
                if (f) focusTrigger();
              }}
              className={cn(floatingPop, className)}
              style={{ '--origin': originFor(pos, side) } as React.CSSProperties}
            />
          )}
        </Floating>
      )}
    </>
  );
}

/**
 * A menu at the pointer on right click, or with the context menu key and Shift F10.
 *   <ContextMenu items={…}><div>Right click here</div></ContextMenu>
 */
export function ContextMenu({ items, header, className, children }: { items: MenuItem[]; header?: React.ReactNode; className?: string; children: React.ReactNode }) {
  const [at, setAt] = useState<{ x: number; y: number } | null>(null);
  const mounted = useMounted();
  const areaRef = useRef<HTMLDivElement>(null);
  const panelRef = useRef<HTMLDivElement>(null);
  const close = useCallback(() => setAt(null), []);
  const refs = useRef([panelRef]).current;
  useDismiss(Boolean(at), close, refs);
  // The last position stays while the menu leaves
  const last = useRef(at);
  if (at) last.current = at;
  const { present, closing } = usePresence(Boolean(at), 150);
  const anchor = useCallback(() => (last.current ? new DOMRect(last.current.x, last.current.y, 0, 0) : null), [at]);

  return (
    <div
      ref={areaRef}
      tabIndex={0}
      className={cn('rounded-xl', className)}
      onContextMenu={(e) => {
        e.preventDefault();
        if (e.clientX === 0 && e.clientY === 0) {
          const r = e.currentTarget.getBoundingClientRect();
          setAt({ x: r.left + 16, y: r.top + 16 });
        } else setAt({ x: e.clientX, y: e.clientY });
      }}
    >
      {children}
      {mounted && present && (
        <Floating anchor={anchor} side="bottom" align="start" offset={2} ref={panelRef}>
          {() => (
            <MenuList
              closing={closing}
              items={items}
              header={header}
              autoFocus="none"
              onClose={(f) => {
                setAt(null);
                if (f) areaRef.current?.focus();
              }}
              className={cn(floatingPop, '[--origin:top_left]')}
            />
          )}
        </Floating>
      )}
    </div>
  );
}