components/chat/ChatMessage.tsx

import { useMemo, useState } from 'react';
import { CodeBlock } from '@/components/ui/CodeBlock';
import {
  IconCheckCircle,
  IconCodeWrap,
  IconCrossCircle,
  IconLandingAI,
  IconListUnordered,
  IconTerminalWindow,
  IconUser,
  IconGlowingDot,
} from '@/components/ui/Icons';
import { MessageUI } from '@/lib/types';
import { ChunkBody, CodeResult, formatStreamLogs } from '@/lib/utils/content';
import {
  Table,
  TableBody,
  TableCell,
  TableHead,
  TableHeader,
  TableRow,
} from '../ui/Table';
import { Button } from '../ui/Button';
import { Dialog, DialogContent, DialogTrigger } from '../ui/Dialog';
import Img from '../ui/Img';
import CodeResultDisplay from '../CodeResultDisplay';
import { useAtom, useSetAtom } from 'jotai';
import { selectedMessageId } from '@/state/chat';
import { Message } from '@prisma/client';
import { Separator } from '../ui/Separator';
import { cn } from '@/lib/utils';

export interface ChatMessageProps {
  message: Message;
  wipAssistantMessage?: MessageUI;
}

export const ChatMessage: React.FC<ChatMessageProps> = ({
  message,
  wipAssistantMessage,
}) => {
  const [messageId, setMessageId] = useAtom(selectedMessageId);
  const { id, mediaUrl, prompt, response, result } = message;
  const [formattedSections, codeResult] = useMemo(
    () => formatStreamLogs(response ?? wipAssistantMessage?.content),
    [response, wipAssistantMessage?.content],
  );
  return (
    <div
      className={cn(
        'rounded-md bg-muted border border-muted p-4 mb-4',
        messageId === id && 'lg:border-primary/50',
        result && 'lg:cursor-pointer',
      )}
      onClick={() => {
        if (result) {
          setMessageId(id);
        }
      }}
    >
      <div className="flex">
        <div className="flex size-8 shrink-0 select-none items-center justify-center rounded-md border shadow bg-background">
          <IconUser />
        </div>
        <div className="flex-1 px-1 ml-4 space-y-2 overflow-hidden">
          <p>{prompt}</p>
          {mediaUrl && (
            <>
              {mediaUrl?.endsWith('.mp4') ? (
                <video src={mediaUrl} controls width={400} height={400} />
              ) : (
                <Dialog>
                  <DialogTrigger asChild>
                    <Img src={mediaUrl} alt={mediaUrl} width={300} />
                  </DialogTrigger>
                  <DialogContent className="max-w-5xl">
                    <Img src={mediaUrl} alt={mediaUrl} quality={100} />
                  </DialogContent>
                </Dialog>
              )}
            </>
          )}
        </div>
      </div>
      {!!formattedSections.length && (
        <>
          <Separator className="bg-primary/30 my-4" />
          <div className="flex">
            <div className="flex size-8 shrink-0 select-none items-center justify-center rounded-md border shadow bg-primary text-primary-foreground">
              <IconLandingAI />
            </div>
            <div className="flex-1 px-1 space-y-4 ml-4 overflow-hidden">
              <Table className="w-[400px]">
                <TableBody>
                  {formattedSections.map(section => (
                    <TableRow
                      className="border-primary/50 h-[56px]"
                      key={section.type}
                    >
                      <TableCell className="text-center text-webkit-center">
                        {ChunkStatusToIconDict[section.status]}
                      </TableCell>
                      <TableCell className="font-medium">
                        {ChunkTypeToTextDict[section.type]}
                      </TableCell>
                      <TableCell className="text-right">
                        <ChunkPayloadAction payload={section.payload} />
                      </TableCell>
                    </TableRow>
                  ))}
                </TableBody>
              </Table>
              {codeResult && (
                <div className="xl:hidden">
                  <CodeResultDisplay codeResult={codeResult} />
                </div>
              )}
              {codeResult && <p>✨ Coding complete</p>}
            </div>
          </div>
        </>
      )}
    </div>
  );
};

const ChunkStatusToIconDict: Record<ChunkBody['status'], React.ReactElement> = {
  started: <IconGlowingDot className="bg-yellow-500/80" />,
  completed: <IconCheckCircle className="text-green-500" />,
  running: <IconGlowingDot className="bg-teal-500/80" />,
  failed: <IconCrossCircle className="text-red-500" />,
};
const ChunkTypeToTextDict: Record<ChunkBody['type'], string> = {
  plans: 'Creating instructions',
  tools: 'Retrieving tools',
  code: 'Generating code',
  final_code: 'Final result',
};
const ChunkPayloadAction: React.FC<{
  payload: ChunkBody['payload'];
}> = ({ payload }) => {
  if (!payload) return null;
  if (Array.isArray(payload)) {
    // [{title: 123, content, 345}, {title: ..., content: ...}] => ['title', 'content']
    const keyArray = Array.from(
      payload.reduce((acc, curr) => {
        Object.keys(curr).forEach(key => acc.add(key));
        return acc;
      }, new Set<string>()),
    );

    return (
      <Dialog>
        <DialogTrigger asChild>
          <Button variant="ghost" size="icon">
            <IconListUnordered />
          </Button>
        </DialogTrigger>
        <DialogContent
          className="max-w-5xl"
          onOpenAutoFocus={e => e.preventDefault()}
        >
          <Table>
            <TableHeader>
              <TableRow className="border-primary/50">
                {keyArray.map(header => (
                  <TableHead key={header}>{header}</TableHead>
                ))}
              </TableRow>
            </TableHeader>
            <TableBody>
              {payload.map((line, index) => (
                <TableRow className="border-primary/50" key={index}>
                  {keyArray.map(header =>
                    header === 'documentation' ? (
                      <TableCell key={header}>
                        <Dialog>
                          <DialogTrigger asChild>
                            <Button
                              variant="ghost"
                              size="icon"
                              className="size-8 ml-[40%]"
                            >
                              <IconTerminalWindow className="text-teal-500 size-4" />
                            </Button>
                          </DialogTrigger>
                          <DialogContent className="max-w-5xl">
                            <CodeBlock language="md" value={line[header]} />
                          </DialogContent>
                        </Dialog>
                      </TableCell>
                    ) : (
                      <TableCell key={header}>{line[header]}</TableCell>
                    ),
                  )}
                </TableRow>
              ))}
            </TableBody>
          </Table>
        </DialogContent>
      </Dialog>
    );
  } else {
    return (
      <Dialog>
        <DialogTrigger asChild>
          <Button variant="ghost" size="icon">
            <IconCodeWrap />
          </Button>
        </DialogTrigger>
        <DialogContent className="max-w-5xl">
          <CodeResultDisplay codeResult={payload as CodeResult} />
        </DialogContent>
      </Dialog>
    );
  }
};

export default ChatMessage;