import type { AriaAttributes, ButtonHTMLAttributes, InputHTMLAttributes, KeyboardEvent, RefObject } from 'react';
import { useEffect, useMemo, useRef } from 'react';
import type { AriaButtonProps, AriaListBoxOptions } from 'react-aria';
import { chain, ListKeyboardDelegate, mergeProps, useId, useOverlayTrigger, useTextField } from 'react-aria';
import { getItemId, listData } from '@react-aria/listbox';
import { ariaHideOutside } from '@react-aria/overlays';
import { useSelectableCollection } from '@react-aria/selection';
import { filterDOMProps, useLabels } from '@react-aria/utils';
import type {
  AriaLabelingProps,
  DOMProps,
  FocusableElement,
  FocusableProps,
  InputBase,
  PressEvent,
} from '@react-types/shared';

import type { MultiSelectState } from './useMultiSelectState';

export interface TriggerOptions extends Omit<InputBase, 'isReadOnly'>, FocusableProps {}

/** Options for useMultiSelect hook. */
export interface UseMultiSelectOptions extends TriggerOptions, AriaLabelingProps, DOMProps {
  /** The ref for the optional list box popup trigger button.  */
  triggerRef?: RefObject<HTMLButtonElement | null>;
  /** The ref for the list box popover. */
  popoverRef: RefObject<Element | null>;
  /** The ref for the input element. */
  filterRef: RefObject<HTMLInputElement | null>;
  /** The ref for the list box. */
  listRef: RefObject<HTMLElement | null>;
}

/** Props that are turned by useMultiSelect hook. */
export interface MultiSelectAria<T> {
  /** Props for the trigger button, to be passed to [useButton](useButton.html). */
  inputProps: AriaButtonProps &
    Pick<ButtonHTMLAttributes<HTMLButtonElement>, 'role'> &
    Pick<AriaAttributes, 'aria-autocomplete'>;
  /** Props for the popover element. */
  popoverProps: DOMProps;
  /** Props for the dialog element. */
  dialogProps: DOMProps & AriaLabelingProps;
  /** Props for the select all button. */
  toggleAllProps: AriaButtonProps & Pick<ButtonHTMLAttributes<HTMLButtonElement>, 'tabIndex'>;
  /** Props for the filter input element. */
  filterProps: InputHTMLAttributes<HTMLInputElement>;
  /** Props for the list box. */
  listProps: AriaListBoxOptions<T>;
}

