import {
  useLiveRef,
  useMergeRefs,
  useSafeLayoutEffect,
  useWrapElement,
  setRef,
  createElement,
  createHook,
  forwardRef,
} from "@ariakit/react-utils";
import type { Options, Props } from "@ariakit/react-utils";
import {
  getDocument,
  isFocusEventOutside,
  disableFocusIn,
  getNextTabbable,
  getPreviousTabbable,
  restoreFocusIn,
} from "@ariakit/utils";
import type { ElementType, MutableRefObject, RefCallback } from "react";
import { useContext, useEffect, useRef, useState } from "react";
import { createPortal } from "react-dom";
import { FocusTrap } from "../focus-trap/focus-trap.tsx";
import { PortalContext } from "./portal-context.tsx";

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

// Returns the best root element for appending portal nodes. When an element
// is in fullscreen mode, portals must be appended inside the fullscreen
// element instead of document.body so they remain visible.
function getRootElement(element?: Element | null) {
  const doc = getDocument(element);
  const { fullscreenElement } = doc;
  const HTMLElementClass = doc.defaultView?.HTMLElement;
  if (HTMLElementClass && fullscreenElement instanceof HTMLElementClass) {
    return fullscreenElement;
  }
  return doc.body;
}

function getPortalElement(
  element: HTMLElement,
  portalElement: PortalOptions["portalElement"],
) {
  if (!portalElement) {
    return getDocument(element).createElement("div");
  }
  if (typeof portalElement === "function") {
    return portalElement(element);
  }
  return portalElement;
}

function getRandomId(prefix = "id") {
  return `${prefix ? `${prefix}-` : ""}${Math.random()
    .toString(36)
    .slice(2, 8)}`;
}

function queueFocus(element?: HTMLElement | null) {
  queueMicrotask(() => {
    element?.focus();
  });
}

interface AttachedPortalRef {
  ref: PortalOptions["portalRef"];
  node: HTMLElement;
  cleanup: void | (() => void);
}

function attachPortalRef(
  ref: PortalOptions["portalRef"],
  node: HTMLElement,
): AttachedPortalRef {
  return { ref, node, cleanup: setRef(ref, node) };
}

function detachPortalRef(attached: AttachedPortalRef) {
  // Preserve React 19 callback ref cleanup semantics. Otherwise, detach the
  // ref with null like any other React ref.
  if (typeof attached.cleanup === "function") {
    attached.cleanup();
  } else {
    setRef(attached.ref, null);
  }
}

/**
 * Returns props to create a `Portal` component.
 * @see https://ariakit.com/components/portal
 * @example
 * ```jsx
 * const props = usePortal();
 * <Role {...props}>Content</Role>
 * ```
 */
