design-system/src/design-system/layout/Sidebar/SidebarTree.tsx

import {
  useState,
  useRef,
  useCallback,
  useEffect,
  useMemo,
  type ReactNode,
  type KeyboardEvent,
  type MouseEvent,
} from 'react'
import { useNavigate } from 'react-router-dom'
import { Star, ChevronRight, ChevronDown } from 'lucide-react'
import styles from './Sidebar.module.css'

// ── Types ────────────────────────────────────────────────────────────────────

export interface SidebarTreeNode {
  id: string
  label: string
  icon?: ReactNode
  badge?: string
  path?: string
  starrable?: boolean
  starKey?: string       // unique key for starring (defaults to id)
  children?: SidebarTreeNode[]
}

export interface SidebarTreeProps {
  nodes: SidebarTreeNode[]
  selectedPath?: string   // current URL path — matches against node.path
  isStarred: (id: string) => boolean
  onToggleStar: (id: string) => void
  className?: string
  filterQuery?: string
  persistKey?: string     // sessionStorage key to persist expand state across remounts
  autoRevealPath?: string | null  // when set, auto-expand the parent of the matching node
  onNavigate?: (path: string) => void
}

// ── Star icons ───────────────────────────────────────────────────────────────

function StarOutline() {
  return <Star size={14} />
}

function StarFilled() {
  return <Star size={14} fill="currentColor" />
}

// ── Persistent expand state ──────────────────────────────────────────────────

function readExpandState(key: string): Set<string> {
  try {
    const raw = sessionStorage.getItem(key)
    if (raw) {
      const arr = JSON.parse(raw)
      if (Array.isArray(arr)) return new Set(arr)
    }
  } catch { /* ignore */ }
  return new Set()
}

function writeExpandState(key: string, ids: Set<string>): void {
  try {
    sessionStorage.setItem(key, JSON.stringify([...ids]))
  } catch { /* ignore */ }
}

// ── Flat node for keyboard nav ───────────────────────────────────────────────

interface FlatNode {
  node: SidebarTreeNode
  depth: number
  parentId: string | null
}

function flattenVisible(
  nodes: SidebarTreeNode[],
  expandedIds: Set<string>,
  depth = 0,
  parentId: string | null = null,
): FlatNode[] {
  const result: FlatNode[] = []
  for (const node of nodes) {
    result.push({ node, depth, parentId })
    if (node.children && node.children.length > 0 && expandedIds.has(node.id)) {
      result.push(...flattenVisible(node.children, expandedIds, depth + 1, node.id))
    }
  }
  return result
}

// ── Filter logic ─────────────────────────────────────────────────────────────

function filterNodes(
  nodes: SidebarTreeNode[],
  query: string,
): { filtered: SidebarTreeNode[]; matchedParentIds: Set<string> } {
  if (!query) return { filtered: nodes, matchedParentIds: new Set() }
  const q = query.toLowerCase()
  const matchedParentIds = new Set<string>()

  function walk(nodeList: SidebarTreeNode[]): SidebarTreeNode[] {
    const result: SidebarTreeNode[] = []
    for (const node of nodeList) {
      const childResults = node.children ? walk(node.children) : []
      const selfMatches = node.label.toLowerCase().includes(q)

      if (selfMatches || childResults.length > 0) {
        if (childResults.length > 0) {
          matchedParentIds.add(node.id)
        }
        result.push({
          ...node,
          children: childResults.length > 0
            ? childResults
            : node.children?.filter((c) => c.label.toLowerCase().includes(q)),
        })
      }
    }
    return result
  }

  return { filtered: walk(nodes), matchedParentIds }
}

// ── Keyboard nav helpers ─────────────────────────────────────────────────────

function handleArrowDown(visibleNodes: FlatNode[], currentIndex: number, focusNode: (id: string) => void) {
  const next = visibleNodes[currentIndex + 1]
  if (next) focusNode(next.node.id)
}

function handleArrowUp(visibleNodes: FlatNode[], currentIndex: number, focusNode: (id: string) => void) {
  const prev = visibleNodes[currentIndex - 1]
  if (prev) focusNode(prev.node.id)
}

function handleArrowRight(
  current: FlatNode | undefined,
  currentIndex: number,
  expandedSet: Set<string>,
  visibleNodes: FlatNode[],
  handleToggle: (id: string) => void,
  focusNode: (id: string) => void,
) {
  if (!current) return
  const hasChildren = current.node.children && current.node.children.length > 0
  if (!hasChildren) return
  if (!expandedSet.has(current.node.id)) {
    handleToggle(current.node.id)
  } else {
    const next = visibleNodes[currentIndex + 1]
    if (next) focusNode(next.node.id)
  }
}

