import { useStoreState } from "@ariakit/react-store";
import { useId, createElement, forwardRef } from "@ariakit/react-utils";
import type { Props } from "@ariakit/react-utils";
import type { ElementType, ReactNode } from "react";
import { useMemo } from "react";
import type {
  CollectionRendererBaseItemProps,
  CollectionRendererItem,
  CollectionRendererItemObject,
  CollectionRendererItemProps,
  CollectionRendererOptions,
} from "../collection/collection-renderer.tsx";
import {
  getCollectionRendererItem,
  getCollectionRendererItemId,
  useCollectionRenderer,
} from "../collection/collection-renderer.tsx";
import type { CollectionStoreItem } from "../collection/collection-store.ts";
import type { SelectStore } from "../select/select-store.ts";
import { useCompositeScopedContext } from "./composite-context.tsx";
import type { CompositeStore, CompositeStoreItem } from "./composite-store.ts";

const TagName = "div" satisfies ElementType;
type TagName = typeof TagName;

interface ItemObject extends CollectionRendererItemObject {
  disabled?: boolean;
}

type Item = ItemObject | CollectionRendererItem;

interface BaseItemProps extends CollectionRendererBaseItemProps {
  "aria-setsize": number;
  "aria-posinset": number;
}

type ItemProps<
  T extends Item,
  P extends BaseItemProps = BaseItemProps,
> = CollectionRendererItemProps<T, P>;

function getItemObject(item: Item): ItemObject {
  if (!item || typeof item !== "object") {
    return { value: item };
  }
  return item;
}

function countItems(items?: number | readonly Item[]): number[] {
  if (!items) return [0];
  if (typeof items === "number") {
    return Array.from({ length: items }, (_, index) => index + 1);
  }
  return items.reduce<number[]>((count, item, index) => {
    const prevCount = count[index - 1] ?? 0;
    const object = getItemObject(item);
    if (!object.items) {
      count[index] = prevCount + 1;
      return count;
    }
    const itemsCount = countItems(object.items)[object.items.length - 1] ?? 0;
    count[index] = prevCount + itemsCount;
    return count;
  }, []);
}

function findFirst(items: readonly Item[], offset = 1): number {
  for (
    let index = offset > 0 ? 0 : items.length - 1;
    index >= 0 && index < items.length;
    index += offset
  ) {
    const item = items[index];
    const object = getItemObject(item);
    if (object.items && findFirst(object.items, offset) !== -1) return index;
    if (!object.disabled) return index;
  }
  return -1;
}

function findLast(items: readonly Item[]) {
  return findFirst(items, -1);
}

function findExactById(
  items: readonly Item[],
  id: string,
  baseId: string,
): number {
  return items.findIndex((item, index) => {
    const itemId = getCollectionRendererItemId(item, index, baseId);
    if (itemId === id) return true;
    const object = getItemObject(item);
    if (object.items?.length) {
      return findExactById(object.items, id, itemId) !== -1;
    }
    return false;
  });
}

interface FindPrefixByIdResult {
  index: number;
  length: number;
}

function findPrefixById(
  items: readonly Item[],
  id: string,
  baseId: string,
): FindPrefixByIdResult {
  const result = { index: -1, length: -1 };

  items.forEach((item, index) => {
    const itemId = getCollectionRendererItemId(item, index, baseId);
    let length = id.startsWith(`${itemId}/`) ? itemId.length : -1;
    const object = getItemObject(item);
    if (object.items?.length) {
      const nestedResult = findPrefixById(object.items, id, itemId);
      length = Math.max(length, nestedResult.length);
    }
    if (length > result.length) {
      result.index = index;
      result.length = length;
    }
  });

  return result;
}

function findById(items: readonly Item[], id: string, baseId: string): number {
  const exactIndex = findExactById(items, id, baseId);
  if (exactIndex !== -1) return exactIndex;
  return findPrefixById(items, id, baseId).index;
}

