components/chat/ChatMessage.tsx

// Inspired by Chatbot-UI and modified to fit the needs of this project
// @see https://github.com/mckaywrigley/chatbot-ui/blob/main/components/Chat/ChatMessage.tsx

import remarkGfm from 'remark-gfm';
import remarkMath from 'remark-math';
import rehypeRaw from 'rehype-raw';

import { useMemo, useState } from 'react';
import { cn } from '@/lib/utils';
import { CodeBlock } from '@/components/ui/CodeBlock';
import { MemoizedReactMarkdown } from '@/components/chat/MemoizedReactMarkdown';
import { IconLandingAI, IconUser } from '@/components/ui/Icons';
import { MessageBase } from '../../lib/types';
import Img from '../ui/Img';
import { getCleanedUpMessages } from '@/lib/messageUtils';
import Loading from '../ui/Loading';
import {
  Table,
  TableCell,
  TableHead,
  TableHeader,
  TableRow,
} from '../ui/Table';
import { Button } from '../ui/Button';
import { Dialog, DialogContent } from '../ui/Dialog';

export interface ChatMessageProps {
  message: MessageBase;
  isLoading: boolean;
}

const Markdown: React.FC<{
  content: string;
  setDetails?: (val: string) => void;
}> = ({ content, setDetails }) => {
  return (
    <>
      <MemoizedReactMarkdown
        className="break-words overflow-auto"
        remarkPlugins={[remarkGfm, remarkMath]}
        rehypePlugins={[rehypeRaw] as any}
        components={{
          table({ children, ...props }) {
            return <Table {...props}>{children}</Table>;
          },
          thead({ children, ...props }) {
            return <TableHeader {...props}>{children}</TableHeader>;
          },
          th({ children, ...props }) {
            return <TableHead {...props}>{children}</TableHead>;
          },
          tr({ children, ...props }) {
            return <TableRow {...props}>{children}</TableRow>;
          },
          td({ children, ...props }) {
            return <TableCell {...props}>{children}</TableCell>;
          },
          button({ children, ...props }) {
            if ('data-details' in props && setDetails) {
              return (
                <Button
                  {...props}
                  onClick={() =>
                    setDetails(decodeURI(props['data-details'] as string))
                  }
                >
                  {children}
                </Button>
              );
            }
            return <Button {...props}>{children}</Button>;
          },
          p({ children, ...props }) {
            if (
              props.node.children.some(
                child => child.type === 'element' && child.tagName === 'img',
              )
            ) {
              return (
                <p className="flex flex-wrap gap-2 items-start">{children}</p>
              );
            }
            return (
              <p className="mb-2 last:mb-0 whitespace-pre-line">{children}</p>
            );
          },
          img(props) {
            if (props.src?.endsWith('.mp4')) {
              return (
                <video src={props.src} controls width={500} height={500} />
              );
            }
            return (
              <Img
                src={props.src ?? '/landing.png'}
                alt={props.alt ?? 'answer-image'}
                quality={100}
                sizes="(min-width: 66em) 15vw,
                        (min-width: 44em) 20vw,
                        100vw"
              />
            );
          },
          code({ node, inline, className, children, ...props }) {
            // if (children.length) {
            //   if (children[0] == '▍') {
            //     return (
            //       <span className="mt-1 cursor-default animate-pulse">▍</span>
            //     );
            //   }

            //   children[0] = (children[0] as string).replace('`▍`', '▍');
            // }

            const match = /language-(\w+)/.exec(className || '');
            // if (inline) {
            //   return (
            //     <code className={className} {...props}>
            //       {children}
            //     </code>
            //   );
            // }

            return (
              <CodeBlock
                key={Math.random()}
                language={(match && match[1]) || ''}
                value={String(children).replace(/\n$/, '')}
                {...props}
              />
            );
          },
        }}
      >
        {content}
      </MemoizedReactMarkdown>
    </>
  );
};

export function ChatMessage({
  message,
  isLoading,
}: ChatMessageProps) {
  const { content } = useMemo(() => {
    return getCleanedUpMessages({
      content: message.content,
      role: message.role,
    });
  }, [message.content, message.role]);
  console.log('[Ming] content:', content);
  console.log('[Ming] raw:', message.content);
  const [details, setDetails] = useState<string>('');
  return (
    <div
      className={cn(
        'group relative mb-6 flex rounded-md bg-muted px-4 py-6 w-3/5',
        message.role === 'user' ? 'ml-auto mr-0 w-3/5' : 'w-4/5',
      )}
    >
      <div
        className={cn(
          'flex size-8 shrink-0 select-none items-center justify-center rounded-md border shadow',
          message.role === 'user'
            ? 'bg-background'
            : 'bg-primary text-primary-foreground',
        )}
      >
        {message.role === 'user' ? <IconUser /> : <IconLandingAI />}
      </div>
      <div className="flex-1 px-1 ml-4 space-y-2 overflow-hidden">
        {content && <Markdown content={content} setDetails={setDetails} />}
        {isLoading && <Loading />}
      </div>
      <Dialog open={!!details} onOpenChange={open => !open && setDetails('')}>
        <DialogContent className="w-11/12">
          <Markdown content={details} />
        </DialogContent>
      </Dialog>
    </div>
  );
}