Tree View

"use client";

import { Button } from "@base-ui/react/button";
import {
  RiArrowRightSLine,
  RiFileLine,
  RiFolder3Line,
  RiFolderOpenLine,
} from "@remixicon/react";
import { cn } from "cn";
import {
  AnimatePresence,
  motion,
  useIsPresent,
  useReducedMotion,
} from "motion/react";
import * as React from "react";

export interface TreeViewNode {
  id: string;
  label: string;
  disabled?: boolean;
  icon?: React.ComponentType<{
    className?: string;
    "aria-hidden"?: boolean;
  }> | null;
  expandedIcon?: React.ComponentType<{
    className?: string;
    "aria-hidden"?: boolean;
  }>;
  trailingContent?: React.ReactNode;
  children?: TreeViewNode[];
}

export interface TreeViewProps extends Omit<
  React.ComponentProps<"div">,
  "children" | "onSelect"
> {
  items: TreeViewNode[];
  size?: "default" | "sm";
  showGuides?: boolean;
  defaultExpanded?: string[];
  selectedId?: string;
  onSelectedIdChange?: (id: string) => void;
  selectionMode?: "single" | "multiple";
  selectedIds?: string[];
  defaultSelectedIds?: string[];
  onSelectedIdsChange?: (ids: string[]) => void;
}

