Tree view

"use client"

import * as React from "react"

import { cn } from "@/lib/utils"
import { TreeBranch, type TreeNode } from "@/components/ui/tree-view/tree-item"

export type { TreeNode }

type FlatNode = { node: TreeNode; level: number; parentId: string | null }

// The focusable set, re-derived from the data and the open branches: a node
// inside a shut branch is not on screen, so the arrow keys must not reach it.
function flattenVisible(
  nodes: TreeNode[],
  expanded: Set<string>,
  level = 1,
  parentId: string | null = null
): FlatNode[] {
  const flat: FlatNode[] = []
  for (const node of nodes) {
    flat.push({ node, level, parentId })
    if (node.children?.length && expanded.has(node.id)) {
      flat.push(...flattenVisible(node.children, expanded, level + 1, node.id))
    }
  }
  return flat
}

// `onSelect` here takes a node id, not a DOM select event, so it replaces the
// div handler of the same name rather than intersecting with it.
export type TreeViewProps = Omit<React.ComponentProps<"div">, "onSelect"> & {
  data: TreeNode[]
  /** Ids open on the first render, when the tree owns its own open set. */
  defaultExpanded?: string[]
  /** Ids open right now. Pass it with onExpandedChange to own the open set yourself. */
  expanded?: string[]
  onExpandedChange?: (ids: string[]) => void
  selected?: string | null
  onSelect?: (id: string) => void
  /** Draws a hairline down each open branch. */
  showLines?: boolean
  /** Pixels of indentation per level. */
  indent?: number
}

/**
 * A file-tree style disclosure list following the WAI-ARIA tree pattern: one
 * tab stop, arrow keys to move and open, Enter or Space to select.
 */
function TreeView({
  className,
  data,
  defaultExpanded,
  expanded,
  onExpandedChange,
  selected = null,
  onSelect,
  showLines = false,
  indent = 16,
  ...props
}: TreeViewProps) {
  const [internalExpanded, setInternalExpanded] = React.useState<string[]>(defaultExpanded ?? [])
  const expandedIds = expanded ?? internalExpanded
  const expandedSet = React.useMemo(() => new Set(expandedIds), [expandedIds])
  const visible = React.useMemo(() => flattenVisible(data, expandedSet), [data, expandedSet])

  const [focusedId, setFocusedId] = React.useState<string | null>(null)
  const itemRefs = React.useRef(new Map<string, HTMLDivElement>())

  // The single tab stop: wherever focus last was, else the selected node, else
  // the first one — and never a node that has since been folded away.
  const activeId = React.useMemo(() => {
    const onScreen = (id: string | null) => id !== null && visible.some((v) => v.node.id === id)
    if (onScreen(focusedId)) return focusedId
    if (onScreen(selected)) return selected
    return visible[0]?.node.id ?? null
  }, [focusedId, selected, visible])

  const setExpandedIds = React.useCallback(
    (next: string[]) => {
      if (expanded === undefined) setInternalExpanded(next)
      onExpandedChange?.(next)
    },
    [expanded, onExpandedChange]
  )

  const toggleExpanded = React.useCallback(
    (id: string) => {
      setExpandedIds(
        expandedSet.has(id) ? expandedIds.filter((each) => each !== id) : [...expandedIds, id]
      )
    },
    [expandedIds, expandedSet, setExpandedIds]
  )

  const focusNode = React.useCallback((id: string) => {
    setFocusedId(id)
    // The element is already in the DOM at tabIndex -1, so it takes focus now
    // rather than one render after the roving index catches up.
    itemRefs.current.get(id)?.focus()
  }, [])

  const activate = React.useCallback(
    (node: TreeNode) => {
      if (node.disabled) return
      onSelect?.(node.id)
      // Every file tree opens a folder when you activate it; the chevron is
      // not meant to be the only way in.
      if (node.children?.length) toggleExpanded(node.id)
    },
    [onSelect, toggleExpanded]
  )

  const handleKeyDown = (event: React.KeyboardEvent<HTMLDivElement>) => {
    // Only keys pressed on a node itself steer the tree; anything a caller put
    // in a badge — a button, a field — keeps its own Space and arrow keys.
    if ((event.target as HTMLElement).getAttribute("role") !== "treeitem") return
    const index = visible.findIndex((v) => v.node.id === activeId)
    if (index === -1) return
    const { node, parentId } = visible[index]
    const isBranch = Boolean(node.children?.length)
    const isExpanded = expandedSet.has(node.id)
    const step = (to: number) => {
      const clamped = Math.min(Math.max(to, 0), visible.length - 1)
      focusNode(visible[clamped].node.id)
    }

    switch (event.key) {
      case "ArrowDown":
        step(index + 1)
        break
      case "ArrowUp":
        step(index - 1)
        break
      case "ArrowRight":
        // Open a shut branch, then step into it on the next press.
        if (isBranch && !isExpanded && !node.disabled) toggleExpanded(node.id)
        else if (isBranch && isExpanded) step(index + 1)
        else return
        break
      case "ArrowLeft":
        // Shut an open branch, then climb out of it on the next press.
        if (isBranch && isExpanded && !node.disabled) toggleExpanded(node.id)
        else if (parentId !== null) focusNode(parentId)
        else return
        break
      case "Home":
        step(0)
        break
      case "End":
        step(visible.length - 1)
        break
      case "Enter":
      case " ":
        activate(node)
        break
      default:
        return
    }
    event.preventDefault()
  }

  return (
    <div
      data-slot="tree-view"
      data-lines={showLines || undefined}
      role="tree"
      onKeyDown={handleKeyDown}
      className={cn("text-sm", className)}
      {...props}
    >
      <TreeBranch
        nodes={data}
        level={1}
        activeId={activeId}
        expandedSet={expandedSet}
        selected={selected}
        showLines={showLines}
        indent={indent}
        itemRefs={itemRefs}
        onActivate={activate}
        onToggle={toggleExpanded}
        onFocusNode={setFocusedId}
      />
    </div>
  )
}

