import { clsx } from '@trail-ui/shared-utils';
import { multiselect, MultiSelectSlots, SlotsToClasses } from '@trail-ui/theme';
import React, {
  useState,
  useRef,
  useEffect,
  forwardRef,
  ForwardedRef,
  ReactElement,
  KeyboardEvent,
  ChangeEvent,
  RefObject,
  useMemo,
} from 'react';
import { Input, Label } from './tw-field';
import { Text } from './tw-text';
import { CheckIcon, ChevronDownIcon, CloseIcon, ErrorIcon } from '@trail-ui/icons';

// Types for handling both string and object options
type BaseOption = {
  label: string;
  value: string;
  [key: string]: any;
};

// Constants
const Keys = {
  Backspace: 'Backspace',
  Clear: 'Clear',
  Down: 'ArrowDown',
  End: 'End',
  Enter: 'Enter',
  Escape: 'Escape',
  Home: 'Home',
  Left: 'ArrowLeft',
  PageDown: 'PageDown',
  PageUp: 'PageUp',
  Right: 'ArrowRight',
  Space: ' ',
  Tab: 'Tab',
  Up: 'ArrowUp',
} as const;

// Define specific action types
const MenuActions = {
  Close: 'Close',
  CloseSelect: 'CloseSelect',
  First: 'First',
  Last: 'Last',
  Next: 'Next',
  Open: 'Open',
  Previous: 'Previous',
  Select: 'Select',
  Space: 'Space',
  Type: 'Type',
} as const;

// Types
type MenuActionType = (typeof MenuActions)[keyof typeof MenuActions];
type NavigationAction =
  | typeof MenuActions.First
  | typeof MenuActions.Last
  | typeof MenuActions.Next
  | typeof MenuActions.Previous;

export interface MultiSelectProps<T extends string | BaseOption> {
  /**
   * @property {T[]} options - The list of options available for selection.
   */
  options: T[];
  /**
   * @property {string} [label] - An optional label for the MultiSelect component.
   * @default 'Select items'
   */
  label?: string;
  /**
   * @property {string} [id] - An optional id for the MultiSelect component.
   */
  id?: string;
  /**
   * @property {string} [labelKey] - The key to use for the label of each option.
   */
  labelKey?: string;
  /**
   * @property {string} [valueKey] - The key to use for the value of each option.
   */
  valueKey?: string;
  /**
   * @param selectedItems  - A callback function that is called when the selected items change.
   * @returns
   */
  onChange?: (selectedItems: T[]) => void;
  /**
   * @property {() => void} [onBlur] - A callback function that is called when the MultiSelect component loses focus.
   */
  onBlur?: () => void;
  /**
   * @property {T[]} [defaultValue] - An optional array of default selected items.
   */
  defaultValue?: T[];
  /**
   * @property {SlotsToClasses<MultiSelectSlots>} [classNames] - An optional object of classes to style the MultiSelect component.
   */
  classNames?: SlotsToClasses<MultiSelectSlots>;
  /**
   * @property {number} [maxTags] - The maximum number of tags to display before showing a '+n more' tag.
   * @default 4
   */
  maxTags?: number;
  /**
   * @property {string} [errorMessage] - An optional error message to display below the MultiSelect component.
   */
  errorMessage?: string;
  /**
   * @property {string} [description] - An optional description to display below the MultiSelect component.
   */
  description?: string;
  /**
   * @property {string} [placeholder] - An optional placeholder to display in the input field.
   */
  placeholder?: string;
  /**
   * @property {React.ReactNode} [errorIcon] - An optional icon to display next to the error message.
   */
  errorIcon?: React.ReactNode;
  /**
   * @property {string} [errorId] - The id of the error message for accessibility.
   */
  errorId?: string;
  /**
   * @property {boolean} [isRequired] - Whether the MultiSelect component is required.
   */
  isRequired?: boolean;
  /**
   * @property {boolean} [isDisabled] - Whether the MultiSelect component is disabled.
   */
  isDisabled?: boolean;
  /**
   * @property {boolean} [isInvalid] - Whether the MultiSelect component is invalid.
   */
  isInvalid?: boolean;
  /**
   * @property {string} [name] - The name of the form field
   */
  name?: string;
  /**
   * @property {T[]} [value] - The controlled value of the component
   */
  value?: T[];
  /**
   * @property {string} [form] - The id of the form this field belongs to
   */
  form?: string;
  /**
   * @property {string} [color] - The color of the tags
   */
  color?: string;
}

