Menu

import {
  FloatingList,
  FloatingNode,
  FloatingTree,
  safePolygon,
  type UseInteractionsReturn,
  useClick,
  useDismiss,
  useFloatingNodeId,
  useFloatingParentNodeId,
  useFloatingTree,
  useHover,
  useInteractions,
  useListItem,
  useListNavigation,
  useMergeRefs,
  useRole,
  useTypeahead,
} from "@floating-ui/react";
import { CaretRightIcon } from "@phosphor-icons/react/dist/ssr";
import {
  createContext,
  Fragment,
  use,
  useCallback,
  useEffect,
  useId,
  useLayoutEffect,
  useMemo,
  useRef,
  useState,
} from "react";

import { Slot } from "@/components/slot";
import { useStableCallback } from "@/foundations/hooks/use-stable-callback";
import { Divider } from "@/components/divider";
import {
  Popover,
  PopoverContext,
  type PopoverOrigin,
  usePopoverContext,
  usePopoverFloating,
} from "@/components/popover";
import { cn, cva } from "@/lib/utils/classnames";

const HOVER_OPEN_DELAY = 75;
const HOVER_CLOSE_DELAY = 150;

type Item = {
  id: string;
  label: string;
  onSelect?: (e: { preventDefault: () => void }) => void;
};

type Items = Record<string, Item>;

interface MenuContextType {
  parent: MenuContextType | null;
  isNested: boolean;
  elementsRef: React.RefObject<(HTMLElement | null)[]>;
  highlightedIndex: number | null;
  setHighlightedIndex: React.Dispatch<React.SetStateAction<number | null>>;
  // Ref so MenuItems can pass it to FloatingFocusManager's `initialFocus` —
  // focusing the input directly avoids FFM's deferred tabbable scan, which
  // on iOS races React's autoFocus and drops the soft keyboard.
  searchInputRef: React.RefObject<HTMLInputElement | null>;
  hasSearchInput: boolean;
  setSearchInputEl: (el: HTMLInputElement | null) => void;
  items: Items;
  registerItem: (item: Item) => () => void;
  getItemProps: UseInteractionsReturn["getItemProps"];
}

const MenuContext = createContext<MenuContextType | null>(null);

const useMenuContext = () => {
  const context = use(MenuContext);

  if (context == null) {
    throw new Error("Menu components must be wrapped in <Menu />");
  }

  return context;
};

interface MenuProps {
  open?: boolean;
  onOpenChange?: (open: boolean) => void;
  placement?: React.ComponentProps<typeof Popover>["placement"];
  offset?: number;
  origin?: PopoverOrigin;
  modal?: boolean;
  children?: React.ReactNode;
}

/**
 * Menu is a list of actions presented in a floating panel. Supports nested
 * submenus by rendering a `<Menu>` inside another menu's items, with the inner
 * menu's trigger being a `<Menu.ItemTrigger>`.
 *
 * @example
 * ```
 * <Menu>
 *   <Menu.Trigger asChild>
 *     <Button>Open</Button>
 *   </Menu.Trigger>
 *   <Menu.Items>
 *     <Menu.Item onSelect={...}>New file</Menu.Item>
 *     <Menu>
 *       <Menu.ItemTrigger>More options</Menu.ItemTrigger>
 *       <Menu.Items>
 *         <Menu.Item onSelect={...}>Sub-action</Menu.Item>
 *       </Menu.Items>
 *     </Menu>
 *   </Menu.Items>
 * </Menu>
 * ```
 */
const Menu = (props: MenuProps) => {
  const parentId = useFloatingParentNodeId();
  // The root menu wraps its descendants in a FloatingTree so submenus can
  // coordinate (sibling close, tree-wide click). Nested menus are already
  // inside the tree from the root, so they render as a Fragment.
  const Container = parentId === null ? FloatingTree : Fragment;

  return (
    <Container>
      <MenuRoot {...props} />
    </Container>
  );
};

