Tree view

'use client';

import { useCallback, useEffect, useId, useMemo, useRef, useState } from 'react';
import { cn, useInView, useTexts } from '@/lib/kit';
import { Icon, type IconName } from '@/components/ui/icon';

/*
 * A tree after the WAI-ARIA APG tree view pattern: one tab stop, arrows move,
 * Right opens or steps into a branch, Left closes or steps to the parent,
 * Home and End jump, * opens all siblings, typing a letter jumps to the next
 * match, Enter or Space selects (or checks). Branches open with a soft
 * height change. Four looks: a file tree with change marks, a docs
 * navigation with links, a permission tree with checkboxes that know their
 * mixed state, and an outline with guide lines and counts.
 */

const TEXTS = {
  en: { tree: 'Tree', expandAll: 'Open all', collapseAll: 'Close all', added: 'added', modified: 'modified', deleted: 'deleted', selected: 'selected', items: 'items' },
  de: { tree: 'Baum', expandAll: 'Alle öffnen', collapseAll: 'Alle schließen', added: 'neu', modified: 'geändert', deleted: 'gelöscht', selected: 'ausgewählt', items: 'Einträge' },
};

export interface TreeNode {
  id: string;
  label: string;
  /** Icon before the label. Files and folders get one by default in the files look. */
  icon?: IconName;
  /** Quiet text at the end, e.g. a size, a count or a role */
  meta?: string;
  /** files: a change mark */
  status?: 'added' | 'modified' | 'deleted';
  /** navigation: the page this node links to */
  href?: string;
  /** checkbox: shown but cannot be changed */
  disabled?: boolean;
  children?: TreeNode[];
}

export type TreeViewVariant = 'files' | 'navigation' | 'checkbox' | 'outline';

export interface TreeViewProps {
  variant?: TreeViewVariant;
  items: TreeNode[];
  /** Accessible name of the tree */
  label?: string;
  /** Ids of open branches at the start */
  defaultExpanded?: string[];
  expanded?: string[];
  onExpandedChange?: (ids: string[]) => void;
  /** The chosen node (files, outline) or the current page (navigation) */
  defaultSelected?: string;
  selected?: string;
  onSelect?: (id: string, node: TreeNode) => void;
  /** checkbox: ids of checked leaves */
  defaultChecked?: string[];
  checked?: string[];
  onCheckedChange?: (ids: string[]) => void;
  /** Lines that connect children to their parent. Default on in outline. */
  guides?: boolean;
  /** A bar above the tree with open all and close all */
  toolbar?: boolean;
  labels?: Partial<(typeof TEXTS)['en']>;
  className?: string;
}

interface Flat {
  node: TreeNode;
  level: number;
  parent: string | null;
  posinset: number;
  setsize: number;
}

/** Visible rows in order: only children of open branches */
function flatten(items: TreeNode[], open: Set<string>, level = 1, parent: string | null = null, out: Flat[] = []) {
  items.forEach((node, i) => {
    out.push({ node, level, parent, posinset: i + 1, setsize: items.length });
    if (node.children?.length && open.has(node.id)) flatten(node.children, open, level + 1, node.id, out);
  });
  return out;
}

const leavesOf = (n: TreeNode): TreeNode[] => (n.children?.length ? n.children.flatMap(leavesOf) : [n]);
const allBranches = (items: TreeNode[]): string[] => items.flatMap((n) => (n.children?.length ? [n.id, ...allBranches(n.children)] : []));
const findPath = (items: TreeNode[], id: string, path: string[] = []): string[] | null => {
  for (const n of items) {
    if (n.id === id) return path;
    if (n.children) {
      const p = findPath(n.children, id, [...path, n.id]);
      if (p) return p;
    }
  }
  return null;
};

const FILE_ICON: Record<string, IconName> = { tsx: 'file-code', ts: 'file-code', js: 'file-code', css: 'file-code', json: 'file-code', md: 'file', png: 'image', jpg: 'image', svg: 'image', pdf: 'file' };
const fileIcon = (n: TreeNode, open: boolean): IconName => n.icon ?? (n.children ? (open ? 'folder-open' : 'folder') : FILE_ICON[n.label.split('.').pop() ?? ''] ?? 'file');

