// TODO: Expose whether the container is expanded through a data attribute.

import { useStoreState } from "@ariakit/react-store";
import {
  useEvent,
  useMergeRefs,
  createElement,
  createHook,
  forwardRef,
} from "@ariakit/react-utils";
import type { Options, Props } from "@ariakit/react-utils";
import {
  isButton,
  isTextField,
  isFocusEventOutside,
  isSelfTarget,
  disableFocusIn,
  getFirstTabbableIn,
  restoreFocusIn,
  removeUndefinedValues,
} from "@ariakit/utils";
import type { ElementType, FocusEvent, KeyboardEvent, MouseEvent } from "react";
import { useEffect, useRef } from "react";
import { useCompositeScopedContext } from "./composite-context.tsx";
import type { CompositeStore } from "./composite-store.ts";
import { selectTextField } from "./utils.ts";

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

function getFirstTabbable(container: HTMLElement) {
  restoreFocusIn(container);
  const tabbable = getFirstTabbableIn(container);
  disableFocusIn(container);
  return tabbable;
}

/**
 * Returns props to create a `CompositeContainer` component. This component
 * renders interactive widgets inside composite items. This should be used in
 * conjunction with the `CompositeItem` component, the `useCompositeItem` hook,
 * or any other component/hook that uses `CompositeItem` underneath.
 * @see https://ariakit.com/components/composite
 * @example
 * ```jsx
 * const store = useCompositeStore();
 * const props = useCompositeContainer({ store });
 * <Composite store={store}>
 *   <CompositeItem {...props}>
 *     <input type="text" />
 *   </CompositeItem>
 * </Composite>
 * ```
 */
export const useCompositeContainer = createHook<
  TagName,
  CompositeContainerOptions