export function useMultiSelect<T>(props: UseMultiSelectOptions, state: MultiSelectState<T>): MultiSelectAria<T> {
  const { popoverRef, filterRef, listRef, autoFocus, isDisabled, id } = props;
  const backupBtnRef = useRef(null);

  const listBoxId = useId();

  const triggerRef = props.triggerRef ?? backupBtnRef;

  const { triggerProps: inputProps, overlayProps: popoverProps } = useOverlayTrigger(
    {
      type: 'listbox',
    },
    state,
    triggerRef
  );

  // Set listbox id so it can be used when calling getItemId later
  listData.set(state, { id: listBoxId });

  const { collection } = state;
  const { disabledKeys } = state.selectionManager;

  const delegate = useMemo(
    () =>
      new ListKeyboardDelegate({
        collection,
        disabledKeys,
        ref: listRef,
      }),
    [collection, disabledKeys, listRef]
  );

  // Use useSelectableCollection to get the keyboard handlers to apply to the textfield
  const { collectionProps } = useSelectableCollection({
    selectionManager: state.selectionManager,
    keyboardDelegate: delegate,
    disallowTypeAhead: false,
    disallowEmptySelection: true,
    shouldFocusWrap: false,
    ref: filterRef,
    // Prevent item scroll behavior from being applied here, should be handled in the user's Popover + ListBox component
    isVirtualized: true,
  });

  const onTriggerKeyDown = (e: KeyboardEvent<FocusableElement>) => {
    if (e.nativeEvent.isComposing) {
      return;
    }

    switch (e.key) {
      case 'Enter': {
        // Prevent form submission if menu is open since we may be selecting a option
        if (state.isOpen && e.key === 'Enter') {
          e.preventDefault();
        }

        break;
      }
      case 'ArrowDown': {
        state.open('first');
        filterRef.current?.focus();
        break;
      }
      case 'ArrowUp': {
        state.open('last');
        filterRef.current?.focus();
        break;
      }
    }
  };

  const onFilterKeyDown = (e: KeyboardEvent<FocusableElement>) => {
    if (e.nativeEvent.isComposing) {
      return;
    }

    switch (e.key) {
      case 'Tab': {
        if (e.shiftKey) {
          e.preventDefault();

          state.close();
        }
        break;
      }
      case 'Enter': {
        e.preventDefault();

        if (state.selectionManager.focusedKey != null) {
          state.selectionManager.toggleSelection(state.selectionManager.focusedKey);
        }

        break;
      }
      case 'Escape': {
        state.close();
        break;
      }
      case 'ArrowDown': {
        if (!state.selectionManager.isFocused) {
          state.selectionManager.setFocused(true);
        }
        break;
      }
      case 'ArrowUp': {
        if (!state.selectionManager.isFocused) {
          state.selectionManager.setFocused(true);
        }
        break;
      }
    }

    if (e.key !== 'ArrowLeft' && e.key !== 'ArrowRight') {
      collectionProps.onKeyDown?.(e);
    }
  };

  const onFilterBlur = () => {
    state.selectionManager.setFocused(false);
    state.setFilterFocused(false);
  };

  const onFilterFocus = () => {
    state.selectionManager.setFocused(true);
    state.setFilterFocused(true);
  };

  const onSelectAllKeyDown = (e: KeyboardEvent<FocusableElement>) => {
    if (e.nativeEvent.isComposing) {
      return;
    }

    switch (e.key) {
      case 'Escape': {
        state.close();
        break;
      }
    }
  };

  const onSelectedAllPress = () => {
    state.toggleFiltered();
  };

  const { inputProps: filterInputProps } = useTextField(
    {
      onChange: state.setSearchQuery,
      onBlur: onFilterBlur,
      value: state.searchQuery,
      onFocus: onFilterFocus,
      onKeyDown: onFilterKeyDown,
      autoComplete: 'off',
      validate: undefined,
      inputElementType: 'input',
      'aria-label': 'Filter options',
    },
    filterRef
  );

  const onInputPress = (e: PressEvent) => {
    if (e.pointerType === 'touch') {
      state.toggle(null);
    }
  };

  const onInputPressStart = (e: PressEvent) => {
    if (e.pointerType !== 'touch') {
      state.toggle(null);
    }
  };

  const inputAriaProps = filterDOMProps(props);

  const inputLabelProps = useLabels({
    id: id ?? inputProps.id,
    'aria-label': props['aria-label'],
    'aria-labelledby': props['aria-labelledby'],
  });

  const dialogLabelProps = useLabels({
    'aria-label': props['aria-label'],
    'aria-labelledby': props['aria-labelledby'],
  });

  const listBoxLabelProps = useLabels({
    id: listBoxId,
    'aria-label': props['aria-label'],
    'aria-labelledby': props['aria-labelledby'],
  });

  const focusedItem =
    state.selectionManager.focusedKey != null && state.isOpen
      ? state.collection.getItem(state.selectionManager.focusedKey)
      : undefined;

  useEffect(() => {
    if (state.isOpen) {
      return ariaHideOutside(
        [triggerRef.current, popoverRef.current].filter((element): element is Element => element != null)
      );
    }
  }, [state.isOpen, triggerRef, popoverRef]);

  return {
    inputProps: {
      ...inputProps,
      ...inputAriaProps,
      ...inputLabelProps,
      onKeyDown: chain(inputProps.onKeyDown, onTriggerKeyDown, props.onKeyDown),
      onKeyUp: props.onKeyUp,
      onPress: onInputPress,
      onPressStart: onInputPressStart,
      onBlur: props.onBlur,
      onFocus: props.onFocus,
      onFocusChange: props.onFocusChange,
      isDisabled,
      role: 'combobox',
      'aria-autocomplete': 'none',
      autoFocus,
    },
    popoverProps,
    dialogProps: dialogLabelProps,
    filterProps: mergeProps(filterInputProps, {
      'aria-activedescendant': focusedItem ? getItemId(state, focusedItem.key) : undefined,
      role: 'combobox',
      autoCorrect: 'off',
      spellCheck: 'false',
      autoFocus: true,
      'aria-autocomplete': 'list',
      'aria-haspopup': 'listbox',
      'aria-expanded': state.isOpen,
      'aria-controls': listBoxId,
    }),
    toggleAllProps: {
      onKeyDown: onSelectAllKeyDown,
      onPress: onSelectedAllPress,
      'aria-pressed': state.selectionManager.isSelectAll,
      tabIndex: 0,
    },
    listProps: mergeProps(listBoxLabelProps, {
      autoFocus: false,
      disallowEmptySelection: true,
      shouldUseVirtualFocus: filterRef.current === document.activeElement,
      shouldSelectOnPressUp: true,
      linkBehavior: 'selection' as const,
    }),
  };
}