const STATUS_MARK: Record<string, { letter: string; tone: string }> = {
  added: { letter: 'A', tone: 'text-success-foreground' },
  modified: { letter: 'M', tone: 'text-warning-foreground' },
  deleted: { letter: 'D', tone: 'text-destructive' },
};

/** The checkbox state of a node: leaves by their id, branches from their leaves */
function checkState(n: TreeNode, checked: Set<string>): 'on' | 'off' | 'mixed' {
  if (!n.children?.length) return checked.has(n.id) ? 'on' : 'off';
  const leaves = leavesOf(n);
  const on = leaves.filter((l) => checked.has(l.id)).length;
  return on === 0 ? 'off' : on === leaves.length ? 'on' : 'mixed';
}

function Box({ state, disabled }: { state: 'on' | 'off' | 'mixed'; disabled?: boolean }) {
  const on = state !== 'off';
  return (
    <span
      aria-hidden
      className={cn(
        'grid size-[17px] shrink-0 place-items-center rounded-[5px] border transition-[background-color,border-color] duration-200 motion-reduce:transition-none',
        on ? 'border-primary bg-primary text-primary-foreground' : 'border-border-strong bg-background',
        disabled && 'opacity-50',
      )}
    >
      <svg viewBox="0 0 24 24" width={12} height={12} fill="none" className="overflow-visible">
        <path d="m5 12.5 4.5 4.5L19 7.5" stroke="currentColor" strokeWidth="3" strokeLinecap="round" strokeLinejoin="round" className={cn('transition-opacity duration-150', state === 'on' ? 'opacity-100' : 'opacity-0')} />
        <path d="M6 12h12" stroke="currentColor" strokeWidth="3" strokeLinecap="round" className={cn('transition-opacity duration-150', state === 'mixed' ? 'opacity-100' : 'opacity-0')} />
      </svg>
    </span>
  );
}

/**
 * A tree of files, pages, permissions or an outline, with full keyboard support.
 *   <TreeView label="Files" items={files} defaultExpanded={['src']} defaultSelected="src/app.tsx" />
 *   <TreeView variant="navigation" label="Docs" items={pages} selected="install" />
 *   <TreeView variant="checkbox" label="Permissions" items={rights} defaultChecked={['read']} />
 *   <TreeView variant="outline" label="Teams" items={org} />
 */
