import {
  batch,
  createStore,
  mergeStore,
  omit,
  setup,
  sync,
} from "@ariakit/store";
import type { Store, StoreOptions, StoreProps } from "@ariakit/store";
import { chain, defaultValue } from "@ariakit/utils";
import type { SetState } from "@ariakit/utils";
import type {
  CollectionStore,
  CollectionStoreItem,
} from "../collection/collection-store.ts";
import { createCollectionStore } from "../collection/collection-store.ts";
import type { ComboboxStore } from "../combobox/combobox-store.ts";
import type {
  CompositeStore,
  CompositeStoreFunctions,
  CompositeStoreItem,
  CompositeStoreOptions,
  CompositeStoreState,
} from "../composite/composite-store.ts";
import { createCompositeStore } from "../composite/composite-store.ts";
import type { SelectStore } from "../select/select-store.ts";

function getFocusedTab(items: TabStoreItem[]) {
  const activeElement = items[0]?.element?.ownerDocument.activeElement;
  if (!activeElement) return;
  return items.find((item) => item.element === activeElement);
}

function getTabById(items: TabStoreItem[], id: string | null | undefined) {
  if (id == null) return;
  return items.find((item) => item.id === id);
}

function isEnabledTab(
  item: TabStoreItem | null | undefined,
): item is TabStoreItem {
  if (!item) return false;
  if (item.disabled) return false;
  if (item.dimmed) return false;
  return true;
}

function createPanelsStore() {
  const panels = createCollectionStore<TabStorePanel>();
  const panelsByTabId = new Map<string, TabStorePanel>();

  sync(panels, ["items"], (state) => {
    panelsByTabId.clear();
    for (const panel of state.items) {
      const { tabId } = panel;
      if (tabId == null) continue;
      if (panelsByTabId.has(tabId)) continue;
      panelsByTabId.set(tabId, panel);
    }
  });

  return {
    panels,
    panel: (tabId) => {
      if (tabId == null) return null;
      return panelsByTabId.get(tabId) || null;
    },
  } satisfies Pick<TabStoreFunctions, "panels" | "panel">;
}

