components/CodeResultDisplay.tsx

import React from 'react';

import { CodeBlock } from './ui/CodeBlock';
import {
  Dialog,
  DialogTrigger,
  DialogContent,
  DialogHeader,
  DialogTitle,
} from './ui/Dialog';
import { Button } from './ui/Button';
import { IconLog, IconTerminalWindow } from './ui/Icons';
import { Separator } from './ui/Separator';
import Img from './ui/Img';
import {
  Carousel,
  CarouselContent,
  CarouselItem,
  CarouselNext,
  CarouselPrevious,
} from './ui/carousel';

export interface CodeResultDisplayProps {}

const CodeResultDisplay: React.FC<{
  codeResult: PrismaJson.FinalCodeBody['payload'];
}> = ({ codeResult }) => {
  const { code, test, result } = codeResult;
  const getDetail = () => {
    if (!result) return {};
    try {
      // IMPORTANT: This is for backwards compatibility with old chat that save result as JSON string
      // updated in https://github.com/landing-ai/vision-agent-ui/pull/86
      const detail =
        typeof result === 'object'
          ? result
          : (JSON.parse(result) as PrismaJson.StructuredResult);
      return {
        results: detail.results,
        stderr: detail.logs.stderr,
        stdout: detail.logs.stdout,
        error: detail.error,
      };
    } catch {
      return {};
    }
  };

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

  const imageResults = results?.filter(_ => !!_.png).map(_ => _.png);
  const videoResults = results?.filter(_ => !!_.mp4).map(_ => _.mp4);
  const finalResult = results?.find(_ => _.is_main_result)?.text;

  return (
    <div
      className="rounded-lg overflow-hidden relative max-w-5xl"
      data-testid="code-result-display-container"
    >
      <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
                data-testid="open-final-test-code-dialog-button"
                variant="ghost"
                size="icon"
                className="size-8"
              >
                <IconTerminalWindow className="text-teal-500 size-4" />
              </Button>
            </DialogTrigger>
            <DialogContent className="max-w-5xl">
              <DialogHeader>
                <DialogTitle>Test code</DialogTitle>
              </DialogHeader>
              <CodeBlock language="python" value={test} />
            </DialogContent>
          </Dialog>
          {Array.isArray(stderr) && !!stderr.join('').trim() && (
            <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()} />
        </>
      )}
      {!!error && (
        <>
          <Separator />
          <CodeBlock
            language="error"
            value={
              error.name +
              '\n' +
              error.value +
              '\n' +
              error.traceback_raw.join('\n')
            }
          />
        </>
      )}
      {!!imageResults.length && (
        <div className="p-4 text-xs lowercase bg-zinc-900 space-y-4 border-t border-muted">
          <div className="flex items-center justify-between">
            <p>image output</p>
            <Dialog>
              <DialogTrigger asChild>
                <Button
                  variant="outline"
                  size="sm"
                  data-testid="view-all-result-images-button"
                >
                  View all
                </Button>
              </DialogTrigger>
              <DialogContent className="max-w-5xl flex justify-center items-center">
                <Carousel className="w-3/4">
                  <CarouselContent>
                    {imageResults.map((png, index) => (
                      <CarouselItem key={'png' + index}>
                        <Img
                          src={png!}
                          width={1200}
                          alt={`detail-result-image-${index}`}
                        />
                      </CarouselItem>
                    ))}
                  </CarouselContent>
                  <CarouselPrevious />
                  <CarouselNext />
                </Carousel>
              </DialogContent>
            </Dialog>
          </div>
          <div className="flex flex-row space-x-4 overflow-auto">
            {imageResults.map((png, index) => (
              <Img
                key={'png' + index}
                src={png!}
                width={200}
                alt={`result-image-${index}`}
              />
            ))}
          </div>
        </div>
      )}
      {!!videoResults.length && (
        <div className="p-4 text-xs lowercase bg-zinc-900 space-y-4 border-t border-muted">
          <p>video output</p>
          <div className="flex flex-row space-x-4 overflow-auto">
            {videoResults.map((mp4, index) => (
              <Dialog key={'png' + index}>
                <DialogTrigger asChild>
                  <video src={mp4} controls width={400} height={400} />
                </DialogTrigger>
                <DialogContent className="max-w-5xl">
                  <video src={mp4} controls width={400} height={400} />
                </DialogContent>
              </Dialog>
            ))}
          </div>
        </div>
      )}
      {!!finalResult && <CodeBlock language="output" value={finalResult} />}
    </div>
  );
};

export default CodeResultDisplay;