import { useStoreState } from "@ariakit/react-store";
import {
  useBooleanEvent,
  useEvent,
  useId,
  useMergeRefs,
  usePortalRef,
  useSafeLayoutEffect,
  useWrapElement,
  createElement,
  createHook,
  forwardRef,
} from "@ariakit/react-utils";
import type { Props } from "@ariakit/react-utils";
import { sync } from "@ariakit/store";
import {
  chain,
  contains,
  getActiveElement,
  getDocument,
  getWindow,
  addGlobalEventListener,
  getFirstTabbableIn,
  isElement,
  isNode,
  isFocusable,
  isSafari,
} from "@ariakit/utils";
import type { BooleanOrCallback } from "@ariakit/utils";
import type {
  ComponentPropsWithRef,
  ElementType,
  FC,
  FocusEvent as ReactFocusEvent,
  ReactElement,
  KeyboardEvent as ReactKeyboardEvent,
  RefObject,
  SyntheticEvent,
} from "react";
import { useCallback, useEffect, useRef, useState } from "react";
import type { DisclosureContentOptions } from "../disclosure/disclosure-content.tsx";
import {
  isHidden,
  useDisclosureContent,
} from "../disclosure/disclosure-content.tsx";
import { useFocusableContainer } from "../focusable/focusable-container.tsx";
import type { FocusableOptions } from "../focusable/focusable.tsx";
import { useFocusable } from "../focusable/focusable.tsx";
import { HeadingLevel } from "../heading/heading-level.tsx";
import type { PortalOptions } from "../portal/portal.tsx";
import { usePortal } from "../portal/portal.tsx";
import { DialogBackdrop } from "./dialog-backdrop.tsx";
import {
  DialogDescriptionContext,
  DialogHeadingContext,
  DialogScopedContextProvider,
  useDialogProviderContext,
} from "./dialog-context.tsx";
import type { DialogStore } from "./dialog-store.ts";
import { useDialogStore } from "./dialog-store.ts";
import {
  disableTree,
  markAndDisableTreeOutside,
} from "./utils/disable-tree.ts";
import {
  isElementInside,
  isElementMarked,
  markTreeInside,
  markTreeOutside,
} from "./utils/mark-tree-outside.ts";
import { prependHiddenDismiss } from "./utils/prepend-hidden-dismiss.ts";
import { supportsInert } from "./utils/supports-inert.ts";
import { useHideOnInteractOutside } from "./utils/use-hide-on-interact-outside.ts";
import { useNestedDialogs } from "./utils/use-nested-dialogs.tsx";
import { usePreventBodyScroll } from "./utils/use-prevent-body-scroll.ts";
import { createWalkTreeSnapshot } from "./utils/walk-tree-outside.ts";

const TagName = "div" satisfies ElementType;
type TagName = typeof TagName;
type HTMLType = HTMLElementTagNameMap[TagName];

const isSafariBrowser = isSafari();
const openModalPortals = new WeakSet<HTMLElement>();

// Elements a dialog captured because nothing named an opener. They stand for
// nothing the application asked for, so a later open may replace them. Tracked
// on the element rather than per dialog or per store, because both the dialog
// and the store it derives are replaced when its component remounts, while the
// store the application owns keeps the value. A mark is never dropped, since
// two dialogs can share an element and one clearing it would leave the other
// reading its own fallback as a name.
// https://github.com/ariakit/ariakit/issues/7095
const capturedDisclosures = new WeakSet<Element>();

function isAlreadyFocusingAnotherElement(dialog?: HTMLElement | null) {
  const activeElement = getActiveElement(dialog);
  if (!activeElement) return false;
  if (dialog && contains(dialog, activeElement)) return false;
  if (isFocusable(activeElement)) return true;
  return false;
}

function getDeepestActiveElement(element: Element) {
  let activeElement = element;
  while (activeElement.shadowRoot) {
    const nextElement = activeElement.shadowRoot.activeElement;
    if (!isElement(nextElement)) break;
    activeElement = nextElement;
  }
  return activeElement;
}

function isElementInDialog(
  element: Element,
  dialog: Element,
  disclosureElement: Element | null,
) {
  if (contains(dialog, element)) return true;
  if (disclosureElement && contains(disclosureElement, element)) return true;
  return isElementInside(element, dialog);
}

function getElementFromProp(
  prop?: HTMLElement | RefObject<HTMLElement | null> | null,
  focusable = false,
) {
  if (!prop) return null;
  const element = "current" in prop ? prop.current : prop;
  if (!element) return null;
  if (focusable) return isFocusable(element) ? element : null;
  return element;
}

/**
 * Returns default modal portals that are later in the same root and still
 * pending in the current layout pass. Portals in the established stack are
 * excluded so this only coordinates the dialogs opening in the same cohort.
 */
function getLaterOpenModalPortals(dialog: HTMLElement) {
  if (!dialog.isConnected) return [];
  const root = dialog.getRootNode() as Document | ShadowRoot;
  const dialogs = root.querySelectorAll<HTMLElement>(
    "[data-dialog][data-dialog-portal][data-open]",
  );
  const portals: HTMLElement[] = [];
  let foundDialog = false;
  for (const currentDialog of dialogs) {
    if (currentDialog === dialog) {
      foundDialog = true;
      continue;
    }
    if (!foundDialog) continue;
    const portalId = currentDialog.getAttribute("data-dialog-portal");
    if (!portalId) continue;
    const portal = root.getElementById(portalId);
    if (!portal || !contains(portal, currentDialog)) continue;
    // Active portals belong to an established stack. DOM order is only used
    // to break ties between dialogs opening in the same layout pass.
    if (openModalPortals.has(portal)) continue;
    portals.push(portal);
  }
  return portals;
}

