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

function queryAllFocusable(el: HTMLElement) {
  const selector = [
    'a[href]',
    'area[href]',
    'input:not([disabled])',
    'select:not([disabled])',
    'textarea:not([disabled])',
    'button:not([disabled])',
    'iframe',
    'object',
    'embed',
    '[tabindex="-1"]',
    '[tabindex="0"]',
    '[contenteditable]',
    'audio[controls]',
    'video[controls]',
    'summary',
  ].join(',');
  return el.querySelectorAll(selector);
}

const focusElement = (el: HTMLElement) => {
  if (typeof el.focus === 'function') {
    el.focus();
  }
};

export const FocusTrap = ({
  autoFocus = true,
  returnFocus = false,
  disabled,
  children,
}: {
  autoFocus?: boolean;
  returnFocus?: boolean;
  disabled?: boolean;
  children: React.ReactNode;
}) => {
  const containerRef = useRef<HTMLDivElement>(null);
  const previouslyFocusedRef = useRef<Element | null>(null);

  function focusFirstChild() {
    if (containerRef.current) {
      const focusable = queryAllFocusable(containerRef.current);
      const fistChild = focusable[0] as HTMLElement | undefined;
      if (fistChild) {
        focusElement(fistChild);
      }
    }
  }

  const handleFocus = useCallback((ev) => {
    if (!containerRef.current || containerRef.current.contains(ev.target)) {
      return;
    }

    ev.stopPropagation();
    ev.preventDefault();
    focusFirstChild();
  }, []);

  useEffect(() => {
    if (disabled) return;

    previouslyFocusedRef.current = document.activeElement;

    if (
      autoFocus &&
      containerRef.current &&
      containerRef.current.contains(document.activeElement) === false
    ) {
      focusFirstChild();
    }

    document.addEventListener('focus', handleFocus, true);
    return () => {
      document.removeEventListener('focus', handleFocus, true);

      if (returnFocus && previouslyFocusedRef.current) {
        focusElement(previouslyFocusedRef.current as HTMLElement);
      }
    };
  }, [handleFocus, autoFocus, returnFocus, disabled]);

  return <div ref={containerRef}>{children}</div>;
};