import { useStoreStateObject } from "@ariakit/react-store";
import {
  useBooleanEvent,
  useEvent,
  useId,
  useMergeRefs,
  useWrapElement,
  createElement,
  createHook,
  forwardRef,
  memo,
} from "@ariakit/react-utils";
import type { Props } from "@ariakit/react-utils";
import { subscribe } from "@ariakit/store";
import {
  getActiveElement,
  getScrollingElement,
  getTextboxSelection,
  getTextboxValue,
  isButton,
  isFocusable,
  isTextbox,
  isTextField,
  isPortalEvent,
  isSelfTarget,
  disabledFromProps,
  warnOnce,
} from "@ariakit/utils";
import type { BooleanOrCallback } from "@ariakit/utils";
import type {
  ElementType,
  FocusEvent,
  KeyboardEvent,
  SyntheticEvent,
} from "react";
import { useCallback, useContext, useMemo, useRef } from "react";
import { withDefaultButtonType } from "../button/utils.ts";
import type { CollectionItemOptions } from "../collection/collection-item.tsx";
import { useCollectionItem } from "../collection/collection-item.tsx";
import type { CommandOptions } from "../command/command.tsx";
import { useCommand } from "../command/command.tsx";
import {
  CompositeItemContext,
  CompositeRowContext,
  useCompositeScopedContext,
} from "./composite-context.tsx";
import type { CompositeStore, CompositeStoreState } from "./composite-store.ts";
import {
  focusSilently,
  getEnabledItem,
  isItem,
  markItemMounted,
  markItemUnmounting,
  selectTextField,
  usePresentItem,
} from "./utils.ts";

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

function isEditableElement(element: HTMLElement) {
  if (isTextbox(element)) return true;
  return element.tagName === "INPUT" && !isButton(element);
}

function getNextPageOffset(scrollingElement: Element, pageUp = false) {
  const height = scrollingElement.clientHeight;
  const { top } = scrollingElement.getBoundingClientRect();
  // Calculates the size of the page based on the scrolling element's height.
  // This is similar to how browsers calculate the scroll position when pressing
  // spacebar, page up, or page down.
  const pageSize = Math.max(height * 0.875, height - 40) * 1.5;
  const pageOffset = pageUp ? height - pageSize + top : pageSize + top;
  if (scrollingElement.tagName === "HTML") {
    return pageOffset + scrollingElement.scrollTop;
  }
  return pageOffset;
}

function getItemOffset(itemElement: Element, pageUp = false) {
  const { top } = itemElement.getBoundingClientRect();
  if (pageUp) {
    // PageUp is always the inverse of PageDown. On PageDown, we consider only
    // the top offset of the element. On PageUp we need to add the height of the
    // element as well so we consider the bottom of it.
    return top + itemElement.clientHeight;
  }
  return top;
}

function findNextPageItemId(
  element: Element,
  store?: CompositeStore,
  next?: CompositeStore["next"],
  pageUp = false,
) {
  if (!store) return;
  if (!next) return;
  const { renderedItems } = store.getState();
  const scrollingElement = getScrollingElement(element);
  if (!scrollingElement) return;
  const nextPageOffset = getNextPageOffset(scrollingElement, pageUp);
  let id: string | null | undefined;
  let prevDifference: number | undefined;
  // We need to loop through the next items to find the one that is closest to
  // the next page offset.
  for (let i = 0; i < renderedItems.length; i += 1) {
    const previousId = id;
    id = next(i);
    if (!id) break;
    if (id === previousId) continue;
    const itemElement = getEnabledItem(store, id)?.element;
    if (!itemElement) continue;
    const itemOffset = getItemOffset(itemElement, pageUp);
    const difference = itemOffset - nextPageOffset;
    const absDifference = Math.abs(difference);
    // On PageUp, the element is at the next page if the difference between its
    // top offset (plus its height) and the next page offset is less than or
    // equal zero. On PageDown, the difference should be greater than or equal
    // zero.
    if ((pageUp && difference <= 0) || (!pageUp && difference >= 0)) {
      // There may be cases when there's a lot of space between the pages, for
      // example, when there's a lot of disabled items. In this case, the first
      // item in the next page might not be the closest one. So we return the
      // previous item id if its difference is less than the current one.
      if (prevDifference !== undefined && prevDifference < absDifference) {
        id = previousId;
      }
      break;
    }
    prevDifference = absDifference;
  }
  return id;
}