export function createTabStore({
  composite: parentComposite,
  combobox,
  ...props
}: TabStoreProps = {}): TabStore {
  const independentKeys = [
    "items",
    "renderedItems",
    "moves",
    "orientation",
    "virtualFocus",
    "compositeElementInFocusOrder",
    "includesBaseElement",
    "compositeElement",
    "baseElement",
    "focusLoop",
    "focusShift",
    "focusWrap",
    "selectOnMove",
  ] as const;

  const store = mergeStore(
    props.store,
    // Combobox also exposes a selectOnMove state, but it controls the selected
    // value rather than tab selection.
    omit(parentComposite as ComboboxStore | undefined, independentKeys),
    omit(combobox, independentKeys),
  );
  const syncState = store?.getState();

  const composite = createCompositeStore({
    ...props,
    store,
    // We need to explicitly set the default value to `false` since we don't
    // want the composite store to default it to `true` when the activeId state
    // is null, which could be the case when rendering combobox with tab.
    compositeElementInFocusOrder: defaultValue(
      props.compositeElementInFocusOrder,
      props.includesBaseElement,
      syncState?.compositeElementInFocusOrder,
      syncState?.includesBaseElement,
      false,
    ),
    orientation: defaultValue(
      props.orientation,
      syncState?.orientation,
      "horizontal" as const,
    ),
    focusLoop: defaultValue(props.focusLoop, syncState?.focusLoop, true),
  });

  const { panels, panel } = createPanelsStore();

  const initialState: TabStoreState = {
    ...composite.getState(),
    selectedId: defaultValue(
      props.selectedId,
      syncState?.selectedId,
      props.defaultSelectedId,
    ),
    selectOnMove: defaultValue(
      props.selectOnMove,
      syncState?.selectOnMove,
      true,
    ),
  };
  const tab = createStore(initialState, composite, store);

  // Selects the active tab when selectOnMove is true. Since we're listening to
  // the moves state, but not the activeId state, this callback will run only
  // when there's a move, which is usually triggered by moving through the tabs
  // using the keyboard.
  setup(tab, () =>
    sync(tab, ["moves"], () => {
      const { activeId, selectOnMove } = tab.getState();
      if (!selectOnMove) return;
      if (!activeId) return;
      const tabItem = composite.item(activeId);
      if (!isEnabledTab(tabItem)) return;
      tab.setState("selectedId", tabItem.id);
    }),
  );

  let pendingRestore = false;
  let restoredSelectedId: TabStoreState["selectedId"];

  // Keep the active and selected tabs in sync.
  setup(tab, () => {
    // A restore armed right before the store was destroyed (for example, when
    // the popover unmounted) must not leak into this init. Reset before
    // registering the listener below since it also runs on registration.
    pendingRestore = false;
    return batch(tab, ["selectedId"], (state, prev) => {
      // There are cases where we don't want to sync activeId with selectedId.
      // For example, restoring the selectedId from a select or combobox
      // selected value. Batch listeners are microtask-coalesced, so a restore
      // may share a flush with a later legitimate change. In that case, the
      // flushed value differs from the restored one and we must still sync.
      if (pendingRestore) {
        pendingRestore = false;
        if (state.selectedId === restoredSelectedId) return;
      }
      // If there's a parent composite widget, we don't need to sync the
      // activeId state with the initial selectedId state. The parent composite
      // widget should handle the initial activeId state.
      if (parentComposite && state.selectedId === prev.selectedId) return;
      const { activeId, renderedItems } = tab.getState();
      if (activeId === state.selectedId) return;
      const focusedTab = getFocusedTab(renderedItems);
      const selectedTab = getTabById(renderedItems, state.selectedId);
      if (focusedTab && isEnabledTab(selectedTab)) {
        composite.move(selectedTab.id);
        return;
      }
      tab.setState("activeId", state.selectedId);
    });
  });

  // Initialize selection from the active or first enabled tab.
  setup(tab, () =>
    sync(tab, ["selectedId", "renderedItems"], (state) => {
      if (state.selectedId !== undefined) return;
      const { activeId, renderedItems } = tab.getState();
      const tabItem = composite.item(activeId);
      if (isEnabledTab(tabItem)) {
        tab.setState("selectedId", tabItem.id);
      } else {
        const tabItem = renderedItems.find(isEnabledTab);
        tab.setState("selectedId", tabItem?.id);
      }
    }),
  );

  // Pair orphan panels with tabs by index.
  setup(tab, () =>
    sync(tab, ["renderedItems"], (state) => {
      const tabs = state.renderedItems;
      if (!tabs.length) return;
      return sync(panels, ["renderedItems"], (state) => {
        const items = state.renderedItems;
        const hasOrphanPanels = items.some((panel) => !panel.tabId);
        if (!hasOrphanPanels) return;
        items.forEach((panel, i) => {
          if (panel.tabId) return;
          const tabItem = tabs[i];
          if (!tabItem) return;
          panels.renderItem({ ...panel, tabId: tabItem.id });
        });
      });
    }),
  );

  // Preserve the selected tab when a select or combobox value is selected
  // within the tab panel.
  let selectedIdFromSelectedValue: string | null | undefined = null;

  setup(tab, () => {
    const backupSelectedId = () => {
      selectedIdFromSelectedValue = tab.getState().selectedId;
    };
    const restoreSelectedId = () => {
      // We suppress the activeId sync to prevent the activeId state from
      // being set to the selectedId state since this is just a restoration of
      // the selectedId state from a select or combobox selected value.
      // setState early-returns on unchanged values and emits no batch event
      // to consume the suppression, so only arm it when the restore will
      // actually change the state.
      const { selectedId } = tab.getState();
      if (selectedId === selectedIdFromSelectedValue) return;
      pendingRestore = true;
      restoredSelectedId = selectedIdFromSelectedValue;
      tab.setState("selectedId", selectedIdFromSelectedValue);
    };
    // The combobox store is also passed as the parent composite store when tabs
    // are rendered inside a combobox, and it exposes setSelectElement as well.
    // Its value state is the input value rather than the selected value, so it
    // is detected by a combobox-only method to make sure it goes through the
    // combobox branch below even when it's only provided as the composite.
    const isParentCombobox =
      !!parentComposite && "setSelectedValue" in parentComposite;
    const parentCombobox = isParentCombobox
      ? (parentComposite as ComboboxStore)
      : combobox;
    if (
      parentComposite &&
      !isParentCombobox &&
      "setSelectElement" in parentComposite
    ) {
      return chain(
        sync(parentComposite, ["value"], backupSelectedId),
        sync(parentComposite, ["mounted"], restoreSelectedId),
      );
    }
    if (!parentCombobox) return;
    return chain(
      sync(parentCombobox, ["selectedValue"], backupSelectedId),
      sync(parentCombobox, ["mounted"], restoreSelectedId),
    );
  });

  return {
    ...composite,
    ...tab,
    panels,
    panel,
    setSelectedId: (id) => tab.setState("selectedId", id),
    select: (id) => {
      tab.setState("selectedId", id);
      composite.move(id);
    },
  };
}

export interface TabStoreItem extends CompositeStoreItem {
  dimmed?: boolean;
}

export interface TabStorePanel extends CollectionStoreItem {
  tabId?: string | null;
}

