Tree table

"use client";

import { useEffect, useMemo, useRef, useState } from "react";
import type { ComponentProps, CSSProperties, FocusEvent, KeyboardEvent, MouseEvent, ReactNode } from "react";
import "./tree-table.css";

export type TreeTableColumn<T> = {
  key: string;
  header: string;
  align?: "start" | "end";
  /** Width in pixels. Columns without one share the remaining space. */
  width?: number;
  /** Renders the cell. Without it the cell shows the row's value for key. */
  render?: (row: T) => ReactNode;
};

export type TreeTableProps<T> = Omit<ComponentProps<"table">, "children"> & {
  /** Names the treegrid for assistive tech. */
  label: string;
  /** The first column holds the indent and the expand chevron. */
  columns: TreeTableColumn<T>[];
  /** The top level rows. */
  rows: T[];
  rowKey: (row: T) => string;
  /** A row's children. Rows that return a non-empty array can be expanded. */
  getChildren: (row: T) => T[] | undefined;
  expanded?: string[];
  defaultExpanded?: string[];
  onExpandedChange?: (keys: string[]) => void;
  /** Caps the height and scrolls the body under a sticky header. */
  maxHeight?: number | string;
  dense?: boolean;
};

type Line<T> = {
  row: T;
  key: string;
  level: number;
  setsize: number;
  posinset: number;
  /** Index of the parent in the visible list, or -1 at the top level. */
  parentIndex: number;
  hasChildren: boolean;
};

type Node<T> = { row: T; parent: string | null };

function valueOf(row: object, key: string): unknown {
  return (row as Record<string, unknown>)[key];
}

function clamp(n: number, min: number, max: number) {
  return Math.min(max, Math.max(min, n));
}