export { TreeView }

"use client"

import * as React from "react"
import { ChevronRightIcon, FileIcon, FolderIcon, FolderOpenIcon } from "lucide-react"

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

export type TreeNode = {
  id: string
  label: React.ReactNode
  /** Replaces the default folder/file icon. */
  icon?: React.ReactNode
  children?: TreeNode[]
  /** Reachable and readable, but not selectable and not openable. */
  disabled?: boolean
  /** Trails the label — a count, a status, a size. */
  badge?: React.ReactNode
}

export type TreeSharedProps = {
  level: number
  /** The one node in the tab order; every other treeitem sits at -1. */
  activeId: string | null
  expandedSet: Set<string>
  selected: string | null
  showLines: boolean
  indent: number
  itemRefs: React.RefObject<Map<string, HTMLDivElement>>
  onActivate: (node: TreeNode) => void
  onToggle: (id: string) => void
  onFocusNode: (id: string) => void
}

export type TreeBranchProps = TreeSharedProps & { nodes: TreeNode[] }

/** Renders one level of nodes; recursion happens through each item's own group. */
function TreeBranch({ nodes, ...shared }: TreeBranchProps) {
  return (
    <>
      {nodes.map((node) => (
        <TreeItem key={node.id} node={node} {...shared} />
      ))}
    </>
  )
}

/** One treeitem: its row, and — when it is open — the group holding its children. */
function TreeItem({ node, ...shared }: TreeSharedProps & { node: TreeNode }) {
  const { level, activeId, expandedSet, selected, showLines, indent, itemRefs } = shared
  const isBranch = Boolean(node.children?.length)
  const isExpanded = isBranch && expandedSet.has(node.id)
  const DefaultIcon = isBranch ? (isExpanded ? FolderOpenIcon : FolderIcon) : FileIcon

  return (
    <div
      ref={(element) => {
        if (element) itemRefs.current.set(node.id, element)
        else itemRefs.current.delete(node.id)
      }}
      role="treeitem"
      // A treeitem holds its own group, so name-from-content would read every
      // descendant along with it. A string label names the item outright.
      aria-label={typeof node.label === "string" ? node.label : undefined}
      aria-expanded={isBranch ? isExpanded : undefined}
      aria-selected={selected === node.id}
      aria-level={level}
      aria-disabled={node.disabled || undefined}
      data-slot="tree-view-item"
      data-branch={isBranch || undefined}
      data-selected={selected === node.id || undefined}
      tabIndex={activeId === node.id ? 0 : -1}
      onFocus={(event) => {
        // Focus bubbles; only the item that actually took it should claim the
        // roving tab stop.
        if (event.target === event.currentTarget) shared.onFocusNode(node.id)
      }}
      className="group/treeitem outline-none"
    >
      {/* The row, not the treeitem, carries the click: a treeitem contains its
          children's rows, and a click on one of those would select the parent too. */}
      <div
        data-slot="tree-view-row"
        onClick={() => {
          shared.onActivate(node)
          shared.onFocusNode(node.id)
          itemRefs.current.get(node.id)?.focus()
        }}
        style={{ paddingInlineStart: `${(level - 1) * indent + 4}px` }}
        className={cn(
          "flex h-7 items-center gap-1.5 rounded-md pe-2 select-none",
          "group-focus-visible/treeitem:focus-outline-inset",
          node.disabled
            ? "cursor-not-allowed text-muted-foreground opacity-60"
            : "cursor-pointer hover:bg-accent",
          selected === node.id && "bg-accent font-medium"
        )}
      >
          <span
          data-slot="tree-view-chevron"
          aria-hidden="true"
          onClick={(event) => {
            // The chevron opens a branch without also selecting it. Keyboard
            // users reach the same behaviour with ArrowRight and ArrowLeft.
            event.stopPropagation()
            if (!node.disabled) shared.onToggle(node.id)
          }}
          className="flex size-4 shrink-0 items-center justify-center text-muted-foreground"
        >
          {isBranch ? (
            <ChevronRightIcon
              className={cn(
                "size-3.5 transition-transform duration-(--duration-fast) ease-(--ease-standard) rtl:-scale-x-100",
                isExpanded && "rotate-90 rtl:-rotate-90"
              )}
            />
          ) : null}
        </span>

        <span
          data-slot="tree-view-icon"
          aria-hidden="true"
          className="flex shrink-0 items-center text-muted-foreground [&_svg]:pointer-events-none [&_svg]:shrink-0 [&_svg:not([class*='size-'])]:size-4"
        >
          {node.icon ?? <DefaultIcon />}
        </span>

        <span data-slot="tree-view-label" className="truncate">
          {node.label}
        </span>

        {node.badge ? (
          <span
            data-slot="tree-view-badge"
            className="ms-auto shrink-0 tabular-nums text-xs text-muted-foreground"
          >
            {node.badge}
          </span>
        ) : null}
      </div>

      {isExpanded ? (
        <div role="group" data-slot="tree-view-group" className="relative">
          {showLines ? (
            <span
              aria-hidden="true"
              className="absolute inset-y-0 w-px bg-border"
              style={{ insetInlineStart: `${(level - 1) * indent + 11}px` }}
            />
          ) : null}
          <TreeBranch {...shared} nodes={node.children ?? []} level={level + 1} />
        </div>
      ) : null}
    </div>
  )
}

export { TreeBranch }