Tree list

"use client";

import * as React from "react";
import { ChevronRightIcon, FileIcon, FolderIcon } from "lucide-react";
import { cn } from "@/lib/utils";

export type TreeNode = {
  id: string;
  label: string;
  /** Child nodes. A node with children is a branch. */
  children?: TreeNode[];
  /** Small mono text on the right, for example a size or a count. */
  meta?: string;
};

type TreeListProps = Omit<React.ComponentProps<"ul">, "children" | "onSelect"> & {
  nodes: TreeNode[];
  /** Branch ids that start open. */
  defaultExpanded?: string[];
  /** Id of the selected node. */
  selected?: string;
  onSelect?: (id: string) => void;
};

type Flat = { node: TreeNode; level: number; parent?: string; index: number; size: number };

function flatten(nodes: TreeNode[], open: Set<string>, level = 1, parent?: string, out: Flat[] = []) {
  nodes.forEach((node, index) => {
    out.push({ node, level, parent, index, size: nodes.length });
    if (node.children?.length && open.has(node.id)) flatten(node.children, open, level + 1, node.id, out);
  });
  return out;
}

/**
 * An expandable tree with the tree pattern's roles and keys. One tab stop.
 * Up and Down move, Right opens or enters, Left closes or goes to the parent,
 * Home and End jump, Enter and Space select.
 */
export function TreeList({
  nodes,
  defaultExpanded = [],
  selected,
  onSelect,
  className,
  "aria-label": ariaLabel = "Tree",
  ...props
}: TreeListProps) {
  const [open, setOpen] = React.useState(() => new Set(defaultExpanded));
  const [active, setActive] = React.useState(selected ?? nodes[0]?.id);
  const refs = React.useRef(new Map<string, HTMLLIElement>());
  const rows = flatten(nodes, open);
  // The active id can sit in a collapsed branch or name no node. The tree still needs its one tab stop.
  const tabStop = rows.some((row) => row.node.id === active) ? active : rows[0]?.node.id;

  function focusId(id: string | undefined) {
    if (!id) return;
    setActive(id);
    refs.current.get(id)?.focus();
  }

  function toggle(id: string, next: boolean) {
    setOpen((prev) => {
      const copy = new Set(prev);
      if (next) copy.add(id);
      else copy.delete(id);
      return copy;
    });
  }

  function onKeyDown(event: React.KeyboardEvent, i: number) {
    const row = rows[i];
    const isBranch = !!row.node.children?.length;
    const expanded = open.has(row.node.id);
    switch (event.key) {
      case "ArrowDown":
        focusId(rows[i + 1]?.node.id);
        break;
      case "ArrowUp":
        focusId(rows[i - 1]?.node.id);
        break;
      case "Home":
        focusId(rows[0].node.id);
        break;
      case "End":
        focusId(rows[rows.length - 1].node.id);
        break;
      case "ArrowRight":
        if (isBranch && !expanded) toggle(row.node.id, true);
        else if (isBranch) focusId(rows[i + 1]?.node.id);
        break;
      case "ArrowLeft":
        if (isBranch && expanded) toggle(row.node.id, false);
        else focusId(row.parent);
        break;
      case "Enter":
      case " ":
        if (isBranch) toggle(row.node.id, !expanded);
        onSelect?.(row.node.id);
        break;
      default:
        return;
    }
    event.preventDefault();
    event.stopPropagation();
  }

  return (
    <ul
      role="tree"
      data-slot="tree-list"
      aria-label={ariaLabel}
      className={cn("m-0 flex w-full list-none flex-col gap-px p-0 text-sm", className)}
      {...props}
    >
      {rows.map(({ node, level, index, size }, i) => {
        const isBranch = !!node.children?.length;
        const expanded = open.has(node.id);
        const isSelected = node.id === selected;
        const Icon = isBranch ? FolderIcon : FileIcon;
        return (
          <li
            key={node.id}
            ref={(el) => {
              if (el) refs.current.set(node.id, el);
              else refs.current.delete(node.id);
            }}
            role="treeitem"
            aria-level={level}
            aria-setsize={size}
            aria-posinset={index + 1}
            aria-expanded={isBranch ? expanded : undefined}
            aria-selected={isSelected}
            tabIndex={node.id === tabStop ? 0 : -1}
            onKeyDown={(e) => e.target === e.currentTarget && onKeyDown(e, i)}
            onClick={(e) => {
              e.stopPropagation();
              setActive(node.id);
              if (isBranch) toggle(node.id, !expanded);
              onSelect?.(node.id);
            }}
            className={cn(
              "flex h-7 cursor-default items-center gap-1.5 rounded-sm pr-2 transition-colors duration-100 ease-out outline-none hover:bg-muted focus-visible:ring-2 focus-visible:ring-ring/50",
              isSelected ? "bg-muted text-foreground" : "text-foreground/80",
            )}
            style={{ paddingLeft: `${(level - 1) * 16 + 4}px` }}
          >
            <ChevronRightIcon
              aria-hidden="true"
              className={cn(
                "size-3.5 shrink-0 text-muted-foreground transition-transform duration-150 ease-out",
                !isBranch && "invisible",
                expanded && "rotate-90",
              )}
            />
            <Icon aria-hidden="true" className="size-3.5 shrink-0 text-muted-foreground" />
            <span className="min-w-0 flex-1 truncate">{node.label}</span>
            {node.meta && <span className="font-mono text-[11px] text-muted-foreground">{node.meta}</span>}
          </li>
        );
      })}
    </ul>
  );
}