import { type RefObject, useEffect, useRef } from "react";

import {
  ensureFocusWithin,
  focusElement,
  getActiveHTMLElement,
  isFocusable,
  isHTMLElement,
  wrapFocus,
} from "../utils/focusManagement";

type FocusTargetRef = {
  current: HTMLElement | null;
};

export type InitialFocusTarget =
  | HTMLElement
  | FocusTargetRef
  | string
  | (() => HTMLElement | null | undefined)
  | null
  | undefined;

export interface UseFocusTrapOptions {
  enabled?: boolean;
  initialFocus?: InitialFocusTarget;
  restoreFocus?: boolean;
  preventScroll?: boolean;
  escapeDeactivates?: boolean;
  onEscapeKey?: (event: KeyboardEvent) => void;
}

function isFocusTargetRef(target: InitialFocusTarget): target is FocusTargetRef {
  return Boolean(target && typeof target === "object" && "current" in target);
}

function resolveInitialFocus(
  initialFocus: InitialFocusTarget,
  container: HTMLElement,
): HTMLElement | null {
  if (!initialFocus) {
    return null;
  }

  if (typeof initialFocus === "function") {
    return initialFocus() ?? null;
  }

  if (typeof initialFocus === "string") {
    try {
      return container.querySelector<HTMLElement>(initialFocus);
    } catch {
      return null;
    }
  }

  if (isFocusTargetRef(initialFocus)) {
    return initialFocus.current;
  }

  return isHTMLElement(initialFocus) ? initialFocus : null;
}

export function useFocusTrap<TContainer extends HTMLElement>(
  containerRef: RefObject<TContainer | null>,
  {
    enabled = true,
    initialFocus,
    restoreFocus = true,
    preventScroll = true,
    escapeDeactivates = true,
    onEscapeKey,
  }: UseFocusTrapOptions = {},
): void {
  const previouslyFocusedElementRef = useRef<HTMLElement | null>(null);

  useEffect(() => {
    if (!enabled || typeof document === "undefined") {
      return undefined;
    }

    const container = containerRef.current;

    if (!container) {
      return undefined;
    }

    const ownerDocument = container.ownerDocument;
    const ownerWindow = ownerDocument.defaultView;
    const focusOptions: FocusOptions = { preventScroll };
    const moveFocusOptions = {
      includeContainer: true,
      includeNegativeTabIndex: true,
      focusOptions,
    };

    previouslyFocusedElementRef.current = getActiveHTMLElement(ownerDocument);

    const focusInitialTarget = () => {
      if (!container.isConnected) {
        return;
      }

      const activeElement = getActiveHTMLElement(ownerDocument);
      if (activeElement && container.contains(activeElement)) {
        return;
      }

      const initialFocusTarget = resolveInitialFocus(initialFocus, container);
      if (
        initialFocusTarget &&
        container.contains(initialFocusTarget) &&
        isFocusable(initialFocusTarget, { includeNegativeTabIndex: true })
      ) {
        focusElement(initialFocusTarget, focusOptions);
        return;
      }

      ensureFocusWithin(container, moveFocusOptions);
    };

    let cancelInitialFocus = () => undefined;

    if (ownerWindow?.requestAnimationFrame) {
      const frame = ownerWindow.requestAnimationFrame(focusInitialTarget);
      cancelInitialFocus = () => ownerWindow.cancelAnimationFrame(frame);
    } else if (ownerWindow?.setTimeout) {
      const timeout = ownerWindow.setTimeout(focusInitialTarget, 0);
      cancelInitialFocus = () => ownerWindow.clearTimeout(timeout);
    } else {
      focusInitialTarget();
    }

    const handleKeyDown = (event: KeyboardEvent) => {
      if (event.key === "Tab") {
        wrapFocus(container, event, moveFocusOptions);
        return;
      }

      if (event.key === "Escape" && escapeDeactivates && onEscapeKey) {
        event.preventDefault();
        event.stopPropagation();
        onEscapeKey(event);
      }
    };

    const handleFocusIn = (event: FocusEvent) => {
      const target = event.target;

      if (!target || typeof target !== "object" || !("nodeType" in target)) {
        return;
      }

      if (!container.contains(target as Node)) {
        ensureFocusWithin(container, moveFocusOptions);
      }
    };

    ownerDocument.addEventListener("keydown", handleKeyDown, true);
    ownerDocument.addEventListener("focusin", handleFocusIn);

    return () => {
      cancelInitialFocus();
      ownerDocument.removeEventListener("keydown", handleKeyDown, true);
      ownerDocument.removeEventListener("focusin", handleFocusIn);

      const previouslyFocusedElement = previouslyFocusedElementRef.current;
      previouslyFocusedElementRef.current = null;

      if (restoreFocus && previouslyFocusedElement?.isConnected) {
        focusElement(previouslyFocusedElement, focusOptions);
      }
    };
  }, [
    containerRef,
    enabled,
    escapeDeactivates,
    initialFocus,
    onEscapeKey,
    preventScroll,
    restoreFocus,
  ]);
}