export function TreeTable<T extends object>({
  label,
  columns,
  rows,
  rowKey,
  getChildren,
  expanded,
  defaultExpanded,
  onExpandedChange,
  maxHeight,
  dense = false,
  className,
  onFocus: onFocusProp,
  onKeyDown: onKeyDownProp,
  ...rest
}: TreeTableProps<T>) {
  const rootRef = useRef<HTMLDivElement>(null);
  const refocus = useRef(false);
  const [innerExpanded, setInnerExpanded] = useState<string[]>(defaultExpanded ?? []);
  // col -1 is the row itself; 0 and up are its cells.
  const [focus, setFocus] = useState<{ key: string | null; col: number }>({ key: null, col: -1 });

  const openKeys = expanded ?? innerExpanded;
  const open = useMemo(() => new Set(openKeys), [openKeys]);

  // Every node, open or not, so a hidden focus can climb to the nearest visible parent.
  const nodes = new Map<string, Node<T>>();
  (function index(list: T[], parent: string | null) {
    for (const row of list) {
      const key = rowKey(row);
      nodes.set(key, { row, parent });
      const kids = getChildren(row);
      if (kids?.length) index(kids, key);
    }
  })(rows, null);

  const visible: Line<T>[] = [];
  (function walk(list: T[], level: number, parentIndex: number) {
    list.forEach((row, i) => {
      const key = rowKey(row);
      const kids = getChildren(row);
      const hasChildren = !!kids && kids.length > 0;
      const at = visible.length;
      visible.push({ row, key, level, setsize: list.length, posinset: i + 1, parentIndex, hasChildren });
      if (hasChildren && open.has(key)) walk(kids, level + 1, at);
    });
  })(rows, 1, -1);

  const lastRow = visible.length - 1;
  const lastCol = columns.length - 1;
  const where = new Map(visible.map((v, i) => [v.key, i]));

  // Resolve the tab stop: the focused row, or its nearest visible ancestor when a branch closed over it.
  let fr = focus.key == null ? 0 : (where.get(focus.key) ?? -1);
  let fc = focus.col;
  if (fr === -1) {
    let up = focus.key == null ? null : (nodes.get(focus.key)?.parent ?? null);
    while (up != null && !where.has(up)) up = nodes.get(up)?.parent ?? null;
    fr = up == null ? 0 : (where.get(up) ?? 0);
    fc = -1;
  }
  fc = clamp(fc, -1, lastCol);

  useEffect(() => {
    if (!refocus.current) return;
    refocus.current = false;
    rootRef.current?.querySelector<HTMLElement>(`[data-tt="${fr}:${fc}"]`)?.focus();
  });

  function commit(next: string[]) {
    if (expanded === undefined) setInnerExpanded(next);
    onExpandedChange?.(next);
  }

  function setOpen(keys: string[], value: boolean) {
    const next = new Set(openKeys);
    for (const k of keys) {
      if (value) next.add(k);
      else next.delete(k);
    }
    commit([...next]);
  }

  function isInside(key: string, ancestor: string) {
    let up = nodes.get(key)?.parent ?? null;
    while (up != null) {
      if (up === ancestor) return true;
      up = nodes.get(up)?.parent ?? null;
    }
    return false;
  }

  function toggle(key: string) {
    const closing = open.has(key);
    // A pointer can close a branch around the focused row; focus then lands on the branch.
    const active = document.activeElement;
    const at = active instanceof HTMLElement && rootRef.current?.contains(active) ? active.closest<HTMLElement>("[data-tt]")?.dataset.tt : undefined;
    const current = at ? visible[Number(at.split(":")[0])]?.key : undefined;
    if (closing && current && isInside(current, key)) {
      refocus.current = true;
      setFocus({ key, col: -1 });
    }
    setOpen([key], !closing);
  }

  function go(r: number, c: number) {
    if (visible.length === 0) return;
    const row = clamp(r, 0, lastRow);
    const col = clamp(c, -1, lastCol);
    setFocus({ key: visible[row].key, col });
    rootRef.current?.querySelector<HTMLElement>(`[data-tt="${row}:${col}"]`)?.focus();
  }

  function onFocus(e: FocusEvent<HTMLTableElement>) {
    onFocusProp?.(e);
    const at = (e.target as HTMLElement).dataset?.tt;
    if (!at) return;
    const [r, c] = at.split(":").map(Number);
    const key = visible[r]?.key;
    if (key && (key !== focus.key || c !== focus.col)) setFocus({ key, col: c });
  }

  function onKeyDown(e: KeyboardEvent<HTMLTableElement>) {
    onKeyDownProp?.(e);
    if (e.defaultPrevented || !(e.target instanceof HTMLElement)) return;
    const at = e.target.dataset.tt;
    if (!at) return;
    const [r, c] = at.split(":").map(Number);
    const line = visible[r];
    if (!line) return;
    const onRow = c === -1;
    const mod = e.ctrlKey || e.metaKey;

    switch (e.key) {
      case "ArrowRight":
        if (!onRow) go(r, c + 1);
        else if (line.hasChildren && !open.has(line.key)) setOpen([line.key], true);
        else go(r, 0);
        break;
      case "ArrowLeft":
        if (!onRow) go(r, c - 1);
        else if (line.hasChildren && open.has(line.key)) setOpen([line.key], false);
        else if (line.parentIndex >= 0) go(line.parentIndex, -1);
        break;
      case "ArrowDown":
        go(r + 1, c);
        break;
      case "ArrowUp":
        go(r - 1, c);
        break;
      case "Home":
        if (mod || onRow) go(0, c);
        else go(r, 0);
        break;
      case "End":
        if (mod || onRow) go(lastRow, c);
        else go(r, lastCol);
        break;
      case "Enter":
        if (!line.hasChildren || c > 0) return;
        toggle(line.key);
        break;
      case "*": {
        const parent = nodes.get(line.key)?.parent ?? null;
        const siblings = parent == null ? rows : (getChildren(nodes.get(parent)!.row) ?? []);
        setOpen(
          siblings.filter((s) => (getChildren(s)?.length ?? 0) > 0).map(rowKey),
          true,
        );
        break;
      }
      default:
        return;
    }
    e.preventDefault();
    e.stopPropagation();
  }

  function onChevron(e: MouseEvent<HTMLSpanElement>, key: string) {
    e.stopPropagation();
    toggle(key);
  }

  const tab = (r: number, c: number) => (r === fr && c === fc ? 0 : -1);

  return (
    <div ref={rootRef} className="eq-tree-table-scroll" data-dense={dense || undefined} style={{ maxHeight }}>
      <table
        role="treegrid"
        aria-label={label}
        aria-rowcount={visible.length + 1}
        aria-colcount={columns.length}
        {...rest}
        className={["eq-tree-table", className].filter(Boolean).join(" ")}
        onFocus={onFocus}
        onKeyDown={onKeyDown}
      >
        <thead className="eq-tree-table-head">
          <tr className="eq-tree-table-row" role="row" aria-rowindex={1}>
            {columns.map((col, i) => (
              <th
                key={col.key}
                className="eq-tree-table-header"
                role="columnheader"
                aria-colindex={i + 1}
                data-align={col.align === "end" && i > 0 ? "end" : undefined}
                style={col.width != null ? { width: col.width } : undefined}
              >
                {col.header}
              </th>
            ))}
          </tr>
        </thead>
        <tbody className="eq-tree-table-body">
          {visible.map((line, r) => (
            <tr
              key={line.key}
              className="eq-tree-table-row"
              role="row"
              aria-rowindex={r + 2}
              aria-level={line.level}
              aria-setsize={line.setsize}
              aria-posinset={line.posinset}
              aria-expanded={line.hasChildren ? open.has(line.key) : undefined}
              data-tt={`${r}:-1`}
              tabIndex={tab(r, -1)}
              style={{ "--_level": line.level - 1 } as CSSProperties}
            >
              {columns.map((col, i) => {
                const value = valueOf(line.row, col.key);
                const content = col.render ? col.render(line.row) : value == null ? "" : String(value);
                return (
                  <td
                    key={col.key}
                    className="eq-tree-table-cell"
                    role="gridcell"
                    aria-colindex={i + 1}
                    data-align={col.align === "end" && i > 0 ? "end" : undefined}
                    data-tt={`${r}:${i}`}
                    tabIndex={tab(r, i)}
                  >
                    {i === 0 ? (
                      <span className="eq-tree-table-content">
                        <span
                          className="eq-tree-table-toggle"
                          aria-hidden="true"
                          onClick={line.hasChildren ? (e) => onChevron(e, line.key) : undefined}
                        >
                          {line.hasChildren && (
                            <svg
                              className="eq-tree-table-chevron"
                              viewBox="0 0 16 16"
                              fill="none"
                              stroke="currentColor"
                              strokeWidth="1.5"
                              strokeLinecap="round"
                              strokeLinejoin="round"
                            >
                              <path d="M6 4l4 4-4 4" />
                            </svg>
                          )}
                        </span>
                        <span className="eq-tree-table-text">{content}</span>
                      </span>
                    ) : (
                      content
                    )}
                  </td>
                );
              })}
            </tr>
          ))}
        </tbody>
      </table>
    </div>
  );
}