import { useEffect, useId, useRef, useState } from 'react';

export type SelectOption = {
  value: string;
  label: string;
};

export type SelectOptionGroup = {
  label: string;
  options: SelectOption[];
};

type SelectFieldProps = {
  id?: string;
  label: string;
  value: string;
  options: SelectOption[];
  groups?: SelectOptionGroup[];
  disabled?: boolean;
  onChange: (value: string) => void;
};

export function SelectField({
  id,
  label,
  value,
  options,
  groups = [],
  disabled = false,
  onChange,
}: SelectFieldProps) {
  const generatedId = useId();
  const fieldId = id ?? generatedId;
  const [isOpen, setIsOpen] = useState(false);
  const rootRef = useRef<HTMLDivElement | null>(null);
  const selectedOption = findOption(value, options, groups);

  useEffect(() => {
    function handlePointerDown(event: PointerEvent) {
      if (!rootRef.current?.contains(event.target as Node)) {
        setIsOpen(false);
      }
    }

    function handleKeyDown(event: KeyboardEvent) {
      if (event.key === 'Escape') {
        setIsOpen(false);
      }
    }

    document.addEventListener('pointerdown', handlePointerDown);
    document.addEventListener('keydown', handleKeyDown);

    return () => {
      document.removeEventListener('pointerdown', handlePointerDown);
      document.removeEventListener('keydown', handleKeyDown);
    };
  }, []);

  function handleSelect(nextValue: string) {
    onChange(nextValue);
    setIsOpen(false);
  }

  return (
    <div className="select-field" ref={rootRef}>
      <label id={`${fieldId}-label`} htmlFor={fieldId}>
        {label}
      </label>
      <button
        id={fieldId}
        type="button"
        className="select-trigger"
        aria-haspopup="listbox"
        aria-expanded={isOpen}
        aria-labelledby={`${fieldId}-label ${fieldId}`}
        disabled={disabled}
        onClick={() => setIsOpen((current) => !current)}
      >
        <span>{selectedOption?.label ?? 'Choose an option'}</span>
        <span className="select-chevron" aria-hidden="true" />
      </button>
      {isOpen ? (
        <div className="select-menu" role="listbox" aria-labelledby={`${fieldId}-label`}>
          {options.map((option) => (
            <SelectItem
              key={option.value}
              option={option}
              isSelected={option.value === value}
              onSelect={handleSelect}
            />
          ))}
          {groups.map((group) => (
            <div className="select-group" key={group.label}>
              <div className="select-group-label">{group.label}</div>
              {group.options.map((option) => (
                <SelectItem
                  key={`${group.label}-${option.value}`}
                  option={option}
                  isSelected={option.value === value}
                  onSelect={handleSelect}
                />
              ))}
            </div>
          ))}
        </div>
      ) : null}
    </div>
  );
}

function SelectItem({
  option,
  isSelected,
  onSelect,
}: {
  option: SelectOption;
  isSelected: boolean;
  onSelect: (value: string) => void;
}) {
  return (
    <button
      type="button"
      className={isSelected ? 'select-option select-option-selected' : 'select-option'}
      role="option"
      aria-selected={isSelected}
      onClick={() => onSelect(option.value)}
    >
      <span>{option.label}</span>
      {isSelected ? <span className="select-check" aria-hidden="true" /> : null}
    </button>
  );
}

function findOption(
  value: string,
  options: SelectOption[],
  groups: SelectOptionGroup[],
): SelectOption | undefined {
  return (
    options.find((option) => option.value === value) ??
    groups.flatMap((group) => group.options).find((option) => option.value === value)
  );
}