useFocusTrap.ts

import { useEffect, useRef } from 'react';

interface UseFocusTrapOptions {
  /** Whether the trap is active */
  isActive: boolean;
  /** Ref to the trigger element that opened the modal (for return focus) */
  triggerRef?: React.RefObject<HTMLElement | null>;
  /** Callback when Escape is pressed */
  onEscape?: () => void;
}

/**
 * Query selector for focusable elements within a container.
 * Includes: buttons, links, inputs, selects, textareas, and elements with tabindex="0"
 */
const FOCUSABLE_SELECTOR = [
  'button:not([disabled]):not([tabindex="-1"])',
  'a[href]:not([tabindex="-1"])',
  'input:not([disabled]):not([tabindex="-1"])',
  'select:not([disabled]):not([tabindex="-1"])',
  'textarea:not([disabled]):not([tabindex="-1"])',
  '[tabindex="0"]',
].join(', ');

/**
 * Trap keyboard focus inside a container while the trap is active.
 *
 * On activation, focus moves to the first focusable element in the
 * container. Tab and Shift+Tab cycle within the container. Escape calls
 * the provided `onEscape` handler. On deactivation, focus returns to
 * `triggerRef` if supplied, otherwise to the element that had focus
 * before the trap was activated.
 *
 * Returns a ref to attach to the container element.
 */
export function useFocusTrap({ isActive, triggerRef, onEscape }: UseFocusTrapOptions) {
  const containerRef = useRef<HTMLDivElement>(null);
  const previousActiveElement = useRef<HTMLElement | null>(null);

  // Store the element that had focus before the modal opened
  useEffect(() => {
    if (isActive) {
      previousActiveElement.current = document.activeElement as HTMLElement;
    }
  }, [isActive]);

  // Handle focus trap and Escape key
  useEffect(() => {
    if (!isActive) return;

    const container = containerRef.current;
    if (!container) return;

    // Get all focusable elements
    const getFocusableElements = (): HTMLElement[] => {
      return Array.from(container.querySelectorAll(FOCUSABLE_SELECTOR));
    };

    // Focus the first focusable element when opened
    const focusableElements = getFocusableElements();
    if (focusableElements.length > 0) {
      // Small delay to ensure DOM is ready
      setTimeout(() => {
        focusableElements[0].focus();
      }, 50);
    }

    // Handle Tab key to trap focus
    const handleKeyDown = (event: KeyboardEvent) => {
      if (event.key !== 'Tab') return;

      const elements = getFocusableElements();
      if (elements.length === 0) return;

      const firstElement = elements[0];
      const lastElement = elements[elements.length - 1];

      // Shift + Tab: moving backwards
      if (event.shiftKey) {
        if (document.activeElement === firstElement) {
          event.preventDefault();
          lastElement.focus();
        }
      } else {
        // Tab: moving forwards
        if (document.activeElement === lastElement) {
          event.preventDefault();
          firstElement.focus();
        }
      }
    };

    // Handle Escape key
    const handleEscape = (event: KeyboardEvent) => {
      if (event.key === 'Escape' && onEscape) {
        event.preventDefault();
        onEscape();
      }
    };

    document.addEventListener('keydown', handleKeyDown);
    document.addEventListener('keydown', handleEscape);

    return () => {
      document.removeEventListener('keydown', handleKeyDown);
      document.removeEventListener('keydown', handleEscape);
    };
  }, [isActive, onEscape]);

  // Return focus to trigger or previous element on close or unmount
  useEffect(() => {
    if (!isActive) return; // only set up return-focus when active

    return () => {
      // Cleanup runs when isActive goes false → true or on unmount
      if (triggerRef?.current) {
        triggerRef.current.focus();
      } else if (previousActiveElement.current) {
        previousActiveElement.current.focus();
      }
    };
  // eslint-disable-next-line react-hooks/exhaustive-deps
  }, [isActive]);

  return containerRef;
}