Tree View

'use client';

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

import { cn } from '@/lib/utils';
import { Collapsible, CollapsibleContent, CollapsibleTrigger } from '@/registry/hirael/bases/radix/ui/collapsible';
import { useControllableState } from '@/registry/hirael/hooks/use-controllable-state';

interface TreeCtx {
  selected: string | undefined;
  setSelected: (value: string) => void;
  tabbable: string | null;
  setTabbable: (value: string) => void;
  registerItem: (value: string, el: HTMLButtonElement) => () => void;
  getVisibleItems: () => HTMLButtonElement[];
}

const TreeViewContext = React.createContext<TreeCtx | null>(null);

const useTreeView = () => {
  const ctx = React.useContext(TreeViewContext);
  if (!ctx) {
    throw new Error('TreeItem must be used inside <TreeView>');
  }

  return ctx;
};

const TreeDepthContext = React.createContext(0);

const TreeParentContext = React.createContext<React.RefObject<HTMLButtonElement | null> | null>(null);

const TREE_INDENT_PER_LEVEL = 14;
const TREE_INDENT_BASE = 8;

export interface TreeViewProps extends Omit<React.ComponentProps<'div'>, 'defaultValue' | 'onChange'> {
  /** Id of the selected leaf. */
  value?: string;
  defaultValue?: string;
  onValueChange?: (value: string) => void;
}

const TreeView = ({ value: valueProp, defaultValue, onValueChange, className, ...props }: TreeViewProps) => {
  const [selected, setSelected] = useControllableState<string | undefined>({
    prop: valueProp,
    defaultProp: defaultValue,
    onChange: onValueChange,
  });

  const itemsRef = React.useRef(new Map<string, HTMLButtonElement>());
  const [tabbable, setTabbable] = React.useState<string | null>(null);

  const getVisibleItems = React.useCallback(
    () =>
      Array.from(itemsRef.current.values())
        .filter((el) => el.isConnected && !el.disabled)
        .sort((a, b) => (a.compareDocumentPosition(b) & Node.DOCUMENT_POSITION_FOLLOWING ? -1 : 1)),
    [],
  );

  const firstVisibleValue = React.useCallback(() => {
    const first = getVisibleItems()[0];
    if (!first) return null;
    for (const [itemValue, itemEl] of itemsRef.current) {
      if (itemEl === first) return itemValue;
    }

    return null;
  }, [getVisibleItems]);

  const registerItem = React.useCallback(
    (value: string, el: HTMLButtonElement) => {
      itemsRef.current.set(value, el);

      return () => {
        itemsRef.current.delete(value);
        setTabbable((current) => (current === value ? firstVisibleValue() : current));
      };
    },
    [firstVisibleValue],
  );

  React.useEffect(() => {
    setTabbable((current) => {
      if (current !== null && itemsRef.current.has(current)) return current;
      if (selected !== undefined && itemsRef.current.has(selected)) {
        return selected;
      }

      return firstVisibleValue();
    });
  }, [selected, firstVisibleValue]);

  const ctx = React.useMemo<TreeCtx>(
    () => ({
      selected,
      setSelected,
      tabbable,
      setTabbable,
      registerItem,
      getVisibleItems,
    }),
    [selected, setSelected, tabbable, setTabbable, registerItem, getVisibleItems],
  );

  return (
    <TreeViewContext.Provider value={ctx}>
      <div data-slot="tree-view" role="tree" className={cn('w-full text-sm select-none', className)} {...props} />
    </TreeViewContext.Provider>
  );
};

export interface TreeItemProps extends Omit<React.ComponentProps<'div'>, 'title'> {
  /** Unique id used for selection. */
  value: string;
  label: React.ReactNode;
  /** Override the leading icon. Pass `null` to hide it entirely. */
  icon?: React.ReactNode;
  /** Controlled expanded state. Pair with `onExpandedChange`. */
  expanded?: boolean;
  defaultExpanded?: boolean;
  onExpandedChange?: (expanded: boolean) => void;
  disabled?: boolean;
}