/**
 * Returns props to create a `Dialog` component.
 * @see https://ariakit.com/components/dialog
 * @example
 * ```jsx
 * const store = useDialogStore();
 * const props = useDialog({ store });
 * <Role {...props}>Dialog</Role>
 * ```
 */
export const useDialog = createHook<TagName, DialogOptions>(function useDialog({
  store: storeProp,
  open: openProp,
  onClose,
  focusable = true,
  modal = true,
  portal = modal,
  backdrop = modal,
  hideOnEscape = true,
  hideOnInteractOutside = true,
  getPersistentElements,
  preventBodyScroll = modal,
  autoFocusOnShow = true,
  autoFocusOnHide = true,
  initialFocus,
  finalFocus,
  unmountOnHide,
  unstable_treeSnapshotKey,
  ...props
}) {
  const context = useDialogProviderContext();
  const ref = useRef<HTMLType>(null);
  const backdropRef = useRef<HTMLDivElement>(null);
  const hasDefaultModalPortal = modal && portal && !props.portalElement;

  const store = useDialogStore({
    store: storeProp || context,
    open: openProp,
    setOpen(open) {
      if (open) return;
      const dialog = ref.current;
      if (!dialog) return;
      const event = new Event("close", { bubbles: false, cancelable: true });
      if (onClose) {
        dialog.addEventListener("close", onClose, { once: true });
      }
      dialog.dispatchEvent(event);
      if (!event.defaultPrevented) return;
      store.setOpen(true);
    },
  });

  // domReady can be also the portal node element so it's updated when the
  // portal node changes (like in between re-renders), triggering effects
  // again.
  const { portalRef, portalNode, domReady } = usePortalRef(
    portal,
    props.portalRef,
  );
  // Sets preserveTabOrder to true only if the dialog is not a modal and is
  // open.
  const preserveTabOrderProp = props.preserveTabOrder;
  const preserveTabOrder = useStoreState(
    store,
    ["mounted"],
    (state) => preserveTabOrderProp && !modal && state.mounted,
  );
  const id = useId(props.id);
  const open = useStoreState(store, "open");
  const mounted = useStoreState(store, "mounted");
  const contentElement = useStoreState(store, "contentElement");
  const hidden = isHidden(mounted, props.hidden, props.alwaysVisible);

  usePreventBodyScroll(contentElement, id, preventBodyScroll && !hidden);

  // Tracks whether focus restoration should be skipped after an outside
  // interaction to match native HTML behavior.
  // Reset when the dialog opens to avoid stale flags from prevented closes
  // (e.g., onClose calling event.preventDefault), async closes with
  // animations, or when autoFocusOnHide is disabled.
  const interactedOutsideRef = useRef(false);
  const focusedStoreRef = useRef<DialogStore | null>(null);
  useSafeLayoutEffect(() => {
    return sync(store, ["open"], (state) => {
      if (!state.open) {
        focusedStoreRef.current = null;
        return;
      }
      interactedOutsideRef.current = false;
      // Native auto-focus may run before layout effects on the initial open
      // render, so preserve focus captured during that commit. Focus history
      // from another store doesn't belong to this open cycle.
      if (focusedStoreRef.current === store) return;
      focusedStoreRef.current = null;
    });
  }, [store]);
  useHideOnInteractOutside({
    store,
    hideOnInteractOutside,
    domReady,
    interactedOutsideRef,
    focusedStoreRef,
  });

  const { wrapElement, nestedDialogs } = useNestedDialogs(store);
  props = useWrapElement(props, wrapElement, [wrapElement]);

  // On Safari, buttons don't receive focus on mousedown unless they have an
  // explicit tabIndex. Non-Ariakit buttons (which don't go through
  // useFocusable) won't have this, so activeElement may still be BODY when the
  // dialog opens. We track the last mousedown target as a fallback.
  const lastMousedownRef = useRef<Element | null>(null);

  if (isSafariBrowser) {
    useEffect(() => {
      if (!domReady) return;
      const dialog = ref.current;
      if (!dialog) return;
      const doc = getDocument(dialog);
      const onMousedown = (event: MouseEvent) => {
        lastMousedownRef.current = event.target as Element;
      };
      doc.addEventListener("mousedown", onMousedown, true);
      return () => {
        doc.removeEventListener("mousedown", onMousedown, true);
      };
    }, [domReady]);
  }

  // Sets disclosure element using the current active element right after the
  // dialog is opened, unless something already named one.
  useSafeLayoutEffect(() => {
    if (!open) return;
    // Native auto-focus can focus the dialog and synchronously move focus
    // outside before layout effects run. In that case, the current element is
    // an escape target rather than the element that disclosed the dialog.
    if (focusedStoreRef.current === store) return;
    const dialog = ref.current;
    // A disclosure element no capture left behind names the opener
    // deliberately, whether it comes from a Disclosure component or from the
    // application assigning it before showing the dialog. The capture is only
    // a fallback for when nothing said otherwise, so it must not overwrite it.
    // https://github.com/ariakit/ariakit/issues/7087
    const hasNamedDisclosure = () => {
      const { disclosureElement } = store.getState();
      if (!disclosureElement) return false;
      if (capturedDisclosures.has(disclosureElement)) return false;
      if (!disclosureElement.isConnected) return false;
      // The disclosure element can't be inside the dialog.
      if (dialog && contains(dialog, disclosureElement)) return false;
      return true;
    };
    if (hasNamedDisclosure()) return;
    const captureDisclosure = (element: HTMLElement) => {
      capturedDisclosures.add(element);
      store.setDisclosureElement(element);
    };
    const activeElement = getActiveElement(dialog, true);
    if (!activeElement) return;
    if (activeElement.tagName === "BODY") {
      // Safari fallback: use the last mousedown target when activeElement is
      // BODY (happens with native buttons that lack an explicit tabIndex).
      const fallback = lastMousedownRef.current;
      lastMousedownRef.current = null;
      if (!fallback?.isConnected) return;
      if (!isFocusable(fallback)) return;
      if (dialog && contains(dialog, fallback)) return;
      captureDisclosure(fallback as HTMLElement);
      return;
    }
    // The disclosure element can't be inside the dialog.
    if (dialog && contains(dialog, activeElement)) return;
    captureDisclosure(activeElement);
  }, [store, open]);

  // Sets --dialog-viewport-height CSS variable to the height of the visual
  // viewport. This allows the dialog to be positioned correctly when the
  // viewport height changes (e.g., when the keyboard is shown on mobile).
  useEffect(() => {
    if (!mounted) return;
    if (!domReady) return;
    const dialog = ref.current;
    if (!dialog) return;
    const win = getWindow(dialog);
    const viewport = win.visualViewport || win;
    const setViewportHeight = () => {
      const height = win.visualViewport?.height ?? win.innerHeight;
      dialog.style.setProperty("--dialog-viewport-height", `${height}px`);
    };
    setViewportHeight();
    viewport.addEventListener("resize", setViewportHeight);
    return () => {
      viewport.removeEventListener("resize", setViewportHeight);
    };
  }, [mounted, domReady]);

  // Renders a hidden dismiss button at the top of the modal dialog element. So
  // that screen reader users aren't trapped in the dialog when there's no
  // visible dismiss button.
  useEffect(() => {
    if (!modal) return;
    if (!mounted) return;
    if (!domReady) return;
    const dialog = ref.current;
    if (!dialog) return;
    // If there's already a DialogDismiss component, it does nothing.
    const existingDismiss = dialog.querySelector("[data-dialog-dismiss]");
    if (existingDismiss) return;
    return prependHiddenDismiss(dialog, store.hide);
  }, [store, modal, mounted, domReady]);

  // TODO: Move this behavior into DisclosureContent.
  // Keep closing animated content inert until its mounted state ends.
  useSafeLayoutEffect(() => {
    if (!supportsInert()) return;
    if (open) return;
    if (!mounted) return;
    if (!domReady) return;
    const dialog = ref.current;
    if (!dialog) return;
    return disableTree(dialog);
  }, [open, mounted, domReady]);

  const canTakeTreeSnapshot = open && domReady;
  const openingCohortRef = useRef<{
    portal: HTMLElement;
    peers: HTMLElement[];
  }>(null);

  // Register this default portal before the tree-marking effect below and
  // capture only later peers still pending in the current open cycle. Retain
  // that cohort through snapshot refreshes and StrictMode effect replay.
  useSafeLayoutEffect(() => {
    if (!id || !hasDefaultModalPortal || !canTakeTreeSnapshot || !portalNode) {
      openingCohortRef.current = null;
      return;
    }
    const dialog = ref.current;
    if (!dialog || !contains(portalNode, dialog)) {
      openingCohortRef.current = null;
      return;
    }
    if (openingCohortRef.current?.portal !== portalNode) {
      openingCohortRef.current = {
        portal: portalNode,
        peers: getLaterOpenModalPortals(dialog),
      };
    }
    openModalPortals.add(portalNode);
    return () => {
      openModalPortals.delete(portalNode);
    };
  }, [id, canTakeTreeSnapshot, hasDefaultModalPortal, portalNode]);

  useSafeLayoutEffect(() => {
    if (!id) return;
    if (!canTakeTreeSnapshot) return;
    const dialog = ref.current;
    // When the dialog opens, we capture a snapshot of the document. This
    // snapshot is then used to disable elements outside the dialog in the
    // subsequent effect. However, the issue arises as this next effect also
    // relies on nested dialogs. Meaning, each time a nested dialog is rendered,
    // we capture a new document snapshot, which might disable third-party
    // dialogs. Hence, we take the snapshot here, independent of any nested
    // dialogs.
    return createWalkTreeSnapshot(id, [dialog]);
  }, [id, canTakeTreeSnapshot, unstable_treeSnapshotKey]);

  const getPersistentElementsProp = useEvent(getPersistentElements);

  // Disables/enables the element tree around the modal dialog element.
  useSafeLayoutEffect(() => {
    if (!id) return;
    if (!canTakeTreeSnapshot) return;
    const { disclosureElement } = store.getState();
    const dialog = contentElement ?? ref.current;
    if (!dialog) return;
    const persistentElements = getPersistentElementsProp() || [];
    const allElements = [
      dialog,
      ...persistentElements,
      ...(openingCohortRef.current?.peers || []),
      ...nestedDialogs.map((dialog) => dialog.getState().contentElement),
    ];
    // Mark known dialog elements before first focus. Re-check disclosure live
    // because hovercards and tooltips may replace it while open.
    // https://github.com/ariakit/ariakit/issues/6344
    const restoreInsideMarks = markTreeInside(dialog, allElements);
    if (modal) {
      return chain(
        restoreInsideMarks,
        markAndDisableTreeOutside(id, allElements),
      );
    }
    return chain(
      restoreInsideMarks,
      markTreeOutside(id, [disclosureElement, ...allElements]),
    );
  }, [
    id,
    store,
    canTakeTreeSnapshot,
    contentElement,
    modal,
    hasDefaultModalPortal,
    getPersistentElementsProp,
    nestedDialogs,
    unstable_treeSnapshotKey,
  ]);

  const mayAutoFocusOnShow = !!autoFocusOnShow;
  const autoFocusOnShowProp = useBooleanEvent(autoFocusOnShow);
  // We have to wait for the dialog to be mounted before allowing focusable
  // elements to be auto focused. Otherwise, there could be unintended scroll
  // jumps. See select-animated browser tests.
  const [autoFocusEnabled, setAutoFocusEnabled] = useState(false);

  // Auto focus on show.
  useEffect(() => {
    if (!open) return;
    if (!mayAutoFocusOnShow) return;
    // Makes sure to wait for the portalNode to be created before moving focus.
    // This is useful for when the Dialog component is unmounted when hidden.
    if (!domReady) return;
    // The dialog element may change for different reasons. For example, when
    // the modal or portal props change, the HTML structure will also change,
    // which will affect the dialog element reference. That's why we're
    // listening to contentElement state here instead of getting the ref.current
    // value. This ensures this effect will re-run when the dialog element
    // reference changes.
    if (!contentElement?.isConnected) return;
    const element =
      getElementFromProp(initialFocus, true) ||
      // If no initial focus is specified, we try to focus the first element
      // with the autofocus attribute. If it's an Ariakit component, the
      // Focusable component will consume the autoFocus prop and add the
      // data-autofocus attribute to the element instead.
      contentElement.querySelector<HTMLElement>(
        "[data-autofocus=true],[autofocus]",
      ) ||
      // We have to fallback to the first focusable element otherwise portaled
      // dialogs with preserveTabOrder set to true will not receive focus
      // properly because the elements aren't tabbable until the dialog receives
      // focus.
      getFirstTabbableIn(contentElement, true, portal && preserveTabOrder) ||
      // Finally, we fallback to the dialog element itself.
      contentElement;
    const isElementFocusable = isFocusable(element);
    if (!autoFocusOnShowProp(isElementFocusable ? element : null)) return;
    setAutoFocusEnabled(true);
    queueMicrotask(() => {
      const { open, disclosureElement } = store.getState();
      // If the dialog was closed between scheduling and executing this
      // microtask, skip the focus call. Otherwise, focusing a now-hidden
      // element could steal focus from the disclosure that focusOnHide already
      // restored.
      if (!open) return;
      // Preserve focus that escaped while placement was pending. Decide
      // ownership in the dialog's document, where iframe focus is represented
      // by its frame element.
      // https://github.com/ariakit/ariakit/pull/7048#discussion_r3712331088
      const documentActiveElement = getDocument(contentElement).activeElement;
      const activeElement = isElement(documentActiveElement)
        ? documentActiveElement
        : null;
      const deepestActiveElement =
        activeElement && getDeepestActiveElement(activeElement);
      if (
        focusedStoreRef.current === store &&
        activeElement &&
        deepestActiveElement &&
        isFocusable(deepestActiveElement) &&
        !isElementInDialog(activeElement, contentElement, disclosureElement) &&
        !isElementInDialog(
          deepestActiveElement,
          contentElement,
          disclosureElement,
        )
      ) {
        return;
      }
      // Safari may drop the browser's focus scroll when virtual focus redirects
      // it, so scroll explicitly before focusing. Re-read focusability here
      // because only scrolling is conditional; a stale `true` would move the
      // page toward an element that no longer accepts focus.
      // https://github.com/ariakit/ariakit/pull/7072#discussion_r3730034764
      if (isFocusable(element)) {
        element.scrollIntoView({ block: "nearest", inline: "nearest" });
      }
      element.focus({ preventScroll: true });
    });
  }, [
    open,
    mayAutoFocusOnShow,
    domReady,
    contentElement,
    initialFocus,
    portal,
    preserveTabOrder,
    store,
    autoFocusOnShowProp,
    focusedStoreRef,
  ]);

  const mayAutoFocusOnHide = !!autoFocusOnHide;
  const autoFocusOnHideProp = useBooleanEvent(autoFocusOnHide);

  // Sets a `hasOpened` flag on an effect so we only auto focus on hide if the
  // dialog was open before.
  const [hasOpened, setHasOpened] = useState(false);

  useSafeLayoutEffect(() => {
    if (!open) return;
    setHasOpened(true);
    return () => setHasOpened(false);
  }, [open]);

  const focusOnHide = useCallback(
    (dialog: HTMLElement | null, retry = true) => {
      // Hide was triggered by an outside interaction that should retain focus
      // on its target, so we skip focus restoration entirely.
      if (interactedOutsideRef.current) return;
      const { disclosureElement } = store.getState();
      // Hide was triggered by a click/focus on a tabbable element outside the
      // dialog. We won't change focus then.
      if (isAlreadyFocusingAnotherElement(dialog)) return;
      let element = getElementFromProp(finalFocus) || disclosureElement;
      if (element?.id) {
        const doc = getDocument(element);
        const selector = `[aria-activedescendant="${element.id}"]`;
        const composite = doc.querySelector<HTMLElement>(selector);
        // If the element is an item in a composite widget that handles focus
        // with the `aria-activedescendant` attribute, we want to focus on the
        // composite element itself.
        if (composite) {
          element = composite;
        }
      }
      // If the element is not focusable by the time the dialog is hidden, it's
      // probably because it's an element inside another popover or menu that
      // also got hidden when this dialog was shown. We'll try to focus on their
      // disclosure element instead.
      if (element && !isFocusable(element)) {
        const maybeParentDialog = element.closest("[data-dialog]");
        if (maybeParentDialog?.id) {
          const doc = getDocument(maybeParentDialog);
          const selector = `[aria-controls~="${maybeParentDialog.id}"]`;
          const control = doc.querySelector<HTMLElement>(selector);
          if (control) {
            element = control;
          }
        }
      }
      const isElementFocusable = element && isFocusable(element);
      if (!isElementFocusable && retry) {
        // If the element is still not focusable by this time, we retry once
        // again on the next frame. This is sometimes necessary because there
        // may be nested dialogs that still need a tick to remove the inert
        // attribute from elements outside.
        requestAnimationFrame(() => focusOnHide(dialog, false));
        return;
      }
      if (!autoFocusOnHideProp(isElementFocusable ? element : null)) return;
      if (!isElementFocusable) return;
      element?.focus();
    },
    [store, finalFocus, autoFocusOnHideProp],
  );

  const focusedOnHideRef = useRef(false);

  // Auto focus on hide with an always rendered dialog.
  useSafeLayoutEffect(() => {
    if (open) return;
    if (!hasOpened) return;
    if (!mayAutoFocusOnHide) return;
    const dialog = ref.current;
    // We don't want to focus on hide twice if the dialog is not unmounted, so
    // we set this flag here that will be checked in the cleanup effect below.
    focusedOnHideRef.current = true;
    focusOnHide(dialog);
  }, [open, hasOpened, domReady, mayAutoFocusOnHide, focusOnHide]);

  // Auto focus on hide with a dialog that gets unmounted when hidden.
  useEffect(() => {
    if (!hasOpened) return;
    if (!mayAutoFocusOnHide) return;
    const dialog = ref.current;
    return () => {
      if (focusedOnHideRef.current) {
        focusedOnHideRef.current = false;
        return;
      }
      focusOnHide(dialog);
    };
  }, [hasOpened, mayAutoFocusOnHide, focusOnHide]);

  const hideOnEscapeProp = useBooleanEvent(hideOnEscape);
  const [escapeEvents] = useState(
    () =>
      new WeakMap<
        KeyboardEvent,
        { accepted: boolean; defaultPrevented: boolean }
      >(),
  );

  const onKeyDownProp = props.onKeyDown;
  const onKeyDownCaptureProp = props.onKeyDownCapture;
  const onFocusCaptureProp = props.onFocusCapture;

  const onFocusCapture = useEvent((event: ReactFocusEvent<HTMLType>) => {
    onFocusCaptureProp?.(event);
    if (!store.getState().open) return;
    const target = event.target;
    if (!isNode(target)) return;
    if (!contains(event.currentTarget, target)) return;
    focusedStoreRef.current = store;
  });

  const acceptEscape = useEvent((event: KeyboardEvent) => {
    if (event.key !== "Escape") return false;
    if (!event.bubbles) return false;
    const result = escapeEvents.get(event);
    if (result) {
      if (event.defaultPrevented && !result.defaultPrevented) return false;
      return result.accepted;
    }
    if (event.defaultPrevented) return false;
    const dialog = ref.current;
    if (!mounted) return false;
    if (!dialog) return false;
    // Ignore the event if the current dialog is marked by another dialog.
    // This guarantees that only the topmost dialog will close on Escape.
    if (isElementMarked(dialog)) return false;
    const accepted = hideOnEscapeProp(event);
    escapeEvents.set(event, {
      accepted,
      defaultPrevented: event.defaultPrevented,
    });
    return accepted;
  });

  const hideOnEscapeEvent = useEvent((event: KeyboardEvent) => {
    const accepted = acceptEscape(event);
    escapeEvents.delete(event);
    if (!accepted) return false;
    store.hide();
    return true;
  });

  const onKeyDown = useEvent((event: ReactKeyboardEvent<HTMLType>) => {
    const nativeEvent = event.nativeEvent;
    const wasPropagationStopped = nativeEvent.cancelBubble;
    onKeyDownProp?.(event);
    if (wasPropagationStopped) {
      escapeEvents.delete(nativeEvent);
      return;
    }
    if (!hideOnEscapeEvent(nativeEvent)) return;
    event.stopPropagation();
  });

  const onKeyDownCapture = useEvent((event: ReactKeyboardEvent<HTMLType>) => {
    const nativeEvent = event.nativeEvent;
    const wasPropagationStopped = nativeEvent.cancelBubble;
    onKeyDownCaptureProp?.(event);
    if (wasPropagationStopped) {
      escapeEvents.delete(nativeEvent);
      return;
    }
    const stoppedAtDialog =
      event.isPropagationStopped() || nativeEvent.cancelBubble;
    if (!stoppedAtDialog) return;
    hideOnEscapeEvent(nativeEvent);
  });

  // Hide on Escape.
  useEffect(() => {
    if (!domReady) return;
    if (!mounted) return;
    const onDocumentKeyDownCapture = (event: KeyboardEvent) => {
      if (event.key !== "Escape") return;
      if (!event.bubbles) return;
      if (event.cancelBubble) {
        escapeEvents.delete(event);
        return;
      }
      if (escapeEvents.has(event)) return;
      const dialog = ref.current;
      if (!dialog) return;
      const target = event.target;
      // Guard against non-node targets (e.g. a synthetic event dispatched on
      // window) so `contains` doesn't throw. `isNode` rather than `isElement`
      // keeps non-element nodes working with `contains`, as before.
      if (!isNode(target)) return;
      const { disclosureElement } = store.getState();
      // This considers valid targets the elements that belong to this dialog
      // tree, including elements marked as outside by this dialog so Escape can
      // close the topmost dialog even when focus is outside.
      const isValidTarget = () => {
        if (isElement(target) && target.tagName === "BODY") return true;
        if (contains(dialog, target)) return true;
        if (!disclosureElement) return true;
        if (contains(disclosureElement, target)) return true;
        if (isElement(target) && isElementMarked(target, dialog.id)) {
          return true;
        }
        return false;
      };
      if (!isValidTarget()) return;
      if (!acceptEscape(event)) return;
      if (!event.cancelBubble) return;
      hideOnEscapeEvent(event);
    };
    const onDocumentKeyDown = (event: KeyboardEvent) => {
      if (!escapeEvents.has(event)) return;
      if (event.cancelBubble) {
        escapeEvents.delete(event);
        return;
      }
      if (!hideOnEscapeEvent(event)) return;
      event.stopPropagation();
    };
    // Listen on the document so Escape works even when the dialog isn't
    // focused. Capture preflights DOM-owned events so hideOnEscape can stop
    // them, while the hide is committed only after descendants handle them.
    const win = contentElement ? getWindow(contentElement) : undefined;
    return chain(
      addGlobalEventListener("keydown", onDocumentKeyDownCapture, true, win),
      addGlobalEventListener("keydown", onDocumentKeyDown, false, win),
    );
  }, [
    store,
    domReady,
    mounted,
    contentElement,
    hideOnEscapeEvent,
    acceptEscape,
    escapeEvents,
  ]);

  // Resets the heading levels inside the modal dialog so they start with h1.
  props = useWrapElement(
    props,
    (element) => (
      <HeadingLevel level={modal ? 1 : undefined}>{element}</HeadingLevel>
    ),
    [modal],
  );

  const hiddenProp = props.hidden;
  const alwaysVisible = props.alwaysVisible;

  // Wraps the dialog with a backdrop element if the backdrop prop is truthy.
  props = useWrapElement(
    props,
    (element) => {
      if (!backdrop) return element;
      return (
        <>
          <DialogBackdrop
            store={store}
            backdrop={backdrop}
            backdropRef={backdropRef}
            hidden={hiddenProp}
            alwaysVisible={alwaysVisible}
          />
          {element}
        </>
      );
    },
    [store, backdrop, hiddenProp, alwaysVisible],
  );

  const [headingId, setHeadingId] = useState<string>();
  const [descriptionId, setDescriptionId] = useState<string>();

  props = useWrapElement(
    props,
    (element) => (
      <DialogScopedContextProvider value={store}>
        <DialogHeadingContext.Provider value={setHeadingId}>
          <DialogDescriptionContext.Provider value={setDescriptionId}>
            {element}
          </DialogDescriptionContext.Provider>
        </DialogHeadingContext.Provider>
      </DialogScopedContextProvider>
    ),
    [store],
  );

  props = {
    "data-dialog": "",
    role: "dialog",
    tabIndex: focusable ? -1 : undefined,
    "aria-labelledby": props["aria-label"] != null ? undefined : headingId,
    "aria-describedby": descriptionId,
    ...props,
    "data-dialog-portal": hasDefaultModalPortal ? portalNode?.id : undefined,
    id,
    ref: useMergeRefs(ref, props.ref),
    onFocusCapture,
    onKeyDown,
    onKeyDownCapture,
  };

  props = useFocusableContainer({
    ...props,
    autoFocusOnShow: autoFocusEnabled,
  });
  props = useDisclosureContent({
    store,
    ...props,
    unstable_otherElementRef: backdropRef,
  });
  props = useFocusable({ ...props, focusable });
  props = usePortal({ portal, ...props, portalRef, preserveTabOrder });

  return props;
});

