import * as listbox from "@zag-js/listbox";
import { normalizeProps, type PropTypes, useMachine } from "@zag-js/react";
import * as React from "react";
import { SuffixContext, useSuffixContext } from "../../contexts";
import { useGetKey, useGetSet } from "../../hooks";

type Context<T = any> = {
  api: listbox.Api<PropTypes, T>;
  items: T[];
};
const ListBoxContext = React.createContext<Context>(null!);
const useListBoxContext = <T,>() => React.use<Context<T>>(ListBoxContext);

type CoreProps<T = any> = {
  children?: React.ReactNode;
  itemToValue: (item: T) => string;
  itemToString: (item: T) => string;
  items: T[];
  loopFocus?: boolean;
  typeahead?: boolean;
  __length?: number;
  "data-id": string;
  "data-label"?: string;
};

type MultipleProps<T = any> = CoreProps<T> & {
  multiple: true;
  onValueChange?: (items: T[]) => void;
  initialValue?: T | null;
};

type SingleProps<T = any> = CoreProps<T> & {
  multiple: false;
  onValueChange?: (item: T[]) => void;
  initialValue?: T | null;
};

const DEFAULT_VALUE = { value: [] };
export function ListBox<T = any>({
  items = [],
  children,
  initialValue,
  itemToValue = defaultItemToValue,
  itemToString = (item: T) => String(item),
  onValueChange,
  multiple = false,
  loopFocus = false,
  typeahead = true,
  __length,
  ...rest
}: MultipleProps<T> | SingleProps<T>) {
  const key = useGetKey(rest);
  const [{ value }, setValue] = useGetSet(
    key,
    initialValue
      ? {
          value: Array.isArray(initialValue)
            ? (initialValue.map(itemToValue) as T[])
            : ([itemToValue(initialValue as T)] as T[]),
        }
      : DEFAULT_VALUE,
  );
  const service = useMachine(listbox.machine as any, {
    id: key,
    collection: listbox.collection({
      items: items || [],
      itemToValue,
      itemToString,
    }),
    typeahead,
    loopFocus,
    value: value as any,
    selectionMode: multiple ? "multiple" : "single",
    onValueChange: (details) => {
      setValue({ value: details.items || [] });
      onValueChange?.(details.items || []);
    },
  }) as listbox.Service<T>;

  const api = listbox.connect(service, normalizeProps);

  return (
    <ListBoxContext value={{ api, items }}>
      <div {...api.getRootProps()}>{children}</div>
    </ListBoxContext>
  );
}

function ListInput<T = any>(props: { "data-id"?: string }) {
  const key = useGetKey(props);
  const initialValue = React.useMemo(() => ({ value: "" }), []);
  const [{ value }, setState] = useGetSet<{ value: string }>(key, initialValue);
  const { api } = useListBoxContext<T>();
  return (
    <input
      {...api.getInputProps({ autoHighlight: true })}
      value={value}
      onChange={(e) => setState({ value: e.target.value })}
    />
  );
}
ListBox.ListInput = ListInput;

function ListContent<T = any>({
  children,
  ...props
}: {
  children?: (value: T, index: number) => React.ReactElement | null;
}) {
  const suffix = useSuffixContext();
  const { api, items } = useListBoxContext<T>();

  return (
    <div {...props} {...api.getContentProps()}>
      {items.map((item, idx) => {
        const itemKey = `${(item as any)?.id ?? idx}`;
        const key = suffix ? `${itemKey}-${suffix}` : itemKey;
        return (
          <SuffixContext.Provider value={key} key={key}>
            <ListBoxItemContext value={item}>
              {children ? children(item, idx) : null}
            </ListBoxItemContext>
          </SuffixContext.Provider>
        );
      })}
    </div>
  );
}

ListBox.ListContent = ListContent as any;

const ListBoxItemContext = React.createContext<any>(null!);
const useListBoxItemContext = <T,>() => React.use<T>(ListBoxItemContext);

function ListItem<T = any>({
  children,
  ...props
}: {
  children?: React.ReactNode;
}) {
  const { api } = useListBoxContext<T>();
  const item = useListBoxItemContext<T>();
  return (
    <div {...props} {...api.getItemProps({ item, highlightOnHover: true })}>
      {children}
    </div>
  );
}
ListContent.ListItem = ListItem;

function ItemText<T = any>(props: any) {
  const { api } = useListBoxContext<T>();
  const item = useListBoxItemContext<T>();
  const label = api.collection.stringifyItem(item);
  return (
    <span {...props} {...api.getItemTextProps({ item })}>
      {label}
    </span>
  );
}
ListItem.ListItemText = ItemText;

function ItemIndicator<T = any>({ children, ...props }) {
  const { api } = useListBoxContext<T>();
  const item = useListBoxItemContext<T>();
  return (
    <span {...props} {...api.getItemIndicatorProps({ item })}>
      {children}
    </span>
  );
}
ListItem.ListItemIndicator = ItemIndicator;

function defaultItemToValue<T>(item: T): string {
  if (!item) return "";
  if (typeof item !== "object") {
    return String(item);
  }
  return (item as any)?.id ?? JSON.stringify(item);
}