ai-chatbot/lib/agents/detail-replacement-batch-runner.ts

import "server-only";

import { and, eq } from "drizzle-orm";
import { db } from "@/lib/db";
import { canvasImageTask, canvasProject, generatedImage, message } from "@/lib/db/schema";
import { getEnabledAiModelsWithProvider } from "@/lib/db/queries";
import { generateCanvasImage } from "@/lib/canvas/generate-image";
import { saveGeneratedImages } from "@/lib/ai/runtime/generated-image-persistence";
import {
  buildAgentToolImagesMarkdown,
  prependImageModelBadge,
} from "@/lib/ai/runtime/agent-media-markdown";
import {
  selectDetailReplacementGenerationInputs,
  shouldIncludeDetailReplacementModelImage,
} from "@/lib/ai/detail-replacement-inputs";
import { isRemoteImageAccessible } from "@/app/(chat)/api/chat/_lib/image-media";
import {
  analyzeReplacementExecutionPlan,
  buildAgentReplacementPrompt,
  isMarketplaceProductDetailReferenceId,
  normalizeReplacementResultImage,
  resolveReplacementReferenceGeometry,
  resolveReplacementScenePlanningRuntime,
} from "@/app/(chat)/api/chat/_lib/detail-replacement";
import { persistAgentDetailReplacementCard } from "@/lib/agents/agent-project-result-persistence";
import { extractAgentProjectStateFromCanvasData } from "@/lib/agents/agent-project-compat";
import {
  completedDetailReplacementCardIds,
  detailReplacementCardMessageId,
  type DetailReplacementBatchCard,
  type DetailReplacementBatchRequest,
  type DetailReplacementBatchResult,
} from "@/lib/agents/detail-replacement-batch";