export function createDialogComponent<T extends DialogOptions>(
  Component: FC<T>,
  useProviderContext = useDialogProviderContext,
) {
  return forwardRef(function DialogComponent(props: T) {
    const context = useProviderContext();
    const store = props.store || context;
    const mounted = useStoreState(
      store,
      ["mounted"],
      (state) => !props.unmountOnHide || state?.mounted || !!props.open,
    );
    if (!mounted) return null;
    return <Component {...props} />;
  });
}

const DialogImpl = forwardRef(function DialogImpl(props: DialogProps) {
  const htmlProps = useDialog(props);
  return createElement(TagName, htmlProps);
});

const DialogWithStore = createDialogComponent(
  DialogImpl,
  useDialogProviderContext,
);

const DialogWithInternalStore = forwardRef(function DialogWithInternalStore(
  props: DialogProps,
) {
  const store = useDialogStore({ open: props.open });
  return <DialogWithStore {...props} store={store} />;
});

/**
 * Renders a dialog similar to the native `dialog` element that's rendered in a
 * [`portal`](https://ariakit.com/reference/dialog#portal) by default.
 *
 * The dialog can be either
 * [`modal`](https://ariakit.com/reference/dialog#modal) or non-modal. The
 * visibility state can be controlled with the
 * [`open`](https://ariakit.com/reference/dialog#open) and
 * [`onClose`](https://ariakit.com/reference/dialog#onclose) props.
 * @see https://ariakit.com/components/dialog
 * @example
 * ```jsx {4-6}
 * const [open, setOpen] = useState(false);
 *
 * <button onClick={() => setOpen(true)}>Open dialog</button>
 * <Dialog open={open} onClose={() => setOpen(false)}>
 *   Dialog
 * </Dialog>
 * ```
 */