>(function useCompositeContainer({ store, ...props }) {
  const context = useCompositeScopedContext();
  store = store || context;

  const ref = useRef<HTMLType>(null);
  const isOpenRef = useRef(false);

  const open = (collapseToEnd = false) => {
    const container = ref.current;
    if (!container) return;
    restoreFocusIn(container);
    const tabbable = getFirstTabbableIn(container);
    if (!tabbable) {
      disableFocusIn(container);
      return;
    }
    isOpenRef.current = true;
    queueMicrotask(() => {
      tabbable.focus();
      if (isTextField(tabbable) || tabbable.isContentEditable) {
        selectTextField(tabbable, collapseToEnd);
      }
    });
  };

  const close = () => {
    const container = ref.current;
    if (!container) return;
    isOpenRef.current = false;
    disableFocusIn(container);
  };

  const renderedItems = useStoreState(store, "renderedItems");

  // Disable focus on the tabbable elements inside the container when the
  // container is mounted.
  useEffect(() => {
    const container = ref.current;
    if (!container) return;
    const isOpen = isOpenRef.current;
    // We need to wait for the items to be populated before we can disable
    // focus, so we consider edge cases where some tabbable elements become
    // disabled after the first render (for example, when rendering nested
    // composite elements).
    if (!isOpen && renderedItems?.length) {
      disableFocusIn(container);
    }
  }, [renderedItems]);

  const onFocusProp = props.onFocus;

  const onFocus = useEvent((event: FocusEvent<HTMLType>) => {
    onFocusProp?.(event);
    if (event.defaultPrevented) return;
    if (!store) return;
    const isOpen = isOpenRef.current;
    if (isSelfTarget(event)) {
      // The container element itself has received focus. Here we make an
      // additional step in case tabbable elements have been added lazily to
      // the DOM. We get all containers in the current composite element and
      // disable all tabbable elements inside them.
      isOpenRef.current = false;
      const { compositeElement } = store.getState();
      const selector = "[data-composite-container]";
      const containers =
        compositeElement?.querySelectorAll<HTMLElement>(selector);
      if (containers) {
        for (const container of containers) {
          disableFocusIn(container);
        }
      }
    } else if (!isOpen) {
      // Otherwise, if any element inside the container has received focus,
      // for example, by a direct user click, we should act as the container
      // has been opened.
      isOpenRef.current = true;
      restoreFocusIn(event.currentTarget);
      // Resets the moves in the store so the composite item will not be
      // focused right after the focusable element inside the container gets
      // focus.
      store?.setState("moves", 0);
    }
  });

  const onBlurProp = props.onBlur;

  const onBlur = useEvent((event: FocusEvent<HTMLType>) => {
    onBlurProp?.(event);
    if (event.defaultPrevented) return;
    if (isFocusEventOutside(event)) {
      close();
    }
  });

  const onKeyDownProp = props.onKeyDown;

  const onKeyDown = useEvent((event: KeyboardEvent<HTMLType>) => {
    onKeyDownProp?.(event);
    // https://github.com/ariakit/ariakit/issues/6579
    if (event.nativeEvent.isComposing) return;
    // Safari may send the confirming Enter after compositionend.
    if (event.keyCode === 229) return;
    if (event.defaultPrevented) return;
    if (event.altKey) return;
    if (event.ctrlKey) return;
    if (event.metaKey) return;
    if (event.shiftKey) return;
    const container = event.currentTarget;
    // Non-space printable and deletion keys enter the editable descendant.
    // Escape returns from any descendant; Enter returns from editable targets.
    if (isSelfTarget(event)) {
      if (event.key.length === 1 && event.key !== " ") {
        const tabbable = getFirstTabbable(container);
        if (!tabbable) return;
        if (isTextField(tabbable) || tabbable.isContentEditable) {
          event.stopPropagation();
          open();
        }
      } else if (event.key === "Delete" || event.key === "Backspace") {
        const tabbable = getFirstTabbable(container);
        if (!tabbable) return;
        if (isTextField(tabbable) || tabbable.isContentEditable) {
          open();
          // Move focus back only when this delete/backspace changes the
          // editable element. If it's already empty, no input event fires;
          // leaving the listener armed would let the next input event steal
          // focus back to the container (#6300).
          //
          // Remove the listener on keyup, before the next key can produce
          // input. The timer is only a fallback for cases where keyup never
          // reaches the container, such as when focus leaves while the key is
          // held; relying on the timer alone can race the next input event.
          const cleanup = () => {
            container.removeEventListener("input", onInput);
            container.removeEventListener("keyup", cleanup);
          };
          const onInput = () => {
            cleanup();
            queueMicrotask(() => container.focus());
          };
          container.addEventListener("input", onInput, { once: true });
          container.addEventListener("keyup", cleanup, { once: true });
          setTimeout(cleanup, 0);
        }
      }
    } else if (event.key === "Escape") {
      queueMicrotask(() => container.focus());
    } else if (event.key === "Enter") {
      const target = event.target as HTMLElement;
      const isInput =
        (target.tagName === "INPUT" && !isButton(target)) ||
        target.tagName === "TEXTAREA";
      if (isInput || target.isContentEditable) {
        event.preventDefault();
        queueMicrotask(() => container.focus());
      }
    }
  });

  const onClickProp = props.onClick;

  const onClick = useEvent((event: MouseEvent<HTMLType>) => {
    onClickProp?.(event);
    if (event.defaultPrevented) return;
    if (isSelfTarget(event) && !event.detail) {
      // A zero-detail click focuses the first tabbable descendant and, when
      // editable, places its caret at the end.
      open(true);
    }
  });

  props = {
    "data-composite-container": "",
    ...props,
    ref: useMergeRefs(ref, props.ref),
    onFocus,
    onBlur,
    onKeyDown,
    onClick,
  };

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

/**
 * Renders a container for interactive widgets inside composite items. This
 * should be used in conjunction with the
 * [`CompositeItem`](https://ariakit.com/reference/composite-item) component or
 * a component that uses
 * [`CompositeItem`](https://ariakit.com/reference/composite-item) underneath.
 * @see https://ariakit.com/components/composite
 * @example
 * ```jsx {3-5}
 * <CompositeProvider>
 *   <Composite>
 *     <CompositeItem render={<CompositeContainer />}>
 *       <input type="text" />
 *     </CompositeItem>
 *   </Composite>
 * </CompositeProvider>
 * ```
 */
export const CompositeContainer = forwardRef(function CompositeContainer(
  props: CompositeContainerProps,
) {
  const htmlProps = useCompositeContainer(props);
  return createElement(TagName, htmlProps);
});

export interface CompositeContainerOptions<
  _T extends ElementType = TagName,
> extends Options {
  /**
   * Object returned by the
   * [`useCompositeStore`](https://ariakit.com/reference/use-composite-store)
   * hook. If not provided, the closest
   * [`Composite`](https://ariakit.com/reference/composite) or
   * [`CompositeProvider`](https://ariakit.com/reference/composite-provider)
   * components' context will be used.
   */
  store?: CompositeStore;
}

export type CompositeContainerProps<T extends ElementType = TagName> = Props<
  T,
  CompositeContainerOptions<T>
>;