const MenuRoot = ({ children, modal = true, placement: propPlacement, ...props }: MenuProps) => {
  const parent = use(MenuContext);
  const parentId = useFloatingParentNodeId();
  const isNested = !!parentId;

  const [highlightedIndex, setHighlightedIndex] = useState<number | null>(null);
  const elementsRef = useRef<(HTMLElement | null)[]>([]);
  const [items, setItems] = useState<Items>({});
  const labelsRef = useRef<string[]>([]);
  const searchInputRef = useRef<HTMLInputElement | null>(null);
  const [hasSearchInput, setHasSearchInput] = useState(false);
  const setSearchInputEl = useCallback((el: HTMLInputElement | null) => {
    searchInputRef.current = el;
    setHasSearchInput(!!el);
  }, []);

  useEffect(() => {
    labelsRef.current = Object.values(items).map((item) => item.label);
  }, [items]);

  const registerItem = useCallback((item: Item) => {
    setItems((prev) => ({ ...prev, [item.id]: item }));

    return () => {
      setItems((prev) => {
        const { [item.id]: _, ...rest } = prev;
        return rest;
      });
    };
  }, []);

  const tree = useFloatingTree();
  const nodeId = useFloatingNodeId();

  const floating = usePopoverFloating({
    nodeId,
    placement: propPlacement ?? (isNested ? "right-start" : "bottom-start"),
    offset: isNested ? -4 : 4,
    // On narrow viewports a 224px-wide submenu can't fit on either side of
    // the trigger. Letting `flip` fall back to vertical placements puts it
    // below (or above) the trigger inside the parent panel — same approach
    // as base-ui's `fallbackAxisSideDirection`.
    flipFallbackPlacements: isNested ? ["left-start", "bottom-start", "top-start"] : undefined,
    ...props,
  });

  // Hover-to-open is only meaningful for nested menus. safePolygon avoids
  // closing the menu when the cursor briefly leaves the trigger to traverse
  // toward the submenu.
  const hover = useHover(floating.context, {
    enabled: isNested,
    delay: { open: HOVER_OPEN_DELAY, close: HOVER_CLOSE_DELAY },
    handleClose: safePolygon({ blockPointerEvents: true }),
    mouseOnly: true,
  });

  const click = useClick(floating.context, {
    event: "click",
    ignoreMouse: isNested,
    toggle: !isNested,
  });

  const role = useRole(floating.context, { role: "menu" });
  const dismiss = useDismiss(floating.context, { bubbles: true });

  const listNavigation = useListNavigation(floating.context, {
    listRef: elementsRef,
    activeIndex: highlightedIndex,
    nested: isNested,
    onNavigate: setHighlightedIndex,
    virtual: hasSearchInput,
  });

  const typeahead = useTypeahead(floating.context, {
    // Disable typeahead when a search input is present (it owns letter keys)
    // or in nested submenus (their parent's typeahead handles letter keys at
    // the root level). Floating UI gates by `open` internally, so a closed
    // sibling menu won't hijack keystrokes.
    enabled: !hasSearchInput && !isNested,
    listRef: labelsRef,
    activeIndex: highlightedIndex,
    onMatch: setHighlightedIndex,
  });

  const interactions = useInteractions([hover, click, role, dismiss, listNavigation, typeahead]);

  // Tree-wide coordination:
  // - 'click' closes every menu in the tree when any item is selected.
  // - 'menuopen' lets siblings (same parent, different node) close themselves
  //   when another submenu opens.
  useEffect(() => {
    if (!tree) return;

    const onTreeClick = () => floating.setOpen(false);
    const onSubMenuOpen = (event: { nodeId: string; parentId: string }) => {
      if (event.nodeId !== nodeId && event.parentId === parentId) {
        floating.setOpen(false);
      }
    };

    tree.events.on("click", onTreeClick);
    tree.events.on("menuopen", onSubMenuOpen);

    return () => {
      tree.events.off("click", onTreeClick);
      tree.events.off("menuopen", onSubMenuOpen);
    };
  }, [tree, nodeId, parentId, floating]);

  useEffect(() => {
    if (floating.open && tree) {
      tree.events.emit("menuopen", { parentId, nodeId });
    }
  }, [tree, nodeId, parentId, floating.open]);

  const popoverContextValue = useMemo(
    () => ({
      ...floating,
      ...interactions,
      modal,
    }),
    [floating, interactions, modal],
  );

  const menuContextValue = useMemo<MenuContextType>(
    () => ({
      parent,
      isNested,
      elementsRef,
      highlightedIndex,
      setHighlightedIndex,
      searchInputRef,
      hasSearchInput,
      setSearchInputEl,
      items,
      registerItem,
      getItemProps: interactions.getItemProps,
    }),
    [
      parent,
      isNested,
      highlightedIndex,
      hasSearchInput,
      setSearchInputEl,
      items,
      registerItem,
      interactions.getItemProps,
    ],
  );

  return (
    <FloatingNode id={nodeId}>
      <PopoverContext value={popoverContextValue}>
        <MenuContext value={menuContextValue}>{children}</MenuContext>
      </PopoverContext>
    </FloatingNode>
  );
};

