Tree View

"use client";

import * as React from "react";
import { AnimatePresence, motion } from "motion/react";

export interface TreeNode {
  /** Unique id across the whole tree. */
  id: string;
  /** Visible label; also used for typeahead. */
  label: string;
  /** Icon shown before the label. Defaults to a folder or file glyph. */
  icon?: React.ReactNode;
  /** Child nodes. Omit and set `hasChildren` to load them lazily with `loadChildren`. */
  children?: TreeNode[];
  /** Marks a parent whose children are loaded on first expand. */
  hasChildren?: boolean;
  /** Focusable but not selectable or checkable. */
  disabled?: boolean;
}

export interface TreeViewProps
  extends Omit<React.HTMLAttributes<HTMLUListElement>, "children" | "onSelect" | "defaultValue" | "defaultChecked"> {
  items: TreeNode[];
  /** Accessible name of the tree. */
  label?: string;
  /** Expanded node ids (controlled). */
  expanded?: string[];
  /** Initially expanded node ids (uncontrolled). */
  defaultExpanded?: string[];
  onExpandedChange?: (ids: string[]) => void;
  /** Selected node ids (controlled). */
  selected?: string[];
  /** Initially selected node ids (uncontrolled). */
  defaultSelected?: string[];
  onSelectedChange?: (ids: string[]) => void;
  /** "single": selection follows focus. "multiple": Ctrl / Shift click and Shift + arrows. */
  selectionMode?: "none" | "single" | "multiple";
  /** Tri-state checkboxes (aria-checked). Replaces selection. */
  checkboxes?: boolean;
  /** Checked ids (controlled). Store leaf ids; parent states are derived. */
  checked?: string[];
  /** Initially checked ids (uncontrolled). */
  defaultChecked?: string[];
  onCheckedChange?: (ids: string[]) => void;
  /** Loads children for nodes with `hasChildren` the first time they expand. */
  loadChildren?: (node: TreeNode) => Promise<TreeNode[]>;
  /** Enter or double-click on a node. */
  onAction?: (node: TreeNode) => void;
  /** Indentation per level in px. */
  indent?: number;
  /** Vertical guide lines for each level. */
  showGuides?: boolean;
  /** Row density. */
  size?: "sm" | "md";
  /** Expand / collapse duration in ms. */
  duration?: number;
  /** Selection, checkbox and focus colour. */
  accent?: string;
}

type CheckState = boolean | "mixed";

interface Row {
  node: TreeNode;
  level: number;
  parentId: string | null;
}

function subscribeReducedMotion(callback: () => void) {
  const mq = window.matchMedia("(prefers-reduced-motion: reduce)");
  mq.addEventListener("change", callback);
  return () => mq.removeEventListener("change", callback);
}

/** SSR-safe: hydrates as "no preference", then follows the media query. */
function usePrefersReducedMotion() {
  return React.useSyncExternalStore(
    subscribeReducedMotion,
    () => window.matchMedia("(prefers-reduced-motion: reduce)").matches,
    () => false,
  );
}

/** Controlled / uncontrolled string-array state. */
function useIdList(value: string[] | undefined, defaultValue: string[] | undefined, onChange?: (ids: string[]) => void) {
  const [internal, setInternal] = React.useState<string[]>(defaultValue ?? []);
  const current = value ?? internal;
  const set = (next: string[]) => {
    if (value === undefined) setInternal(next);
    onChange?.(next);
  };
  return [current, set] as const;
}

/**
 * A file-explorer tree implementing the WAI-ARIA tree pattern: roving focus, arrow / Home / End navigation,
 * typeahead, "*" to expand siblings, single or multiple selection, optional tri-state checkboxes, lazily
 * loaded children with a spinner, and animated expand / collapse.
 */