export const Dialog = forwardRef(function Dialog(props: DialogProps) {
  const context = useDialogProviderContext();
  // The hoisted store is only needed for the unmountOnHide mounted-state gate.
  if (props.store || context || !props.unmountOnHide) {
    return <DialogWithStore {...props} />;
  }
  return <DialogWithInternalStore {...props} />;
});

export interface DialogOptions<T extends ElementType = TagName>
  extends FocusableOptions<T>, PortalOptions<T>, DisclosureContentOptions<T> {
  /**
   * Object returned by the
   * [`useDialogStore`](https://ariakit.com/reference/use-dialog-store) hook. If
   * not provided, the closest
   * [`DialogProvider`](https://ariakit.com/reference/dialog-provider)
   * component's context will be used. Otherwise, an internal store will be
   * created.
   */
  store?: DialogStore;
  /**
   * Controls the open state of the dialog. This is similar to the
   * [`open`](https://developer.mozilla.org/en-US/docs/Web/API/HTMLDialogElement/open)
   * attribute on native dialog elements.
   *
   * Live examples:
   * - [Dialog with scrollable
   *   backdrop](https://ariakit.com/examples/dialog-backdrop-scrollable)
   * - [Dialog with details &
   *   summary](https://ariakit.com/examples/dialog-details)
   * - [Warning on Dialog
   *   hide](https://ariakit.com/examples/dialog-hide-warning)
   * - [Dialog with Menu](https://ariakit.com/examples/dialog-menu)
   */
  open?: boolean;
  /**
   * This is an event handler prop triggered when the dialog's `close` event is
   * dispatched. The `close` event is similar to the native dialog
   * [`close`](https://developer.mozilla.org/en-US/docs/Web/API/HTMLDialogElement/close_event)
   * event. The only difference is that this event can be canceled with
   * `event.preventDefault()`, which will prevent the dialog from hiding.
   *
   * It's important to note that this event only fires when the dialog store's
   * [`open`](https://ariakit.com/reference/use-dialog-store#open) state is set
   * to `false`. If the controlled
   * [`open`](https://ariakit.com/reference/dialog#open) prop value changes, or
   * if the dialog's visibility is altered in any other way (such as unmounting
   * the dialog without adjusting the open state), this event won't be
   * triggered.
   *
   * Live examples:
   * - [Dialog with scrollable
   *   backdrop](https://ariakit.com/examples/dialog-backdrop-scrollable)
   * - [Dialog with details &
   *   summary](https://ariakit.com/examples/dialog-details)
   * - [Warning on Dialog
   *   hide](https://ariakit.com/examples/dialog-hide-warning)
   * - [Dialog with Menu](https://ariakit.com/examples/dialog-menu)
   */
  onClose?: (event: Event) => void;
  /**
   * Determines whether the dialog is modal. Modal dialogs have distinct states
   * and behaviors:
   * - The [`portal`](https://ariakit.com/reference/dialog#portal) and
   *   [`preventBodyScroll`](https://ariakit.com/reference/dialog#preventbodyscroll)
   *   props are set to `true`. They can still be manually set to `false`.
   * - When using the [`Heading`](https://ariakit.com/reference/heading) or
   *   [`DialogHeading`](https://ariakit.com/reference/dialog-heading)
   *   components within the dialog, their level will be reset so they start
   *   with `h1`.
   * - A visually hidden dismiss button will be rendered if the
   *   [`DialogDismiss`](https://ariakit.com/reference/dialog-dismiss) component
   *   hasn't been used. This allows screen reader users to close the dialog.
   * - When the dialog is open, element tree outside it will be inert.
   *
   * Live examples:
   * - [Combobox with Tabs](https://ariakit.com/examples/combobox-tabs)
   * - [Dialog with details &
   *   summary](https://ariakit.com/examples/dialog-details)
   * - [Form with Select](https://ariakit.com/examples/form-select)
   * - [Context menu](https://ariakit.com/examples/menu-context-menu)
   * - [Responsive Popover](https://ariakit.com/examples/popover-responsive)
   * @default true
   */
  modal?: boolean;
  /**
   * Determines whether there will be a backdrop behind the dialog. On modal
   * dialogs, this is `true` by default. Besides a `boolean`, this prop can also
   * be a React component or JSX element that will be rendered as the backdrop.
   *
   * **Note**: If a custom component is used, it must [accept ref and spread all
   * props to its underlying DOM
   * element](https://ariakit.com/guide/composition#custom-components-must-be-open-for-extension),
   * the same way a native element would.
   *
   * Live examples:
   * - [Animated Dialog](https://ariakit.com/examples/dialog-animated)
   * - [Dialog with scrollable
   *   backdrop](https://ariakit.com/examples/dialog-backdrop-scrollable)
   * - [Dialog with
   *   Motion](https://ariakit.com/examples/dialog-framer-motion)
   * - [Dialog with Menu](https://ariakit.com/examples/dialog-menu)
   * - [Nested Dialog](https://ariakit.com/examples/dialog-nested)
   * - [Dialog with Next.js App
   *   Router](https://ariakit.com/examples/dialog-next-router)
   * @example
   * ```jsx
   * <Dialog backdrop={<div className="backdrop" />} />
   * ```
   */
  backdrop?:
    | boolean
    | ReactElement<ComponentPropsWithRef<"div">>
    | ElementType<ComponentPropsWithRef<"div">>;
  /**
   * Determines if the dialog will hide when the user presses the Escape key.
   *
   * This prop can be either a boolean or a function that accepts an event as an
   * argument and returns a boolean. The event object represents the keydown
   * event that initiated the hide action, which could be either a native
   * keyboard event or a React synthetic event.
   *
   * **Note**: The dialog stops handled Escape events from its React subtree
   * before they reach ancestor React bubble handlers. An ancestor capture
   * handler can own the event by stopping its propagation. If this function
   * runs before such a handler, it can call `event.stopPropagation()` to
   * prevent the handler from receiving the event.
   * @default true
   */
  hideOnEscape?: BooleanOrCallback<KeyboardEvent | ReactKeyboardEvent>;
  /**
   * Determines if the dialog should hide when the user clicks or focuses on an
   * element outside the dialog.
   *
   * This prop can be either a boolean or a function that takes an event as an
   * argument and returns a boolean. The event object represents the event that
   * triggered the action, which could be a native event or a React synthetic
   * event of various types.
   *
   * Live examples:
   * - [Selection Popover](https://ariakit.com/examples/popover-selection)
   * @default true
   */
  hideOnInteractOutside?: BooleanOrCallback<Event | SyntheticEvent>;
  /**
   * When a dialog is open, the elements outside of it are disabled to prevent
   * interaction if the dialog is
   * [`modal`](https://ariakit.com/reference/dialog#modal). For non-modal
   * dialogs, interacting with elements outside the dialog prompts it to close.
   *
   * This function allows you to return an iterable collection of elements that
   * will be considered as part of the dialog, thus excluding them from this
   * behavior.
   *
   * **Note**: The elements returned by this function must exist in the DOM when
   * the dialog opens.
   *
   * Live examples:
   * - [Dialog with
   *   React-Toastify](https://ariakit.com/examples/dialog-react-toastify)
   */
  getPersistentElements?: () => Iterable<Element>;
  /**
   * Determines whether the body scrolling will be prevented when the dialog is
   * shown. This is automatically set to `true` when the dialog is
   * [`modal`](https://ariakit.com/reference/dialog#modal). You can disable this
   * prop if you want to implement your own logic.
   */
  preventBodyScroll?: boolean;
  /**
   * Determines whether an element inside the dialog will receive focus when the
   * dialog is shown. By default, this is usually the first tabbable element in
   * the dialog or the dialog itself. The
   * [`initialFocus`](https://ariakit.com/reference/dialog#initialfocus) prop
   * can be used to set a different element to receive focus.
   *
   * Live examples:
   * - [Warning on Dialog
   *   hide](https://ariakit.com/examples/dialog-hide-warning)
   * - [Sliding Menu](https://ariakit.com/examples/menu-slide)
   * - [Selection Popover](https://ariakit.com/examples/popover-selection)
   * @default true
   */
  autoFocusOnShow?: BooleanOrCallback<HTMLElement | null>;
  /**
   * Determines whether an element outside of the dialog will be focused when
   * the dialog is hidden if another element hasn't been focused in the action
   * of hiding the dialog (for example, by clicking or tabbing into another
   * tabbable element outside of the dialog).
   *
   * By default, this is usually the disclosure element. The
   * [`finalFocus`](https://ariakit.com/reference/dialog#finalfocus) prop can be
   * used to define a different element to be focused.
   *
   * Live examples:
   * - [Dialog with Next.js App
   *   Router](https://ariakit.com/examples/dialog-next-router)
   * - [Sliding menu](https://ariakit.com/examples/menu-slide)
   * @default true
   */
  autoFocusOnHide?: BooleanOrCallback<HTMLElement | null>;
  /**
   * Specifies the element that will receive focus when the dialog is first
   * opened. It can be an `HTMLElement` or a `React.RefObject` with an
   * `HTMLElement`.
   *
   * If
   * [`autoFocusOnShow`](https://ariakit.com/reference/dialog#autofocusonshow)
   * is set to `false`, this prop will have no effect. If left unset, the dialog
   * will attempt to determine the initial focus element in the following order:
   * 1. A [Focusable](https://ariakit.com/components/focusable) element with an
   *    [`autoFocus`](https://ariakit.com/reference/focusable#autofocus) prop.
   * 2. The first tabbable element inside the dialog.
   * 3. The first focusable element inside the dialog.
   * 4. The dialog element itself.
   */
  initialFocus?: HTMLElement | RefObject<HTMLElement | null> | null;
  /**
   * Determines the element that will receive focus once the dialog is closed,
   * provided that no other element has been focused while the dialog was being
   * hidden (e.g., by clicking or tabbing into another tabbable element outside
   * of the dialog).
   * - If
   *   [`autoFocusOnHide`](https://ariakit.com/reference/dialog#autofocusonhide)
   *   is set to `false`, this prop will have no effect.
   * - If left unset, the element that was focused before the dialog was opened
   *   will be focused again.
   */
  finalFocus?: HTMLElement | RefObject<HTMLElement | null> | null;
  /**
   * @private
   */
  unstable_treeSnapshotKey?: string | number | boolean | null;
}

export type DialogProps<T extends ElementType = TagName> = Props<
  T,
  DialogOptions<T>
>;