export function TreeView({
  items,
  size = "default",
  showGuides = true,
  defaultExpanded = [],
  selectedId,
  onSelectedIdChange,
  selectionMode = "single",
  selectedIds,
  defaultSelectedIds = [],
  onSelectedIdsChange,
  ref,
  className,
  onKeyDown,
  ...props
}: TreeViewProps) {
  const [expanded, setExpanded] = React.useState(
    () => new Set(defaultExpanded)
  );
  const [internalSelected, setInternalSelected] = React.useState<string>();
  const [internalSelectedIds, setInternalSelectedIds] =
    React.useState(defaultSelectedIds);
  const anchor = React.useRef<string | null>(null);
  const search = React.useRef({ text: "", time: 0 });
  const [focused, setFocused] = React.useState<string>();
  const root = React.useRef<HTMLDivElement>(null);
  const selected = selectedId ?? internalSelected;
  const multiSelected = selectedIds ?? internalSelectedIds;
  const rows: {
    node: TreeViewNode;
    level: number;
    parent?: string;
    last: boolean;
    // Whether the ancestor at each level (1-based, index 0 is level 1) still
    // has a sibling to come, so its trunk keeps running past this row.
    trail: boolean[];
  }[] = [];
  function collect(
    nodes: TreeViewNode[],
    level: number,
    parent?: string,
    trail: boolean[] = []
  ) {
    for (const [index, node] of nodes.entries()) {
      const last = index === nodes.length - 1;
      rows.push({ node, level, parent, last, trail });
      if (expanded.has(node.id) && node.children) {
        collect(node.children, level + 1, node.id, [...trail, !last]);
      }
    }
  }
  collect(items, 1);
  const enabled = rows.filter(({ node }) => !node.disabled);
  const active = enabled.some(({ node }) => node.id === focused)
    ? focused
    : enabled[0]?.node.id;
  function focus(id?: string) {
    if (!id) {
      return;
    }
    setFocused(id);
    for (const el of root.current?.querySelectorAll<HTMLElement>(
      '[role="treeitem"]'
    ) ?? []) {
      if (el.dataset.id === id) {
        el.focus();
      }
    }
  }
  function toggle(id: string) {
    setExpanded((previous) => {
      const next = new Set(previous);
      if (next.has(id)) {
        next.delete(id);
      } else {
        next.add(id);
      }
      return next;
    });
  }
  function select(id: string, range = false) {
    if (selectionMode === "multiple") {
      let next: string[];
      const start = enabled.findIndex(({ node }) => node.id === anchor.current);
      const end = enabled.findIndex(({ node }) => node.id === id);
      if (range && start !== -1 && end !== -1) {
        next = [
          ...new Set([
            ...multiSelected,
            ...enabled
              .slice(Math.min(start, end), Math.max(start, end) + 1)
              .map(({ node }) => node.id),
          ]),
        ];
      } else {
        next = multiSelected.includes(id)
          ? multiSelected.filter((value) => value !== id)
          : [...multiSelected, id];
        anchor.current = id;
      }
      setInternalSelectedIds(next);
      onSelectedIdsChange?.(next);
    } else {
      setInternalSelected(id);
      onSelectedIdChange?.(id);
    }
  }
  return (
    <div
      {...props}
      ref={(element) => {
        root.current = element;
        if (typeof ref === "function") {
          ref(element);
        } else if (ref) {
          ref.current = element;
        }
      }}
      role="tree"
      aria-multiselectable={selectionMode === "multiple" || undefined}
      data-slot="tree-view"
      className={cn(
        "flex w-full min-w-0 flex-col rounded-[24px] border border-border bg-[var(--surface-raised,var(--background))] p-2 outline-none",
        className
      )}
      onKeyDown={(event) => {
        onKeyDown?.(event);
        if (
          event.defaultPrevented ||
          event.altKey ||
          ((event.ctrlKey || event.metaKey) &&
            !(selectionMode === "multiple" && event.key.toLowerCase() === "a"))
        ) {
          return;
        }
        const index = enabled.findIndex(({ node }) => node.id === active);
        const row = enabled[index];
        if (!row) {
          return;
        }
        const rtl = getComputedStyle(event.currentTarget).direction === "rtl";
        const forward = rtl ? "ArrowLeft" : "ArrowRight";
        const back = rtl ? "ArrowRight" : "ArrowLeft";
        if (
          selectionMode === "multiple" &&
          (event.ctrlKey || event.metaKey) &&
          event.key.toLowerCase() === "a"
        ) {
          const next = enabled.map(({ node }) => node.id);
          setInternalSelectedIds(next);
          onSelectedIdsChange?.(next);
        } else if (event.key === "ArrowDown") {
          focus(enabled[Math.min(index + 1, enabled.length - 1)]?.node.id);
        } else if (event.key === "ArrowUp") {
          focus(enabled[Math.max(index - 1, 0)]?.node.id);
        } else if (event.key === "Home") {
          focus(enabled[0]?.node.id);
        } else if (event.key === "End") {
          focus(enabled.at(-1)?.node.id);
        } else if (event.key === forward && row.node.children?.length) {
          if (!expanded.has(row.node.id)) {
            toggle(row.node.id);
          } else {
            focus(enabled[index + 1]?.node.id);
          }
        } else if (event.key === back) {
          if (expanded.has(row.node.id)) {
            toggle(row.node.id);
          } else {
            focus(row.parent);
          }
        } else if (event.key === "Enter" || event.key === " ") {
          select(row.node.id, event.shiftKey);
        } else if (event.key.length === 1 && !event.ctrlKey && !event.metaKey) {
          const now = Date.now();
          const previous =
            now - search.current.time < 500 ? search.current.text : "";
          const query = (previous + event.key).toLocaleLowerCase();
          search.current = { text: query, time: now };
          const matching = (value: string) =>
            enabled.find(({ node }) =>
              node.label.toLocaleLowerCase().startsWith(value)
            );
          const ordered = [
            ...enabled.slice(index + 1),
            ...enabled.slice(0, index + 1),
          ];
          const match =
            ordered.find(({ node }) =>
              node.label.toLocaleLowerCase().startsWith(query)
            ) ?? matching(query);
          if (match) {
            focus(match.node.id);
          }
        } else {
          return;
        }
        event.preventDefault();
      }}
    >
      {rows.length === 0 && (
        <div className="px-2 py-6 text-center text-sm text-muted-foreground">
          No items.
        </div>
      )}
      <AnimatePresence initial={false}>
        {rows.map(({ node, level, last, trail }) => {
          const branch = !!node.children?.length;
          const open = expanded.has(node.id);
          const Icon =
            node.icon === null
              ? null
              : (open && node.expandedIcon) ||
                node.icon ||
                (branch
                  ? open
                    ? RiFolderOpenLine
                    : RiFolder3Line
                  : RiFileLine);
          return (
            <AnimatedTreeRow key={node.id}>
              <div
                role="treeitem"
                data-slot="tree-view-item"
                data-id={node.id}
                aria-level={level}
                aria-expanded={branch ? open : undefined}
                aria-selected={
                  selectionMode === "multiple"
                    ? multiSelected.includes(node.id)
                    : selected === node.id
                }
                aria-disabled={node.disabled || undefined}
                tabIndex={active === node.id ? 0 : -1}
                onFocus={() => setFocused(node.id)}
                onKeyDown={() => {}}
                onClick={(event) => {
                  if (!node.disabled) {
                    focus(node.id);
                    select(node.id, event.shiftKey);
                  }
                }}
                style={{
                  paddingInlineStart: `calc(0.5rem + ${level - 1} * 1.25rem)`,
                }}
                className={cn(
                  "relative isolate flex min-h-9 min-w-0 cursor-pointer items-center gap-2 rounded-[12px] pe-2 before:absolute before:inset-0 before:-z-10 before:rounded-[inherit] text-sm leading-5 text-foreground outline-none transition-[background-color,color,box-shadow] duration-150 motion-reduce:transition-none hover:bg-[var(--surface-hover,var(--muted))] focus-visible:bg-[var(--surface-hover,var(--muted))] focus-visible:ring-2 focus-visible:ring-ring focus-visible:ring-inset aria-selected:transition-none aria-selected:not-hover:not-focus-visible:before:bg-[var(--surface-hover,var(--muted))] aria-disabled:cursor-default aria-disabled:text-muted-foreground aria-disabled:hover:bg-transparent",
                  size === "sm" && "min-h-8 text-[13px] leading-[18px]"
                )}
              >
                {showGuides && (level > 1 || (branch && open)) && (
                  <TreeGuides
                    level={level}
                    last={last}
                    trail={trail}
                    branch={branch}
                    open={open}
                  />
                )}
                <span className="flex shrink-0 items-center gap-1">
                  {branch ? (
                    <Button
                      tabIndex={-1}
                      disabled={node.disabled}
                      aria-label={`${open ? "Collapse" : "Expand"} ${node.label}`}
                      onClick={(event) => {
                        event.stopPropagation();
                        toggle(node.id);
                        focus(node.id);
                      }}
                      className="flex size-5 shrink-0 cursor-pointer items-center justify-center rounded-[6px] text-muted-foreground outline-none hover:text-foreground focus-visible:ring-2 focus-visible:ring-ring"
                    >
                      <RiArrowRightSLine
                        aria-hidden
                        className={cn(
                          "size-4 opacity-50 transition-transform duration-150 motion-reduce:transition-none rtl:rotate-180",
                          open && "rotate-90 rtl:rotate-90"
                        )}
                      />
                    </Button>
                  ) : (
                    <span
                      aria-hidden
                      className={cn("h-5 shrink-0", level > 1 ? "w-2" : "w-5")}
                    />
                  )}
                  {Icon && (
                    <Icon aria-hidden className="size-4 shrink-0 opacity-50" />
                  )}
                </span>
                <span className="min-w-0 flex-1 truncate text-start">
                  <bdi>{node.label}</bdi>
                </span>
                {node.trailingContent != null && (
                  <span className="ms-auto shrink-0 text-xs leading-4 tracking-[0.15px] text-muted-foreground">
                    {node.trailingContent}
                  </span>
                )}
              </div>
            </AnimatedTreeRow>
          );
        })}
      </AnimatePresence>
    </div>
  );
}

