import React, { useMemo } from 'react';
import { Listbox } from '@headlessui/react';

import { forwardRef } from '../utils/react';
import {
  SelectButton as BaseSelectButton,
  SelectButtonProps as BaseSelectButtonProps,
} from '../select-button';

/*-- Types --*/

type ListboxButtonContext = {
  open: boolean;
  disabled: boolean;
};

type ListboxButtonProps = {
  children?:
    | React.ReactNode
    | ((context: ListboxButtonContext) => React.ReactElement);
  className?: string | ((context: ListboxButtonContext) => string);
};

export interface SelectButtonProps
  extends ListboxButtonProps,
    BaseSelectButtonProps {}

/*-- Main --*/

export const SelectButton = forwardRef<
  typeof BaseSelectButton,
  SelectButtonProps
>(function SelectButton({ as, ...rest }, ref) {
  const WrappedBaseSelectButton = useMemo(
    () =>
      React.forwardRef(function WrappedBaseSelectButton(props: any, ref) {
        return <BaseSelectButton ref={ref} {...props} as={as} />;
      }),
    [as],
  );

  return <Listbox.Button ref={ref} as={WrappedBaseSelectButton} {...rest} />;
});