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';

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);

/** Horizontal inset per nesting level, plus the base inset, in px. */
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 [internal, setInternal] = React.useState<string | undefined>(defaultValue);
  const selected = valueProp ?? internal;

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

  const setSelected = React.useCallback(
    (next: string) => {
      if (valueProp === undefined) setInternal(next);
      onValueChange?.(next);
    },
    [valueProp, onValueChange],
  );

  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 select-none text-sm', 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;
  defaultExpanded?: boolean;
  disabled?: boolean;
}

const TreeItem = ({
  value,
  label,
  icon,
  defaultExpanded = false,
  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] = React.useState(defaultExpanded);

  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]);

  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 triggerStyle: React.CSSProperties = {
    paddingInlineStart: depth * TREE_INDENT_PER_LEVEL + TREE_INDENT_BASE,
  };
  const triggerClassName = cn(
    'group flex h-7 w-full items-center gap-1.5 rounded-sm pe-2 text-start outline-none transition-colors',
    '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',
    className,
  );
  const triggerProps = {
    ref: triggerRef,
    type: 'button' as const,
    disabled,
    tabIndex: tabbable === value ? 0 : -1,
    'data-slot': 'tree-item-trigger',
    style: triggerStyle,
    className: triggerClassName,
    onKeyDown: onTriggerKeyDown,
    onFocus: () => setTabbable(value),
  };

  // Folder open/closed and the chevron rotation are driven off the trigger's
  // data-[state] (set by Collapsible) instead of tracked React state.
  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
        aria-hidden
        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 className="flex shrink-0 items-center [&_svg]:size-4">{leadingIcon}</span>}
      <span 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={isSelected} aria-expanded={open} {...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" role="treeitem" aria-selected={isSelected} {...props}>
      <button {...triggerProps} data-state={isSelected ? 'selected' : undefined} onClick={() => setSelected(value)}>
        {labelRow}
      </button>
    </div>
  );
};

export { TreeView, TreeItem };