lovart3/components/elements/response.tsx

"use client";

import { type ComponentProps, memo } from "react";
import { Streamdown } from "streamdown";
import { cn } from "@/lib/utils";
import { ImagePreview } from "@/components/image-preview";
import { VideoPreview } from "@/components/video-preview";

type ResponseProps = ComponentProps<typeof Streamdown> & {
  onImageEdit?: (imageUrl: string) => void;
  onDrawingComplete?: (annotatedImageUrl: string) => void;
  onImageEnhance?: (imageUrl: string) => void;
  onImageRmbg?: (imageUrl: string) => void;
  onImageExtend?: (imageUrl: string) => void;
};

export const Response = memo(
  ({ className, onImageEdit, onDrawingComplete, onImageEnhance, onImageRmbg, onImageExtend, ...props }: ResponseProps) => (
    <Streamdown
      className={cn(
        "size-full [&>*:first-child]:mt-0 [&>*:last-child]:mb-0 [&_code]:whitespace-pre-wrap [&_code]:break-words [&_pre]:max-w-full [&_pre]:overflow-x-auto",
        className
      )}
      components={{
        img: ({ src, alt }: { src?: string; alt?: string }) => {
          if (!src) {
            return null;
          }
          return <ImagePreview 
            src={src} 
            alt={alt} 
            onImageEdit={onImageEdit} 
            onDrawingComplete={onDrawingComplete}
            onImageEnhance={onImageEnhance}
            onImageRmbg={onImageRmbg}
            onImageExtend={onImageExtend}
          />;
        },
        video: ({ src, alt }: { src?: string; alt?: string }) => {
          if (!src) {
            return null;
          }
          return <VideoPreview src={src} alt={alt} />;
        },
      }}
      {...props}
    />
  ),
  (prevProps: ResponseProps, nextProps: ResponseProps) => 
    prevProps.children === nextProps.children && 
    prevProps.onImageEdit === nextProps.onImageEdit && 
    prevProps.onDrawingComplete === nextProps.onDrawingComplete &&
    prevProps.onImageEnhance === nextProps.onImageEnhance &&
    prevProps.onImageRmbg === nextProps.onImageRmbg &&
    prevProps.onImageExtend === nextProps.onImageExtend
);

Response.displayName = "Response";