Tree View

"use client";

import * as React from "react";
import { motion, useReducedMotion } from "motion/react";
import { ChevronRight, File, Folder, FolderOpen } from "lucide-react";
import { cn } from "@/lib/utils";

export interface TreeNode {
  id: string;
  name: string;
  /** Custom leading icon; folders and files get sensible defaults. */
  icon?: React.ReactNode;
  /** Present (even empty) marks the node as a folder. */
  children?: TreeNode[];
}

interface FlatNode {
  node: TreeNode;
  depth: number;
  parent: TreeNode | null;
}

export interface TreeViewProps
  extends Omit<React.ComponentProps<"ul">, "onSelect"> {
  data: TreeNode[];
  /** Folder ids expanded on first render. */
  defaultExpandedIds?: string[];
  /**
   * Auto-expand every folder up to (and including) this depth on first
   * render. Combined with (union of) `defaultExpandedIds`. `1` = top level.
   */
  defaultExpandedDepth?: number;
  /** Node id selected on first render. */
  defaultSelectedId?: string;
  /** Draws the vertical indentation rails linking each level. */
  showGuides?: boolean;
  /** Fires whenever a node is activated (click, Enter, Space). */
  onSelect?: (node: TreeNode) => void;
  /** Disables the expand/collapse and highlight motion. */
  static?: boolean;
}

/**
 * A file tree. Folders reveal children through the sanctioned
 * `grid-template-rows: 0fr ↔ 1fr` transition with the chevron rotating in
 * sync; the selected row's highlight is a single shared element that
 * slides between rows (motion layout) instead of each row restyling.
 * Keyboard follows the ARIA tree pattern with a roving tabindex:
 * arrows move/expand/collapse, Home/End jump, Enter and Space activate.
 * Collapsed branches stay mounted for the animation but are `inert`.
 */