interface MenuTriggerProps extends React.ComponentPropsWithRef<"button"> {
  asChild?: boolean;
}

/**
 * The trigger element that opens the menu. For root menus this is the entry
 * point users click. For nested menus, prefer `<Menu.ItemTrigger>` which
 * styles it as a menu item with a chevron.
 */
const MenuTrigger = ({
  ref: refProp,
  asChild,
  children,
  className,
  ...props
}: MenuTriggerProps) => {
  const popover = usePopoverContext();
  const { isNested, parent } = useMenuContext();
  const item = useListItem();

  const ref = useMergeRefs([popover.refs.setReference, item.ref, refProp]);

  const isHighlighted = parent?.highlightedIndex === item.index;
  const Comp = asChild ? Slot : "button";

  // When nested, the trigger is also an item of the parent menu — so it must
  // pick up parent's getItemProps for keyboard nav and registration.
  const referenceProps = popover.getReferenceProps(isNested ? parent?.getItemProps(props) : props);

  return (
    <Comp
      ref={ref}
      type={asChild ? undefined : "button"}
      className={cn(!asChild && "disabled:opacity-40", className)}
      tabIndex={isNested ? (isHighlighted ? 0 : -1) : 0}
      data-state={popover.open ? "open" : "closed"}
      data-highlighted={isHighlighted || undefined}
      {...referenceProps}
    >
      {children}
    </Comp>
  );
};

const itemTriggerStyle = cva({
  base: [
    "relative mx-(--inset) flex w-[calc(100%-calc(var(--inset)*2))] cursor-pointer select-none items-center gap-1.5 rounded-lg px-3 py-1.5",
    "font-medium text-base text-foreground/80 outline-none",
    "first-of-type:mt-(--inset) last-of-type:mb-(--inset)",
    "data-disabled:pointer-events-none data-disabled:opacity-50",
    "active:bg-background-secondary data-[state=open]:bg-background-secondary data-highlighted:bg-background-secondary",
  ],
});

/**
 * A menu item that opens a nested submenu. Must be used inside a nested
 * `<Menu>` (i.e. one rendered as a child of `<Menu.Items>`).
 */
const MenuItemTrigger = ({ children, className, ...props }: MenuTriggerProps) => {
  const { parent } = useMenuContext();

  if (!parent) {
    throw new Error("<Menu.ItemTrigger> must be rendered inside a nested <Menu>.");
  }

  return (
    <MenuTrigger asChild {...props}>
      <button type="button" role="menuitem" className={itemTriggerStyle({ className })}>
        {children}
        <CaretRightIcon className="ml-auto text-foreground-secondary" />
      </button>
    </MenuTrigger>
  );
};

interface MenuItemsProps extends React.ComponentPropsWithRef<"div"> {
  /**
   * Drop the floating panel and just render the items inline. Use this when
   * something else owns the surface — a Dialog for a command palette, a
   * Drawer for a mobile menu, a card for embedded actions. You keep all of
   * Menu's keyboard nav, search, and item registration; you just don't get
   * the popover.
   */
  inline?: boolean;
}

