Tree

"use client";

import {
  createContext,
  useContext,
  useEffect,
  useRef,
  useState,
  type ComponentProps,
  type FocusEvent,
  type KeyboardEvent,
  type ReactNode,
} from "react";
import "./tree.css";

export type TreeProps = Omit<ComponentProps<"ul">, "children" | "onSelect"> & {
  /** Names the tree for assistive tech, e.g. "Project files". */
  label: string;
  /** Id of the selected item, when you control it. */
  selected?: string | null;
  defaultSelected?: string;
  onSelectedChange?: (id: string) => void;
  /** Ids of the items that start open. */
  defaultExpanded?: string[];
  /** TreeItem elements. */
  children: ReactNode;
};

type TreeContext = {
  expanded: Set<string>;
  selected: string | null;
  tabStop: string | null;
  toggle: (id: string) => void;
  select: (id: string) => void;
};

const Context = createContext<TreeContext>({
  expanded: new Set(),
  selected: null,
  tabStop: null,
  toggle: () => {},
  select: () => {},
});

const Level = createContext(1);

function itemOf(node: EventTarget | null) {
  return node instanceof Element ? node.closest<HTMLElement>('[role="treeitem"]') : null;
}

export function Tree({
  label,
  selected: selectedProp,
  defaultSelected,
  onSelectedChange,
  defaultExpanded,
  className,
  children,
  onKeyDown,
  onFocus,
  ...rest
}: TreeProps) {
  const ref = useRef<HTMLUListElement>(null);
  const [expanded, setExpanded] = useState(() => new Set(defaultExpanded));
  const [internal, setInternal] = useState<string | null>(defaultSelected ?? null);
  const [focused, setFocused] = useState<string | null>(null);
  const selected = selectedProp !== undefined ? selectedProp : internal;

  function visible() {
    return Array.from(ref.current?.querySelectorAll<HTMLElement>('[role="treeitem"]') ?? []).filter(
      (item) => !item.parentElement?.closest('[role="treeitem"][aria-expanded="false"]'),
    );
  }

  // One item is always the tab stop: the focused one, else the selected one, else the first.
  // If a collapse hides it, the stop moves up to the nearest visible ancestor.
  const [tabStop, setTabStop] = useState<string | null>(null);
  useEffect(() => {
    const items = visible();
    const want = focused ?? selected;
    let item = items.find((node) => node.dataset.id === want);
    if (!item && want) {
      const hidden = ref.current?.querySelector<HTMLElement>(`[role="treeitem"][data-id="${CSS.escape(want)}"]`);
      let parent = itemOf(hidden?.parentElement ?? null);
      while (parent && !items.includes(parent)) parent = itemOf(parent.parentElement);
      item = parent ?? undefined;
    }
    setTabStop((item ?? items[0])?.dataset.id ?? null);
  }, [focused, selected, expanded]);

  function toggle(id: string, open?: boolean) {
    setExpanded((current) => {
      const next = new Set(current);
      if (open ?? !next.has(id)) next.add(id);
      else next.delete(id);
      return next;
    });
  }

  function select(id: string) {
    if (selectedProp === undefined) setInternal(id);
    if (id !== selected) onSelectedChange?.(id);
  }

  function handleFocus(event: FocusEvent<HTMLUListElement>) {
    onFocus?.(event);
    const item = itemOf(event.target);
    if (item?.dataset.id) setFocused(item.dataset.id);
  }

  function handleKeyDown(event: KeyboardEvent<HTMLUListElement>) {
    onKeyDown?.(event);
    if (event.defaultPrevented) return;
    const item = itemOf(event.target);
    const id = item?.dataset.id;
    if (!item || id === undefined) return;
    const items = visible();
    const index = items.indexOf(item);
    const parent = itemOf(item.parentElement);
    const isParent = item.hasAttribute("aria-expanded");
    const isOpen = item.getAttribute("aria-expanded") === "true";
    const usable = item.getAttribute("aria-disabled") !== "true";
    let next: HTMLElement | null | undefined;

    switch (event.key) {
      case "ArrowDown":
        next = items[index + 1];
        break;
      case "ArrowUp":
        next = items[index - 1];
        break;
      case "Home":
        next = items[0];
        break;
      case "End":
        next = items[items.length - 1];
        break;
      case "ArrowRight":
        if (!isParent) break;
        if (!isOpen) {
          if (usable) toggle(id, true);
        } else {
          next = item.querySelector<HTMLElement>('[role="group"] > [role="treeitem"]');
        }
        event.preventDefault();
        break;
      case "ArrowLeft":
        if (isParent && isOpen && usable) {
          toggle(id, false);
          event.preventDefault();
        } else {
          next = parent;
        }
        break;
      case "Enter":
      case " ":
        event.preventDefault();
        if (usable) select(id);
        return;
      case "*": {
        const siblings = Array.from(item.parentElement?.children ?? []).filter(
          (node): node is HTMLElement =>
            node instanceof HTMLElement && node.getAttribute("aria-expanded") === "false" && node.getAttribute("aria-disabled") !== "true",
        );
        setExpanded((current) => new Set([...current, ...siblings.map((node) => node.dataset.id ?? "")]));
        event.preventDefault();
        return;
      }
      default: {
        if (event.key.length !== 1 || event.ctrlKey || event.metaKey || event.altKey) return;
        const letter = event.key.toLowerCase();
        for (let step = 1; step <= items.length; step++) {
          const candidate = items[(index + step) % items.length];
          const text = candidate.querySelector(":scope > .eq-tree-row .eq-tree-label")?.textContent ?? "";
          if (text.trim().toLowerCase().startsWith(letter)) {
            next = candidate;
            break;
          }
        }
      }
    }

    if (!next) return;
    event.preventDefault();
    next.focus();
  }

  return (
    <ul
      {...rest}
      ref={ref}
      role="tree"
      aria-label={label}
      className={["eq-tree", className].filter(Boolean).join(" ")}
      onFocus={handleFocus}
      onKeyDown={handleKeyDown}
    >
      <Context.Provider value={{ expanded, selected, tabStop, toggle, select }}>{children}</Context.Provider>
    </ul>
  );
}

