components/chat/ChatMessage.tsx

import { useEffect, useMemo, useRef, useState } from 'react';
import { CodeBlock } from '@/components/ui/CodeBlock';
import {
  IconCheckCircle,
  IconCodeWrap,
  IconCrossCircle,
  IconLandingAI,
  IconListUnordered,
  IconTerminalWindow,
  IconUser,
  IconGlowingDot,
} from '@/components/ui/Icons';
import { WIPChunkBodyGroup, 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 } from 'jotai';
import { selectedMessageId } from '@/state/chat';
import { Message } from '@prisma/client';
import { Separator } from '../ui/Separator';
import { cn } from '@/lib/utils';
import toast from 'react-hot-toast';
import {
  EXECUTE_CODE_FAILURE_TITLE,
  EXECUTE_CODE_SUCCESS_TITLE,
  EXECUTE_CODE_TITLE,
  GENERATE_CODE_TITLE,
  PLAN_TITLE,
  TOOLS_TITLE,
} from '@/lib/constants';

export interface ChatMessageProps {
  message: Message;
  loading?: boolean;
  wipAssistantMessage?: PrismaJson.MessageBody[];
}

const getParsedStreamLogs = (content: string) => {
  const streamLogs = content.split('\n').filter(log => !!log);

  const buffer = streamLogs.pop();
  const parsedStreamLogs: WIPChunkBodyGroup[] = [];
  try {
    streamLogs.forEach(streamLog =>
      parsedStreamLogs.push(JSON.parse(streamLog)),
    );
  } catch {
    toast.error('Error parsing stream logs');
  }

  if (buffer) {
    try {
      const lastLog = JSON.parse(buffer);
      parsedStreamLogs.push(lastLog);
    } catch {
      console.log(buffer);
    }
  }
  return parsedStreamLogs;
};

export const ChatMessage: React.FC<ChatMessageProps> = ({
  message,
  wipAssistantMessage,
  loading,
}) => {
  const [messageId, setMessageId] = useAtom(selectedMessageId);
  const { id, mediaUrl, prompt, response, result, responseBody } = message;
  const { formattedSections, finalResult, finalError } = useMemo(
    () =>
      formatStreamLogs(
        wipAssistantMessage ??
          responseBody ??
          (response ? getParsedStreamLogs(response) : []),
        result,
      ),
    [wipAssistantMessage, responseBody, response, result],
  );
  return (
    <div
      className={cn(
        'rounded-md bg-muted border border-muted p-4 pb-5 mb-4 relative',
        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, index) => (
                    <TableRow
                      className="border-primary/50 h-[56px]"
                      key={index}
                    >
                      <TableCell className="text-center text-webkit-center">
                        {ChunkStatusToIconDict[section.status]}
                      </TableCell>
                      <TableCell className="font-medium">
                        <ChunkTypeToText
                          useTimer={!finalResult && !finalError}
                          chunk={section}
                        />
                      </TableCell>
                      <TableCell className="text-right">
                        <ChunkPayloadAction payload={section.payload} />
                      </TableCell>
                    </TableRow>
                  ))}
                </TableBody>
              </Table>
              {finalResult && (
                <>
                  <div className="xl:hidden">
                    <CodeResultDisplay codeResult={finalResult} />
                  </div>
                  <p>✨ Coding complete</p>
                </>
              )}
              {!finalResult && finalError && (
                <>
                  <p>❌ {finalError.name}</p>
                  <div>
                    <CodeBlock
                      language="error"
                      value={
                        finalError.value +
                        '\n' +
                        finalError.traceback_raw.join('\n')
                      }
                    />
                  </div>
                </>
              )}
            </div>
          </div>
        </>
      )}
      <div
        className={cn(
          'w-1/3 h-1 rounded-full overflow-hidden bg-zinc-700 absolute left-1/2 -translate-x-1/2 bottom-2',
          loading ? 'opacity-100' : 'opacity-0',
        )}
      >
        <div className="h-full bg-primary animate-progress origin-left-right" />
      </div>
    </div>
  );
};

const ChunkStatusToIconDict: Record<
  WIPChunkBodyGroup['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 ChunkTypeToText: React.FC<{
  chunk: WIPChunkBodyGroup;
  useTimer: boolean;
}> = ({ chunk, useTimer }) => {
  const { status, type, timestamp, duration } = chunk;

  const [mSeconds, setMSeconds] = useState(0);
  const isExecuting = !['completed', 'failed'].includes(status);

  useEffect(() => {
    if (isExecuting && timestamp && useTimer) {
      const timerId = setInterval(() => {
        setMSeconds(Date.now() - Date.parse(timestamp));
      }, 200);
      return () => clearInterval(timerId);
    }
  }, [isExecuting, timestamp, useTimer]);

  const displayMs = isExecuting && useTimer ? mSeconds : duration;

  const durationDisplay = displayMs
    ? `(${Math.round(displayMs / 100) / 10}s)`
    : '';

  if (type === 'plans')
    return (
      <p>
        {PLAN_TITLE} {durationDisplay}
      </p>
    );
  if (type === 'tools')
    return (
      <p>
        {TOOLS_TITLE} {durationDisplay}
      </p>
    );
  if (type === 'code' && status === 'started')
    return (
      <p>
        {GENERATE_CODE_TITLE} {durationDisplay}
      </p>
    );
  if (type === 'code' && status === 'running')
    return (
      <p>
        {EXECUTE_CODE_TITLE} {durationDisplay}
      </p>
    );
  if (type === 'code' && status === 'completed')
    return (
      <p>
        {EXECUTE_CODE_SUCCESS_TITLE} {durationDisplay}
      </p>
    );
  if (type === 'code' && status === 'failed')
    return (
      <p>
        {EXECUTE_CODE_FAILURE_TITLE} {durationDisplay}
      </p>
    );

  return null;
};
const ChunkPayloadAction: React.FC<{
  payload: WIPChunkBodyGroup['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 if ((payload as PrismaJson.FinalCodeBody['payload']).code) {
    return (
      <Dialog>
        <DialogTrigger asChild>
          <Button variant="ghost" size="icon">
            <IconCodeWrap />
          </Button>
        </DialogTrigger>
        <DialogContent className="max-w-5xl">
          <CodeResultDisplay
            codeResult={payload as PrismaJson.FinalCodeBody['payload']}
          />
        </DialogContent>
      </Dialog>
    );
  }
  return null;
};

export default ChatMessage;