function AnimatedTreeRow({ children }: { children: React.ReactNode }) {
  const present = useIsPresent();
  const reduced = useReducedMotion();
  return (
    <motion.div
      role="presentation"
      className="relative min-w-0 shrink-0 overflow-hidden"
      aria-hidden={!present || undefined}
      inert={!present}
      initial={
        reduced
          ? false
          : {
              height: 0,
              paddingTop: 0,
              paddingBottom: 0,
              opacity: 0,
              y: -4,
              filter: "blur(4px)",
            }
      }
      animate={{
        height: "auto",
        paddingTop: 1,
        paddingBottom: 1,
        opacity: 1,
        y: 0,
        filter: "blur(0px)",
      }}
      exit={{
        height: 0,
        paddingTop: 0,
        paddingBottom: 0,
        opacity: 0,
        y: reduced ? 0 : -4,
        filter: reduced ? "blur(0px)" : "blur(4px)",
      }}
      transition={{ duration: reduced ? 0 : 0.28, ease: [0.22, 1, 0.36, 1] }}
    >
      {children}
    </motion.div>
  );
}

const axis = (depth: number) => `calc(1.125rem + ${depth} * 1.25rem)`;

// Trunks for ancestors that continue, a curved elbow into this row, a tail
// when a sibling follows, and a bridge down from an open folder's icon.
function TreeGuides({
  level,
  last,
  trail,
  branch,
  open,
}: {
  level: number;
  last: boolean;
  trail: boolean[];
  branch: boolean;
  open: boolean;
}) {
  const width = branch ? 12 : 16;
  return (
    <span
      aria-hidden
      data-tree-view-guides
      className="pointer-events-none absolute inset-0 text-[color:var(--border-strong,var(--border))]"
    >
      {trail.map(
        (continues, index) =>
          index >= 1 &&
          continues && (
            <span
              key={index}
              className="absolute -inset-y-px w-px bg-current"
              style={{ insetInlineStart: axis(index - 1) }}
            />
          )
      )}
      {level > 1 && (
        <span
          className="absolute -inset-y-px"
          style={{ insetInlineStart: axis(level - 2), width }}
        >
          <span
            className="absolute start-0 top-0 w-px bg-current"
            style={{ bottom: "calc(50% + 6px)" }}
          />
          <svg
            width={width}
            height="7"
            viewBox={`0 0 ${width} 7`}
            fill="none"
            className="absolute start-0 rtl:-scale-x-100"
            style={{ top: "calc(50% - 6px)" }}
          >
            <path
              d={`M0.5 0 Q0.5 6 6.5 6 H${width - 0.5}`}
              stroke="currentColor"
              strokeWidth="1"
            />
          </svg>
          {!last && (
            <span
              className="absolute start-0 bottom-0 w-px bg-current"
              style={{ top: "calc(50% - 6px)" }}
            />
          )}
        </span>
      )}
      {branch && open && (
        <span
          className="absolute -bottom-px w-px bg-current"
          style={{ insetInlineStart: axis(level - 1), top: "calc(50% + 8px)" }}
        />
      )}
    </span>
  );
}