"use client";

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

const FOCUSABLE_SELECTOR =
  'button, [href], input, select, textarea, [tabindex]:not([tabindex="-1"])';

let lastInteractedFocusable: HTMLElement | null = null;
let interactionTrackingBound = false;

function closestFocusable(target: EventTarget | null): HTMLElement | null {
  if (!(target instanceof Element)) return null;
  const match = target.closest(FOCUSABLE_SELECTOR);
  return match instanceof HTMLElement ? match : null;
}

/** Menu items unmount on select; restore to the control that opened the menu. */
function menuTriggerFor(item: HTMLElement): HTMLElement | null {
  if (item.getAttribute("role") !== "menuitem") return null;
  const menu = item.closest("[role='menu']");
  const menuId = menu?.getAttribute("id");
  if (menuId) {
    const trigger = document.querySelector(
      `[aria-controls="${CSS.escape(menuId)}"]`,
    );
    if (trigger instanceof HTMLElement) return trigger;
  }
  const expanded = document.querySelector(
    '[aria-haspopup="menu"][aria-expanded="true"]',
  );
  return expanded instanceof HTMLElement ? expanded : null;
}

function stableFocusableFrom(target: EventTarget | null): HTMLElement | null {
  const focusable = closestFocusable(target);
  if (!focusable) return null;
  return menuTriggerFor(focusable) ?? focusable;
}

function isRestorable(
  node: HTMLElement | null,
  dialog: HTMLElement | null,
): node is HTMLElement {
  if (!node?.isConnected) return false;
  if (node === document.body || node === document.documentElement) return false;
  if (dialog?.contains(node)) return false;
  return true;
}

function retainLastInteracted(): void {
  if (interactionTrackingBound || typeof document === "undefined") return;
  interactionTrackingBound = true;
  const save = (event: Event) => {
    const el = stableFocusableFrom(event.target);
    if (!el) return;
    lastInteractedFocusable = el;
  };
  document.addEventListener("pointerdown", save, true);
  document.addEventListener("focusin", save);
}

retainLastInteracted();

function snapshotTrigger(dialog: HTMLElement | null): HTMLElement | null {
  if (lastInteractedFocusable && !lastInteractedFocusable.isConnected) {
    lastInteractedFocusable = null;
  }
  const active =
    document.activeElement instanceof HTMLElement
      ? document.activeElement
      : null;
  const fromActive = isRestorable(active, dialog)
    ? (menuTriggerFor(active) ?? active)
    : null;
  if (fromActive && isRestorable(fromActive, dialog)) return fromActive;
  if (isRestorable(lastInteractedFocusable, dialog)) {
    return lastInteractedFocusable;
  }
  return null;
}

function restoreFocus(node: HTMLElement | null): void {
  if (!node?.isConnected) return;
  node.focus();
}

/**
 * Escape-to-close, body scroll lock, focus move-in, tab trap, and restore
 * focus to the control that opened a Create-shell modal.
 */
export function useCreateModalA11y(
  isOpen: boolean,
  onClose: () => void,
  dialogRef: RefObject<HTMLDivElement | null>,
): void {
  const previousActiveElementRef = useRef<HTMLElement | null>(null);

  useEffect(() => {
    if (!isOpen) return;

    const handleEscape = (e: KeyboardEvent) => {
      if (e.key === "Escape") {
        onClose();
      }
    };

    document.addEventListener("keydown", handleEscape);
    return () => {
      document.removeEventListener("keydown", handleEscape);
    };
  }, [isOpen, onClose]);

  useLayoutEffect(() => {
    if (!isOpen) return;

    previousActiveElementRef.current = snapshotTrigger(dialogRef.current);
    document.body.style.overflow = "hidden";

    if (dialogRef.current) {
      const focusableElements = dialogRef.current.querySelectorAll(
        FOCUSABLE_SELECTOR,
      );
      const firstElement = focusableElements[0] as HTMLElement | undefined;
      if (firstElement) {
        firstElement.focus();
      } else {
        dialogRef.current.setAttribute("tabindex", "-1");
        dialogRef.current.focus();
      }
    }

    const handleTab = (e: KeyboardEvent) => {
      if (e.key !== "Tab" || !dialogRef.current) return;

      const focusableElements = dialogRef.current.querySelectorAll(
        FOCUSABLE_SELECTOR,
      );
      const firstElement = focusableElements[0] as HTMLElement | undefined;
      const lastElement = focusableElements[
        focusableElements.length - 1
      ] as HTMLElement | undefined;

      if (e.shiftKey) {
        if (document.activeElement === firstElement) {
          e.preventDefault();
          lastElement?.focus();
        }
      } else if (document.activeElement === lastElement) {
        e.preventDefault();
        firstElement?.focus();
      }
    };

    document.addEventListener("keydown", handleTab);

    return () => {
      document.body.style.overflow = "";
      document.removeEventListener("keydown", handleTab);
      restoreFocus(previousActiveElementRef.current);
    };
  }, [dialogRef, isOpen]);
}