const MenuItems = ({ ref: refProp, children, className, inline, ...props }: MenuItemsProps) => {
  const popover = usePopoverContext();
  const { isNested, elementsRef, searchInputRef, hasSearchInput } = useMenuContext();

  const ref = useMergeRefs([popover.refs.setFloating, refProp]);

  if (inline) {
    return (
      <FloatingList elementsRef={elementsRef}>
        <div
          ref={ref}
          className={cn("font-medium text-foreground outline-none", className)}
          {...popover.getFloatingProps(props)}
        >
          {children}
        </div>
      </FloatingList>
    );
  }

  return (
    <FloatingList elementsRef={elementsRef}>
      <Popover.Panel
        ref={ref}
        context={popover.context}
        modal={popover.modal}
        isPositioned={popover.isPositioned}
        // Point FFM at the search input so it skips the tabbable scan and
        // focuses the input directly — keeps iOS gesture context intact.
        initialFocus={hasSearchInput ? searchInputRef : isNested ? -1 : 0}
        returnFocus={!isNested}
        animate={!isNested}
        className={cn(
          "z-50 max-h-(--max-height) w-56 scroll-py-(--inset) overflow-auto rounded-xl border border-border bg-background font-medium text-foreground shadow-lg outline-none",
          className,
        )}
        {...popover.getFloatingProps(props)}
      >
        {children}
      </Popover.Panel>
    </FloatingList>
  );
};

const itemStyle = cva({
  base: [
    "relative mx-(--inset) flex w-[calc(100%-calc(var(--inset)*2))] cursor-pointer select-none items-center gap-1.5 rounded-lg px-3 py-1.5",
    "font-medium text-base outline-none",
    "first-of-type:mt-(--inset) last-of-type:mb-(--inset)",
    "data-disabled:pointer-events-none data-disabled:opacity-50",
  ],
  variants: {
    variant: {
      default:
        "text-foreground/80 active:bg-background-secondary data-highlighted:bg-background-secondary",
      destructive: "text-error active:bg-error/10 data-highlighted:bg-error/10",
    },
  },
  defaultVariants: { variant: "default" },
});

interface MenuItemProps extends React.ComponentPropsWithRef<"button"> {
  onSelect?: Item["onSelect"];
  asChild?: boolean;
  variant?: "default" | "destructive";
}

const MenuItem = ({
  ref: refProp,
  children,
  className,
  disabled,
  variant,
  onClick,
  onSelect,
  onKeyDown,
  asChild,
  ...props
}: MenuItemProps) => {
  const itemId = useId();
  const innerRef = useRef<HTMLButtonElement | null>(null);
  const { registerItem, highlightedIndex, getItemProps, searchInputRef } = useMenuContext();
  const popoverCtx = usePopoverContext();
  const tree = useFloatingTree();
  const stableOnSelect = useStableCallback(onSelect);

  const { ref: listItemRef, index } = useListItem();
  const ref = useMergeRefs([listItemRef, refProp, innerRef]);

  const isHighlighted = highlightedIndex === index;
  const Comp = asChild ? Slot : "button";

  useLayoutEffect(() => {
    const text = innerRef.current?.textContent;
    if (!text) return;

    return registerItem({
      id: itemId,
      label: text,
      onSelect: stableOnSelect,
    });
  }, [registerItem, itemId, stableOnSelect]);

  const handleClick = (e: React.MouseEvent<HTMLButtonElement>) => {
    onClick?.(e);
    if (e.defaultPrevented) return;

    stableOnSelect?.(e);
    if (e.defaultPrevented) return;

    // Close every menu in the tree (root + any open submenus).
    if (tree) {
      tree.events.emit("click");
    } else {
      popoverCtx.setOpen(false);
    }
  };

  const handleKeyDown = (e: React.KeyboardEvent<HTMLButtonElement>) => {
    onKeyDown?.(e);
    if (e.defaultPrevented) return;

    // Native button Enter already triggers click → handleClick → onSelect, so
    // we don't handle Enter here. If an item ends up focused while a search
    // input is present (rare — usually focus stays on the input in virtual
    // mode), forward subsequent keystrokes back to the input so typing keeps
    // working.
    if (searchInputRef.current && e.key !== "Enter") {
      searchInputRef.current.focus();
    }
  };

  return (
    <Comp
      ref={ref}
      type={asChild ? undefined : "button"}
      role="menuitem"
      data-item-id={itemId}
      data-highlighted={isHighlighted || undefined}
      tabIndex={isHighlighted ? 0 : -1}
      disabled={disabled || undefined}
      data-disabled={disabled || undefined}
      className={itemStyle({ variant, className })}
      {...getItemProps({
        ...props,
        onKeyDown: handleKeyDown,
        onClick: handleClick,
      })}
    >
      {children}
    </Comp>
  );
};