export const usePortal = createHook<TagName, PortalOptions>(function usePortal({
  preserveTabOrder,
  preserveTabOrderAnchor,
  portalElement,
  portalRef,
  portal = true,
  ...props
}) {
  const ref = useRef<HTMLType>(null);
  const refProp = useMergeRefs(ref, props.ref);
  const context = useContext(PortalContext);
  const [portalNode, setPortalNode] = useState<HTMLElement | null>(null);
  const [anchorPortalNode, setAnchorPortalNode] = useState<HTMLElement | null>(
    null,
  );

  const outerBeforeRef = useRef<HTMLSpanElement>(null);
  const innerBeforeRef = useRef<HTMLSpanElement>(null);
  const innerAfterRef = useRef<HTMLSpanElement>(null);
  const outerAfterRef = useRef<HTMLSpanElement>(null);

  const portalRefProp = useLiveRef(portalRef);
  // Tracks the currently attached portalRef so the two effects below can
  // detach and re-attach it without sharing dependencies.
  const attachedPortalRefRef = useRef<AttachedPortalRef | null>(null);

  // Create the portal node and attach it to the DOM.
  useSafeLayoutEffect(() => {
    const element = ref.current;
    if (!element || !portal) {
      setPortalNode(null);
      return;
    }
    const portalEl = getPortalElement(element, portalElement);
    // TODO: Warn when document.body is used as the portal element.
    if (!portalEl) {
      setPortalNode(null);
      return;
    }
    const isPortalInDocument = portalEl.isConnected;
    if (!isPortalInDocument) {
      const rootElement = context || getRootElement(element);
      rootElement.appendChild(portalEl);
    }
    // If the portal element doesn't have an id already, set one.
    if (!portalEl.id) {
      // Use the element's id so rendering <Portal id="some-id" /> will
      // produce predictable results.
      portalEl.id = element.id ? `portal/${element.id}` : getRandomId();
    }
    // Set the internal portal node state and attach the portalRef prop. The
    // ref is read through a live ref so its identity is not a dependency of
    // this effect: a portalRef identity change must not recreate the portal
    // node. The effect below re-fires the ref in that case.
    setPortalNode(portalEl);
    attachedPortalRefRef.current = attachPortalRef(
      portalRefProp.current,
      portalEl,
    );
    return () => {
      const attached = attachedPortalRefRef.current;
      // Detach the portalRef first so ref cleanups still observe a connected
      // portal node.
      if (attached) {
        attachedPortalRefRef.current = null;
        detachPortalRef(attached);
      }
      // Connected portals keep their DOM node.
      if (!isPortalInDocument) {
        portalEl.remove();
      }
    };
  }, [portal, portalElement, context]);

  // Re-fire the portalRef against the same portal node when only its identity
  // changes (e.g. an inline callback on a parent re-render), mirroring how
  // React re-fires element refs without recreating the DOM node.
  useSafeLayoutEffect(() => {
    const attached = attachedPortalRefRef.current;
    if (!attached) return;
    if (attached.ref === portalRef) return;
    detachPortalRef(attached);
    attachedPortalRefRef.current = attachPortalRef(portalRef, attached.node);
  }, [portalRef]);

  // Move the portal node when fullscreen state changes so it stays visible.
  useEffect(() => {
    if (!portalNode) return;
    if (context) return;
    if (portalElement) return;
    const doc = getDocument(portalNode);
    const onFullscreenChange = () => {
      const rootElement = getRootElement(portalNode);
      if (portalNode.parentElement !== rootElement) {
        rootElement.appendChild(portalNode);
      }
    };
    // Sync immediately in case fullscreen was entered before this effect
    // ran, which can happen if the portal mounts while already in
    // fullscreen mode. Skip when the captured node is already disconnected,
    // which happens for a StrictMode cleanup node whose layout cleanup
    // already removed it.
    if (portalNode.isConnected) {
      onFullscreenChange();
    }
    doc.addEventListener("fullscreenchange", onFullscreenChange);
    return () => {
      doc.removeEventListener("fullscreenchange", onFullscreenChange);
    };
  }, [portalNode, context, portalElement]);

  // Create the anchor portal node and attach it to the DOM.
  useSafeLayoutEffect(() => {
    if (!portal) return;
    if (!preserveTabOrder) return;
    if (!preserveTabOrderAnchor) return;
    const doc = getDocument(preserveTabOrderAnchor);
    const element = doc.createElement("span");
    element.style.position = "fixed";
    preserveTabOrderAnchor.insertAdjacentElement("afterend", element);
    setAnchorPortalNode(element);
    return () => {
      element.remove();
      setAnchorPortalNode(null);
    };
  }, [portal, preserveTabOrder, preserveTabOrderAnchor]);

  // When preserveTabOrder is true, make sure elements inside the portal
  // element are tabbable only when the portal has already been focused,
  // either by tabbing into a focus trap element outside or using the mouse.
  useEffect(() => {
    if (!portalNode) return;
    if (!preserveTabOrder) return;
    let raf = 0;
    const onFocus = (event: FocusEvent) => {
      if (!isFocusEventOutside(event)) return;
      const focusing = event.type === "focusin";
      cancelAnimationFrame(raf);
      if (focusing) {
        return restoreFocusIn(portalNode);
      }
      // Wait for the next frame to allow tabindex changes after the focus
      // event.
      raf = requestAnimationFrame(() => {
        disableFocusIn(portalNode, true);
      });
    };
    // Listen to the event on the capture phase so they run before the focus
    // trap elements onFocus prop is called.
    portalNode.addEventListener("focusin", onFocus, true);
    portalNode.addEventListener("focusout", onFocus, true);
    return () => {
      cancelAnimationFrame(raf);
      portalNode.removeEventListener("focusin", onFocus, true);
      portalNode.removeEventListener("focusout", onFocus, true);
    };
  }, [portalNode, preserveTabOrder]);

  props = useWrapElement(
    props,
    (element) => {
      element = (
        // While the portal node is not in the DOM, we need to pass the
        // current context to the portal context, otherwise it's going to
        // reset to the body element on nested portals.
        <PortalContext.Provider value={portalNode || context}>
          {element}
        </PortalContext.Provider>
      );

      if (!portal) return element;

      if (!portalNode) {
        // If the element should be rendered within a portal, but the portal
        // node is not yet in the DOM, we'll return an empty div element. We
        // assign the id to the element so we can use it to set the portal id
        // later on. We're using position: fixed here so that the browser
        // doesn't add margin to the element when setting gap on a parent
        // element.
        return (
          <span
            ref={refProp}
            id={props.id}
            style={{ position: "fixed" }}
            hidden
          />
        );
      }

      element = (
        <>
          {preserveTabOrder && (
            <FocusTrap
              ref={innerBeforeRef}
              data-focus-trap={props.id}
              className="__focus-trap-inner-before"
              onFocus={(event) => {
                if (isFocusEventOutside(event, portalNode)) {
                  queueFocus(getNextTabbable());
                } else {
                  queueFocus(outerBeforeRef.current);
                }
              }}
            />
          )}
          {element}
          {preserveTabOrder && (
            <FocusTrap
              ref={innerAfterRef}
              data-focus-trap={props.id}
              className="__focus-trap-inner-after"
              onFocus={(event) => {
                if (isFocusEventOutside(event, portalNode)) {
                  queueFocus(getPreviousTabbable());
                } else {
                  queueFocus(outerAfterRef.current);
                }
              }}
            />
          )}
        </>
      );

      element = createPortal(element, portalNode);

      let preserveTabOrderElement = (
        <>
          {preserveTabOrder && (
            <FocusTrap
              ref={outerBeforeRef}
              data-focus-trap={props.id}
              className="__focus-trap-outer-before"
              onFocus={(event) => {
                // If the event is coming from the outer after focus trap, it
                // means there's no tabbable element inside the portal. In
                // this case, we don't focus the inner before focus trap, but
                // the previous tabbable element outside the portal.
                const fromOuter = event.relatedTarget === outerAfterRef.current;
                if (!fromOuter && isFocusEventOutside(event, portalNode)) {
                  queueFocus(innerBeforeRef.current);
                } else {
                  queueFocus(getPreviousTabbable());
                }
              }}
            />
          )}
          {preserveTabOrder && (
            // We're using position: fixed here so that the browser doesn't
            // add margin to the element when setting gap on a parent element.
            <span aria-owns={portalNode.id} style={{ position: "fixed" }} />
          )}
          {preserveTabOrder && (
            <FocusTrap
              ref={outerAfterRef}
              data-focus-trap={props.id}
              className="__focus-trap-outer-after"
              onFocus={(event) => {
                if (isFocusEventOutside(event, portalNode)) {
                  queueFocus(innerAfterRef.current);
                } else {
                  const nextTabbable = getNextTabbable();
                  // If the next tabbable element is the inner before focus
                  // trap, this means we're at the end of the document or the
                  // portal was placed right after the original spot in the
                  // React tree. We need to wait for the next frame so the
                  // preserveTabOrder effect can run and disable the inner
                  // before focus trap. If there's no tabbable element after
                  // that, the focus will stay on this element.
                  if (nextTabbable === innerBeforeRef.current) {
                    requestAnimationFrame(() => getNextTabbable()?.focus());
                    return;
                  }
                  queueFocus(nextTabbable);
                }
              }}
            />
          )}
        </>
      );

      if (anchorPortalNode && preserveTabOrder) {
        preserveTabOrderElement = createPortal(
          preserveTabOrderElement,
          anchorPortalNode,
        );
      }

      return (
        <>
          {preserveTabOrderElement}
          {element}
        </>
      );
    },
    [portalNode, context, portal, props.id, preserveTabOrder, anchorPortalNode],
  );

  props = {
    ...props,
    ref: refProp,
  };

  return props;
});