function handleArrowLeft(
  current: FlatNode | undefined,
  expandedSet: Set<string>,
  handleToggle: (id: string) => void,
  focusNode: (id: string) => void,
) {
  if (!current) return
  const hasChildren = current.node.children && current.node.children.length > 0
  if (hasChildren && expandedSet.has(current.node.id)) {
    handleToggle(current.node.id)
  } else if (current.parentId !== null) {
    focusNode(current.parentId)
  }
}

// ── SidebarTree ──────────────────────────────────────────────────────────────

export function SidebarTree({
  nodes,
  selectedPath,
  isStarred,
  onToggleStar,
  className,
  filterQuery,
  persistKey,
  autoRevealPath,
  onNavigate,
}: SidebarTreeProps) {
  const routerNavigate = useNavigate()
  const navigate = onNavigate ?? routerNavigate

  // Expand/collapse state — optionally persisted to sessionStorage
  const [userExpandedIds, setUserExpandedIds] = useState<Set<string>>(
    () => persistKey ? readExpandState(persistKey) : new Set(),
  )

  // Auto-expand parent when autoRevealPath changes (e.g. from Cmd-K navigation)
  useEffect(() => {
    if (!autoRevealPath) return
    for (const node of nodes) {
      // Check if a child of this node matches the reveal path
      if (node.children?.some((child) => child.path === autoRevealPath)) {
        if (!userExpandedIds.has(node.id)) {
          setUserExpandedIds((prev) => {
            const next = new Set(prev)
            next.add(node.id)
            if (persistKey) writeExpandState(persistKey, next)
            return next
          })
        }
        break
      }
      // Also check if the node itself matches (top-level node, no parent to expand)
      if (node.path === autoRevealPath) break
    }
  }, [autoRevealPath]) // eslint-disable-line react-hooks/exhaustive-deps

  // Filter
  const { filtered, matchedParentIds } = useMemo(
    () => filterNodes(nodes, filterQuery ?? ''),
    [nodes, filterQuery],
  )

  // Effective expanded set: user toggles + auto-expanded from search
  const expandedSet = useMemo(() => {
    if (filterQuery) {
      return new Set([...userExpandedIds, ...matchedParentIds])
    }
    return userExpandedIds
  }, [userExpandedIds, matchedParentIds, filterQuery])

  function handleToggle(id: string) {
    setUserExpandedIds((prev) => {
      const next = new Set(prev)
      if (next.has(id)) {
        next.delete(id)
      } else {
        next.add(id)
      }
      if (persistKey) writeExpandState(persistKey, next)
      return next
    })
  }

  // Keyboard navigation
  const [focusedId, setFocusedId] = useState<string | null>(null)
  const treeRef = useRef<HTMLUListElement>(null)

  const visibleNodes = useMemo(
    () => flattenVisible(filtered, expandedSet),
    [filtered, expandedSet],
  )

  const getFocusedIndex = useCallback(() => {
    if (focusedId === null) return -1
    return visibleNodes.findIndex((fn) => fn.node.id === focusedId)
  }, [focusedId, visibleNodes])

  function focusNode(id: string) {
    const el = treeRef.current?.querySelector(`[data-nodeid="${CSS.escape(id)}"]`) as HTMLElement | null
    if (el) {
      el.focus()
    } else {
      setFocusedId(id)
    }
  }

  const handleKeyDown = useCallback(
    (e: KeyboardEvent<HTMLUListElement>) => {
      const currentIndex = getFocusedIndex()
      const current = visibleNodes[currentIndex]

      switch (e.key) {
        case 'ArrowDown': { e.preventDefault(); handleArrowDown(visibleNodes, currentIndex, focusNode); break }
        case 'ArrowUp': { e.preventDefault(); handleArrowUp(visibleNodes, currentIndex, focusNode); break }
        case 'ArrowRight': { e.preventDefault(); handleArrowRight(current, currentIndex, expandedSet, visibleNodes, handleToggle, focusNode); break }
        case 'ArrowLeft': { e.preventDefault(); handleArrowLeft(current, expandedSet, handleToggle, focusNode); break }
        case 'Enter': { e.preventDefault(); if (current?.node.path) navigate(current.node.path); break }
        case 'Home': { e.preventDefault(); if (visibleNodes.length > 0) focusNode(visibleNodes[0].node.id); break }
        case 'End': { e.preventDefault(); if (visibleNodes.length > 0) focusNode(visibleNodes[visibleNodes.length - 1].node.id); break }
      }
    },
    // eslint-disable-next-line react-hooks/exhaustive-deps
    [visibleNodes, expandedSet, focusedId],
  )

  return (
    <ul
      ref={treeRef}
      role="tree"
      className={`${styles.tree} ${className ?? ''}`}
      onKeyDown={handleKeyDown}
    >
      {filtered.map((node) => (
        <SidebarTreeRow
          key={node.id}
          node={node}
          depth={0}
          expandedSet={expandedSet}
          selectedPath={selectedPath}
          focusedId={focusedId}
          isStarred={isStarred}
          onToggle={handleToggle}
          onToggleStar={onToggleStar}
          onFocus={setFocusedId}
          navigate={navigate}
        />
      ))}
    </ul>
  )
}

