import { useStoreState, useStoreStateObject } from "@ariakit/react-store";
import {
  useId,
  useMergeRefs,
  useSafeLayoutEffect,
  useWrapElement,
  createElement,
  createHook,
  forwardRef,
} from "@ariakit/react-utils";
import type { Options, Props } from "@ariakit/react-utils";
import { afterPaint, invariant, removeUndefinedValues } from "@ariakit/utils";
import type { ElementType, RefObject } from "react";
import { useMemo, useRef, useState } from "react";
import { flushSync } from "react-dom";
import { DialogScopedContextProvider } from "../dialog/dialog-context.tsx";
import { useDisclosureProviderContext } from "./disclosure-context.tsx";
import type { DisclosureStore } from "./disclosure-store.ts";

const TagName = "div" satisfies ElementType;
type TagName = typeof TagName;
type HTMLType = HTMLElementTagNameMap[TagName];
type TransitionState = "enter" | "leave" | null;

function afterTimeout(timeoutMs: number, cb: () => void) {
  const timeoutId = setTimeout(cb, timeoutMs);
  return () => clearTimeout(timeoutId);
}

function parseCSSTime(time: string | undefined) {
  const value = time?.trim() || "0s";
  const multiplier = value.endsWith("ms") ? 1 : 1000;
  const parsed = Number.parseFloat(value) * multiplier;
  // Non-numeric values such as `animation-duration: auto` parse to `NaN`; treat
  // them as 0 so they don't poison the `Math.max` that combines end times.
  return Number.isNaN(parsed) ? 0 : parsed;
}

// CSS cycles shorter delay/duration lists to match the property/name list.
// Pair each item's delay and duration; independent maxima can overestimate.
function getEndTime(names: string, delays: string, durations: string) {
  const nameList = names.split(",");
  const delayList = delays.split(",");
  const durationList = durations.split(",");
  let endTime = 0;
  for (const [index, name] of nameList.entries()) {
    // `transition-property: none` and `animation-name: none` mean nothing runs
    // for that item, so it doesn't contribute to the end time even if a delay
    // or duration is still set.
    if (name.trim() === "none") continue;
    const delay = parseCSSTime(delayList[index % delayList.length]);
    const duration = parseCSSTime(durationList[index % durationList.length]);
    endTime = Math.max(endTime, delay + duration);
  }
  return endTime;
}

// Take the longest transition or animation end time for this element.
function getElementEndTime(element: Element) {
  const {
    transitionProperty,
    transitionDuration,
    transitionDelay,
    animationName,
    animationDuration,
    animationDelay,
  } = getComputedStyle(element);
  return Math.max(
    getEndTime(transitionProperty, transitionDelay, transitionDuration),
    getEndTime(animationName, animationDelay, animationDuration),
  );
}

export function isHidden(
  mounted: boolean,
  hidden?: boolean | null,
  alwaysVisible?: boolean | null,
) {
  return !alwaysVisible && hidden !== false && (!mounted || !!hidden);
}

/**
 * Returns props to create a `DislosureContent` component.
 * @see https://ariakit.com/components/disclosure
 * @example
 * ```jsx
 * const store = useDisclosureStore();
 * const props = useDisclosureContent({ store });
 * <Disclosure store={store}>Disclosure</Disclosure>
 * <Role {...props}>Content</Role>
 * ```
 */
export const useDisclosureContent = createHook<
  TagName,
  DisclosureContentOptions