export function TreeView({
  variant = 'files', items, label, defaultExpanded, expanded: expandedProp, onExpandedChange, defaultSelected, selected: selectedProp, onSelect,
  defaultChecked, checked: checkedProp, onCheckedChange, guides, toolbar, labels, className,
}: TreeViewProps) {
  const t = useTexts(TEXTS, labels);
  const baseId = useId();
  // The current page opens its parents by itself
  const [innerOpen, setInnerOpen] = useState<string[]>(() => {
    const start = new Set(defaultExpanded ?? []);
    const sel = selectedProp ?? defaultSelected;
    if (sel) findPath(items, sel)?.forEach((id) => start.add(id));
    return [...start];
  });
  const openIds = expandedProp ?? innerOpen;
  const open = useMemo(() => new Set(openIds), [openIds]);
  const [innerSel, setInnerSel] = useState(defaultSelected);
  const selected = selectedProp ?? innerSel;
  const [innerChecked, setInnerChecked] = useState<string[]>(defaultChecked ?? []);
  const checked = useMemo(() => new Set(checkedProp ?? innerChecked), [checkedProp, innerChecked]);
  const rows = useMemo(() => flatten(items, open), [items, open]);
  const [focusId, setFocusId] = useState<string | null>(null);
  const tabId = focusId && rows.some((r) => r.node.id === focusId) ? focusId : selected && rows.some((r) => r.node.id === selected) ? selected : rows[0]?.node.id;
  const treeRef = useRef<HTMLUListElement>(null);
  const typed = useRef({ text: '', at: 0 });
  const [said, setSaid] = useState('');
  const [ref, inView] = useInView<HTMLDivElement>();
  const showGuides = guides ?? variant === 'outline';
  const [settled, setSettled] = useState(false);
  useEffect(() => {
    if (!inView) return;
    const id = window.setTimeout(() => setSettled(true), 900);
    return () => window.clearTimeout(id);
  }, [inView]);

  const setOpen = useCallback(
    (next: string[]) => {
      setInnerOpen(next);
      onExpandedChange?.(next);
    },
    [onExpandedChange],
  );
  const toggleOpen = (id: string, to?: boolean) => {
    const isOpen = open.has(id);
    const want = to ?? !isOpen;
    if (want === isOpen) return;
    setOpen(want ? [...openIds, id] : openIds.filter((x) => x !== id));
  };
  const focusRow = (id: string) => {
    setFocusId(id);
    requestAnimationFrame(() => document.getElementById(`${baseId}-${id}`)?.focus());
  };
  const select = (n: TreeNode) => {
    setInnerSel(n.id);
    onSelect?.(n.id, n);
  };
  const toggleCheck = (n: TreeNode) => {
    if (n.disabled) return;
    const leaves = leavesOf(n).filter((l) => !l.disabled);
    const state = checkState(n, checked);
    const next = new Set(checked);
    leaves.forEach((l) => (state === 'on' ? next.delete(l.id) : next.add(l.id)));
    const list = [...next];
    setInnerChecked(list);
    onCheckedChange?.(list);
    setSaid(`${n.label}: ${state === 'on' ? '0' : leaves.length}/${leaves.length}`);
  };
  const activate = (n: TreeNode) => {
    if (variant === 'checkbox') toggleCheck(n);
    else if (variant === 'navigation' && n.children?.length) toggleOpen(n.id);
    else {
      select(n);
      if (n.children?.length) toggleOpen(n.id);
    }
  };

  const onKeyDown = (e: React.KeyboardEvent, r: Flat, i: number) => {
    const n = r.node;
    const branch = Boolean(n.children?.length);
    const go = (j: number) => {
      e.preventDefault();
      const row = rows[Math.max(0, Math.min(rows.length - 1, j))];
      if (row) focusRow(row.node.id);
    };
    switch (e.key) {
      case 'ArrowDown': return go(i + 1);
      case 'ArrowUp': return go(i - 1);
      case 'Home': return go(0);
      case 'End': return go(rows.length - 1);
      case 'ArrowRight':
        e.preventDefault();
        if (branch && !open.has(n.id)) toggleOpen(n.id, true);
        else if (branch) go(i + 1);
        return;
      case 'ArrowLeft':
        e.preventDefault();
        if (branch && open.has(n.id)) toggleOpen(n.id, false);
        else if (r.parent) focusRow(r.parent);
        return;
      case '*': {
        e.preventDefault();
        const siblings = rows.filter((x) => x.parent === r.parent && x.node.children?.length).map((x) => x.node.id);
        setOpen([...new Set([...openIds, ...siblings])]);
        return;
      }
      case 'Enter':
      case ' ':
        e.preventDefault();
        // A page of the navigation opens its link
        if (variant === 'navigation' && n.href && !branch) {
          select(n);
          return (e.currentTarget as HTMLElement).querySelector<HTMLAnchorElement>(':scope > a')?.click();
        }
        return activate(n);
    }
    // Type to jump: the next row that starts with the typed letters
    if (e.key.length === 1 && /\S/.test(e.key) && !e.ctrlKey && !e.metaKey && !e.altKey) {
      const now = Date.now();
      typed.current = { text: now - typed.current.at < 600 ? typed.current.text + e.key.toLowerCase() : e.key.toLowerCase(), at: now };
      const q = typed.current.text;
      const order = [...rows.slice(i + 1), ...rows.slice(0, i + 1)];
      const hit = order.find((x) => x.node.label.toLowerCase().startsWith(q));
      if (hit) focusRow(hit.node.id);
    }
  };

  /* ---------- one row ---------- */
  const renderNode = (node: TreeNode, level: number, parent: string | null, posinset: number, setsize: number, index: { i: number }): React.ReactNode => {
    const branch = Boolean(node.children?.length);
    const isOpen = branch && open.has(node.id);
    const i = rows.findIndex((r) => r.node.id === node.id);
    const flat = rows[i];
    const isSel = selected === node.id;
    const state = variant === 'checkbox' ? checkState(node, checked) : null;
    // The first entrance is staggered, branches opened later come in at once
    const delay = settled ? 0 : Math.min(index.i++, 14) * 28;
    const indent = variant === 'outline' ? 20 : 16;
    const pad = 8 + (level - 1) * indent;
    const Tag = variant === 'navigation' && node.href && !branch ? 'a' : 'div';

    const content = (
      <>
        {/* Chevron for branches, a spacer for leaves, so labels line up */}
        <span aria-hidden className="grid size-4 shrink-0 place-items-center text-subtle-foreground">
          {branch && <Icon name="chevron-right" size={14} strokeWidth={2} className={cn('transition-[rotate] duration-300 ease-[var(--ease-kit)] motion-reduce:transition-none', isOpen && 'rotate-90')} />}
        </span>
        {state && <Box state={state} disabled={node.disabled} />}
        {variant === 'files' && <Icon name={fileIcon(node, isOpen)} size={16} className={cn('shrink-0', branch ? 'text-[color-mix(in_oklab,var(--primary)_70%,var(--muted-foreground))]' : 'text-muted-foreground')} />}
        {variant !== 'files' && node.icon && <Icon name={node.icon} size={16} className="shrink-0 text-muted-foreground" />}
        <span className={cn('min-w-0 flex-1 truncate', node.status === 'deleted' && 'text-muted-foreground line-through decoration-[color-mix(in_oklab,var(--destructive)_60%,transparent)]', variant === 'outline' && level === 1 && 'font-medium')}>{node.label}</span>
        {node.meta && <span className="shrink-0 text-xs tabular-nums text-subtle-foreground">{node.meta}</span>}
        {node.status && (
          <span className={cn('w-3 shrink-0 text-center text-[11px] font-semibold', STATUS_MARK[node.status].tone)}>
            {STATUS_MARK[node.status].letter}
            <span className="sr-only">, {t[node.status]}</span>
          </span>
        )}
      </>
    );

    return (
      <li
        key={node.id}
        id={`${baseId}-${node.id}`}
        role="treeitem"
        aria-level={level}
        aria-setsize={setsize}
        aria-posinset={posinset}
        aria-expanded={branch ? isOpen : undefined}
        aria-selected={variant === 'checkbox' ? undefined : variant === 'navigation' ? undefined : isSel}
        aria-current={variant === 'navigation' && isSel ? 'page' : undefined}
        aria-checked={state ? (state === 'mixed' ? 'mixed' : state === 'on') : undefined}
        aria-disabled={node.disabled || undefined}
        tabIndex={node.id === tabId ? 0 : -1}
        onFocus={(e) => {
          if (e.target === e.currentTarget) setFocusId(node.id);
        }}
        onKeyDown={(e) => {
          if (e.target !== e.currentTarget && !(e.target as HTMLElement).matches('a')) return;
          e.stopPropagation();
          if (flat) onKeyDown(e, flat, i);
        }}
        className="group/node relative outline-none"
      >
        <Tag
          {...(Tag === 'a' ? { href: node.href, tabIndex: -1 } : {})}
          onClick={(e: React.MouseEvent) => {
            e.stopPropagation();
            if (Tag === 'a') {
              select(node);
              return;
            }
            setFocusId(node.id);
            activate(node);
          }}
          style={{ paddingInlineStart: pad, '--d': `${delay}ms` } as React.CSSProperties}
          className={cn(
            'kitfx-tree-row relative flex min-h-11 cursor-default select-none items-center gap-2 rounded-md pe-2.5 text-sm text-foreground transition-[background-color,color] duration-150 sm:min-h-8',
            'hover:bg-accent group-focus-visible/node:shadow-[inset_0_0_0_2px_color-mix(in_oklab,var(--primary)_55%,transparent)]',
            variant === 'navigation' && 'text-muted-foreground hover:text-foreground',
            variant === 'navigation' && isSel && 'font-medium text-foreground',
            variant !== 'navigation' && variant !== 'checkbox' && isSel && 'bg-[color-mix(in_oklab,var(--primary)_12%,transparent)] hover:bg-[color-mix(in_oklab,var(--primary)_16%,transparent)]',
            node.disabled && 'cursor-not-allowed opacity-60',
            Tag === 'a' && 'cursor-pointer',
            inView ? '' : 'kitfx-tree-wait',
          )}
        >
          {/* navigation: the current page gets a bar on the guide line */}
          {variant === 'navigation' && isSel && <span aria-hidden className="absolute inset-y-1.5 w-[2px] rounded-full bg-primary transition-transform duration-300 starting:scale-y-0" style={{ insetInlineStart: pad - 8 }} />}
          {content}
        </Tag>
        {branch && (
          // Grid rows from 0fr to 1fr give the branch a soft height without measuring
          <div inert={!isOpen || undefined} className={cn('grid transition-[grid-template-rows,opacity] duration-300 ease-[var(--ease-kit)] motion-reduce:transition-none', isOpen ? 'grid-rows-[1fr] opacity-100' : 'grid-rows-[0fr] opacity-0')}>
            <ul role="group" className="relative min-h-0 overflow-hidden">
              {(showGuides || variant === 'navigation') && (
                <span aria-hidden className="pointer-events-none absolute bottom-1 top-0 w-px bg-border" style={{ insetInlineStart: pad + 7.5 }} />
              )}
              {node.children!.map((c, k) => renderNode(c, level + 1, node.id, k + 1, node.children!.length, index))}
            </ul>
          </div>
        )}
      </li>
    );
  };

  const counter = { i: 0 };
  const total = variant === 'checkbox' ? items.flatMap(leavesOf).length : 0;

  return (
    <div ref={ref} className={cn('w-full', className)}>
      <style href="kitfx-tree" precedence="default">{TREE_CSS}</style>
      {toolbar && (
        <div className="mb-2 flex items-center justify-between gap-2 px-1">
          <span className="text-xs text-subtle-foreground tabular-nums">
            {variant === 'checkbox' ? `${[...checked].filter((id) => items.flatMap(leavesOf).some((l) => l.id === id)).length}/${total} ${t.selected}` : `${rows.length} ${t.items}`}
          </span>
          <span className="flex gap-1">
            <button type="button" onClick={() => setOpen(allBranches(items))} className="inline-flex h-11 items-center rounded-md px-2 text-xs font-medium text-muted-foreground outline-none transition-colors hover:bg-accent hover:text-foreground focus-visible:bg-accent focus-visible:shadow-[inset_0_0_0_2px_color-mix(in_oklab,var(--primary)_55%,transparent)] sm:h-7">
              {t.expandAll}
            </button>
            <button type="button" onClick={() => setOpen([])} className="inline-flex h-11 items-center rounded-md px-2 text-xs font-medium text-muted-foreground outline-none transition-colors hover:bg-accent hover:text-foreground focus-visible:bg-accent focus-visible:shadow-[inset_0_0_0_2px_color-mix(in_oklab,var(--primary)_55%,transparent)] sm:h-7">
              {t.collapseAll}
            </button>
          </span>
        </div>
      )}
      <ul ref={treeRef} role="tree" aria-label={label ?? t.tree} aria-multiselectable={variant === 'checkbox' || undefined} className="relative">
        {items.map((n, k) => renderNode(n, 1, null, k + 1, items.length, counter))}
      </ul>
      <span aria-live="polite" className="sr-only">
        {said}
      </span>
    </div>
  );
}

/* Rows rise in one after another the first time the tree scrolls into view, and when a branch opens */
const TREE_CSS = `
@keyframes kitfx-tree-in { from { opacity: 0; transform: translateX(-6px); } }
@media (prefers-reduced-motion: no-preference) {
  .kitfx-tree-row { animation: kitfx-tree-in .42s var(--ease-kit) both; animation-delay: var(--d, 0ms); }
  .kitfx-tree-wait { animation-play-state: paused; }
}
@media (scripting: none) { .kitfx-tree-row { animation: none; } }
`;