interface MenuSectionContextType {
  setTitleId: (id: string) => void;
}

const MenuSectionContext = createContext<MenuSectionContextType | null>(null);

const MenuSection = ({ children, className, ...props }: React.ComponentPropsWithRef<"div">) => {
  const [titleId, setTitleId] = useState<string | undefined>(undefined);

  return (
    <MenuSectionContext value={{ setTitleId }}>
      {/** biome-ignore lint/a11y/useSemanticElements: maintain div */}
      <div
        role="group"
        aria-labelledby={titleId}
        className={cn("flex flex-col items-stretch border-border not-first:border-t", className)}
        {...props}
      >
        {children}
      </div>
    </MenuSectionContext>
  );
};

const MenuHeading = ({
  children,
  id: propsId,
  className,
  ...props
}: React.ComponentPropsWithRef<"div">) => {
  const generatedId = useId();
  const id = propsId ?? generatedId;
  const ctx = use(MenuSectionContext);

  useLayoutEffect(() => {
    if (ctx) ctx.setTitleId(id);
  }, [ctx, id]);

  return (
    <div
      id={id}
      className={cn("px-3.5 pt-3 pb-1 font-medium text-foreground-secondary text-sm", className)}
      {...props}
    >
      {children}
    </div>
  );
};

const MenuDivider = ({ className, ...props }: React.ComponentPropsWithRef<"div">) => {
  return <Divider className={cn("my-(--inset)", className)} {...props} />;
};

interface MenuSearchInputProps extends React.ComponentPropsWithRef<"input"> {
  isLoading?: boolean;
}

const MenuSearchInput = ({
  ref: refProp,
  onChange,
  onKeyDown,
  isLoading,
  ...props
}: MenuSearchInputProps) => {
  const internalRef = useRef<HTMLInputElement | null>(null);
  const { highlightedIndex, setHighlightedIndex, items, setSearchInputEl, elementsRef } =
    useMenuContext();
  const popoverCtx = usePopoverContext();
  const tree = useFloatingTree();

  // Callback ref keeps the menu's searchInputRef in sync as the input mounts
  // and unmounts. Using a callback ref (not a useEffect) means the menu's ref
  // is set during commit, before FloatingFocusManager's rAF fires — so its
  // `initialFocus` ref dereferences cleanly.
  const ref = useMergeRefs([refProp, internalRef, setSearchInputEl]);

  const handleChange = (e: React.ChangeEvent<HTMLInputElement>) => {
    onChange?.(e);
    setHighlightedIndex(0);
  };

  const handleKeyDown = (e: React.KeyboardEvent<HTMLInputElement>) => {
    if (e.key === "Enter" && highlightedIndex !== null) {
      const id = elementsRef.current[highlightedIndex]?.dataset.itemId;
      if (id && items[id]) {
        items[id].onSelect?.(e);
        // Mirror MenuItem.handleClick: close the whole tree on selection
        // unless onSelect explicitly preventDefault'd.
        if (!e.defaultPrevented) {
          if (tree) {
            tree.events.emit("click");
          } else {
            popoverCtx.setOpen(false);
          }
        }
      }
    }
    onKeyDown?.(e);
  };

  return (
    <Popover.SearchInput
      ref={ref}
      onChange={handleChange}
      onKeyDown={handleKeyDown}
      isLoading={isLoading}
      {...props}
    />
  );
};

const MenuEmpty = Popover.Empty;

const CompoundMenu = Object.assign(Menu, {
  Trigger: MenuTrigger,
  ItemTrigger: MenuItemTrigger,
  Items: MenuItems,
  Item: MenuItem,
  Section: MenuSection,
  Heading: MenuHeading,
  Divider: MenuDivider,
  SearchInput: MenuSearchInput,
  Empty: MenuEmpty,
});

const useMenuPopoverContext = usePopoverContext;

export { CompoundMenu as Menu, useMenuContext, useMenuPopoverContext };