// ── Row ──────────────────────────────────────────────────────────────────────

interface SidebarTreeRowProps {
  node: SidebarTreeNode
  depth: number
  expandedSet: Set<string>
  selectedPath?: string
  focusedId: string | null
  isStarred: (id: string) => boolean
  onToggle: (id: string) => void
  onToggleStar: (id: string) => void
  onFocus: (id: string) => void
  navigate: (path: string) => void
}

function SidebarTreeRow({
  node,
  depth,
  expandedSet,
  selectedPath,
  focusedId,
  isStarred,
  onToggle,
  onToggleStar,
  onFocus,
  navigate,
}: SidebarTreeRowProps) {
  const hasChildren = node.children && node.children.length > 0
  const isExpanded = expandedSet.has(node.id)
  const isSelected = Boolean(node.path && selectedPath === node.path)
  const isFocused = focusedId === node.id
  const effectiveStarKey = node.starKey ?? node.id
  const starred = isStarred(effectiveStarKey)

  function handleRowClick() {
    if (node.path) {
      navigate(node.path)
    }
  }

  function handleChevronClick(e: MouseEvent) {
    e.stopPropagation()
    onToggle(node.id)
  }

  function handleStarClick(e: MouseEvent) {
    e.stopPropagation()
    onToggleStar(effectiveStarKey)
  }

  const rowClass = [
    styles.treeRow,
    isSelected ? styles.treeRowActive : '',
  ]
    .filter(Boolean)
    .join(' ')

  return (
    <li role="none">
      {/* S1082: No onKeyDown here by design — the parent <ul role="tree"> carries
          onKeyDown={handleKeyDown} which handles Enter (navigate) and all arrow keys
          per the WAI-ARIA tree widget pattern. Adding a duplicate handler here would
          fire the action twice. */}
      <div
        role="treeitem"
        aria-expanded={hasChildren ? isExpanded : undefined}
        aria-selected={isSelected}
        tabIndex={isFocused ? 0 : -1}
        data-nodeid={node.id}
        className={rowClass}
        style={{ paddingLeft: `${depth * 16 + 8}px` }}
        onClick={handleRowClick}
        onFocus={() => onFocus(node.id)}
      >
        {/* Chevron */}
        <span className={styles.treeChevronSlot}>
          {hasChildren ? (
            <button
              className={styles.treeChevron}
              onClick={handleChevronClick}
              tabIndex={-1}
              aria-label={isExpanded ? 'Collapse' : 'Expand'}
            >
              {isExpanded ? <ChevronDown size={12} /> : <ChevronRight size={12} />}
            </button>
          ) : null}
        </span>

        {/* Icon (health dot, arrow, etc.) */}
        {node.icon && (
          <span className={styles.treeIcon} aria-hidden="true">
            {node.icon}
          </span>
        )}

        {/* Label */}
        <span className={styles.treeLabel}>{node.label}</span>

        {/* Badge */}
        {node.badge && (
          <span className={styles.treeBadge}>{node.badge}</span>
        )}

        {/* Star */}
        {node.starrable && (
          <button
            className={`${styles.treeStar} ${starred ? styles.treeStarActive : ''}`}
            onClick={handleStarClick}
            tabIndex={-1}
            aria-label={starred ? 'Remove from starred' : 'Add to starred'}
          >
            {starred ? <StarFilled /> : <StarOutline />}
          </button>
        )}
      </div>

      {/* Children */}
      {hasChildren && isExpanded && (
        <ul role="group" className={styles.treeGroup}>
          {node.children!.map((child) => (
            <SidebarTreeRow
              key={child.id}
              node={child}
              depth={depth + 1}
              expandedSet={expandedSet}
              selectedPath={selectedPath}
              focusedId={focusedId}
              isStarred={isStarred}
              onToggle={onToggle}
              onToggleStar={onToggleStar}
              onFocus={onFocus}
              navigate={navigate}
            />
          ))}
        </ul>
      )}
    </li>
  )
}