>(function useDisclosureContent({
  store,
  alwaysVisible,
  unstable_otherElementRef: otherElementRef,
  ...props
}) {
  const context = useDisclosureProviderContext();
  store = store || context;

  invariant(
    store,
    process.env.NODE_ENV !== "production" &&
      "DisclosureContent must receive a `store` prop or be wrapped in a DisclosureProvider component.",
  );

  const ref = useRef<HTMLType>(null);
  const id = useId(props.id);
  const [transition, setTransition] = useState<TransitionState>(null);
  const { open, mounted, animated, contentElement } = useStoreStateObject(
    store,
    {
      open: "open",
      mounted: "mounted",
      animated: "animated",
      contentElement: "contentElement",
    },
  );
  const otherElement = useStoreState(store.disclosure, "contentElement");
  const hasClosedRef = useRef(false);

  // This is a workaround to avoid the content element from being reset to null
  // on fast refresh.
  useSafeLayoutEffect(() => {
    if (!ref.current) return;
    store?.setContentElement(ref.current);
  }, [store]);

  // When the disclosure content element is rendered, we automatically set the
  // animated state to true. If there's no enter animation, the animated state
  // will be set to false later on.
  useSafeLayoutEffect(() => {
    let previousAnimated: boolean | number | undefined;
    store?.setState("animated", (animated) => {
      previousAnimated = animated;
      return true;
    });
    return () => {
      if (previousAnimated === undefined) return;
      store?.setState("animated", previousAnimated);
    };
  }, [store]);

  useSafeLayoutEffect(() => {
    if (!animated) {
      // When animation detection has disabled animations (e.g., no CSS
      // transition was detected on the content element), manage data-enter
      // so wrapper elements using :has([data-enter]) work correctly even
      // when the content element itself has no transitions.
      if (!open) {
        hasClosedRef.current = true;
        setTransition(null);
      } else if (hasClosedRef.current) {
        // On reopen after animation detection disabled animations, restore
        // data-enter so enter styles are re-applied.
        hasClosedRef.current = false;
        setTransition("enter");
      }
      return;
    }
    // When the disclosure content element is rendered in a portal, we need to
    // wait for the portal to be mounted and connected to the DOM before we
    // can start the animation.
    if (!contentElement?.isConnected) {
      setTransition(null);
      return;
    }
    // Double requestAnimationFrame is necessary here to avoid potential bugs
    // when the data attribute is added before the element is fully rendered
    // in the DOM, which wouldn't trigger the animation.
    return afterPaint(() => {
      setTransition(open ? "enter" : mounted ? "leave" : null);
    });
  }, [animated, contentElement, open, mounted]);

  useSafeLayoutEffect(() => {
    if (!store) return;
    if (!animated) return;
    if (!transition) return;
    if (!contentElement) return;
    const stopAnimation = () => store?.setState("animating", false);
    const stopAnimationSync = () => flushSync(stopAnimation);
    // Ignore transition states that don't match the current open state. This
    // may happen because transitions are updated asynchronously using
    // requestAnimationFrame.
    if (transition === "leave" && open) return;
    if (transition === "enter" && !open) return;
    // When the animated state is a number, the user has manually set the
    // animation timeout, so we just respect it.
    if (typeof animated === "number") {
      const timeout = animated;
      return afterTimeout(timeout, stopAnimationSync);
    }
    // Derive a timeout because transition and animation end events may never
    // fire. Use the longest timeline when multiple transitions or animations
    // run.
    const elements = [contentElement];
    // If we're rendering a dialog backdrop, otherElement will be the dialog
    // element itself. We need to consider both the backdrop and the dialog
    // animation/transition durations and delays because the dialog may be
    // animated while the backdrop is not.
    if (otherElement) {
      elements.push(otherElement);
    }
    // Conversely, if we're rendering a dialog, its backdrop may be animated
    // while the dialog itself is not. The backdrop element isn't tracked in
    // the store, so the dialog passes it in through otherElementRef.
    const relatedElement = otherElementRef?.current;
    if (relatedElement) {
      elements.push(relatedElement);
    }
    // The timeout is the longest end time among the content element and the
    // related elements, so none of them gets hidden before its own transition
    // or animation ends.
    const timeout = Math.max(...elements.map(getElementEndTime));
    // If the timeout is zero, there's no animation or transition, either
    // because they weren't defined in the CSS or the duration was explicitly
    // set to zero. In this scenario, we can halt the animation right away
    // and, if we're entering, we can set the animated state to false so
    // the element unmounts immediately on close without waiting for a
    // leave animation.
    if (!timeout) {
      if (transition === "enter") {
        store.setState("animated", false);
      }
      stopAnimation();
      return;
    }
    // Since setTimeout may be delayed by a few milliseconds, we need to
    // subtract a frame duration from the timeout to make sure the animation is
    // stopped right after it ends, preventing flickering.
    const frameRate = 1000 / 60;
    const maxTimeout = Math.max(timeout - frameRate, 0);
    return afterTimeout(maxTimeout, stopAnimationSync);
  }, [
    store,
    animated,
    contentElement,
    otherElement,
    otherElementRef,
    open,
    transition,
  ]);

  props = useWrapElement(
    props,
    (element) => (
      <DialogScopedContextProvider value={store}>
        {element}
      </DialogScopedContextProvider>
    ),
    [store],
  );

  const hidden = isHidden(mounted, props.hidden, alwaysVisible);
  const styleProp = props.style;
  const style = useMemo(() => {
    if (hidden) {
      return { ...styleProp, display: "none" };
    }
    return styleProp;
  }, [hidden, styleProp]);

  props = {
    "data-open": open || undefined,
    "data-enter": transition === "enter" || undefined,
    "data-leave": transition === "leave" || undefined,
    hidden,
    ...props,
    id,
    ref: useMergeRefs(id ? store.setContentElement : null, ref, props.ref),
    style,
  };

  // useDialog computes focusable defaults after this hook. Remove sentinels
  // before they can overwrite those later defaults.
  return removeUndefinedValues(props);
});