/**
 * Renders an element using [React
 * Portal](https://react.dev/reference/react-dom/createPortal).
 *
 * By default, the portal element is a `div` element appended to the
 * `document.body` element. You can customize this with the
 * [`portalElement`](https://ariakit.com/reference/portal#portalelement) prop.
 *
 * The
 * [`preserveTabOrder`](https://ariakit.com/reference/portal#preservetaborder)
 * prop allows this component to manage the tab order of the elements. It
 * ensures the tab order remains consistent with the original location where the
 * portal was rendered in the React tree, instead of the final location in the
 * DOM. The
 * [`preserveTabOrderAnchor`](https://ariakit.com/reference/portal#preservetaborderanchor)
 * prop can specify a different location from which the tab order is preserved.
 * @see https://ariakit.com/components/portal
 * @example
 * ```jsx
 * <Portal>Content</Portal>
 * ```
 */
export const Portal = forwardRef(function Portal(props: PortalProps) {
  const htmlProps = usePortal(props);
  return createElement(TagName, htmlProps);
});

export interface PortalOptions<
  _T extends ElementType = TagName,
> extends Options {
  /**
   * When enabled, `preserveTabOrder` will keep the DOM element's tab order the
   * same as the order in which the underlying
   * [`Portal`](https://ariakit.com/reference/portal) component was mounted in
   * the React tree.
   *
   * If the
   * [`preserveTabOrderAnchor`](https://ariakit.com/reference/portal#preservetaborderanchor)
   * prop is provided, the tab order will be preserved relative to that element.
   * @default false
   */
  preserveTabOrder?: boolean;
  /**
   * An anchor element for maintaining the tab order when
   * [`preserveTabOrder`](https://ariakit.com/reference/portal#preservetaborder)
   * prop is enabled. The tab order will be kept relative to this element.
   *
   * By default, the tab order is kept relative to the original location in the
   * React tree where the underlying
   * [`Portal`](https://ariakit.com/reference/portal) component was mounted.
   * @example
   * ```jsx {18-20}
   * const [anchor, setAnchor] = useState(null);
   *
   * <button ref={setAnchor}>Order 0</button>
   * <button>Order 2</button>
   *
   * // Rendered at the end of the document.
   * <Portal>
   *   <button>Order 5</button>
   * </Portal>
   *
   * // Rendered at the end of the document, but the tab order is preserved.
   * <Portal preserveTabOrder>
   *   <button>Order 3</button>
   * </Portal>
   *
   * // Rendered at the end of the document, but the tab order is preserved
   * // relative to the anchor element.
   * <Portal preserveTabOrder preserveTabOrderAnchor={anchor}>
   *   <button>Order 1</button>
   * </Portal>
   *
   * <button>Order 4</button>
   * ```
   */
  preserveTabOrderAnchor?: Element | null;
  /**
   * `portalRef` is similar to `ref` but is scoped to the portal node. It's
   * useful when you need to be informed when the portal element is appended to
   * the DOM or removed from the DOM.
   *
   * Live examples:
   * - [Form with Select](https://ariakit.com/examples/form-select)
   * @example
   * ```jsx
   * const [portalElement, setPortalElement] = useState(null);
   *
   * <Portal portalRef={setPortalElement} />
   * ```
   */
  portalRef?: RefCallback<HTMLElement> | MutableRefObject<HTMLElement | null>;
  /**
   * Determines whether the element should be rendered as a React Portal.
   *
   * Live examples:
   * - [Combobox with integrated
   *   filter](https://ariakit.com/examples/combobox-filtering-integrated)
   * - [Dialog with Menu](https://ariakit.com/examples/dialog-menu)
   * - [Hovercard with keyboard
   *   support](https://ariakit.com/examples/hovercard-disclosure)
   * - [Menubar](https://ariakit.com/components/menubar)
   * - [Standalone Popover](https://ariakit.com/examples/popover-standalone)
   * - [Animated Select](https://ariakit.com/examples/select-animated)
   * @default true
   */
  portal?: boolean;
  /**
   * An HTML element or a memoized callback function that returns an HTML
   * element to be used as the portal element. By default, the portal element
   * will be a `div` element appended to the `document.body`.
   *
   * Live examples:
   * - [Navigation Menubar](https://ariakit.com/examples/menubar-navigation)
   * @example
   * ```jsx
   * const [portal, setPortal] = useState(null);
   *
   * <Portal portalElement={portal} />
   * <div ref={setPortal} />
   * ```
   * @example
   * ```jsx
   * const getPortalElement = useCallback(() => {
   *   const div = document.createElement("div");
   *   const portalRoot = document.getElementById("portal-root");
   *   portalRoot.appendChild(div);
   *   return div;
   * }, []);
   *
   * <Portal portalElement={getPortalElement} />
   * ```
   */
  portalElement?:
    | ((element: HTMLElement) => HTMLElement | null)
    | HTMLElement
    | null;
}

export type PortalProps<T extends ElementType = TagName> = Props<
  T,
  PortalOptions<T>
>;