import React from 'react';
import { Menu as BaseMenu } from '@headlessui/react';

import { useCombinedRef } from '../utils/hooks';
import { forwardRef } from '../utils/react';
import { useDropdownMenuContext } from './dropdown-menu.context';

/*-- Types --*/

type BaseMenuButtonContext = {
  open: boolean;
};

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

export interface DropdownMenuButtonProps extends BaseMenuButtonProps {}

/*-- Main --*/

export const DropdownMenuButton = forwardRef<'button', DropdownMenuButtonProps>(
  function DropdownMenuButton(props, ref) {
    const setReference = useDropdownMenuContext(
      ({ refs }) => refs.setReference,
    );

    const combinedRef = useCombinedRef(ref, setReference);

    return <BaseMenu.Button ref={combinedRef} {...props} />;
  },
);