function getActionFromKey(key: string, menuOpen: boolean): MenuActionType | undefined {
  if (!menuOpen && key === Keys.Down) return MenuActions.Open;
  if (!menuOpen && key === Keys.Enter) return MenuActions.Open;
  if (key === Keys.Down) return MenuActions.Next;
  if (key === Keys.Up) return MenuActions.Previous;
  if (key === Keys.Home) return MenuActions.First;
  if (key === Keys.End) return MenuActions.Last;
  if (key === Keys.Escape) return MenuActions.Close;
  if (key === Keys.Enter) return MenuActions.CloseSelect;
  if (key === Keys.Backspace || key === Keys.Clear || key.length === 1) return MenuActions.Type;
  return undefined;
}

function getUpdatedIndex(current: number, max: number, action: NavigationAction): number {
  switch (action) {
    case MenuActions.First:
      return 0;
    case MenuActions.Last:
      return max;
    case MenuActions.Previous:
      return Math.max(0, current - 1);
    case MenuActions.Next:
      return Math.min(max, current + 1);
    default:
      return current;
  }
}

const getTagColor = (value: string): string => {
  const firstChar = value?.toLowerCase()[0];
  switch (true) {
    case 'a' <= firstChar && firstChar <= 'd':
      return 'green';
    case 'e' <= firstChar && firstChar <= 'h':
      return 'purple';
    case 'i' <= firstChar && firstChar <= 'l':
      return 'red';
    case 'm' <= firstChar && firstChar <= 'p':
      return 'yellow';
    case 'q' <= firstChar && firstChar <= 't':
      return 'blue';
    default:
      return 'default';
  }
};

const handleColor = (color: string) => {
  switch (color) {
    case 'blue':
      return 'bg-blue-100 ';
    case 'green':
      return 'bg-green-100';
    case 'purple':
      return 'bg-purple-100';
    case 'red':
      return 'bg-red-100';
    case 'yellow':
      return 'bg-yellow-50';
    default:
      return 'bg-neutral-200';
  }
};

const MAX_DISPLAYED_TAGS = 50;