export function TreeView({
  items,
  label = "Tree",
  expanded: expandedProp,
  defaultExpanded,
  onExpandedChange,
  selected: selectedProp,
  defaultSelected,
  onSelectedChange,
  selectionMode = "single",
  checkboxes = false,
  checked: checkedProp,
  defaultChecked,
  onCheckedChange,
  loadChildren,
  onAction,
  indent = 16,
  showGuides = true,
  size = "md",
  duration = 200,
  accent = "#FF7448",
  className = "",
  style,
  onKeyDown,
  ...props
}: TreeViewProps) {
  const uid = React.useId();
  const reduceMotion = usePrefersReducedMotion();
  const rootRef = React.useRef<HTMLUListElement>(null);
  const itemRefs = React.useRef(new Map<string, HTMLLIElement>());
  const typeaheadRef = React.useRef({ buffer: "", timer: 0 });
  const anchorRef = React.useRef<string | null>(null);

  const [expandedIds, setExpandedIds] = useIdList(expandedProp, defaultExpanded, onExpandedChange);
  const [selectedIds, setSelectedIds] = useIdList(selectedProp, defaultSelected, onSelectedChange);
  const [checkedIds, setCheckedIds] = useIdList(checkedProp, defaultChecked, onCheckedChange);
  const [loaded, setLoaded] = React.useState<Record<string, TreeNode[]>>({});
  const [loading, setLoading] = React.useState<string[]>([]);
  const [focusedId, setFocusedId] = React.useState<string | null>(null);
  const [announcement, setAnnouncement] = React.useState("");

  // Latest lists for async completions (a load can resolve after several re-renders).
  const latestRef = React.useRef({ expandedIds, checkedIds });
  React.useEffect(() => {
    latestRef.current = { expandedIds, checkedIds };
  });

  const expandedSet = new Set(expandedIds);
  const selectedSet = new Set(selectedIds);
  const checkedSet = new Set(checkedIds);
  const selectable = !checkboxes && selectionMode !== "none";
  const multiple = selectable && selectionMode === "multiple";

  const childrenOf = (node: TreeNode): TreeNode[] | undefined => node.children ?? loaded[node.id];
  const isParent = (node: TreeNode) => Boolean(childrenOf(node)?.length) || Boolean(node.hasChildren && !childrenOf(node));

  // Every known node by id (loaded or not visible), plus the visible rows in order.
  const byId = new Map<string, Row>();
  const indexAll = (nodes: TreeNode[], level: number, parentId: string | null) => {
    for (const node of nodes) {
      byId.set(node.id, { node, level, parentId });
      const kids = childrenOf(node);
      if (kids) indexAll(kids, level + 1, node.id);
    }
  };
  indexAll(items, 1, null);

  const visible: Row[] = [];
  const flatten = (nodes: TreeNode[]) => {
    for (const node of nodes) {
      visible.push(byId.get(node.id)!);
      const kids = childrenOf(node);
      if (kids && expandedSet.has(node.id)) flatten(kids);
    }
  };
  flatten(items);
  const visibleIndex = (id: string | null) => visible.findIndex((row) => row.node.id === id);

  // Tab stop: the focused row if visible, else the first selected visible row, else the first row.
  const tabStop =
    (focusedId && visibleIndex(focusedId) >= 0 ? focusedId : null) ??
    visible.find((row) => selectedSet.has(row.node.id))?.node.id ??
    visible[0]?.node.id ??
    null;

  /* --------------------------------- checks --------------------------------- */

  const checkCache = new Map<string, CheckState>();
  const checkState = (node: TreeNode): CheckState => {
    const cached = checkCache.get(node.id);
    if (cached !== undefined) return cached;
    const kids = childrenOf(node);
    let state: CheckState;
    if (!kids?.length) state = checkedSet.has(node.id);
    else {
      let all = true;
      let none = true;
      for (const kid of kids) {
        const s = checkState(kid);
        if (s !== true) all = false;
        if (s !== false) none = false;
      }
      state = all ? true : none ? false : "mixed";
    }
    checkCache.set(node.id, state);
    return state;
  };

  /** Leaf ids under a node (a node without known children counts as its own leaf). */
  const leavesOf = (node: TreeNode): string[] => {
    const kids = childrenOf(node);
    if (!kids?.length) return node.disabled ? [] : [node.id];
    return kids.flatMap(leavesOf);
  };

  const toggleChecked = (node: TreeNode) => {
    if (node.disabled) return;
    const next = new Set(checkedIds);
    const turnOn = checkState(node) !== true;
    for (const id of leavesOf(node)) {
      if (turnOn) next.add(id);
      else next.delete(id);
    }
    setCheckedIds([...next]);
  };

  /* ------------------------------ expand / load ------------------------------ */

  const load = (node: TreeNode) => {
    if (!loadChildren || loading.includes(node.id)) return;
    setLoading((prev) => [...prev, node.id]);
    setAnnouncement(`Loading ${node.label}…`);
    loadChildren(node).then(
      (kids) => {
        setLoaded((prev) => ({ ...prev, [node.id]: kids }));
        setLoading((prev) => prev.filter((id) => id !== node.id));
        setAnnouncement(`${node.label}: ${kids.length} ${kids.length === 1 ? "item" : "items"} loaded`);
        // A checked placeholder hands its state down to the newly loaded children.
        const latestChecked = latestRef.current.checkedIds;
        if (latestChecked.includes(node.id)) {
          const next = new Set(latestChecked.filter((id) => id !== node.id));
          kids.forEach((kid) => !kid.disabled && next.add(kid.id));
          setCheckedIds([...next]);
        }
      },
      () => {
        setLoading((prev) => prev.filter((id) => id !== node.id));
        setExpandedIds(latestRef.current.expandedIds.filter((id) => id !== node.id));
        setAnnouncement(`Could not load ${node.label}`);
      },
    );
  };

  const setExpanded = (node: TreeNode, open: boolean) => {
    if (!isParent(node) || expandedSet.has(node.id) === open) return;
    setExpandedIds(open ? [...expandedIds, node.id] : expandedIds.filter((id) => id !== node.id));
    if (open && !childrenOf(node) && node.hasChildren) load(node);
  };

  /* -------------------------------- selection -------------------------------- */

  const rangeIds = (fromId: string, toId: string) => {
    const a = visibleIndex(fromId);
    const b = visibleIndex(toId);
    if (a < 0 || b < 0) return [toId];
    const [start, end] = a < b ? [a, b] : [b, a];
    return visible.slice(start, end + 1).filter((row) => !row.node.disabled).map((row) => row.node.id);
  };

  const selectOnly = (node: TreeNode) => {
    if (!selectable || node.disabled) return;
    anchorRef.current = node.id;
    if (selectedIds.length !== 1 || selectedIds[0] !== node.id) setSelectedIds([node.id]);
  };

  const toggleSelected = (node: TreeNode) => {
    if (!multiple || node.disabled) return;
    anchorRef.current = node.id;
    setSelectedIds(selectedSet.has(node.id) ? selectedIds.filter((id) => id !== node.id) : [...selectedIds, node.id]);
  };

  const selectRange = (toId: string, additive = false) => {
    if (!multiple) return;
    const anchor = anchorRef.current ?? tabStop ?? toId;
    const range = rangeIds(anchor, toId);
    setSelectedIds(additive ? [...new Set([...selectedIds, ...range])] : range);
  };

  /* ---------------------------------- focus ---------------------------------- */

  /** Keeps the row inside the tree's own scroll area without scrolling the page. */
  const reveal = (id: string) => {
    const root = rootRef.current;
    const row = itemRefs.current.get(id)?.querySelector<HTMLElement>("[data-row]");
    if (!root || !row || root.scrollHeight <= root.clientHeight) return;
    const rootRect = root.getBoundingClientRect();
    const rect = row.getBoundingClientRect();
    if (rect.top < rootRect.top) root.scrollTop += rect.top - rootRect.top - 4;
    else if (rect.bottom > rootRect.bottom) root.scrollTop += rect.bottom - rootRect.bottom + 4;
  };

  const focusRow = (id: string) => {
    setFocusedId(id);
    itemRefs.current.get(id)?.focus({ preventScroll: true });
    reveal(id);
  };

  /* -------------------------------- keyboard -------------------------------- */

  const typeahead = (char: string, fromIndex: number) => {
    const state = typeaheadRef.current;
    window.clearTimeout(state.timer);
    state.buffer += char.toLowerCase();
    state.timer = window.setTimeout(() => (state.buffer = ""), 500);
    const start = state.buffer.length === 1 ? fromIndex + 1 : fromIndex;
    for (let i = 0; i < visible.length; i++) {
      const row = visible[(start + i) % visible.length];
      if (row.node.label.toLowerCase().startsWith(state.buffer)) return row.node.id;
    }
    return null;
  };

  const handleKeyDown = (event: React.KeyboardEvent<HTMLUListElement>) => {
    onKeyDown?.(event);
    if (event.defaultPrevented) return;
    const li = (event.target as HTMLElement).closest<HTMLLIElement>("[role=treeitem]");
    const id = li?.dataset.id;
    if (!id || li !== event.target) return;
    const row = byId.get(id);
    const index = visibleIndex(id);
    if (!row || index < 0) return;
    const { node } = row;
    const mod = event.ctrlKey || event.metaKey;

    const move = (targetIndex: number) => {
      const target = visible[Math.max(0, Math.min(visible.length - 1, targetIndex))];
      if (!target) return;
      focusRow(target.node.id);
      if (multiple && event.shiftKey) {
        if (mod) selectRange(target.node.id, true);
        else toggleSelected(target.node);
      } else if (selectionMode === "single" && selectable) selectOnly(target.node);
    };

    let handled = true;
    switch (event.key) {
      case "ArrowDown":
        move(index + 1);
        break;
      case "ArrowUp":
        move(index - 1);
        break;
      case "ArrowRight":
        if (isParent(node)) {
          if (!expandedSet.has(node.id)) setExpanded(node, true);
          else if (childrenOf(node)?.length) move(index + 1);
        }
        break;
      case "ArrowLeft":
        if (isParent(node) && expandedSet.has(node.id)) setExpanded(node, false);
        else if (row.parentId) move(visibleIndex(row.parentId));
        break;
      case "Home":
        if (multiple && event.shiftKey && mod) selectRange(visible[0].node.id, true);
        move(0);
        break;
      case "End":
        if (multiple && event.shiftKey && mod) selectRange(visible[visible.length - 1].node.id, true);
        move(visible.length - 1);
        break;
      case "Enter":
        if (isParent(node)) setExpanded(node, !expandedSet.has(node.id));
        if (!multiple) selectOnly(node);
        onAction?.(node);
        break;
      case " ":
        if (checkboxes) toggleChecked(node);
        else if (multiple && event.shiftKey) selectRange(node.id);
        else if (multiple) toggleSelected(node);
        else selectOnly(node);
        break;
      case "*": {
        const siblings = row.parentId ? (childrenOf(byId.get(row.parentId)!.node) ?? []) : items;
        const toOpen = siblings.filter((sibling) => isParent(sibling) && !expandedSet.has(sibling.id));
        if (toOpen.length) {
          setExpandedIds([...expandedIds, ...toOpen.map((sibling) => sibling.id)]);
          toOpen.forEach((sibling) => !childrenOf(sibling) && sibling.hasChildren && load(sibling));
        }
        break;
      }
      default:
        if (multiple && mod && event.key.toLowerCase() === "a") {
          setSelectedIds(visible.filter((r) => !r.node.disabled).map((r) => r.node.id));
        } else if (event.key.length === 1 && !mod && !event.altKey && event.key.trim()) {
          const match = typeahead(event.key, index);
          if (match) focusRow(match);
          if (match && selectionMode === "single" && selectable) selectOnly(byId.get(match)!.node);
        } else handled = false;
    }
    if (handled) event.preventDefault();
  };

  /* --------------------------------- pointer --------------------------------- */

  const handleRowClick = (event: React.MouseEvent, node: TreeNode) => {
    focusRow(node.id);
    if (checkboxes) {
      if (isParent(node)) setExpanded(node, !expandedSet.has(node.id));
      else toggleChecked(node);
      return;
    }
    if (multiple && event.shiftKey) {
      selectRange(node.id, event.ctrlKey || event.metaKey);
      return;
    }
    if (multiple && (event.ctrlKey || event.metaKey)) {
      toggleSelected(node);
      return;
    }
    selectOnly(node);
    if (isParent(node)) setExpanded(node, !expandedSet.has(node.id));
  };

  /* ---------------------------------- render --------------------------------- */

  const rowHeight = size === "sm" ? "h-7 text-[12.5px]" : "h-8 text-[13px]";
  const transition = { duration: reduceMotion ? 0 : duration / 1000, ease: [0.22, 1, 0.36, 1] as const };
  const GUTTER = 8; // px before the first chevron

  const renderNodes = (nodes: TreeNode[], level: number): React.ReactNode =>
    nodes.map((node, i) => {
      const kids = childrenOf(node);
      const parent = isParent(node);
      const isExpanded = parent && expandedSet.has(node.id);
      const isLoading = loading.includes(node.id);
      const isSelected = selectable && selectedSet.has(node.id);
      const check = checkboxes ? checkState(node) : false;
      const labelId = `${uid}-${node.id}-label`;
      const offset = GUTTER + (level - 1) * indent;

      return (
        <li
          key={node.id}
          ref={(el) => {
            if (el) itemRefs.current.set(node.id, el);
            else itemRefs.current.delete(node.id);
          }}
          role="treeitem"
          data-id={node.id}
          aria-labelledby={labelId}
          aria-level={level}
          aria-setsize={nodes.length}
          aria-posinset={i + 1}
          aria-expanded={parent ? isExpanded : undefined}
          aria-selected={selectable ? isSelected : undefined}
          aria-checked={checkboxes ? check : undefined}
          aria-busy={isLoading || undefined}
          aria-disabled={node.disabled || undefined}
          tabIndex={node.id === tabStop ? 0 : -1}
          onFocus={(event) => {
            if (event.target === event.currentTarget) setFocusedId(node.id);
          }}
          className="outline-hidden [&:focus-visible>[data-row]]:outline-2 [&:focus-visible>[data-row]]:-outline-offset-2 [&:focus-visible>[data-row]]:outline-solid [&:focus-visible>[data-row]]:outline-[var(--tv-accent)]"
        >
          <div
            data-row=""
            onClick={(event) => handleRowClick(event, node)}
            onDoubleClick={() => onAction?.(node)}
            className={`relative flex cursor-default select-none items-center gap-1.5 rounded-md pr-2 transition-colors duration-150 ${rowHeight} ${
              isSelected
                ? "bg-[color-mix(in_oklab,var(--tv-accent)_14%,transparent)] text-zinc-950 dark:bg-[color-mix(in_oklab,var(--tv-accent)_20%,transparent)] dark:text-white"
                : "text-zinc-700 hover:bg-zinc-900/[0.045] dark:text-zinc-300 dark:hover:bg-white/[0.05]"
            } ${node.disabled ? "opacity-45" : ""}`}
            style={{ paddingLeft: offset }}
          >
            {showGuides &&
              Array.from({ length: level - 1 }, (_, l) => (
                <span
                  key={l}
                  aria-hidden="true"
                  className="absolute inset-y-0 w-px bg-zinc-900/10 dark:bg-white/10"
                  style={{ left: GUTTER + l * indent + 7.5 }}
                />
              ))}
            {parent ? (
              <span
                aria-hidden="true"
                onClick={(event) => {
                  // The chevron only toggles, it never changes the selection.
                  event.stopPropagation();
                  focusRow(node.id);
                  setExpanded(node, !isExpanded);
                }}
                className="flex size-4 shrink-0 items-center justify-center rounded text-zinc-400 hover:text-zinc-700 dark:text-zinc-500 dark:hover:text-zinc-200"
              >
                <svg
                  viewBox="0 0 16 16"
                  className={`size-3 ${isExpanded ? "rotate-90" : ""} ${reduceMotion ? "" : "transition-transform duration-200"}`}
                  fill="none"
                >
                  <path d="M6 3.5 10.5 8 6 12.5" stroke="currentColor" strokeWidth={1.75} strokeLinecap="round" strokeLinejoin="round" />
                </svg>
              </span>
            ) : (
              <span aria-hidden="true" className="size-4 shrink-0" />
            )}
            {checkboxes && (
              <span
                aria-hidden="true"
                onClick={(event) => {
                  event.stopPropagation();
                  focusRow(node.id);
                  toggleChecked(node);
                }}
                className={`flex size-[15px] shrink-0 items-center justify-center rounded-[4px] border transition-colors ${
                  check
                    ? "border-[var(--tv-accent)] bg-[var(--tv-accent)] text-white"
                    : "border-zinc-900/25 bg-white dark:border-white/25 dark:bg-zinc-900"
                }`}
              >
                {check === true && (
                  <svg viewBox="0 0 12 12" className="size-2.5" fill="none">
                    <path d="m2.5 6.2 2.2 2.2 4.8-4.9" stroke="currentColor" strokeWidth={1.8} strokeLinecap="round" strokeLinejoin="round" />
                  </svg>
                )}
                {check === "mixed" && <span className="h-[1.75px] w-2 rounded-full bg-current" />}
              </span>
            )}
            <span
              aria-hidden="true"
              className={`flex size-4 shrink-0 items-center justify-center [&_svg]:size-4 ${
                isSelected ? "text-[var(--tv-accent)]" : "text-zinc-400 dark:text-zinc-500"
              }`}
            >
              {node.icon ?? (parent ? <FolderGlyph open={isExpanded} /> : <FileGlyph />)}
            </span>
            <span id={labelId} className="min-w-0 flex-1 truncate tracking-tight">
              {node.label}
            </span>
            {isLoading && (
              <svg aria-hidden="true" viewBox="0 0 16 16" className={`size-3.5 shrink-0 text-[var(--tv-accent)] ${reduceMotion ? "" : "animate-spin"}`} fill="none">
                <circle cx="8" cy="8" r="6" stroke="currentColor" strokeOpacity=".25" strokeWidth="2" />
                <path d="M14 8a6 6 0 0 0-6-6" stroke="currentColor" strokeWidth="2" strokeLinecap="round" />
              </svg>
            )}
          </div>
          <AnimatePresence initial={false}>
            {isExpanded && kids && kids.length > 0 && (
              <motion.ul
                key="group"
                role="group"
                initial={{ height: 0, opacity: 0 }}
                animate={{ height: "auto", opacity: 1 }}
                exit={{ height: 0, opacity: 0 }}
                transition={transition}
                className="overflow-hidden"
              >
                {renderNodes(kids, level + 1)}
              </motion.ul>
            )}
          </AnimatePresence>
        </li>
      );
    });

  return (
    <>
      <ul
        ref={rootRef}
        role="tree"
        aria-label={label}
        aria-multiselectable={multiple || undefined}
        onKeyDown={handleKeyDown}
        className={`relative overflow-auto overscroll-contain ${className}`}
        style={{ ["--tv-accent" as string]: accent, ...style }}
        {...props}
      >
        {renderNodes(items, 1)}
      </ul>
      <div role="status" aria-live="polite" className="sr-only">
        {announcement}
      </div>
    </>
  );
}