export function useCompositeRenderer<T extends Item = any>({
  store,
  orientation: orientationProp,
  persistentIndices: persistentIndicesProp,
  children: renderItem,
  "aria-setsize": ariaSetSize,
  "aria-posinset": ariaPosInSet = 1,
  ...props
}: CompositeRendererProps<T>) {
  const context = useCompositeScopedContext();
  store = store || (context as typeof store);

  const orientation = useStoreState(
    store,
    ["orientation"],
    (state) =>
      orientationProp ??
      (state?.orientation === "both" ? "vertical" : state?.orientation),
  );

  // SelectRenderer passes a SelectStore through the base CompositeStore type.
  const stateStore = store as typeof store | SelectStore;
  const items = useStoreState(stateStore, ["mounted", "items"], (state) => {
    if (!state) return props.items;
    if ("mounted" in state && !state.mounted) return 0;
    return props.items ?? (state.items as T[]);
  });

  const id = useId(props.id);

  const itemsCount = useMemo(() => countItems(items), [items]);

  const setSize = useMemo(
    () => ariaSetSize ?? itemsCount[itemsCount.length - 1] ?? 0,
    [ariaSetSize, itemsCount],
  );

  const firstIndex = useMemo(() => {
    if (!items) return -1;
    if (typeof items === "number") return 0;
    if (!items.length) return -1;
    return findFirst(items);
  }, [items]);

  const lastIndex = useMemo(() => {
    if (!items) return -1;
    if (typeof items === "number") return items - 1;
    if (!items.length) return -1;
    return findLast(items);
  }, [items]);

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

  const activeIndex = useMemo(() => {
    if (!id) return -1;
    if (!items) return -1;
    if (activeId == null) return -1;
    if (typeof items === "number") return -1;
    if (!items.length) return -1;
    return findById(items, activeId, id);
  }, [id, items, activeId]);

  const persistentIndices = useMemo(() => {
    const indices = [firstIndex, activeIndex, lastIndex].filter(
      (index) => index >= 0,
    );
    if (persistentIndicesProp) {
      return [...persistentIndicesProp, ...indices];
    }
    return indices;
  }, [firstIndex, activeIndex, lastIndex, persistentIndicesProp]);

  return useCollectionRenderer({
    id,
    store,
    orientation,
    persistentIndices,
    ...props,
    children: (item) => {
      const nextItem = Object.assign({}, item, {
        "aria-setsize": setSize,
        "aria-posinset": ariaPosInSet + (itemsCount[item.index - 1] ?? 0),
      });
      return renderItem?.(nextItem as ItemProps<T>);
    },
  });
}

export {
  getCollectionRendererItem as getCompositeRendererItem,
  getCollectionRendererItemId as getCompositeRendererItemId,
};

export type CompositeRendererItemObject = ItemObject;
export type CompositeRendererItem = Item;
export type CompositeRendererBaseItemProps = BaseItemProps;
export type CompositeRendererItemProps<
  T extends Item,
  P extends BaseItemProps = BaseItemProps,
> = ItemProps<T, P>;

export const CompositeRenderer = forwardRef(function CompositeRenderer<
  T extends Item = any,
>(props: CompositeRendererProps<T>) {
  const htmlProps = useCompositeRenderer(props);
  return createElement(TagName, htmlProps);
});

export interface CompositeRendererOptions<T extends Item = any> extends Omit<
  CollectionRendererOptions<T>,
  "store" | "children"
> {
  /**
   * Object returned by the
   * [`useCompositeStore`](https://ariakit.com/reference/use-composite-store)
   * hook. If not provided, the closest
   * [Composite](https://ariakit.com/components/composite) component's context
   * will be used.
   *
   * The store
   * [`items`](https://ariakit.com/reference/use-composite-store#items) state
   * will be used to render the items if the
   * [`items`](https://ariakit.com/reference/composite-items#items) prop is not
   * provided.
   */
  store?: CompositeStore<
    T extends CollectionStoreItem ? T : CompositeStoreItem
  >;
  /**
   * The `children` should be a function that receives item props and returns a
   * React element. The item props should be spread onto the element that
   * renders the item.
   */
  children?: (item: ItemProps<T>) => ReactNode;
}

export interface CompositeRendererProps<T extends Item = any> extends Props<
  TagName,
  CompositeRendererOptions<T>
> {}