components/chat/ChatMessage.tsx

import { useMemo, useState } from 'react';
import { CodeBlock } from '@/components/ui/CodeBlock';
import {
  IconCheckCircle,
  IconCodeWrap,
  IconCrossCircle,
  IconLandingAI,
  IconListUnordered,
  IconTerminalWindow,
  IconUser,
  IconOutput,
  IconLog,
  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 { Separator } from '../ui/Separator';
import {
  Tooltip,
  TooltipContent,
  TooltipTrigger,
} from '@/components/ui/Tooltip';

import { Dialog, DialogContent, DialogTrigger } from '../ui/Dialog';
import { Markdown } from './MemoizedReactMarkdown';
import Img from '../ui/Img';

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

export function ChatMessage({ message, isLoading }: ChatMessageProps) {
  const { role, content, mediaUrl } = message;

  return role === 'user' ? (
    <UserChatMessage content={content} mediaUrl={mediaUrl} />
  ) : (
    <AssistantChatMessage content={content} />
  );
}

const UserChatMessage: React.FC<{
  content: string;
  mediaUrl?: string;
}> = ({ content, mediaUrl }) => {
  return (
    <div className="group relative mb-6 flex rounded-md bg-muted p-6 ml-auto mr-0 w-3/5">
      <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-3 overflow-hidden">
        <p>{content}</p>
        {mediaUrl && (
          <>
            {mediaUrl?.endsWith('.mp4') ? (
              <video src={mediaUrl} controls width={500} height={500} />
            ) : (
              <Img
                src={mediaUrl}
                alt={mediaUrl}
                quality={100}
                sizes="(min-width: 66em) 15vw,
                        (min-width: 44em) 20vw,
                        100vw"
              />
            )}
          </>
        )}
      </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>
    );
  }
};

const CodeResultDisplay: React.FC<{
  codeResult: CodeResult;
}> = ({ codeResult }) => {
  const { code, test, result } = codeResult;
  const getDetail = () => {
    if (!result) return {};
    try {
      const detail = JSON.parse(result);
      return {
        results: detail.results,
        stderr: detail.logs.stderr,
        stdout: detail.logs.stdout,
      };
    } catch {
      return {};
    }
  };

  const { results, stderr, stdout } = getDetail();

  return (
    <div className="rounded-lg overflow-hidden relative max-w-5xl">
      <CodeBlock language="python" value={code} />
      <div className="rounded-lg relative">
        <div className="absolute left-1/2 -translate-x-1/2 -top-4 z-10">
          <Dialog>
            <DialogTrigger asChild>
              <Button variant="ghost" size="icon" className="size-8">
                <IconTerminalWindow className="text-teal-500 size-4" />
              </Button>
            </DialogTrigger>
            <DialogContent className="max-w-5xl">
              <CodeBlock language="python" value={test} />
            </DialogContent>
          </Dialog>
          {Array.isArray(stderr) && (
            <Dialog>
              <DialogTrigger asChild>
                <Button variant="ghost" size="icon" className="size-8">
                  <IconLog className="text-gray-500 size-4" />
                </Button>
              </DialogTrigger>
              <DialogContent className="max-w-5xl">
                <CodeBlock language="vim" value={stderr.join('').trim()} />
              </DialogContent>
            </Dialog>
          )}
        </div>
      </div>
      {Array.isArray(stdout) && !!stdout.join('').trim() && (
        <>
          <Separator />
          <CodeBlock language="print" value={stdout.join('').trim()} />
        </>
      )}
      {!!results.length && (
        <>
          <Separator />
          <CodeBlock language="output" value={results} />
        </>
      )}
      <Separator />
    </div>
  );
};

const AssistantChatMessage: React.FC<{
  content: string;
}> = ({ content }) => {
  const [formattedSections, codeResult] = useMemo(
    () => formatStreamLogs(content),
    [content],
  );

  return (
    <div className="group relative mb-6 flex rounded-md bg-muted p-6 w-full">
      <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 && <CodeResultDisplay codeResult={codeResult} />}
      </div>
    </div>
  );
};

export default UserChatMessage;