Implementing Keyboard Navigation in JSON Tree Views

import React, { useEffect, useMemo, useRef, useState } from 'react';

type TreeNode = {
  id: string;
  label: string;
  level: number;
  parentId: string | null;
  childIds: string[];
  valuePreview: string;
};

type Props = {
  nodes: Map<string, TreeNode>;
  rootIds: string[];
};

export function JsonTree({ nodes, rootIds }: Props) {
  const [expandedIds, setExpandedIds] = useState<Set<string>>(() => new Set(rootIds));
  const [activeId, setActiveId] = useState<string | null>(rootIds[0] ?? null);
  const [selectedId, setSelectedId] = useState<string | null>(null);
  const itemRefs = useRef(new Map<string, HTMLLIElement>());
  const typeaheadRef = useRef<{ value: string; timeout?: number }>({ value: '' });

  const visibleIds = useMemo(
    () => buildVisibleIds(rootIds, nodes, expandedIds),
    [rootIds, nodes, expandedIds],
  );

  useEffect(() => {
    if (!activeId) return;
    const element = itemRefs.current.get(activeId);
    element?.focus();
    element?.scrollIntoView({ block: 'nearest' });
  }, [activeId]);

  const isExpanded = (id: string) => expandedIds.has(id);

  const toggleExpanded = (id: string) => {
    setExpandedIds((previous) => {
      const next = new Set(previous);
      next.has(id) ? next.delete(id) : next.add(id);
      return next;
    });
  };

  const moveByOffset = (offset: number) => {
    if (!activeId) return;
    const currentIndex = visibleIds.indexOf(activeId);
    const nextId = visibleIds[currentIndex + offset];
    if (nextId) setActiveId(nextId);
  };

  const runTypeahead = (character: string) => {
    const buffer = (typeaheadRef.current.value + character).toLowerCase();
    typeaheadRef.current.value = buffer;
    window.clearTimeout(typeaheadRef.current.timeout);
    typeaheadRef.current.timeout = window.setTimeout(() => {
      typeaheadRef.current.value = '';
    }, 500);

    const startIndex = activeId ? visibleIds.indexOf(activeId) + 1 : 0;
    const searchOrder = [...visibleIds.slice(startIndex), ...visibleIds.slice(0, startIndex)];
    const match = searchOrder.find((id) => nodes.get(id)?.label.toLowerCase().startsWith(buffer));
    if (match) setActiveId(match);
  };

  const onKeyDown = (event: React.KeyboardEvent<HTMLUListElement>) => {
    if (!activeId) return;

    const node = nodes.get(activeId)!;
    const parentId = node.parentId;
    const firstChildId = node.childIds[0];
    const hasChildren = node.childIds.length > 0;

    switch (event.key) {
      case 'ArrowDown':
        event.preventDefault();
        moveByOffset(1);
        break;
      case 'ArrowUp':
        event.preventDefault();
        moveByOffset(-1);
        break;
      case 'ArrowRight':
        event.preventDefault();
        if (!hasChildren) break;
        if (!isExpanded(node.id)) toggleExpanded(node.id);
        else if (firstChildId) setActiveId(firstChildId);
        break;
      case 'ArrowLeft':
        event.preventDefault();
        if (hasChildren && isExpanded(node.id)) toggleExpanded(node.id);
        else if (parentId) setActiveId(parentId);
        break;
      case 'Home':
        event.preventDefault();
        setActiveId(visibleIds[0] ?? activeId);
        break;
      case 'End':
        event.preventDefault();
        setActiveId(visibleIds.at(-1) ?? activeId);
        break;
      case 'Enter':
      case ' ':
        event.preventDefault();
        if (hasChildren) toggleExpanded(node.id);
        else setSelectedId(node.id);
        break;
      default:
        if (event.key.length === 1 && /\S/.test(event.key)) runTypeahead(event.key);
    }
  };

  const renderNode = (id: string): React.ReactNode => {
    const node = nodes.get(id)!;
    const hasChildren = node.childIds.length > 0;
    const expanded = isExpanded(id);

    return (
      <li
        key={id}
        ref={(element) => {
          if (element) itemRefs.current.set(id, element);
          else itemRefs.current.delete(id);
        }}
        role='treeitem'
        tabIndex={activeId === id ? 0 : -1}
        aria-level={node.level}
        aria-selected={selectedId === id ? true : undefined}
        aria-expanded={hasChildren ? expanded : undefined}
        onClick={() => setActiveId(id)}
        className={activeId === id ? 'outline outline-2 outline-sky-600' : undefined}
      >
        <span className='font-medium'>{node.label}</span>{' '}
        <span className='text-gray-500'>{node.valuePreview}</span>

        {hasChildren && expanded ? (
          <ul role='group' className='pl-4'>
            {node.childIds.map(renderNode)}
          </ul>
        ) : null}
      </li>
    );
  };

  return (
    <ul role='tree' aria-label='JSON document structure' onKeyDown={onKeyDown}>
      {rootIds.map(renderNode)}
    </ul>
  );
}

function buildVisibleIds(
  rootIds: string[],
  nodes: Map<string, TreeNode>,
  expandedIds: Set<string>,
) {
  const visible: string[] = [];

  const visit = (id: string) => {
    visible.push(id);
    if (!expandedIds.has(id)) return;
    nodes.get(id)?.childIds.forEach(visit);
  };

  rootIds.forEach(visit);
  return visible;
}