lovart3/components/elements/response.tsx

"use client";

import type { ComponentProps } 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;
  onImageAngleEdit?: (imageUrl: string) => void;
  onImageWatermarkRemove?: (imageUrl: string) => void;
  isProcessingEnhance?: boolean;
  isProcessingRmbg?: boolean;
  isProcessingExtend?: boolean;
  isProcessingAngleEdit?: boolean;
  isProcessingWatermarkRemove?: boolean;
};

function splitMarkdownImages(content: string) {
  const segments: Array<
    | { type: "text"; value: string }
    | { type: "image"; alt?: string; src: string }
  > = [];

  const imageRegex = /!\[([^\]]*)\]\((https?:\/\/[^\s)]+)\)/g;
  let lastIndex = 0;
  let match: RegExpExecArray | null;

  while ((match = imageRegex.exec(content)) !== null) {
    if (match.index > lastIndex) {
      segments.push({
        type: "text",
        value: content.slice(lastIndex, match.index),
      });
    }

    segments.push({
      type: "image",
      alt: match[1],
      src: match[2],
    });

    lastIndex = imageRegex.lastIndex;
  }

  if (lastIndex < content.length) {
    segments.push({
      type: "text",
      value: content.slice(lastIndex),
    });
  }

  return segments;
}

// 注意：移除了 memo 的自定义比较函数，确保 isProcessing* 状态变化时组件能正确重新渲染
// 之前的 memo 比较函数可能导致 Streamdown 内部的 ImagePreview 组件没有正确更新
export const Response = ({ 
  className, 
  onImageEdit, 
  onDrawingComplete, 
  onImageEnhance, 
  onImageRmbg, 
  onImageExtend, 
  onImageAngleEdit, 
  onImageWatermarkRemove,
  isProcessingEnhance, 
  isProcessingRmbg, 
  isProcessingExtend, 
  isProcessingAngleEdit, 
  isProcessingWatermarkRemove,
  ...props 
}: ResponseProps) => {
  const streamdownClassName = 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
  );
  const processingKey = `${isProcessingEnhance}-${isProcessingRmbg}-${isProcessingExtend}-${isProcessingAngleEdit}-${isProcessingWatermarkRemove}`;

  const streamdownComponents = {
    img: ({ src, alt }: { src?: string; alt?: string }) => {
      if (!src) {
        return null;
      }

      return (
        <ImagePreview
          key={`${src}-${processingKey}`}
          src={src}
          alt={alt}
          onImageEdit={onImageEdit}
          onDrawingComplete={onDrawingComplete}
          onImageEnhance={onImageEnhance}
          onImageRmbg={onImageRmbg}
          onImageExtend={onImageExtend}
          onImageAngleEdit={onImageAngleEdit}
          onImageWatermarkRemove={onImageWatermarkRemove}
          isProcessingEnhance={isProcessingEnhance}
          isProcessingRmbg={isProcessingRmbg}
          isProcessingExtend={isProcessingExtend}
          isProcessingAngleEdit={isProcessingAngleEdit}
          isProcessingWatermarkRemove={isProcessingWatermarkRemove}
        />
      );
    },
    video: ({ src, alt }: { src?: string; alt?: string }) => {
      if (!src) {
        return null;
      }
      return <VideoPreview src={src} alt={alt} />;
    },
  };

  const rawText = typeof props.children === "string" ? props.children : null;
  if (rawText && rawText.includes("![") && rawText.includes("](")) {
    const segments = splitMarkdownImages(rawText);
    const hasImageSegments = segments.some((segment) => segment.type === "image");

    if (hasImageSegments) {
      return (
        <div className={streamdownClassName}>
          {segments.map((segment, index) => {
            if (segment.type === "image") {
              return (
                <ImagePreview
                  key={`${segment.src}-${processingKey}-${index}`}
                  src={segment.src}
                  alt={segment.alt}
                  onImageEdit={onImageEdit}
                  onDrawingComplete={onDrawingComplete}
                  onImageEnhance={onImageEnhance}
                  onImageRmbg={onImageRmbg}
                  onImageExtend={onImageExtend}
                  onImageAngleEdit={onImageAngleEdit}
                  onImageWatermarkRemove={onImageWatermarkRemove}
                  isProcessingEnhance={isProcessingEnhance}
                  isProcessingRmbg={isProcessingRmbg}
                  isProcessingExtend={isProcessingExtend}
                  isProcessingAngleEdit={isProcessingAngleEdit}
                  isProcessingWatermarkRemove={isProcessingWatermarkRemove}
                />
              );
            }

            if (!segment.value.trim()) {
              return null;
            }

            return (
              <Streamdown
                className={streamdownClassName}
                components={streamdownComponents}
                key={`text-${index}`}
              >
                {segment.value}
              </Streamdown>
            );
          })}
        </div>
      );
    }
  }

  return (
    <Streamdown
      className={streamdownClassName}
      components={streamdownComponents}
      {...props}
    />
  );
};

Response.displayName = "Response";