export interface TabStoreState extends CompositeStoreState<TabStoreItem> {
  /** @default "horizontal" */
  orientation: CompositeStoreState<TabStoreItem>["orientation"];
  /** @default true */
  focusLoop: CompositeStoreState<TabStoreItem>["focusLoop"];
  /**
   * The id of the tab whose panel is currently visible. If it's `undefined`, it
   * will be automatically set to the first enabled tab.
   *
   * Live examples:
   * - [Tab with React Router](https://ariakit.com/examples/tab-react-router)
   * - [Combobox with Tabs](https://ariakit.com/examples/combobox-tabs)
   * - [Select with Combobox and
   *   Tabs](https://ariakit.com/examples/select-combobox-tab)
   * - [Command Menu with
   *   Tabs](https://ariakit.com/examples/dialog-combobox-tab-command-menu)
   */
  selectedId: TabStoreState["activeId"];
  /**
   * Determines if the tab should be selected when it receives focus. If set to
   * `false`, the tab will only be selected upon clicking, not when using arrow
   * keys to shift focus.
   *
   * Live examples:
   * - [Tab with React Router](https://ariakit.com/examples/tab-react-router)
   * - [Select with Combobox and
   *   Tabs](https://ariakit.com/examples/select-combobox-tab)
   * @default true
   */
  selectOnMove?: boolean;
}

export interface TabStoreFunctions extends CompositeStoreFunctions<TabStoreItem> {
  /**
   * Sets the
   * [`selectedId`](https://ariakit.com/reference/tab-provider#selectedid) state.
   * If another tab has DOM focus and the selected tab is enabled, focus will
   * move to the selected tab. If you want to always move focus, use the
   * [`select`](https://ariakit.com/reference/use-tab-store#select) function
   * instead.
   * @example
   * // Selects the tab with id "tab-1"
   * store.setSelectedId("tab-1");
   * // Toggles between "tab-1" and "tab-2"
   * store.setSelectedId((id) => id === "tab-1" ? "tab-2" : "tab-1"));
   * // Selects the first tab
   * store.setSelectedId(store.first());
   * // Selects the next tab
   * store.setSelectedId(store.next());
   */
  setSelectedId: SetState<TabStoreState["selectedId"]>;
  /**
   * A collection store containing the tab panels.
   *
   * Live examples:
   * - [Animated TabPanel](https://ariakit.com/examples/tab-panel-animated)
   */
  panels: CollectionStore<TabStorePanel>;
  /**
   * Gets the panel associated with the given tab id.
   * @example
   * const panel = store.panel("tab-1");
   */
  panel: (tabId: string | null | undefined) => TabStorePanel | null;
  /**
   * Selects the tab for the given id and moves focus to it. If you want to set
   * the [`selectedId`](https://ariakit.com/reference/tab-provider#selectedid)
   * state without always moving focus, use the
   * [`setSelectedId`](https://ariakit.com/reference/use-tab-store#setselectedid-1)
   * function instead.
   * @example
   * // Selects the tab with id "tab-1"
   * store.select("tab-1");
   * // Selects the first tab
   * store.select(store.first());
   * // Selects the next tab
   * store.select(store.next());
   */
  select: TabStore["move"];
}

export interface TabStoreOptions
  extends
    StoreOptions<
      TabStoreState,
      "orientation" | "focusLoop" | "selectedId" | "selectOnMove"
    >,
    CompositeStoreOptions<TabStoreItem> {
  /**
   * A reference to another [composite
   * store](https://ariakit.com/reference/use-composite-store). This is used when
   * rendering tabs as part of another composite widget such as
   * [Combobox](https://ariakit.com/components/combobox) or
   * [Select](https://ariakit.com/components/select). The stores will share the
   * same state.
   */
  composite?: CompositeStore | SelectStore | null;
  /**
   * A reference to a [combobox
   * store](https://ariakit.com/reference/use-combobox-store). This is used when
   * rendering tabs inside a
   * [Combobox](https://ariakit.com/components/combobox).
   */
  combobox?: ComboboxStore | null;
  /**
   * The id of the tab whose panel is currently visible. If it's `undefined`, it
   * will be automatically set to the first enabled tab.
   *
   * Live examples:
   * - [Combobox with Tabs](https://ariakit.com/examples/combobox-tabs)
   * - [Animated TabPanel](https://ariakit.com/examples/tab-panel-animated)
   * - [Select with Combobox and
   *   Tabs](https://ariakit.com/examples/select-combobox-tab)
   * - [Command Menu with
   *   Tabs](https://ariakit.com/examples/dialog-combobox-tab-command-menu)
   */
  defaultSelectedId?: TabStoreState["selectedId"];
}

export interface TabStoreProps
  extends TabStoreOptions, StoreProps<TabStoreState> {}

export interface TabStore extends TabStoreFunctions, Store<TabStoreState> {}