import { useState, useRef, useCallback } from 'react';

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

export interface TreeNode {
  id: string;
  label: string;
  icon?: React.ReactNode;
  children?: TreeNode[];
}

export interface UseTreeViewProps {
  /** The tree data to manage. */
  data: TreeNode[];
  /** Node IDs that are expanded by default (uncontrolled). */
  defaultExpanded?: string[];
  /**
   * Controlled selected node ID.
   * When provided, selection state is managed externally.
   */
  selectedNodeId?: string;
  /** Called when a node is clicked or activated via keyboard. */
  onNodeClick?: (node: TreeNode) => void;
  /** Called when a node becomes the selected item. */
  onNodeSelect?: (node: TreeNode) => void;
}

export interface UseTreeViewReturn {
  // ── State ─────────────────────────────────────────────────────────────────
  /** Set of currently expanded node IDs. */
  expanded: Set<string>;
  /** The currently selected node ID (controlled or internal). */
  effectiveSelectedId: string | undefined;

  // ── Refs ──────────────────────────────────────────────────────────────────
  /**
   * Map of node ID → button element.
   * Attach via the `setRef` helper returned by `getNodeRef`.
   */
  nodeRefs: React.MutableRefObject<Map<string, HTMLButtonElement>>;

  // ── Helpers ───────────────────────────────────────────────────────────────
  /** Returns a ref-setter callback for a given node ID. */
  getNodeRef: (nodeId: string) => (el: HTMLButtonElement | null) => void;

  // ── Handlers ──────────────────────────────────────────────────────────────
  /** Toggle the expanded state of a node. */
  toggleExpand: (nodeId: string) => void;
  /** Select a node (updates internal state if uncontrolled, fires callbacks). */
  handleSelect: (node: TreeNode) => void;
  /**
   * Full WAI-ARIA keyboard handler.
   * Attach to `onKeyDown` of each tree-item button.
   */
  handleKeyDown: (e: React.KeyboardEvent, node: TreeNode) => void;

  // ── Computed ──────────────────────────────────────────────────────────────
  /** Returns the ordered list of all currently visible nodes (respects expand state). */
  getVisibleNodes: () => TreeNode[];
}

// ─────────────────────────────────────────────────────────────────────────────
// Hook
// ─────────────────────────────────────────────────────────────────────────────

/**
 * Headless hook for hierarchical tree-view logic.
 *
 * @description
 * Manages expand/collapse state, selection (controlled or uncontrolled),
 * WAI-ARIA keyboard navigation (Arrow keys, Home, End, Enter/Space),
 * and focus management via node refs. Pair with any custom tree UI.
 *
 * @example
 * ```tsx
 * const { expanded, effectiveSelectedId, toggleExpand, handleSelect, handleKeyDown, getNodeRef } =
 *   useTreeView({ data, onNodeSelect });
 * ```
 */
export function useTreeView({
  data,
  defaultExpanded = [],
  selectedNodeId,
  onNodeClick,
  onNodeSelect,
}: UseTreeViewProps): UseTreeViewReturn {
  const [expanded, setExpanded] = useState<Set<string>>(new Set(defaultExpanded));
  const [internalSelectedId, setInternalSelectedId] = useState<string | undefined>(undefined);
  const nodeRefs = useRef<Map<string, HTMLButtonElement>>(new Map());

  const isControlled = selectedNodeId !== undefined;
  const effectiveSelectedId = isControlled ? selectedNodeId : internalSelectedId;

  // ── Helpers ───────────────────────────────────────────────────────────────

  const getNodeRef = useCallback(
    (nodeId: string) => (el: HTMLButtonElement | null) => {
      if (el) nodeRefs.current.set(nodeId, el);
      else nodeRefs.current.delete(nodeId);
    },
    []
  );

  const getVisibleNodes = useCallback((): TreeNode[] => {
    const result: TreeNode[] = [];
    const traverse = (nodes: TreeNode[]) => {
      for (const node of nodes) {
        result.push(node);
        if (node.children?.length && expanded.has(node.id)) traverse(node.children);
      }
    };
    traverse(data);
    return result;
  }, [data, expanded]);

  const findParent = useCallback(
    (nodes: TreeNode[], targetId: string): TreeNode | null => {
      for (const n of nodes) {
        if (n.children?.some(c => c.id === targetId)) return n;
        if (n.children) {
          const found = findParent(n.children, targetId);
          if (found) return found;
        }
      }
      return null;
    },
    // findParent is a pure recursive function over `data` — no reactive deps needed
    // eslint-disable-next-line react-hooks/exhaustive-deps
    []
  );

  // ── Handlers ──────────────────────────────────────────────────────────────

  const toggleExpand = useCallback((nodeId: string) => {
    setExpanded(prev => {
      const next = new Set(prev);
      next.has(nodeId) ? next.delete(nodeId) : next.add(nodeId);
      return next;
    });
  }, []);

  const handleSelect = useCallback(
    (node: TreeNode) => {
      if (!isControlled) setInternalSelectedId(node.id);
      onNodeSelect?.(node);
      onNodeClick?.(node);
    },
    [isControlled, onNodeClick, onNodeSelect]
  );

  const handleKeyDown = useCallback(
    (e: React.KeyboardEvent, node: TreeNode) => {
      const visibleNodes = getVisibleNodes();
      const idx = visibleNodes.findIndex(n => n.id === node.id);
      const hasChildren = !!node.children?.length;
      const isExpanded = expanded.has(node.id);

      const focusNode = (id: string) => {
        nodeRefs.current.get(id)?.focus();
      };

      switch (e.key) {
        case 'ArrowDown': {
          e.preventDefault();
          const next = visibleNodes[idx + 1];
          if (next) focusNode(next.id);
          break;
        }
        case 'ArrowUp': {
          e.preventDefault();
          const prev = visibleNodes[idx - 1];
          if (prev) focusNode(prev.id);
          break;
        }
        case 'ArrowRight': {
          e.preventDefault();
          if (hasChildren && !isExpanded) {
            toggleExpand(node.id);
          } else if (hasChildren && isExpanded && node.children?.length) {
            focusNode(node.children[0].id);
          }
          break;
        }
        case 'ArrowLeft': {
          e.preventDefault();
          if (hasChildren && isExpanded) {
            toggleExpand(node.id);
          } else {
            const parent = findParent(data, node.id);
            if (parent) focusNode(parent.id);
          }
          break;
        }
        case 'Home': {
          e.preventDefault();
          if (visibleNodes.length > 0) focusNode(visibleNodes[0].id);
          break;
        }
        case 'End': {
          e.preventDefault();
          const last = visibleNodes[visibleNodes.length - 1];
          if (last) focusNode(last.id);
          break;
        }
        case ' ': {
          e.preventDefault();
          if (hasChildren) {
            toggleExpand(node.id);
          } else {
            handleSelect(node);
          }
          break;
        }
        case 'Enter': {
          e.preventDefault();
          if (hasChildren) toggleExpand(node.id);
          handleSelect(node);
          break;
        }
      }
    },
    [data, expanded, findParent, getVisibleNodes, handleSelect, toggleExpand]
  );

  return {
    expanded,
    effectiveSelectedId,
    nodeRefs,
    getNodeRef,
    toggleExpand,
    handleSelect,
    handleKeyDown,
    getVisibleNodes,
  };
}