const TreeItem = ({
  value,
  label,
  icon,
  expanded: expandedProp,
  defaultExpanded = false,
  onExpandedChange,
  disabled = false,
  className,
  children,
  ...props
}: TreeItemProps) => {
  const { selected, setSelected, tabbable, setTabbable, registerItem, getVisibleItems } = useTreeView();
  const depth = React.useContext(TreeDepthContext);
  const parentTriggerRef = React.useContext(TreeParentContext);
  const triggerRef = React.useRef<HTMLButtonElement>(null);
  const [open, setOpen] = useControllableState({
    prop: expandedProp,
    defaultProp: defaultExpanded,
    onChange: onExpandedChange,
  });

  const hasChildren = React.Children.count(children) > 0;
  const isSelected = !hasChildren && selected === value;

  React.useLayoutEffect(() => {
    const el = triggerRef.current;
    if (!el) return;

    return registerItem(value, el);
  }, [registerItem, value]);

  // Hands focus to the parent trigger on unmount (e.g. an ancestor collapsed). Separate from the
  // registration effect, which re-runs on `value` changes while the item stays mounted.
  React.useLayoutEffect(() => {
    const el = triggerRef.current;
    const parentTrigger = parentTriggerRef?.current ?? null;

    return () => {
      const holdsFocus = el?.closest('[data-slot="tree-item"]')?.contains(document.activeElement) ?? false;
      if (holdsFocus) parentTrigger?.focus();
    };
  }, [parentTriggerRef]);

  const onTriggerKeyDown = (event: React.KeyboardEvent<HTMLButtonElement>) => {
    const { key } = event;
    if (
      key !== 'ArrowDown' &&
      key !== 'ArrowUp' &&
      key !== 'ArrowLeft' &&
      key !== 'ArrowRight' &&
      key !== 'Home' &&
      key !== 'End'
    ) {
      return;
    }
    event.preventDefault();
    const items = getVisibleItems();
    const index = items.indexOf(event.currentTarget);
    const rtl = getComputedStyle(event.currentTarget).direction === 'rtl';
    const expandKey = rtl ? 'ArrowLeft' : 'ArrowRight';

    if (key === 'ArrowDown') {
      items[index + 1]?.focus();
    } else if (key === 'ArrowUp') {
      items[index - 1]?.focus();
    } else if (key === 'Home') {
      items[0]?.focus();
    } else if (key === 'End') {
      items[items.length - 1]?.focus();
    } else if (key === expandKey) {
      if (!hasChildren) return;
      if (open) items[index + 1]?.focus();
      else setOpen(true);
    } else if (hasChildren && open) {
      setOpen(false);
    } else {
      parentTriggerRef?.current?.focus();
    }
  };

  const triggerProps = {
    ref: triggerRef,
    type: 'button' as const,
    disabled,
    tabIndex: tabbable === value ? 0 : -1,
    'data-slot': 'tree-item-trigger',
    style: { paddingInlineStart: depth * TREE_INDENT_PER_LEVEL + TREE_INDENT_BASE },
    className: cn(
      'group flex h-7 w-full items-center gap-1.5 rounded-sm pe-2 text-start transition-colors outline-none',
      'hover:bg-accent focus-visible:bg-accent focus-visible:ring-2 focus-visible:ring-ring',
      'disabled:cursor-not-allowed disabled:opacity-50',
      isSelected ? 'bg-accent font-medium text-foreground' : 'text-foreground/80',
    ),
    onKeyDown: onTriggerKeyDown,
    onFocus: () => setTabbable(value),
  };

  const leadingIcon =
    icon !== undefined ? (
      icon
    ) : hasChildren ? (
      <>
        <Folder className="size-4 text-muted-foreground group-data-[state=open]:hidden" />
        <FolderOpen className="hidden size-4 text-muted-foreground group-data-[state=open]:block" />
      </>
    ) : (
      <File className="size-4 text-muted-foreground" />
    );

  const labelRow = (
    <>
      <ChevronRight
        className={cn(
          'size-3.5 shrink-0 text-muted-foreground transition-transform duration-150',
          hasChildren ? 'group-data-[state=open]:rotate-90 rtl:group-data-[state=closed]:rotate-180' : 'invisible',
        )}
      />
      {leadingIcon != null && (
        <span data-slot="tree-item-icon" className="flex shrink-0 items-center [&_svg]:size-4">
          {leadingIcon}
        </span>
      )}
      <span data-slot="tree-item-label" className="min-w-0 truncate">
        {label}
      </span>
    </>
  );

  if (hasChildren) {
    return (
      <Collapsible asChild open={open} onOpenChange={setOpen}>
        <div
          data-slot="tree-item"
          role="treeitem"
          aria-selected={false}
          aria-expanded={open}
          aria-level={depth + 1}
          className={className}
          {...props}
        >
          <CollapsibleTrigger asChild>
            <button {...triggerProps}>{labelRow}</button>
          </CollapsibleTrigger>
          <CollapsibleContent asChild>
            <div role="group">
              <TreeDepthContext.Provider value={depth + 1}>
                <TreeParentContext.Provider value={triggerRef}>{children}</TreeParentContext.Provider>
              </TreeDepthContext.Provider>
            </div>
          </CollapsibleContent>
        </div>
      </Collapsible>
    );
  }

  return (
    <div
      data-slot="tree-item"
      data-selected={isSelected ? '' : undefined}
      role="treeitem"
      aria-selected={isSelected}
      aria-level={depth + 1}
      className={className}
      {...props}
    >
      <button {...triggerProps} data-selected={isSelected ? '' : undefined} onClick={() => setSelected(value)}>
        {labelRow}
      </button>
    </div>
  );
};

export { TreeView, TreeItem };