export function TreeView({
  data,
  defaultExpandedIds = [],
  defaultExpandedDepth = 0,
  defaultSelectedId,
  showGuides = true,
  onSelect,
  static: isStatic = false,
  className,
  ...props
}: TreeViewProps) {
  const [expanded, setExpanded] = React.useState<Set<string>>(() => {
    const ids = new Set(defaultExpandedIds);
    if (defaultExpandedDepth > 0) {
      const walk = (nodes: TreeNode[], depth: number) => {
        for (const node of nodes) {
          if (node.children) {
            if (depth <= defaultExpandedDepth) ids.add(node.id);
            walk(node.children, depth + 1);
          }
        }
      };
      walk(data, 1);
    }
    return ids;
  });
  const [selected, setSelected] = React.useState<string | null>(
    defaultSelectedId ?? null,
  );
  const [focused, setFocused] = React.useState<string | null>(null);
  const rowRefs = React.useRef(new Map<string, HTMLDivElement>());
  const treeId = React.useId();
  const reducedMotion = useReducedMotion() ?? false;
  const noMotion = reducedMotion || isStatic;

  // Visible rows in document order — the keyboard navigation space.
  const flat = React.useMemo<FlatNode[]>(() => {
    const list: FlatNode[] = [];
    const walk = (nodes: TreeNode[], depth: number, parent: TreeNode | null) => {
      for (const node of nodes) {
        list.push({ node, depth, parent });
        if (node.children && expanded.has(node.id)) {
          walk(node.children, depth + 1, node);
        }
      }
    };
    walk(data, 1, null);
    return list;
  }, [data, expanded]);

  const focusTarget = focused ?? selected ?? flat[0]?.node.id ?? null;

  const toggle = (id: string) => {
    setExpanded((prev) => {
      const next = new Set(prev);
      if (next.has(id)) next.delete(id);
      else next.add(id);
      return next;
    });
  };

  const activate = (node: TreeNode) => {
    setSelected(node.id);
    setFocused(node.id);
    if (node.children) toggle(node.id);
    onSelect?.(node);
  };

  const focusId = (id: string | undefined | null) => {
    if (!id) return;
    setFocused(id);
    rowRefs.current.get(id)?.focus();
  };

  const handleKeyDown = (event: React.KeyboardEvent, node: TreeNode) => {
    const ids = flat.map((f) => f.node.id);
    const i = ids.indexOf(node.id);
    const isFolder = !!node.children;
    switch (event.key) {
      case "ArrowDown":
        event.preventDefault();
        focusId(ids[i + 1]);
        break;
      case "ArrowUp":
        event.preventDefault();
        focusId(ids[i - 1]);
        break;
      case "ArrowRight":
        event.preventDefault();
        if (!isFolder) break;
        if (!expanded.has(node.id)) toggle(node.id);
        else focusId(node.children?.[0]?.id);
        break;
      case "ArrowLeft":
        event.preventDefault();
        if (isFolder && expanded.has(node.id)) toggle(node.id);
        else focusId(flat[i]?.parent?.id);
        break;
      case "Home":
        event.preventDefault();
        focusId(ids[0]);
        break;
      case "End":
        event.preventDefault();
        focusId(ids[ids.length - 1]);
        break;
      case "Enter":
      case " ":
        event.preventDefault();
        activate(node);
        break;
    }
  };

  const renderNodes = (nodes: TreeNode[], depth: number): React.ReactNode =>
    nodes.map((node) => {
      const isFolder = !!node.children;
      const isExpanded = expanded.has(node.id);
      const isSelected = selected === node.id;
      return (
        <li key={node.id} role="none">
          <div
            role="treeitem"
            aria-level={depth}
            aria-expanded={isFolder ? isExpanded : undefined}
            aria-selected={isSelected}
            tabIndex={focusTarget === node.id ? 0 : -1}
            ref={(el) => {
              if (el) rowRefs.current.set(node.id, el);
              else rowRefs.current.delete(node.id);
            }}
            onClick={() => activate(node)}
            onKeyDown={(event) => handleKeyDown(event, node)}
            onFocus={() => setFocused(node.id)}
            className={cn(
              "relative isolate flex h-7 cursor-default items-center gap-1.5 rounded-md px-2 text-sm select-none",
              "transition-colors duration-(--duration-fast)",
              isSelected
                ? "text-foreground"
                : "text-muted-foreground hover:bg-muted/40 hover:text-foreground",
            )}
          >
            {isSelected &&
              (noMotion ? (
                <span
                  aria-hidden
                  className="pointer-events-none absolute inset-0 -z-10 rounded-md bg-accent"
                />
              ) : (
                <motion.span
                  aria-hidden
                  layoutId={`${treeId}-highlight`}
                  transition={{ type: "spring", duration: 0.3, bounce: 0 }}
                  className="pointer-events-none absolute inset-0 -z-10 rounded-md bg-accent"
                />
              ))}
            {isFolder ? (
              <ChevronRight
                aria-hidden
                className={cn(
                  "size-3.5 shrink-0 text-muted-foreground/70",
                  isExpanded && "rotate-90",
                  !isStatic &&
                    "transition-[rotate] duration-(--duration-quick) ease-(--ease-out) motion-reduce:transition-none",
                )}
              />
            ) : (
              <span aria-hidden className="w-3.5 shrink-0" />
            )}
            <span
              aria-hidden
              className="shrink-0 text-muted-foreground [&_svg]:size-4"
            >
              {node.icon ??
                (isFolder ? (
                  isExpanded ? (
                    <FolderOpen className="size-4" />
                  ) : (
                    <Folder className="size-4" />
                  )
                ) : (
                  <File className="size-4" />
                ))}
            </span>
            <span className="truncate">{node.name}</span>
          </div>
          {isFolder && (
            <div
              className={cn(
                "grid",
                isExpanded ? "grid-rows-[1fr]" : "grid-rows-[0fr]",
                !isStatic &&
                  "transition-[grid-template-rows] motion-reduce:transition-none",
                !isStatic &&
                  (isExpanded
                    ? "duration-(--duration-base) ease-(--ease-out)"
                    : "duration-(--duration-quick) ease-(--ease-exit)"),
              )}
            >
              <div inert={!isExpanded} className="min-h-0 overflow-hidden">
                <ul
                  role="group"
                  className={cn(
                    // ml-[15px] centers the rail under the parent chevron
                    // (px-2 + half of the 3.5 chevron ≈ 15px); pl-2 keeps
                    // child labels clear of the guide.
                    "mt-0.5 ml-[15px] flex flex-col gap-0.5 pl-2",
                    showGuides && "border-l border-border",
                  )}
                >
                  {renderNodes(node.children ?? [], depth + 1)}
                </ul>
              </div>
            </div>
          )}
        </li>
      );
    });

  return (
    <ul
      role="tree"
      data-slot="tree-view"
      className={cn("flex w-full flex-col gap-0.5", className)}
      {...props}
    >
      {renderNodes(data, 1)}
    </ul>
  );
}