MessageActions

import React, { memo, useRef, useState } from 'react';
import { CopyButton, Group, Popover, Text } from '@mantine/core';
import {
  IconArrowBackUp,
  IconCheck,
  IconCopy,
  IconGitBranch,
  IconPencil,
  IconRefresh,
  IconThumbDown,
  IconThumbDownFilled,
  IconThumbUp,
  IconThumbUpFilled,
} from '@tabler/icons-react';
import { cx } from '../../utils/cx';
import { getErrorMessage } from '../../utils/error-message';
import { FeedbackForm, type FeedbackFormLabels } from '../FeedbackForm/FeedbackForm';
import type { FeedbackDetails, FeedbackReason, MessageFeedbackValue } from '../types';
import { usePendingActions } from '../../hooks/use-pending-actions';
import { MessageActionButton } from '../MessageActionButton/MessageActionButton';
import classes from './MessageActions.module.css';

export interface MessageActionsLabels {
  toolbar: string;
  copy: string;
  copied: string;
  edit: string;
  retry: string;
  rewind: string;
  branch: string;
  good: string;
  bad: string;
  error: string;
  /** Labels of the feedback form */
  feedbackForm: Partial<FeedbackFormLabels>;
}

export const DEFAULT_MESSAGE_ACTIONS_LABELS: MessageActionsLabels = {
  toolbar: 'Message actions',
  copy: 'Copy',
  copied: 'Copied',
  edit: 'Edit',
  retry: 'Retry',
  rewind: 'Rewind to here',
  branch: 'Branch from here',
  good: 'Good response',
  bad: 'Bad response',
  error: 'Action failed',
  feedbackForm: {},
};

export interface MessageActionsProps {
  /** Author of the message: user messages get edit and rewind, assistant messages get retry, feedback and branch */
  messageRole: 'user' | 'assistant';
  /** Text copied by the copy action, the action is hidden when empty */
  text?: string;
  /** Time label rendered before the actions */
  timestamp?: string;
  /** Opens inline editing of a user message */
  onEdit?: () => void;
  /** Regenerates an assistant message; the button shows a loader until the promise settles */
  onRetry?: () => void | Promise<void>;
  /** Returns the conversation to this message */
  onRewind?: () => void | Promise<void>;
  /** Starts a new branch from this message */
  onBranch?: () => void | Promise<void>;
  /** Rates an assistant message, adds the thumbs up and thumbs down buttons */
  onFeedback?: (value: MessageFeedbackValue, details?: FeedbackDetails) => void | Promise<void>;
  /** Controlled rating */
  feedback?: MessageFeedbackValue | null;
  /** Reasons offered after a thumbs down */
  feedbackReasons?: FeedbackReason[];
  /** Disables every action except copy, for example while the agent is responding */
  disabled?: boolean;
  /** `hover` reveals the actions on hover and focus of the toolbar or of an ancestor with `data-message-actions-host`, touch devices always see them; `always` keeps them visible, `hover` by default */
  visibility?: 'hover' | 'always';
  /** Horizontal placement of the actions, `start` by default */
  align?: 'start' | 'end';
  /** Host buttons after the built-in actions, usually `MessageActionButton`s */
  actions?: React.ReactNode;
  /** Overrides of the default English labels */
  labels?: Partial<MessageActionsLabels>;
  /** Class name added to the root element */
  className?: string;
  /** Inline styles added to the root element */
  style?: React.CSSProperties;
}

const ICON_SIZE = 16;

