import React, { forwardRef } from 'react';
import { ContextMenu as BaseContextMenu } from '@base-ui/react/context-menu';
import cs from 'classnames';
import classes from './ContextMenu.module.scss';

type TriggerBaseProps = Omit<
  React.ComponentPropsWithoutRef<typeof BaseContextMenu.Trigger>,
  'children' | 'render'
>;

type TriggerProps = TriggerBaseProps & {
  asChild?: boolean;
  children: React.ReactNode;
};
type TriggerComponent = React.ForwardRefExoticComponent<
  TriggerProps & React.RefAttributes<HTMLDivElement>
>;

const Trigger: TriggerComponent = forwardRef<HTMLDivElement, TriggerProps>((props, ref) => {
  const { asChild = false, children, ...rest } = props;

  if (asChild && !React.isValidElement(children)) {
    if (process.env.NODE_ENV !== 'production') {
      throw new Error('ContextMenu.Trigger with asChild expects a single valid React element.');
    }
    return null;
  }

  const render = asChild ? (children as React.ReactElement) : undefined;

  return (
    <BaseContextMenu.Trigger
      ref={ref}
      render={render}
      {...rest}
    >
      {asChild ? undefined : children}
    </BaseContextMenu.Trigger>
  );
});

Trigger.displayName = 'ContextMenu.Trigger';

type PositionerProps = React.ComponentPropsWithoutRef<typeof BaseContextMenu.Positioner>;
type PositionerComponent = React.ForwardRefExoticComponent<
  PositionerProps & React.RefAttributes<HTMLDivElement>
>;

const Positioner: PositionerComponent = forwardRef<HTMLDivElement, PositionerProps>((props, ref) => {
  const { className, ...rest } = props;
  return (
    <BaseContextMenu.Positioner
      ref={ref}
      className={cs(classes['chat-context-menu__positioner'], className)}
      {...rest}
    />
  );
});

Positioner.displayName = 'ContextMenu.Positioner';

type PopupProps = React.ComponentPropsWithoutRef<typeof BaseContextMenu.Popup>;
type PopupComponent = React.ForwardRefExoticComponent<
  PopupProps & React.RefAttributes<HTMLDivElement>
>;

const Popup: PopupComponent = forwardRef<HTMLDivElement, PopupProps>((props, ref) => {
  const { className, ...rest } = props;
  return (
    <BaseContextMenu.Popup
      ref={ref}
      className={cs(classes['chat-context-menu__popup'], className)}
      finalFocus={false}
      {...rest}
    />
  );
});

Popup.displayName = 'ContextMenu.Popup';

interface PortalProps extends React.ComponentPropsWithoutRef<typeof BaseContextMenu.Portal> {
  children?: React.ReactNode;
}

const Portal = (props: PortalProps) => {
  const { children, ...rest } = props;
  return (
    <BaseContextMenu.Portal {...rest}>
      <BaseContextMenu.Backdrop className={classes['chat-context-menu__backdrop']} />
      {children}
    </BaseContextMenu.Portal>
  );
};

Portal.displayName = 'ContextMenu.Portal';

interface ContextMenuItemProps extends React.ComponentPropsWithoutRef<typeof BaseContextMenu.Item> {
  destructive?: boolean;
  icon?: React.ReactNode;
}
type ItemComponent = React.ForwardRefExoticComponent<
  ContextMenuItemProps & React.RefAttributes<HTMLDivElement>
>;

const Item: ItemComponent = forwardRef<HTMLDivElement, ContextMenuItemProps>((props, ref) => {
  const { destructive = false, icon, className, children, ...rest } = props;
  return (
    <BaseContextMenu.Item
      ref={ref}
      className={cs(
        classes['chat-context-menu__item'],
        { [classes['chat-context-menu__item--destructive']]: destructive },
        className,
      )}
      {...rest}
    >
      {icon != null && <span className={classes['chat-context-menu__item-icon']}>{icon}</span>}
      {children}
    </BaseContextMenu.Item>
  );
});

Item.displayName = 'ContextMenu.Item';

type SeparatorProps = React.ComponentPropsWithoutRef<typeof BaseContextMenu.Separator>;
type SeparatorComponent = React.ForwardRefExoticComponent<
  SeparatorProps & React.RefAttributes<HTMLDivElement>
>;

const Separator: SeparatorComponent = forwardRef<HTMLDivElement, SeparatorProps>((props, ref) => {
  const { className, ...rest } = props;
  return (
    <BaseContextMenu.Separator
      ref={ref}
      className={cs(classes['chat-context-menu__separator'], className)}
      {...rest}
    />
  );
});

Separator.displayName = 'ContextMenu.Separator';

interface GroupProps extends React.ComponentPropsWithoutRef<typeof BaseContextMenu.Group> {
  label?: string;
}
type GroupComponent = React.ForwardRefExoticComponent<
  GroupProps & React.RefAttributes<HTMLDivElement>
>;

const Group: GroupComponent = forwardRef<HTMLDivElement, GroupProps>((props, ref) => {
  const { label, className, children, ...rest } = props;
  return (
    <BaseContextMenu.Group
      ref={ref}
      className={cs(classes['chat-context-menu__group'], className)}
      {...rest}
    >
      {label && (
        <BaseContextMenu.GroupLabel className={classes['chat-context-menu__group-label']}>
          {label}
        </BaseContextMenu.GroupLabel>
      )}
      {children}
    </BaseContextMenu.Group>
  );
});

Group.displayName = 'ContextMenu.Group';

interface ContextMenuComponent {
  Root: typeof BaseContextMenu.Root;
  Trigger: typeof Trigger;
  Portal: typeof Portal;
  Positioner: typeof Positioner;
  Popup: typeof Popup;
  Item: typeof Item;
  Separator: typeof Separator;
  Group: typeof Group;
}

export const ContextMenu: ContextMenuComponent = {
  Root: BaseContextMenu.Root,
  Trigger,
  Portal,
  Positioner,
  Popup,
  Item,
  Separator,
  Group,
};