function targetIsAnotherItem(event: SyntheticEvent, store: CompositeStore) {
  if (isSelfTarget(event)) return false;
  return isItem(store, event.target as HTMLElement);
}

/**
 * Returns props to create a `CompositeItem` component.
 * @see https://ariakit.com/components/composite
 * @example
 * ```jsx
 * const store = useCompositeStore();
 * const props = useCompositeItem({ store });
 * <Role {...props}>Item 1</Role>
 * ```
 */
export const useCompositeItem = createHook<TagName, CompositeItemOptions>(
  function useCompositeItem({
    store,
    rowId: rowIdProp,
    preventScrollOnKeyDown = false,
    moveOnKeyPress = true,
    tabbable = false,
    getItem: getItemProp,
    typeaheadText,
    "aria-setsize": ariaSetSizeProp,
    "aria-posinset": ariaPosInSetProp,
    unstable_scrollIntoView: scrollIntoView,
    ...props
  }) {
    const context = useCompositeScopedContext();
    store = store || context;

    const id = useId(props.id);
    const ref = useRef<HTMLType>(null);
    const mountedElementRef = useRef<HTMLType>(null);
    const markUnmountingRef = useCallback((element: HTMLType | null) => {
      const mountedElement = mountedElementRef.current;
      if (!element && mountedElement) markItemUnmounting(mountedElement);
      if (element) markItemMounted(element);
      mountedElementRef.current = element;
    }, []);
    const row = useContext(CompositeRowContext);
    const disabled = disabledFromProps(props);
    const trulyDisabled = disabled && !props.accessibleWhenDisabled;
    // Snapshot before the props object is replaced below (useCollectionItem
    // consumes this prop), so the onFocus handler can read it at event time.
    const shouldRegisterItem = props.shouldRegisterItem;

    // Sibling selectors below (ariaPosInSet) also need the row id during
    // render. They can't read the destructured rowId const since it's declared
    // by the same statement they're arguments to, which would hit the temporal
    // dead zone. See https://github.com/ariakit/ariakit/issues/6334
    const getRowId = (
      state?: Pick<CompositeStoreState, "compositeElement">,
    ) => {
      if (rowIdProp) return rowIdProp;
      if (!state) return;
      if (!row?.compositeElement) return;
      if (row.compositeElement !== state.compositeElement) return;
      return row.id;
    };

    const { rowId, compositeElement, ariaSetSize, ariaPosInSet } =
      useStoreStateObject(store, ["compositeElement", "renderedItems"], {
        rowId: getRowId,
        compositeElement(state) {
          return state?.compositeElement || undefined;
        },
        ariaSetSize(state) {
          if (ariaSetSizeProp != null) return ariaSetSizeProp;
          if (!state) return;
          if (!row?.ariaSetSize) return;
          if (row.compositeElement !== state.compositeElement) return;
          return row.ariaSetSize;
        },
        ariaPosInSet(state) {
          if (ariaPosInSetProp != null) return ariaPosInSetProp;
          if (!state) return;
          if (!row?.ariaPosInSet) return;
          if (row.compositeElement !== state.compositeElement) return;
          const rowId = getRowId(state);
          const itemsInRow = state.renderedItems.filter(
            (item) => item.rowId === rowId,
          );
          return (
            row.ariaPosInSet + itemsInRow.findIndex((item) => item.id === id)
          );
        },
      });

    const { isActiveItem, isTabbable } = useStoreStateObject(
      store,
      [
        "activeId",
        "compositeElement",
        "renderedItems",
        "virtualFocus",
        "items",
      ],
      {
        isActiveItem(state) {
          return !!state && state.activeId === id;
        },
        isTabbable(state) {
          if (!state) return true;
          // The composite element is published in a layout effect, one commit
          // after the items mount, while rendered items are only published on
          // the next animation frame. Both are empty before hydration and on
          // the first commit, which is when items must keep their native tab
          // order. Both conditions are necessary: a composite store may never
          // get a composite element, and its items must still roam.
          if (!state.compositeElement && !state.renderedItems.length) {
            return true;
          }
          if (state.virtualFocus) return false;
          if (!state.renderedItems.length) return true;
          if (tabbable) return true;
          if (state.activeId === null) return false;
          // If activeId refers to an item that's disabled or not connected to the
          // DOM, we make all items tabbable so users can tab into the composite
          // widget. Once the activeId is valid, we restore the roving tabindex. See
          // https://github.com/ariakit/ariakit/issues/3232
          // https://github.com/ariakit/ariakit/issues/4129
          const item = store?.item(state.activeId);
          if (item?.disabled) return true;
          if (!item?.element) return true;
          return state.activeId === id;
        },
      },
    );

    const getItem = useCallback<NonNullable<CollectionItemOptions["getItem"]>>(
      (item) => {
        const nextItem = {
          ...item,
          id: id || item.id,
          rowId,
          disabled: trulyDisabled,
          children: item.element?.textContent,
          typeaheadText,
        };
        if (getItemProp) {
          return getItemProp(nextItem);
        }
        return nextItem;
      },
      [id, rowId, trulyDisabled, typeaheadText, getItemProp],
    );

    const onFocusProp = props.onFocus;
    const hasFocusedComposite = useRef(false);
    // This event-created subscription cannot use effect cleanup: StrictMode's
    // simulated unmount would cancel it permanently. It self-cancels when focus
    // leaves or item unregistration updates the store.
    const cancelScheduledFocusRedirectRef = useRef<(() => void) | null>(null);
    const present = usePresentItem(store);

    const onFocus = useEvent((event: FocusEvent<HTMLType>) => {
      onFocusProp?.(event);
      if (event.defaultPrevented) return;
      if (isPortalEvent(event)) return;
      if (!id) return;
      if (!store) return;
      // If the target is another item, this probably means that composite items
      // are nested. This is okay when building, for example, tree or treegrid
      // elements. In this case, we just ignore the focus event on this parent
      // item.
      if (targetIsAnotherItem(event, store)) return;
      const { virtualFocus, compositeElement } = store.getState();
      store.setActiveId(id);
      // If the composite item is a text field, we'll select its content when
      // focused. This guarantees that pressing arrow keys will move to the
      // previous/next composite items instead of moving the cursor inside the
      // text field.
      if (isTextbox(event.currentTarget)) {
        selectTextField(event.currentTarget);
      }
      // When using aria-activedescendant, we want to make sure that the
      // composite container receives focus, not the composite item.
      if (!virtualFocus) {
        // DOM focus moved with preventScroll, so a marked item still needs the
        // widget's presentation even without the virtual-focus handoff below.
        if (isSelfTarget(event) && store.item(id)) {
          present({
            id,
            markedOnly: true,
            requireFocus: true,
            scrollIntoView,
          });
        }
        return;
      }
      // But we'll only do this if the focused element is the composite item
      // itself
      if (!isSelfTarget(event)) return;
      // and the composite item is not a text field or contenteditable element.
      if (isEditableElement(event.currentTarget)) return;

      const redirectFocusToCompositeElement = (
        relatedTarget: Element | null,
        compositeElement: HTMLElement,
      ) => {
        if (!isFocusable(compositeElement)) {
          if (process.env.NODE_ENV !== "production") {
            warnOnce(
              "A composite widget with `virtualFocus` enabled requires a " +
                "focusable composite element. Set the `focusable` prop to " +
                "`true` or the `virtualFocus` option to `false`.",
              compositeElement,
            );
          }
          return;
        }
        const fromComposite =
          relatedTarget === compositeElement || isItem(store, relatedTarget);
        // Virtual-focus handoff prevents the browser from scrolling the item,
        // so present registered, marked targets after popup placement. The
        // internal item-to-composite move does not abandon the request.
        if (store.item(id)) {
          present({
            id,
            markedOnly: true,
            requireFocus: true,
            scrollIntoView,
          });
        }
        hasFocusedComposite.current = true;
        // If the previously focused element is a composite or composite item
        // component, we'll transfer focus silently to the composite element.
        // That's because this is just a transition event, the composite
        // element was likely already focused, so we're just immediately
        // returning focus to it when navigating through the items.
        if (fromComposite) {
          focusSilently(compositeElement);
        }

        // Otherwise, the composite element is likely not focused, so we need
        // this focus event to propagate so consumers can use the onFocus prop
        // on <Composite>. Scrolling is left to whoever is presenting this item:
        // the composite element can be inside a popup that hasn't been
        // positioned yet, and letting the browser scroll to it moves the page.
        else {
          compositeElement.focus({ preventScroll: true });
        }
      };

      if (compositeElement?.isConnected) {
        redirectFocusToCompositeElement(event.relatedTarget, compositeElement);
        return;
      }

      // The composite element may be published a commit later or be temporarily
      // disconnected. Wait for a connected element before redirecting focus.
      // https://github.com/ariakit/ariakit/issues/6623

      // Items that opt out of registering themselves in the store never
      // produce the unregister store update that the scheduled redirect below
      // relies on to self-clean, so they keep the previous behavior of
      // dropping the redirect.
      if (shouldRegisterItem === false) return;
      const { currentTarget, relatedTarget } = event;
      const cancelScheduledFocusRedirect = () => {
        cancelScheduledFocusRedirectRef.current?.();
        cancelScheduledFocusRedirectRef.current = null;
      };
      cancelScheduledFocusRedirect();
      // Subscribe to every store update, not just compositeElement changes, so
      // the pending redirect is also discarded when the item unmounts without
      // a composite element ever arriving.
      cancelScheduledFocusRedirectRef.current = subscribe(store, null, () => {
        // The redirect is no longer relevant if the item lost DOM focus in
        // the meantime, including when it was unmounted.
        if (getActiveElement(currentTarget) !== currentTarget) {
          cancelScheduledFocusRedirect();
          return;
        }
        const state = store.getState();
        const nextCompositeElement = state.compositeElement;
        // Keep waiting until a connected composite element is stored.
        if (!nextCompositeElement?.isConnected) return;
        cancelScheduledFocusRedirect();
        if (!state.virtualFocus) return;
        redirectFocusToCompositeElement(relatedTarget, nextCompositeElement);
      });
    });

    const onBlurCaptureProp = props.onBlurCapture;

    const onBlurCapture = useEvent((event: FocusEvent<HTMLType>) => {
      onBlurCaptureProp?.(event);
      if (event.defaultPrevented) return;
      const state = store?.getState();
      if (state?.virtualFocus && hasFocusedComposite.current) {
        // When hasFocusedComposite is true, composite has been focused right
        // after focusing on this item. This is an intermediate blur event, so
        // we ignore it.
        hasFocusedComposite.current = false;
        event.preventDefault();
        event.stopPropagation();
      }
    });

    const onKeyDownProp = props.onKeyDown;
    const preventScrollOnKeyDownProp = useBooleanEvent(preventScrollOnKeyDown);
    const moveOnKeyPressProp = useBooleanEvent(moveOnKeyPress);

    const onKeyDown = useEvent((event: KeyboardEvent<HTMLType>) => {
      onKeyDownProp?.(event);
      if (event.defaultPrevented) return;
      if (!isSelfTarget(event)) return;
      if (!store) return;
      const { currentTarget } = event;
      const state = store.getState();
      const item = store.item(id);
      const isGrid = !!item?.rowId;
      const isVertical = state.orientation !== "horizontal";
      const isHorizontal = state.orientation !== "vertical";
      // If the composite element is a text field, the Home and End keys should
      // be performed on the text field, not the composite item, unless the
      // composite is a grid or has a horizontal orientation.
      const canHomeEnd = () => {
        if (isGrid) return true;
        if (isHorizontal) return true;
        if (!state.compositeElement) return true;
        if (!isTextField(state.compositeElement)) return true;
        return false;
      };
      const keyMap = {
        ArrowUp: (isGrid || isVertical) && store.up,
        ArrowRight: (isGrid || isHorizontal) && store.next,
        ArrowDown: (isGrid || isVertical) && store.down,
        ArrowLeft: (isGrid || isHorizontal) && store.previous,
        Home: () => {
          if (!canHomeEnd()) return;
          if (!isGrid || event.ctrlKey) {
            return store?.first();
          }
          return store?.previous(-1);
        },
        End: () => {
          if (!canHomeEnd()) return;
          if (!isGrid || event.ctrlKey) {
            return store?.last();
          }
          return store?.next(-1);
        },
        PageUp: () => {
          return findNextPageItemId(currentTarget, store, store?.up, true);
        },
        PageDown: () => {
          return findNextPageItemId(currentTarget, store, store?.down);
        },
      };
      const action = keyMap[event.key as keyof typeof keyMap];
      if (action) {
        // If the composite item is a textbox, we'll only move focus to the
        // previous/next composite items when the cursor is at the beginning or
        // end of the text. This is to avoid moving focus when the user is
        // navigating through the text.
        if (isTextbox(currentTarget)) {
          const selection = getTextboxSelection(currentTarget);
          const isLeft = isHorizontal && event.key === "ArrowLeft";
          const isRight = isHorizontal && event.key === "ArrowRight";
          const isUp = isVertical && event.key === "ArrowUp";
          const isDown = isVertical && event.key === "ArrowDown";
          if (isRight || isDown) {
            const { length: valueLength } = getTextboxValue(currentTarget);
            if (selection.end !== valueLength) return;
          } else if ((isLeft || isUp) && selection.start !== 0) return;
        }
        const nextId = action();
        if (preventScrollOnKeyDownProp(event) || nextId !== undefined) {
          if (!moveOnKeyPressProp(event)) return;
          event.preventDefault();
          store.move(nextId);
        }
      }
    });

    const providerValue = useMemo(
      () => ({ id, compositeElement }),
      [id, compositeElement],
    );

    props = useWrapElement(
      props,
      (element) => (
        <CompositeItemContext.Provider value={providerValue}>
          {element}
        </CompositeItemContext.Provider>
      ),
      [providerValue],
    );

    props = {
      "data-active-item": isActiveItem || undefined,
      ...props,
      id,
      ref: useMergeRefs(ref, markUnmountingRef, props.ref),
      tabIndex: isTabbable ? props.tabIndex : -1,
      onFocus,
      onBlurCapture,
      onKeyDown,
    };

    props = useCommand(props);
    props = useCollectionItem<TagName>({
      store,
      ...props,
      getItem,
      shouldRegisterItem: id ? shouldRegisterItem : false,
    });

    return {
      ...props,
      "aria-setsize": ariaSetSize,
      "aria-posinset": ariaPosInSet,
    };
  },
);

