import React, { useMemo } from 'react';
import { Menu as HeadlessMenu } from '@headlessui/react';
import { clsx } from 'clsx';

import { useCombinedRef } from '../utils/hooks';
import { applyClassName } from '../utils/misc';
import { forwardRef } from '../utils/react';
import { Menu } from '../menu';
import { useDropdownMenuContext } from './dropdown-menu.context';
import * as css from './dropdown-menu-items.css';

/*-- Types --*/

type HeadlessMenuItemsContext = {
  open: boolean;
};

type HeadlessMenuItemsProps = {
  children?:
    | React.ReactNode
    | ((context: HeadlessMenuItemsContext) => React.ReactNode);
  className?: string | ((context: HeadlessMenuItemsContext) => string);
};

export type DropdownMenuItemsProps = Omit<
  HeadlessMenuItemsProps,
  'className'
> & {
  className?:
    | HeadlessMenuItemsProps['className']
    | {
        base?: HeadlessMenuItemsProps['className'];
        arrow?: HeadlessMenuItemsProps['className'];
      };
  compact?: boolean;
};

/*-- Main --*/

export const DropdownMenuItems = forwardRef<'ul', DropdownMenuItemsProps>(
  function DropdownMenuItems(
    { style: userStyle, children, className, compact, ...rest },
    ref,
  ) {
    const baseAttributes = useDropdownMenuContext(
      ({ attributes }) => attributes.floating,
    );
    const baseStyle = useDropdownMenuContext(({ styles }) => styles.floating);
    const baseArrowStyle = useDropdownMenuContext(({ styles }) => styles.arrow);
    const withArrow = useDropdownMenuContext(({ withArrow }) => withArrow);
    const setFloating = useDropdownMenuContext(({ refs }) => refs.setFloating);
    const setArrow = useDropdownMenuContext(({ refs }) => refs.setArrow);
    const keepMounted = useDropdownMenuContext(
      ({ keepMounted }) => keepMounted,
    );

    const style = useMemo(
      () => ({ ...userStyle, ...baseStyle }),
      [baseStyle, userStyle],
    );

    const arrowStyle = useMemo(
      () => ({
        ...baseArrowStyle,
        transform: `${baseArrowStyle.transform} rotate(45deg)`,
      }),
      [baseArrowStyle],
    );

    const cn = useMemo(() => {
      const classNameObject =
        typeof className === 'string' || typeof className === 'function'
          ? { base: className }
          : className || {};

      return {
        base: (itemsContext: HeadlessMenuItemsContext) =>
          clsx(
            applyClassName(classNameObject.base, itemsContext),
            css.items({ open: itemsContext.open, withArrow }),
          ),
        arrow: (itemsContext: HeadlessMenuItemsContext) =>
          clsx(applyClassName(classNameObject.arrow, itemsContext), css.arrow),
      };
    }, [className, withArrow]);

    const combinedRef = useCombinedRef(ref, setFloating);

    return (
      <HeadlessMenu.Items
        as={Menu}
        // @ts-ignore `as` doesn't work well in `@headlessui/react`
        ref={combinedRef}
        style={style}
        className={cn.base}
        compact={compact}
        {...baseAttributes}
        {...rest}
        static={keepMounted}
      >
        {itemsContext => (
          <>
            {withArrow && (
              <div
                ref={setArrow}
                style={arrowStyle}
                className={cn.arrow(itemsContext)}
              />
            )}
            <>
              {typeof children === 'function'
                ? children(itemsContext)
                : children}
            </>
          </>
        )}
      </HeadlessMenu.Items>
    );
  },
);