import { useState, useRef, useCallback, type ReactNode, type KeyboardEvent } from 'react'
import styles from './TreeView.module.css'

export interface TreeNode {
  id: string
  label: string
  icon?: ReactNode
  children?: TreeNode[]
  meta?: string
}

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

function flattenVisibleNodes(
  nodes: TreeNode[],
  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(...flattenVisibleNodes(node.children, expandedIds, depth + 1, node.id))
    }
  }
  return result
}

// ── 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)
  }
}

interface TreeViewProps {
  nodes: TreeNode[]
  onSelect?: (id: string) => void
  selectedId?: string
  expandedIds?: string[]
  onToggle?: (id: string) => void
  className?: string
}

export function TreeView({
  nodes,
  onSelect,
  selectedId,
  expandedIds: controlledExpandedIds,
  onToggle,
  className,
}: TreeViewProps) {
  // Controlled vs uncontrolled expansion
  const isControlled = controlledExpandedIds !== undefined && onToggle !== undefined
  const [internalExpandedIds, setInternalExpandedIds] = useState<Set<string>>(new Set())

  const expandedSet = isControlled
    ? new Set(controlledExpandedIds)
    : internalExpandedIds

  function handleToggle(id: string) {
    if (isControlled) {
      onToggle!(id)
    } else {
      setInternalExpandedIds((prev) => {
        const next = new Set(prev)
        if (next.has(id)) {
          next.delete(id)
        } else {
          next.add(id)
        }
        return next
      })
    }
  }

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

  const visibleNodes = flattenVisibleNodes(nodes, expandedSet)

  function getFocusedIndex() {
    if (focusedId === null) return -1
    return visibleNodes.findIndex((fn) => fn.node.id === focusedId)
  }

  function focusNode(id: string) {
    // We focus the element directly. All tree items have tabIndex={-1} by default
    // which means programmatic focus works even without tabIndex=0.
    // The element's onFocus handler will fire and call setFocusedId — but that
    // happens synchronously within the React event so it's properly batched.
    const el = treeRef.current?.querySelector(`[data-nodeid="${id}"]`) as HTMLElement | null
    if (el) {
      // Temporarily make focusable if not already
      el.focus()
    } else {
      // Element not in DOM yet (e.g. after expand); update state so it renders
      // with tabIndex=0 and the browser's next focus movement will work.
      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) onSelect?.(current.node.id); 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.root} ${className ?? ''}`}
      onKeyDown={handleKeyDown}
    >
      {nodes.map((node) => (
        <TreeNodeRow
          key={node.id}
          node={node}
          depth={0}
          expandedSet={expandedSet}
          selectedId={selectedId}
          focusedId={focusedId}
          onToggle={handleToggle}
          onSelect={onSelect}
          onFocus={setFocusedId}
        />
      ))}
    </ul>
  )
}

interface TreeNodeRowProps {
  node: TreeNode
  depth: number
  expandedSet: Set<string>
  selectedId?: string
  focusedId: string | null
  onToggle: (id: string) => void
  onSelect?: (id: string) => void
  onFocus: (id: string) => void
}

function TreeNodeRow({
  node,
  depth,
  expandedSet,
  selectedId,
  focusedId,
  onToggle,
  onSelect,
  onFocus,
}: TreeNodeRowProps) {
  const hasChildren = node.children && node.children.length > 0
  const isExpanded = expandedSet.has(node.id)
  const isSelected = selectedId === node.id
  const isFocused = focusedId === node.id

  function handleClick() {
    if (hasChildren) {
      onToggle(node.id)
    }
    onSelect?.(node.id)
  }

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

  return (
    <li role="none">
      {/* S1082: No onKeyDown here by design — the parent <ul role="tree"> carries
          onKeyDown={handleKeyDown} which handles Enter (select) 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 || (focusedId === null && depth === 0 && node === node) ? 0 : -1}
        data-nodeid={node.id}
        className={rowClass}
        style={{ paddingLeft: `${depth * 16 + 8}px` }}
        onClick={handleClick}
        onFocus={() => onFocus(node.id)}
      >
        <span className={styles.chevronSlot}>
          {hasChildren ? (
            <span className={styles.chevron} aria-hidden="true">
              {isExpanded ? '▾' : '▸'}
            </span>
          ) : null}
        </span>
        {node.icon && (
          <span className={styles.icon} aria-hidden="true">
            {node.icon}
          </span>
        )}
        <span className={styles.label}>{node.label}</span>
        {node.meta && (
          <span className={styles.meta}>{node.meta}</span>
        )}
      </div>
      {hasChildren && isExpanded && (
        <ul role="group" className={styles.children}>
          {node.children!.map((child) => (
            <TreeNodeRow
              key={child.id}
              node={child}
              depth={depth + 1}
              expandedSet={expandedSet}
              selectedId={selectedId}
              focusedId={focusedId}
              onToggle={onToggle}
              onSelect={onSelect}
              onFocus={onFocus}
            />
          ))}
        </ul>
      )}
    </li>
  )
}