function MultiSelect<T extends string | BaseOption>(
  {
    options,
    label,
    labelKey = 'value',
    valueKey = 'id',
    id = 'multiselect',
    onChange,
    onBlur,
    defaultValue = [],
    classNames,
    maxTags = MAX_DISPLAYED_TAGS,
    errorMessage,
    description,
    errorIcon = (
      <ErrorIcon
        className="h-4 w-4 text-red-800 dark:text-red-600"
        role="img"
        aria-label="Error"
        aria-hidden={false}
      />
    ),
    errorId,
    isRequired,
    isDisabled,
    isInvalid,
    placeholder = 'Select items',
    value,
    form,
    name,
    color,
    ...otherprops
  }: MultiSelectProps<T>,
  ref: ForwardedRef<HTMLDivElement>,
) {
  // State
  const [isOpen, setIsOpen] = useState(false);
  const [selectedItems, setSelectedItems] = useState<T[]>(
    Array.isArray(value) ? value : Array.isArray(defaultValue) ? defaultValue : [],
  );
  const [searchValue, setSearchValue] = useState('');
  const [activeIndex, setActiveIndex] = useState(-1);

  // Refs
  const listboxRef = useRef<HTMLUListElement>(null);
  const inputRef = useRef<HTMLInputElement>(null);
  const localRef = useRef<HTMLDivElement>(null);
  const optionRefs = useRef<Array<RefObject<HTMLLIElement>>>([]);
  const hiddenInputRef = useRef<HTMLInputElement>(null);

  // Sync the forwarded ref with our local ref
  useEffect(() => {
    if (!ref) return;

    if (typeof ref === 'function') {
      ref(localRef.current);
    } else {
      ref.current = localRef.current;
    }
  }, [ref]);

  // Styles
  const slots = useMemo(() => multiselect(), []);

  // Helpers
  const getOptionLabel = (option: T): string => {
    if (typeof option === 'string') return option;
    return (option as BaseOption)[labelKey] || (option as BaseOption).label || '';
  };

  const getOptionValue = (option: T): string => {
    if (typeof option === 'string') return option;
    return (option as BaseOption)[valueKey] || (option as BaseOption).value || '';
  };

  // Update selectedItems when value prop changes
  useEffect(() => {
    if (value !== undefined) {
      setSelectedItems(Array.isArray(value) ? value : []);
    }
  }, [value]);

  // Filter options based on search input
  const filteredOptions = options.filter((option) =>
    getOptionLabel(option).toLowerCase().includes(searchValue.toLowerCase()),
  );

  // Initialize refs array
  useEffect(() => {
    optionRefs.current = filteredOptions.map(() => React.createRef());
  }, [filteredOptions]);

  // Click outside handler
  useEffect(() => {
    const handleClickOutside = (event: MouseEvent) => {
      if (localRef.current && !localRef.current.contains(event.target as Node)) {
        setIsOpen(false);
      }
    };

    document.addEventListener('mousedown', handleClickOutside);
    return () => document.removeEventListener('mousedown', handleClickOutside);
  }, []);

  // Reset active index when filtered options change
  useEffect(() => {
    if (activeIndex === -1 && filteredOptions.length > 0) {
      setActiveIndex(0);
    }
  }, [filteredOptions, activeIndex]);

  useEffect(() => {
    if (!isOpen) {
      setActiveIndex(-1);
    }
  }, [isOpen]);

  const toggleItem = (option: T) => {
    if (isDisabled) return;

    const isSelected = selectedItems.some(
      (item) => getOptionValue(item) === getOptionValue(option),
    );

    const newItems = isSelected
      ? selectedItems.filter((item) => getOptionValue(item) !== getOptionValue(option))
      : [...selectedItems, option];

    setSelectedItems(newItems);
    onChange?.(newItems);
    setSearchValue('');
    inputRef.current?.focus();

    // Update hidden input
    if (hiddenInputRef.current) {
      const values = newItems.map(getOptionValue);
      hiddenInputRef.current.value = values.join(',');
    }
  };

  // Event handlers
  const handleInputChange = (e: ChangeEvent<HTMLInputElement>) => {
    if (isDisabled) return;
    setSearchValue(e.target.value);
    setIsOpen(true);
    if (activeIndex === -1) {
      setActiveIndex(0);
    }
  };

  const handleInputKeyDown = (e: KeyboardEvent<HTMLInputElement | HTMLOrSVGElement>) => {
    const max = filteredOptions.length - 1;
    const action = getActionFromKey(e.key, isOpen);

    if (isDisabled) {
      e.preventDefault();
      e.stopPropagation();
      return;
    }

    if (e.key === Keys.Tab) {
      if (isOpen) {
        setIsOpen(false);
        setActiveIndex(-1);
        setSearchValue('');
      }
      return;
    }

    if (!action) return;

    switch (action) {
      case MenuActions.Next:
      case MenuActions.Last:
      case MenuActions.First:
      case MenuActions.Previous:
        e.preventDefault();
        const nextIndex = getUpdatedIndex(activeIndex, max, action);
        setActiveIndex(nextIndex);
        break;
      case MenuActions.CloseSelect:
        e.preventDefault();
        if (activeIndex >= 0 && filteredOptions[activeIndex]) {
          toggleItem(filteredOptions[activeIndex]);
        }
        break;
      case MenuActions.Close:
        e.preventDefault();
        setIsOpen(false);
        setSearchValue('');
        break;
      case MenuActions.Open:
        e.preventDefault();
        setIsOpen(true);
        break;
    }
  };

  const handleOptionKeyDown = (e: KeyboardEvent<HTMLLIElement>, index: number, option: T) => {
    if (isDisabled) {
      e.preventDefault();
      e.stopPropagation();
      return;
    }

    switch (e.key) {
      case Keys.Enter:
      case Keys.Space:
        e.preventDefault();
        toggleItem(option);
        break;
      case Keys.Escape:
        e.preventDefault();
        setIsOpen(false);
        inputRef.current?.focus();
        break;
      case Keys.Down:
        e.preventDefault();
        if (index < filteredOptions.length - 1) {
          setActiveIndex(index + 1);
          optionRefs.current[index + 1]?.current?.focus();
        }
        break;
      case Keys.Up:
        e.preventDefault();
        if (index > 0) {
          setActiveIndex(index - 1);
          optionRefs.current[index - 1]?.current?.focus();
        } else {
          inputRef.current?.focus();
          setActiveIndex(-1);
        }
        break;
    }
  };

  const handleTagKeyDown = (e: KeyboardEvent<HTMLSpanElement>, index: number) => {
    if (isDisabled) {
      e.preventDefault();
      e.stopPropagation();
      return;
    }

    switch (e.key) {
      case Keys.Enter:
      case Keys.Space:
        e.preventDefault();
        toggleItem(selectedItems[index]);
        inputRef.current?.focus();
        break;
      case Keys.Escape:
        e.preventDefault();
        inputRef.current?.focus();
        break;
    }
  };

  // Render tags
  const renderTags = () => {
    if (!Array.isArray(selectedItems) || selectedItems.length === 0) return null;

    const displayedTags = selectedItems.slice(0, maxTags);
    const remainingCount = selectedItems.length - maxTags;

    return (
      <ul className="flex w-full flex-wrap items-center gap-1">
        {displayedTags.map((item, index) => (
          <li key={getOptionValue(item)}>
            <span
              id={`${id}-tag-${index}`}
              className={`${handleColor(
                color ?? getTagColor(getOptionLabel(item)),
              )} focus-visible:outline-focus inline-flex animate-[fadeIn_0.2s_ease-in-out] cursor-pointer items-center rounded px-2 py-0.5 pr-1 text-sm text-neutral-950 focus-visible:outline`}
              role="button"
              aria-label={`Remove ${getOptionLabel(item)}`}
              tabIndex={0}
              onClick={(e) => {
                e.stopPropagation();
                toggleItem(item);
              }}
              onKeyDown={(e) => handleTagKeyDown(e, index)}
            >
              {getOptionLabel(item)}
              <CloseIcon className="ml-1 h-4 w-4 text-neutral-800" />
            </span>
          </li>
        ))}
        {remainingCount > 0 && (
          <li>
            <span
              role="status"
              aria-label={`${remainingCount} more items selected`}
              className="inline-flex items-center rounded-md bg-neutral-200 px-2 py-0.5 text-sm text-neutral-800"
            >
              +{remainingCount} more
            </span>
          </li>
        )}
      </ul>
    );
  };

  return (
    <div
      ref={localRef}
      className="group relative flex w-full max-w-xs flex-col transition-opacity duration-200 data-[disabled=true]:pointer-events-none data-[disabled=true]:opacity-50 motion-reduce:transition-none"
    >
      <Label
        id={`${id}-label`}
        htmlFor={`${id}-input`}
        isDisabled={isDisabled}
        requiredHint={isRequired}
        className={slots.label({ class: classNames?.label })}
      >
        {label}
      </Label>

      <div className="relative">
        <input
          ref={hiddenInputRef}
          type="hidden"
          name={name}
          value={Array.isArray(selectedItems) ? selectedItems.map(getOptionValue).join(',') : ''}
          form={form}
          required={isRequired}
        />

        <div className={slots.control({ class: classNames?.control })}>
          <div className="flex flex-1 flex-wrap items-center gap-2">
            {renderTags()}
            <Input
              ref={inputRef}
              id={`${id}-input`}
              autoComplete="off"
              aria-autocomplete="list"
              onBlur={onBlur}
              className={slots.input({ class: classNames?.input })}
              value={searchValue}
              onChange={handleInputChange}
              onKeyDown={handleInputKeyDown}
              onClick={() => !isDisabled && setIsOpen(true)}
              placeholder={
                Array.isArray(selectedItems) && selectedItems.length === 0 ? `${placeholder}` : ''
              }
              aria-expanded={isOpen}
              aria-controls={isOpen ? 'multiselect-listbox' : undefined}
              aria-activedescendant={
                isOpen && activeIndex >= 0 ? `${id}-option-${activeIndex}` : undefined
              }
              disabled={isDisabled}
              required={
                isRequired &&
                (!Array.isArray(selectedItems) || selectedItems.length === 0) &&
                !searchValue
              }
              aria-invalid={isInvalid}
              aria-required={isRequired}
              role="combobox"
              aria-describedby={`${id}-selection-status ${errorId || ''}`}
              aria-haspopup="listbox"
              {...otherprops}
            />
            <ChevronDownIcon
              height={24}
              width={24}
              onKeyDown={handleInputKeyDown}
              onClick={() => !isDisabled && setIsOpen(!isOpen)}
              className={`${isOpen ? 'rotate-180' : ''} cursor-pointer`}
            />
          </div>
        </div>

        {errorMessage ? (
          <Text
            id={errorId}
            slot="errorMessage"
            aria-live="polite"
            elementType="div"
            className={slots.errorMessage({ class: classNames?.errorMessage })}
          >
            {errorIcon}
            <span>{errorMessage}</span>
          </Text>
        ) : description ? (
          <Text
            slot="description"
            elementType="div"
            className={slots.description({ class: classNames?.description })}
          >
            <span>{description}</span>
          </Text>
        ) : null}

        <div id={`${id}-selection-status`} className="sr-only" aria-live="polite">
          {Array.isArray(selectedItems) && selectedItems.length > 0
            ? `${selectedItems.length} items selected: ${selectedItems.map(getOptionLabel).join(', ')}`
            : 'No items selected'}
        </div>

        {isOpen && filteredOptions.length > 0 && (
          <ul
            ref={listboxRef}
            role="listbox"
            id="multiselect-listbox"
            aria-label={label}
            aria-multiselectable="true"
            className="absolute z-10 mt-2 max-h-60 w-full overflow-auto rounded-md bg-neutral-50 py-2 text-base shadow-lg ring-1 ring-black ring-opacity-5 focus:outline-none"
          >
            {filteredOptions.map((option, index) => {
              const isSelected =
                Array.isArray(selectedItems) &&
                selectedItems.some((item) => getOptionValue(item) === getOptionValue(option));
              const isActive = index === activeIndex;

              return (
                <li
                  key={getOptionValue(option)}
                  ref={optionRefs.current[index]}
                  role="option"
                  id={`${id}-option-${index}`}
                  aria-selected={isSelected}
                  tabIndex={isActive ? 0 : -1}
                  className={`relative flex cursor-pointer justify-between border-l-[3px] p-2 hover:bg-purple-50 focus:bg-purple-50 ${isActive ? 'border-purple-600 bg-purple-50' : 'border-transparent'} ${isSelected ? 'font-medium' : ''}`}
                  onClick={() => toggleItem(option)}
                  onKeyDown={(e) => handleOptionKeyDown(e, index, option)}
                  onMouseEnter={() => setActiveIndex(index)}
                >
                  <span className={`block truncate text-sm text-neutral-900`}>
                    {getOptionLabel(option)}
                  </span>
                  {isSelected && (
                    <span className="flex items-center pl-1.5">
                      <CheckIcon className="h-5 w-5 text-purple-600" />
                    </span>
                  )}
                </li>
              );
            })}
          </ul>
        )}
      </div>
    </div>
  );
}

const _MultiSelect = forwardRef(MultiSelect) as <T extends string | BaseOption>(
  props: MultiSelectProps<T> & { ref?: ForwardedRef<HTMLDivElement> },
) => ReactElement;

export { _MultiSelect as MultiSelect };