const DisclosureContentImpl = forwardRef(function DisclosureContentImpl(
  props: DisclosureContentProps,
) {
  const htmlProps = useDisclosureContent(props);
  return createElement(TagName, htmlProps);
});

/**
 * Renders an element that can be shown or hidden by a
 * [`Disclosure`](https://ariakit.com/components/disclosure) component.
 * @see https://ariakit.com/components/disclosure
 * @example
 * ```jsx {3}
 * <DisclosureProvider>
 *   <Disclosure>Disclosure</Disclosure>
 *   <DisclosureContent>Content</DisclosureContent>
 * </DisclosureProvider>
 * ```
 */
export const DisclosureContent = forwardRef(function DisclosureContent({
  unmountOnHide,
  ...props
}: DisclosureContentProps) {
  const context = useDisclosureProviderContext();
  const store = props.store || context;
  const mounted = useStoreState(
    store,
    ["mounted"],
    (state) => !unmountOnHide || state?.mounted,
  );
  if (mounted === false) return null;
  return <DisclosureContentImpl {...props} />;
});

export interface DisclosureContentOptions<
  _T extends ElementType = TagName,
> extends Options {
  /**
   * Object returned by the
   * [`useDisclosureStore`](https://ariakit.com/reference/use-disclosure-store)
   * hook. If not provided, the closest
   * [`DisclosureProvider`](https://ariakit.com/reference/disclosure-provider)
   * component's context will be used.
   */
  store?: DisclosureStore;
  /**
   * A ref to another element whose CSS transitions and animations should be
   * taken into account when computing the animation timeout, such as the
   * dialog's backdrop element, which may be animated while the dialog itself
   * is not.
   * @private
   */
  unstable_otherElementRef?: RefObject<HTMLElement | null>;
  /**
   * Determines whether the content element should remain visible even when the
   * [`open`](https://ariakit.com/reference/disclosure-provider#open) state is
   * `false`. If this prop is set to `true`, the `hidden` prop and the `display:
   * none` style will not be applied, unless explicitly set otherwise.
   *
   * This prop is particularly useful when using third-party animation libraries
   * such as Motion or React Spring, where the element needs to be
   * visible for exit animations to work.
   *
   * Live examples:
   * - [Dialog with
   *   Motion](https://ariakit.com/examples/dialog-framer-motion)
   * - [Menu with
   *   Motion](https://ariakit.com/examples/menu-framer-motion)
   * - [Tooltip with
   *   Motion](https://ariakit.com/examples/tooltip-framer-motion)
   * - [Dialog with details &
   *   summary](https://ariakit.com/examples/dialog-details)
   * @default false
   */
  alwaysVisible?: boolean;
  /**
   * When set to `true`, the content element will be unmounted and removed from
   * the DOM when it's hidden.
   *
   * Live examples:
   * - [Navigation Menubar](https://ariakit.com/examples/menubar-navigation)
   * - [Combobox with integrated
   *   filter](https://ariakit.com/examples/combobox-filtering-integrated)
   * - [Textarea with inline
   *   Combobox](https://ariakit.com/examples/combobox-textarea)
   * - [Standalone Popover](https://ariakit.com/examples/popover-standalone)
   * - [Animated Select](https://ariakit.com/examples/select-animated)
   * - [Multi-Select](https://ariakit.com/examples/select-multiple)
   * @default false
   */
  unmountOnHide?: boolean;
}

export type DisclosureContentProps<T extends ElementType = TagName> = Props<
  T,
  DisclosureContentOptions<T>
>;