export async function executeDetailReplacementBatch(params: {
  taskId: string;
  userId: string;
  request: DetailReplacementBatchRequest;
}): Promise<DetailReplacementBatchResult> {
  const { taskId, userId, request } = params;
  const [storedTask] = await db.select({ result: canvasImageTask.result })
    .from(canvasImageTask)
    .where(and(eq(canvasImageTask.id, taskId), eq(canvasImageTask.userId, userId)))
    .limit(1);
  const storedResult = storedTask?.result as DetailReplacementBatchResult | null;
  const cards = Array.isArray(storedResult?.cards) ? [...storedResult.cards] : [];
  const completedIds = completedDetailReplacementCardIds(storedResult);
  const models = await getEnabledAiModelsWithProvider();
  const planningRuntime = await resolveReplacementScenePlanningRuntime({ models });

  const publish = async (currentCardId: string | null) => {
    const result: DetailReplacementBatchResult = {
      taskKind: "agent-detail-replacement-batch",
      cards: [...cards],
      currentCardId,
    };
    await db.update(canvasImageTask).set({ result, updatedAt: new Date() })
      .where(and(eq(canvasImageTask.id, taskId), eq(canvasImageTask.status, "running")));
    return result;
  };

  const persistAssistantAndAlbum = async (params: {
    card: DetailReplacementBatchCard;
    text: string;
    prompt: string;
    sourceImageUrls: string[];
  }) => {
    const messageId = detailReplacementCardMessageId(taskId, params.card.cardId);
    await db.insert(message).values({
      id: messageId,
      chatId: request.chatId,
      role: "assistant",
      parts: [{ type: "text", text: params.text }],
      attachments: [],
      createdAt: new Date(),
    }).onConflictDoUpdate({
      target: message.id,
      set: { parts: [{ type: "text", text: params.text }], createdAt: new Date() },
    });
    if (params.card.status !== "completed") return;
    const [existingImage] = await db.select({ id: generatedImage.id })
      .from(generatedImage)
      .where(and(eq(generatedImage.messageId, messageId), eq(generatedImage.userId, userId)))
      .limit(1);
    if (existingImage) return;
    await saveGeneratedImages(
      params.text, userId, request.chatId, messageId, params.prompt, request.model,
      {
        sourceImageUrls: params.sourceImageUrls,
        relationType: "edited_from",
        sourceType: "detail_replacement",
        projectId: request.projectId,
        userPrompt: request.userPrompt,
        optimizedPrompt: params.prompt,
        requestedImageCount: 1,
        aspectRatio: "auto",
        imageSize: request.imageSize,
      },
    ).catch((error) => {
      console.error("[Detail Replacement Batch] Album persistence failed", {
        taskId, cardId: params.card.cardId, error,
      });
    });
  };

  for (const [index, reference] of request.references.entries()) {
    if (completedIds.has(reference.id)) continue;
    await publish(reference.id);
    const startedAt = Date.now();
    let card: DetailReplacementBatchCard;
    let sourceImageUrls = [...request.productImageUrls, reference.imageUrl];
    let prompt = request.prompt;
    const [storedProject] = await db.select({ canvasData: canvasProject.canvasData })
      .from(canvasProject)
      .where(and(eq(canvasProject.id, request.projectId), eq(canvasProject.userId, userId)))
      .limit(1);
    const recoveredRecord = extractAgentProjectStateFromCanvasData(storedProject?.canvasData)
      .replacementResultState?.detail_replacement?.[reference.id];
    if (
      recoveredRecord?.url &&
      recoveredRecord.status !== "fallback" &&
      recoveredRecord.text.includes(`<!-- AGENT_REPLACEMENT_RUN_ID:${taskId} -->`)
    ) {
      card = {
        cardId: reference.id,
        status: recoveredRecord.status,
        imageUrl: recoveredRecord.url,
        thumbnailUrl: recoveredRecord.url,
        text: recoveredRecord.text,
      };
      await persistAssistantAndAlbum({
        card, text: recoveredRecord.text, prompt, sourceImageUrls,
      });
      const previousIndex = cards.findIndex((item) => item.cardId === card.cardId);
      if (previousIndex >= 0) cards.splice(previousIndex, 1);
      cards.push(card);
      completedIds.add(reference.id);
      await publish(null);
      continue;
    }
    try {
      if (reference.passthrough) {
        card = {
          cardId: reference.id,
          status: "passthrough",
          imageUrl: reference.imageUrl,
          thumbnailUrl: reference.imageUrl,
          text: "已按人工标记沿用参考图，直接应用当前切片，不参与继续生成。",
        };
      } else {
        const plan = await analyzeReplacementExecutionPlan({
          modeType: "detail_replacement",
          replacementStrategy: request.replacementStrategy,
          referenceImageUrl: reference.imageUrl,
          referenceLabel: reference.label,
          productImageUrls: request.productImageUrls,
          modelImageUrl: request.modelImageUrl,
          workflowPrompt: request.replacementStrategy === "product_only" ? null : request.prompt,
          planningRuntime,
          models,
        });
        const geometry = await resolveReplacementReferenceGeometry(reference.imageUrl);
        const selectedIndexes = plan.scenePlan.selectedProductImageIndexes
          .map((value) => value - 1)
          .filter((value) => value >= 0 && value < request.productImageUrls.length);
        const selectedProducts = selectedIndexes.length > 0
          ? selectedIndexes.map((value) => request.productImageUrls[value]!)
          : request.productImageUrls.slice(0, 2);
        const inputs = selectDetailReplacementGenerationInputs({
          productImageUrls: selectedProducts,
          modelImageUrl: request.modelImageUrl,
          includeModelImage: Boolean(
            request.replacementStrategy !== "product_only" &&
            request.modelImageUrl &&
            (shouldIncludeDetailReplacementModelImage(plan.scenePlan) ||
              selectedProducts.length === 0),
          ),
          rerunReferenceImageUrls: reference.rerunReferenceImages?.map((image) => image.url),
          referenceImageUrl: reference.imageUrl,
        });
        sourceImageUrls = inputs.imageUrls;
        prompt = buildAgentReplacementPrompt({
          basePrompt: request.prompt,
          modeType: "detail_replacement",
          replacementStrategy: request.replacementStrategy,
          marketplaceProductDetailReference:
            isMarketplaceProductDetailReferenceId(reference.id) ||
            isMarketplaceProductDetailReferenceId(reference.candidateId),
          referenceLabel: reference.label,
          referenceIndex: index,
          totalReferences: request.references.length,
          hasModelImage: Boolean(inputs.modelImageUrl),
          productImageCount: inputs.productImageUrls.length,
          originalProductImageCount: request.productImageUrls.length,
          rerunReferenceCount: inputs.supplementImageUrls.length,
          templateReferenceCount: 0,
          scenePlan: plan.scenePlan,
          referenceGeometry: geometry,
        });
        const generated = await generateCanvasImage({
          prompt,
          model: request.model,
          ownerKey: userId,
          aspectRatio: geometry?.aspectRatio || "auto",
          imageSize: request.imageSize,
          imageCount: 1,
          referenceImages: inputs.imageUrls,
          promptExtend: false,
          providerImageParams: { inputImageRoleByUrl: inputs.inputImageRoleByUrl },
          debugLabel: "agent-detail-replacement-batch",
          debugMeta: { projectId: request.projectId, cardId: reference.id },
        });
        const image = generated.images[0];
        if (!image?.originalUrl) throw new Error("模型未返回结果图片");
        let finalImage = image;
        try {
          finalImage = await normalizeReplacementResultImage({
            generatedImageUrl: image.originalUrl,
            generatedThumbnailUrl: image.thumbnailUrl,
            referenceImageUrl: reference.imageUrl,
            referenceGeometry: geometry,
            scenePlan: plan.scenePlan,
            preserveReferenceOutsideSubject: request.replacementStrategy === "product_only",
            preserveReferenceText: true,
            uploadPrefix: "agent-detail-replacement",
          });
        } catch (error) {
          console.warn("[Detail Replacement Batch] Normalize failed, retaining generated image", {
            taskId, cardId: reference.id, error,
          });
          if (!await isRemoteImageAccessible(image.originalUrl).catch(() => false)) {
            throw error;
          }
        }
        card = {
          cardId: reference.id,
          status: "completed",
          imageUrl: finalImage.originalUrl,
          thumbnailUrl: finalImage.thumbnailUrl || finalImage.originalUrl,
          text: "",
        };
      }
    } catch (error) {
      const message = error instanceof Error ? error.message : "生成失败";
      console.error("[Detail Replacement Batch] Card failed", {
        taskId, cardId: reference.id, elapsedMs: Date.now() - startedAt, error,
      });
      card = {
        cardId: reference.id,
        status: "fallback",
        imageUrl: reference.imageUrl,
        thumbnailUrl: reference.imageUrl,
        text: `生成失败，当前先沿用参考图：${message}`,
        error: message,
      };
    }

    const assistantText = prependImageModelBadge(
      [
        `<!-- AGENT_REPLACEMENT_CARD_ID:${reference.id} -->`,
        `<!-- AGENT_REPLACEMENT_RUN_ID:${taskId} -->`,
        `<!-- AGENT_REPLACEMENT_STATUS:${card.status === "completed" ? "completed" : "failed"} -->`,
        `### ${reference.label}`,
        card.text,
        buildAgentToolImagesMarkdown([{
          originalUrl: card.imageUrl,
          thumbnailUrl: card.thumbnailUrl,
          imageSize: request.imageSize,
        }]),
      ].filter(Boolean).join("\n\n"),
      card.status === "completed" ? request.model : null,
    );

    const persisted = await persistAgentDetailReplacementCard({
      projectId: request.projectId,
      userId,
      cardId: reference.id,
      imageUrl: card.imageUrl,
      referenceImageUrl: reference.imageUrl,
      label: reference.label,
      status: card.status,
      text: assistantText,
    });
    if (!persisted) throw new Error(`无法保存详情替换结果：${reference.id}`);

    await persistAssistantAndAlbum({ card, text: assistantText, prompt, sourceImageUrls });
    const previousIndex = cards.findIndex((item) => item.cardId === card.cardId);
    if (previousIndex >= 0) cards.splice(previousIndex, 1);
    cards.push(card);
    completedIds.add(reference.id);
    await publish(null);
  }
  return publish(null);
}