Tree View

"use client";

import * as React from "react";
import { ChevronRight, File, Folder } from "lucide-react";

import { cn } from "@/lib/utils";

export interface TreeNode {
  id: string;
  label: string;
  icon?: React.ReactNode;
  children?: TreeNode[];
}

export interface TreeViewProps {
  data: TreeNode[];
  defaultExpanded?: string[];
  selected?: string;
  onSelect?: (id: string) => void;
  className?: string;
  "aria-label"?: string;
}

interface FlatNode {
  node: TreeNode;
  depth: number;
  posinset: number;
  setsize: number;
}

function flattenVisible(nodes: TreeNode[], expanded: Set<string>, depth = 0): FlatNode[] {
  const result: FlatNode[] = [];
  nodes.forEach((node, index) => {
    result.push({ node, depth, posinset: index + 1, setsize: nodes.length });
    if (node.children && node.children.length > 0 && expanded.has(node.id)) {
      result.push(...flattenVisible(node.children, expanded, depth + 1));
    }
  });
  return result;
}

/**
 * WAI-ARIA tree pattern: role="tree" of role="treeitem"s with roving
 * tabindex over currently-visible (ancestor-expanded) items only.
 * ArrowRight expands / descends into an already-expanded node;
 * ArrowLeft collapses / ascends to the parent.
 */
export function TreeView({
  data,
  defaultExpanded = [],
  selected,
  onSelect,
  className,
  "aria-label": ariaLabel = "Tree",
}: TreeViewProps) {
  const [expanded, setExpanded] = React.useState<Set<string>>(() => new Set(defaultExpanded));
  const [focusedId, setFocusedId] = React.useState<string | null>(() => data[0]?.id ?? null);
  const rootRef = React.useRef<HTMLDivElement>(null);

  const visible = flattenVisible(data, expanded);
  const visibleIds = visible.map((v) => v.node.id);

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

  const focusId = (id: string) => {
    setFocusedId(id);
    requestAnimationFrame(() => {
      rootRef.current?.querySelector<HTMLElement>(`[data-tree-id="${id}"]`)?.focus();
    });
  };

  const handleKeyDown = (event: React.KeyboardEvent, flat: FlatNode) => {
    const { node, depth } = flat;
    const hasChildren = Boolean(node.children?.length);
    const currentIndex = visibleIds.indexOf(node.id);

    if (event.key === "ArrowDown") {
      event.preventDefault();
      const next = visible[currentIndex + 1];
      if (next) focusId(next.node.id);
    } else if (event.key === "ArrowUp") {
      event.preventDefault();
      const prev = visible[currentIndex - 1];
      if (prev) focusId(prev.node.id);
    } else if (event.key === "ArrowRight") {
      event.preventDefault();
      if (hasChildren && !expanded.has(node.id)) {
        toggleExpanded(node.id);
      } else if (hasChildren) {
        const next = visible[currentIndex + 1];
        if (next) focusId(next.node.id);
      }
    } else if (event.key === "ArrowLeft") {
      event.preventDefault();
      if (hasChildren && expanded.has(node.id)) {
        toggleExpanded(node.id);
      } else if (depth > 0) {
        for (let i = currentIndex - 1; i >= 0; i--) {
          if (visible[i].depth < depth) {
            focusId(visible[i].node.id);
            break;
          }
        }
      }
    } else if (event.key === "Home") {
      event.preventDefault();
      if (visible[0]) focusId(visible[0].node.id);
    } else if (event.key === "End") {
      event.preventDefault();
      if (visible[visible.length - 1]) focusId(visible[visible.length - 1].node.id);
    } else if (event.key === "Enter" || event.key === " ") {
      event.preventDefault();
      onSelect?.(node.id);
      if (hasChildren) toggleExpanded(node.id);
    }
  };

  return (
    <div ref={rootRef} role="tree" aria-label={ariaLabel} className={cn("flex flex-col text-sm", className)}>
      {visible.map((flat) => {
        const { node, depth, posinset, setsize } = flat;
        const hasChildren = Boolean(node.children?.length);
        const isExpanded = expanded.has(node.id);
        const isFocusTarget = focusedId === node.id;
        return (
          <div
            key={node.id}
            data-tree-id={node.id}
            role="treeitem"
            aria-expanded={hasChildren ? isExpanded : undefined}
            aria-selected={selected === node.id}
            aria-level={depth + 1}
            aria-setsize={setsize}
            aria-posinset={posinset}
            tabIndex={isFocusTarget ? 0 : -1}
            onFocus={() => setFocusedId(node.id)}
            onClick={() => {
              onSelect?.(node.id);
              if (hasChildren) toggleExpanded(node.id);
              setFocusedId(node.id);
            }}
            onKeyDown={(event) => handleKeyDown(event, flat)}
            style={{ paddingLeft: `calc(var(--tree-indent-step) * ${depth} + var(--tree-indent-base))` }}
            className={cn(
              "flex cursor-pointer items-center gap-1.5 rounded-md py-1.5 pr-2 outline-none",
              "focus-visible:ring-2 focus-visible:ring-ring focus-visible:ring-inset",
              selected === node.id ? "bg-accent text-accent-foreground" : "text-foreground hover:bg-accent/50",
            )}
          >
            <span
              aria-hidden="true"
              className={cn(
                "flex size-4 shrink-0 items-center justify-center text-muted-foreground transition-transform",
                hasChildren && isExpanded && "rotate-90",
              )}
            >
              {hasChildren ? <ChevronRight className="size-3.5" /> : null}
            </span>
            <span className="text-muted-foreground" aria-hidden="true">
              {node.icon ?? (hasChildren ? <Folder className="size-4" /> : <File className="size-4" />)}
            </span>
            <span className="truncate">{node.label}</span>
          </div>
        );
      })}
    </div>
  );
}