import { ListboxButton } from '@headlessui/react';
import mergeProps from 'merge-props';
import { useContext } from 'react';
import { PolymorphicWithOverrides } from '../../../common/polymorphicWithOverrides/PolymorphicWithOverrides';
import { Merge } from '../../../utils';
import { SelectInputTriggerButtonPropsContext } from '../SelectInput.contexts';

export type SelectInputTriggerButtonElementType = 'button' | React.ComponentType;

export type SelectInputTriggerButtonProps<
  T extends SelectInputTriggerButtonElementType = 'button',
> = Merge<React.ComponentPropsWithoutRef<T>, { as?: T }>;

/**
 * The trigger button component for SelectInput.
 * Uses Headless UI's ListboxButton with polymorphic support to allow
 * rendering as different element types.
 */
export function SelectInputTriggerButton<T extends SelectInputTriggerButtonElementType = 'button'>({
  as = 'button' as T,
  ...restProps
}: SelectInputTriggerButtonProps<T>) {
  const { ref, onClick, onKeyDown, size, ...interactionProps } = useContext(
    SelectInputTriggerButtonPropsContext,
  );

  return (
    <ListboxButton
      ref={ref}
      as={PolymorphicWithOverrides}
      role="combobox"
      __overrides={{ as, size, ...interactionProps } as Record<string, unknown>}
      {...mergeProps({ onClick, onKeyDown }, restProps)}
    />
  );
}