export type TreeItemProps = Omit<ComponentProps<"li">, "id" | "children"> & {
  /** Identifies the item for selection and expansion. Not used as the DOM id. */
  id: string;
  label: ReactNode;
  /** An icon before the label, e.g. a folder or file glyph. */
  icon?: ReactNode;
  /** Nested TreeItem elements. Their presence makes the item expandable. */
  children?: ReactNode;
  disabled?: boolean;
};

export function TreeItem({ id, label, icon, children, disabled, className, ...rest }: TreeItemProps) {
  const { expanded, selected, tabStop, toggle, select } = useContext(Context);
  const level = useContext(Level);
  const isParent = children != null && children !== false;
  const isOpen = isParent && expanded.has(id);

  return (
    <li
      {...rest}
      role="treeitem"
      className={["eq-tree-item", className].filter(Boolean).join(" ")}
      data-id={id}
      tabIndex={tabStop === id ? 0 : -1}
      aria-level={level}
      aria-expanded={isParent ? isOpen : undefined}
      aria-selected={disabled ? undefined : selected === id}
      aria-disabled={disabled || undefined}
    >
      <div
        className="eq-tree-row"
        onClick={() => {
          if (disabled) return;
          select(id);
          if (isParent) toggle(id);
        }}
      >
        <span className="eq-tree-toggle" aria-hidden="true">
          {isParent && (
            <svg className="eq-tree-chevron" viewBox="0 0 16 16" fill="none" stroke="currentColor" strokeWidth="1.5" strokeLinecap="round" strokeLinejoin="round">
              <path d="M6 4l4 4-4 4" />
            </svg>
          )}
        </span>
        {icon != null && (
          <span className="eq-tree-icon" aria-hidden="true">
            {icon}
          </span>
        )}
        <span className="eq-tree-label">{label}</span>
      </div>
      {isParent && (
        <div className="eq-tree-panel">
          <ul role="group" className="eq-tree-group">
            <Level.Provider value={level + 1}>{children}</Level.Provider>
          </ul>
        </div>
      )}
    </li>
  );
}