/** Toolbar under a chat message: copy, edit, retry, rewind, branch and feedback, each shown when its callback is set */
export const MessageActions = memo(function MessageActions({
  messageRole,
  text,
  timestamp,
  onEdit,
  onRetry,
  onRewind,
  onBranch,
  onFeedback,
  feedback,
  feedbackReasons,
  disabled = false,
  visibility = 'hover',
  align = 'start',
  actions,
  labels: labelsProp,
  className,
  style,
}: MessageActionsProps) {
  const labels = { ...DEFAULT_MESSAGE_ACTIONS_LABELS, ...labelsProp };
  const { isPending: actionPending, run } = usePendingActions(labels.error);
  const [internalFeedback, setInternalFeedback] = useState<MessageFeedbackValue | null>(null);
  const [feedbackOpen, setFeedbackOpen] = useState(false);
  const [feedbackError, setFeedbackError] = useState<string | null>(null);
  const downSentRef = useRef(false);
  const rating =
    feedback === undefined || (feedbackOpen && internalFeedback) ? internalFeedback : feedback;
  const isUser = messageRole === 'user';
  const locked = disabled || actionPending();

  const rate = async (value: MessageFeedbackValue) => {
    if (!onFeedback) {
      return;
    }
    setFeedbackError(null);
    setInternalFeedback(value);
    setFeedbackOpen(true);
    if (value === 'down') {
      downSentRef.current = false;
      return;
    }
    try {
      await onFeedback('up');
    } catch (error) {
      setInternalFeedback(null);
      setFeedbackError(getErrorMessage(error, labels.error));
    }
  };

  const copy = text ? (
    <CopyButton value={text} timeout={2000}>
      {({ copied, copy: copyText }) => (
        <MessageActionButton
          label={copied ? labels.copied : labels.copy}
          onClick={copyText}
          icon={copied ? <IconCheck size={ICON_SIZE} /> : <IconCopy size={ICON_SIZE} />}
        />
      )}
    </CopyButton>
  ) : null;

  const feedbackButtons = !isUser && onFeedback && (
    <Popover
      opened={feedbackOpen}
      onChange={(opened) => {
        setFeedbackOpen(opened);
        if (!opened && internalFeedback === 'down' && !downSentRef.current) {
          setInternalFeedback(null);
        }
      }}
      position="bottom-start"
      width={300}
      shadow="md"
      withArrow
      trapFocus={rating === 'down'}
    >
      <Popover.Target>
        <Group gap={2} wrap="nowrap">
          <MessageActionButton
            label={labels.good}
            active={rating === 'up'}
            disabled={disabled}
            onClick={() => rate('up')}
            icon={
              rating === 'up' ? (
                <IconThumbUpFilled size={ICON_SIZE} />
              ) : (
                <IconThumbUp size={ICON_SIZE} />
              )
            }
          />
          <MessageActionButton
            label={labels.bad}
            active={rating === 'down'}
            disabled={disabled}
            onClick={() => rate('down')}
            icon={
              rating === 'down' ? (
                <IconThumbDownFilled size={ICON_SIZE} />
              ) : (
                <IconThumbDown size={ICON_SIZE} />
              )
            }
          />
        </Group>
      </Popover.Target>
      <Popover.Dropdown>
        {feedbackError ? (
          <Text size="sm" c="var(--ae-danger)">
            {feedbackError}
          </Text>
        ) : (
          <FeedbackForm
            value={internalFeedback ?? 'up'}
            reasons={feedbackReasons}
            labels={labels.feedbackForm}
            onSubmit={async (details) => {
              await onFeedback('down', details);
              downSentRef.current = true;
            }}
            onSkip={async () => {
              await onFeedback('down');
              downSentRef.current = true;
            }}
          />
        )}
      </Popover.Dropdown>
    </Popover>
  );

  return (
    <Group
      role="toolbar"
      aria-label={labels.toolbar}
      gap={2}
      wrap="nowrap"
      justify={align === 'end' ? 'flex-end' : 'flex-start'}
      className={cx(classes.root, className)}
      style={style}
      data-part="message-actions"
      data-visibility={visibility}
      data-active={feedbackOpen || actionPending() || undefined}
    >
      {timestamp && (
        <Text component="span" size="xs" c="dimmed" px={4}>
          {timestamp}
        </Text>
      )}
      {isUser && onRewind && (
        <MessageActionButton
          label={labels.rewind}
          disabled={locked}
          loading={actionPending('rewind')}
          onClick={() => run('rewind', onRewind, { exclusive: true })}
          icon={<IconArrowBackUp size={ICON_SIZE} />}
        />
      )}
      {isUser && onEdit && (
        <MessageActionButton
          label={labels.edit}
          action="edit"
          disabled={locked}
          onClick={onEdit}
          icon={<IconPencil size={ICON_SIZE} />}
        />
      )}
      {copy}
      {!isUser && onRetry && (
        <MessageActionButton
          label={labels.retry}
          disabled={locked}
          loading={actionPending('retry')}
          onClick={() => run('retry', onRetry, { exclusive: true })}
          icon={<IconRefresh size={ICON_SIZE} />}
        />
      )}
      {feedbackButtons}
      {!isUser && onBranch && (
        <MessageActionButton
          label={labels.branch}
          disabled={locked}
          loading={actionPending('branch')}
          onClick={() => run('branch', onBranch, { exclusive: true })}
          icon={<IconGitBranch size={ICON_SIZE} />}
        />
      )}
      {actions}
    </Group>
  );
});

MessageActions.displayName = 'MessageActions';