function FolderGlyph({ open }: { open: boolean }) {
  return open ? (
    <svg viewBox="0 0 16 16" fill="none">
      <path d="M1.75 4.25c0-.69.56-1.25 1.25-1.25h3l1.5 1.5H13c.69 0 1.25.56 1.25 1.25V6" stroke="currentColor" strokeWidth="1.25" strokeLinejoin="round" />
      <path d="M1.9 12.4 3.4 7.1c.15-.5.6-.85 1.13-.85h9.6c.75 0 1.28.72 1.07 1.43l-1.33 4.5c-.15.5-.6.82-1.12.82H3a1.1 1.1 0 0 1-1.1-.6Z" stroke="currentColor" strokeWidth="1.25" strokeLinejoin="round" />
    </svg>
  ) : (
    <svg viewBox="0 0 16 16" fill="none">
      <path d="M1.75 4.25c0-.69.56-1.25 1.25-1.25h3l1.5 1.5H13c.69 0 1.25.56 1.25 1.25v6c0 .69-.56 1.25-1.25 1.25H3c-.69 0-1.25-.56-1.25-1.25v-7.5Z" stroke="currentColor" strokeWidth="1.25" strokeLinejoin="round" />
    </svg>
  );
}

function FileGlyph() {
  return (
    <svg viewBox="0 0 16 16" fill="none">
      <path d="M4 1.75h5l3.25 3.25v8c0 .69-.56 1.25-1.25 1.25H4c-.69 0-1.25-.56-1.25-1.25V3c0-.69.56-1.25 1.25-1.25Z" stroke="currentColor" strokeWidth="1.25" strokeLinejoin="round" />
      <path d="M9 1.75V5h3.25" stroke="currentColor" strokeWidth="1.25" strokeLinejoin="round" />
    </svg>
  );
}