/**
 * Renders a focusable item as part of a composite widget. The `tabindex`
 * attribute is automatically managed by this component based on the
 * [`virtualFocus`](https://ariakit.com/reference/composite-provider#virtualfocus)
 * option.
 *
 * When this component receives DOM focus or is virtually focused (when the
 * [`virtualFocus`](https://ariakit.com/reference/composite-provider#virtualfocus)
 * option is set to `true`), the element will automatically receive the
 * [`data-active-item`](https://ariakit.com/guide/styling#data-active-item)
 * attribute. This can be used to style the focused item, no matter the focus
 * approach employed.
 * @see https://ariakit.com/components/composite
 * @example
 * ```jsx {3-5}
 * <CompositeProvider>
 *   <Composite>
 *     <CompositeItem>Item 1</CompositeItem>
 *     <CompositeItem>Item 2</CompositeItem>
 *     <CompositeItem>Item 3</CompositeItem>
 *   </Composite>
 * </CompositeProvider>
 * ```
 */
export const CompositeItem = memo(
  forwardRef(function CompositeItem(props: CompositeItemProps) {
    const htmlProps = useCompositeItem(withDefaultButtonType(props));
    return createElement(TagName, htmlProps);
  }),
);

export interface CompositeItemOptions<T extends ElementType = TagName>
  extends CommandOptions<T>, CollectionItemOptions<T> {
  /**
   * 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;
  /**
   * Determines how the item is scrolled into view when it's presented.
   * @private
   */
  unstable_scrollIntoView?: (element: HTMLElement) => void;
  /**
   * Determines if the item should be registered as part of the collection. If
   * this is set to `false`, the item won't be accessible via arrow keys.
   */
  shouldRegisterItem?: CollectionItemOptions<T>["shouldRegisterItem"];
  /**
   * The id that will be used to group items in the same row. This is usually
   * retrieved by the
   * [`CompositeRow`](https://ariakit.com/reference/composite-row) component
   * through context so in most cases you don't need to set it manually.
   *
   * Live examples:
   * - [Command Menu with
   *   Tabs](https://ariakit.com/examples/dialog-combobox-tab-command-menu)
   */
  rowId?: string;
  /**
   * Whether the scroll behavior should be prevented when pressing arrow keys on
   * the first or the last items.
   * @deprecated Use CSS
   * [`scroll-margin`](https://developer.mozilla.org/en-US/docs/Web/CSS/scroll-margin)
   * instead.
   * @default false
   */
  preventScrollOnKeyDown?: BooleanOrCallback<KeyboardEvent<HTMLElement>>;
  /**
   * Determines if pressing arrow keys while this item is in focus should move
   * focus to a different item.
   *
   * **Note**: To entirely disable focus moving within a composite widget, you
   * can use the
   * [`focusOnMove`](https://ariakit.com/reference/composite#focusonmove) prop
   * on the composite component instead. If you want to control the behavior
   * _only when arrow keys are pressed_, where
   * [`focusOnMove`](https://ariakit.com/reference/composite#focusonmove) may
   * not be applicable, this prop must be set on all composite items because
   * they each manage their own key presses, as well as on the composite
   * component itself.
   * @default true
   * @example
   * ```jsx
   * <Composite moveOnKeyPress={false}>
   *   <CompositeItem moveOnKeyPress={false} />
   *   <CompositeItem moveOnKeyPress={false} />
   * </Composite>
   * ```
   */
  moveOnKeyPress?: BooleanOrCallback<KeyboardEvent<HTMLElement>>;
  /**
   * When the `tabbable` prop is set to `true`, the [roving
   * tabindex](https://www.w3.org/WAI/ARIA/apg/practices/keyboard-interface/#kbd_roving_tabindex)
   * method is partially disabled for this element. This means that the
   * `tabIndex` prop won't be assigned `-1` when the item is inactive. In
   * addition to using arrow keys, users will be able to tab to this element,
   * leading to the composite widget no longer existing as a single tab stop.
   *
   * As per the [ARIA spec](https://w3c.github.io/aria/#composite):
   *
   * > Authors **SHOULD** ensure that a composite widget exists as a single
   * > navigation stop within the larger navigation system of the web page.
   *
   * Additionally, as stated in
   * [RFC-2119](https://www.rfc-editor.org/rfc/rfc2119.txt):
   *
   * > **SHOULD** This word, or the adjective "RECOMMENDED", mean that there may
   * > exist valid reasons in particular circumstances to ignore a particular
   * > item, but the full implications must be understood and carefully weighed
   * > before choosing a different course.
   *
   * Therefore, while this may be allowed, you should think carefully about the
   * implications of using this prop.
   *
   * **Note**: This prop has no effect when the
   * [`virtualFocus`](https://ariakit.com/reference/composite-provider#virtualfocus)
   * option is enabled.
   *
   * Live examples:
   * - [Navigation Menubar](https://ariakit.com/examples/menubar-navigation)
   */
  tabbable?: boolean;
  /**
   * The text used by typeahead to match this item. Use this when the rendered
   * text starts with decorative or custom content.
   *
   * Set this to an empty string to exclude the item from typeahead matching.
   * @example
   * ```jsx
   * <CompositeItem typeaheadText="Canada">🇨🇦 Canada</CompositeItem>
   * ```
   */
  typeaheadText?: string;
}

export type CompositeItemProps<T extends ElementType = TagName> = Props<
  T,
  CompositeItemOptions<T>
>;