import React, { useMemo, useState } from 'react';
import type { MenuTriggerProps, MenuTriggerState } from 'react-stately';
import { Item, Section, useMenuTriggerState } from 'react-stately';

import type { MultiSelectListProps, MultiSelectListState } from './useMultiSelectListState';
import { useMultiSelectListState } from './useMultiSelectListState';

/** MultiSelect option selectors. */
export type MultiSelectSelectors<T> = {
  /** Gets the key for the option. */
  getOptionValue: (option: T) => string;
  /** Gets the text label for the option. */
  getOptionLabel: (option: T) => string;
  /** Renders the label for the option. If not present, `getOptionLabel` is used. */
  renderLabel?: (option: T) => React.ReactNode | undefined;
  /** Gets the children for the option (meant for section). */
  getGroupOptions?: (option: T) => T[] | undefined;
};

/** Props for the useMultiSelectState hook. */
export interface MultiSelectStateProps<T>
  extends Omit<MultiSelectListProps<T>, 'children' | 'onSelectionChange'>, MultiSelectSelectors<T>, MenuTriggerProps {
  /** Handler that is called when the selection changes. */
  onSelectionChange?: (selection: string[]) => void;
}

export interface MultiSelectState<T> extends MultiSelectListState<T>, MenuTriggerState {
  /** Whether the select is currently focused. */
  isFilterFocused: boolean;

  /** Sets whether the select is focused. */
  setFilterFocused(isFocused: boolean): void;
}

export function useMultiSelectState<T extends object>({
  getOptionValue,
  getOptionLabel,
  renderLabel,
  getGroupOptions,
  ...props
}: MultiSelectStateProps<T>): MultiSelectState<T> {
  const [isFilterFocused, setFilterFocused] = useState(false);

  const children = useMemo(
    () =>
      makeChildrenSelector({
        getOptionValue,
        getOptionLabel,
        renderLabel,
        getGroupOptions,
      }),
    [getOptionValue, getOptionLabel, renderLabel, getGroupOptions]
  );

  const listState = useMultiSelectListState({
    ...props,
    children,
    disallowEmptySelection: false,
    onSelectionChange: keys => {
      if (props.onSelectionChange) {
        if (keys === 'all') {
          const disabled = new Set(props.disabledKeys ?? []);
          props.onSelectionChange(
            extractAllItemKeys([...(props.items ?? [])], getOptionValue, getGroupOptions).filter(
              key => !disabled.has(key)
            )
          );
        } else {
          props.onSelectionChange([...keys] as string[]);
        }
      }
    },
  });

  const triggerState = useMenuTriggerState({
    ...props,
    onOpenChange: isOpen => {
      props.onOpenChange?.(isOpen);

      if (isOpen) {
        listState.setSearchQuery('');
      }
    },
  });

  return {
    ...listState,
    ...triggerState,
    isFilterFocused,
    setFilterFocused,
  };
}

function makeChildrenSelector<T extends object>(selectors: MultiSelectSelectors<T>) {
  return function selectChild(option: T) {
    const { getOptionValue, getOptionLabel, renderLabel = getOptionLabel, getGroupOptions } = selectors;

    const children = getGroupOptions?.(option);

    if (children) {
      return (
        <Section key={getOptionValue(option)} items={children} title={getOptionLabel(option)}>
          {item => (
            <Item key={getOptionValue(item)} textValue={getOptionLabel(item)}>
              {renderLabel(item)}
            </Item>
          )}
        </Section>
      );
    }

    return (
      <Item key={getOptionValue(option)} textValue={getOptionLabel(option)}>
        {renderLabel(option)}
      </Item>
    );
  };
}

export function extractAllItemKeys<T extends object>(
  items: T[],
  getValue: MultiSelectSelectors<T>['getOptionValue'],
  getGroupOptions: MultiSelectSelectors<T>['getGroupOptions']
): string[] {
  const keys: string[] = [];

  items.forEach(item => {
    const children = getGroupOptions?.(item);

    if (children) {
      keys.push(...extractAllItemKeys(children, getValue, getGroupOptions));
    } else {
      keys.push(getValue(item));
    }
  });

  return keys;
}