import { RefObject, useEffect, useRef } from 'react';

interface UseDialogFocusTrapParams {
  active: boolean;
  containerRef: RefObject<HTMLElement | null>;
}

const focusableSelector = [
  'a[href]',
  'button:not([disabled])',
  'input:not([disabled])',
  'select:not([disabled])',
  'textarea:not([disabled])',
  '[tabindex]:not([tabindex="-1"])',
].join(',');

const isElementHiddenFromAT = (node: HTMLElement): boolean => {
  if (node.getAttribute('aria-hidden') === 'true') return true;
  if (node.closest('[aria-hidden="true"]')) return true;
  return false;
};

const isInert = (node: HTMLElement): boolean => {
  // `inert` is not in TS lib dom everywhere.
  // eslint-disable-next-line @typescript-eslint/no-explicit-any
  if ((node as any).inert) return true;
  if (node.closest('[inert]')) return true;
  return false;
};

const getFocusableElements = (root: HTMLElement): HTMLElement[] => {
  const nodes = Array.from(root.querySelectorAll<HTMLElement>(focusableSelector));
  return nodes.filter((node) => !isElementHiddenFromAT(node) && !isInert(node));
};

export const useDialogFocusTrap = ({ active, containerRef }: UseDialogFocusTrapParams): void => {
  const previouslyFocusedRef = useRef<HTMLElement | null>(null);

  useEffect(() => {
    if (!active) return undefined;

    const root = containerRef.current;
    if (!root) return undefined;

    previouslyFocusedRef.current = document.activeElement as HTMLElement | null;

    const focusInitial = () => {
      const focusables = getFocusableElements(root);
      (focusables[0] ?? root).focus();
    };

    const raf = window.requestAnimationFrame(focusInitial);

    const onKeyDown = (e: KeyboardEvent) => {
      if (e.key !== 'Tab') return;

      const focusables = getFocusableElements(root);
      if (focusables.length === 0) {
        e.preventDefault();
        root.focus();
        return;
      }

      const activeEl = document.activeElement as HTMLElement | null;
      const currentIndex = activeEl ? focusables.indexOf(activeEl) : -1;
      const lastIndex = focusables.length - 1;

      if (e.shiftKey) {
        if (currentIndex <= 0) {
          e.preventDefault();
          focusables[lastIndex].focus();
        }
      } else if (currentIndex === -1 || currentIndex >= lastIndex) {
        e.preventDefault();
        focusables[0].focus();
      }
    };

    const onFocusIn = (e: FocusEvent) => {
      const target = e.target as HTMLElement | null;

      // If focus went to an element outside the dialog, redirect it back
      if (target && !root.contains(target)) {
        e.preventDefault();
        e.stopPropagation();

        // Focus back to the first focusable element in the dialog
        const focusables = getFocusableElements(root);

        if (focusables.length > 0) {
          focusables[0].focus();
        } else {
          root.focus();
        }
      }
    };

    document.addEventListener('keydown', onKeyDown);
    document.addEventListener('focusin', onFocusIn, true);

    return () => {
      window.cancelAnimationFrame(raf);
      document.removeEventListener('keydown', onKeyDown);
      document.removeEventListener('focusin', onFocusIn, true);

      const prev = previouslyFocusedRef.current;
      if (prev && document.contains(prev)) {
        prev.focus();
      }
    };
  }, [active, containerRef]);
};

export default useDialogFocusTrap;