ai-chatbot/lib/canvas/layout-replace-task.ts

import "server-only";

import { and, eq } from "drizzle-orm";
import {
  normalizeReplacementResultImage,
  resolveReplacementReferenceGeometry,
} from "@/app/(chat)/api/chat/_lib/detail-replacement";
import { saveGeneratedImages } from "@/lib/ai/runtime/generated-image-persistence";
import {
  generateCanvasImage,
  type CanvasGenerateImageResult,
} from "@/lib/canvas/generate-image";
import {
  buildCanvasLayoutReplacePrompt,
  rankCanvasLayoutReplaceModels,
} from "@/lib/canvas/layout-replace-model-selection";
import { resolveCanvasLayoutReplaceSourceProfile } from "@/lib/canvas/layout-replace-source-profile";
import { db } from "@/lib/db";
import { getEnabledAiModelsWithProvider } from "@/lib/db/queries";
import { generatedImage } from "@/lib/db/schema";

export interface CanvasLayoutReplaceTaskRequest {
  taskKind: "canvas-layout-replace";
  sourceImageUrl: string;
  referenceImageUrl: string;
  supplementalPrompt?: string;
  model?: string;
  imageSize?: string;
  sourceNodeId?: string;
  agentProjectId?: string;
  referenceLabel?: string;
  debugMeta?: Record<string, unknown>;
  batchId?: string;
  batchItemId?: string;
  batchPosition?: number;
  batchTotal?: number;
  canvasProjectId?: string;
}

function normalizeHttpUrl(value: unknown) {
  const url = String(value || "").trim();
  return /^https?:\/\//i.test(url) ? url : "";
}

export function isCanvasLayoutReplaceTaskRequest(
  request: unknown,
): request is CanvasLayoutReplaceTaskRequest {
  return Boolean(
    request &&
      typeof request === "object" &&
      !Array.isArray(request) &&
      (request as { taskKind?: unknown }).taskKind === "canvas-layout-replace",
  );
}

function normalizeLayoutReplacePromptSupplement(value: unknown) {
  return String(value || "")
    .trim()
    .replace(/\s+/g, " ")
    .slice(0, 600);
}

async function saveLayoutReplaceResultToGallery(params: {
  request: CanvasLayoutReplaceTaskRequest;
  result: CanvasGenerateImageResult;
  userId: string;
}) {
  const image = params.result.images[0];
  const originalUrl = normalizeHttpUrl(image?.originalUrl || image?.thumbnailUrl);
  const thumbnailUrl = normalizeHttpUrl(image?.thumbnailUrl || image?.originalUrl);
  if (!originalUrl || !thumbnailUrl) return;

  const [existing] = await db
    .select({ id: generatedImage.id })
    .from(generatedImage)
    .where(
      and(
        eq(generatedImage.userId, params.userId),
        eq(generatedImage.originalUrl, originalUrl),
      ),
    )
    .limit(1);
  if (existing) return;

  const imageSize = params.request.imageSize || "2K";
  const prompt = params.result.prompt || "AI 模版替换结果";
  const modelId =
    params.result.modelRef || params.result.model || "canvas-layout-replace";
  await saveGeneratedImages(
    `![thumb|||${originalUrl}|||${imageSize}](${thumbnailUrl})`,
    params.userId,
    null,
    null,
    prompt,
    modelId,
    {
      sourceImageUrls: [params.request.sourceImageUrl],
      relationType: "edited_from",
      sourceType: "canvas_layout_replace",
      projectId: params.request.canvasProjectId || null,
      userPrompt: params.request.supplementalPrompt || "AI 模版替换",
      optimizedPrompt: prompt,
      requestedImageCount: 1,
      imageSize,
      skipBilling: true,
    },
  );
}

export async function executeCanvasLayoutReplaceTask(params: {
  request: CanvasLayoutReplaceTaskRequest;
  userId: string;
}): Promise<CanvasGenerateImageResult> {
  const sourceImageUrl = normalizeHttpUrl(params.request.sourceImageUrl);
  const referenceImageUrl = normalizeHttpUrl(params.request.referenceImageUrl);
  if (!sourceImageUrl) {
    throw new Error("缺少当前产品图片");
  }
  if (!referenceImageUrl) {
    throw new Error("缺少排版参考图");
  }

  const models = await getEnabledAiModelsWithProvider().catch((error) => {
    console.warn("[Canvas Layout Replace] Failed to load models:", error);
    return undefined;
  });
  const rankedModels = rankCanvasLayoutReplaceModels(
    models || [],
    params.request.model,
  );
  const selectedModel = rankedModels[0];
  if (!selectedModel) {
    throw new Error("AI 模版替换未配置可用的非 FAL 多图编辑模型");
  }

  const [referenceGeometry, sourceProfile] = await Promise.all([
    resolveReplacementReferenceGeometry(referenceImageUrl),
    resolveCanvasLayoutReplaceSourceProfile(sourceImageUrl),
  ]);
  const finalPrompt = buildCanvasLayoutReplacePrompt({
    referenceHeight: referenceGeometry?.height,
    referenceWidth: referenceGeometry?.width,
    sourceProductAspectRatio: sourceProfile?.silhouetteAspectRatio,
    supplementalPrompt: normalizeLayoutReplacePromptSupplement(
      params.request.supplementalPrompt,
    ),
  });
  const retryVariants = rankedModels.slice(1, 3).map((candidate, index) => ({
    allowFailover: false,
    label: `non-fal-model-${index + 2}`,
    model: candidate.model.id,
    providerImageParams: { quality: "high" },
  }));

  const generated = await generateCanvasImage({
    prompt: finalPrompt,
    displayPrompt: "排版替换结果",
    model: selectedModel.model.id,
    ownerKey: params.userId,
    aspectRatio: referenceGeometry?.aspectRatio || "auto",
    imageSize: params.request.imageSize || "2K",
    imageCount: 1,
    promptExtend: false,
    providerImageParams: { quality: "high" },
    referenceImages: [sourceImageUrl, referenceImageUrl],
    allowFailover: false,
    retryVariants,
    debugLabel: "canvas-layout-replace",
    debugMeta: {
      ...(params.request.debugMeta || {}),
      operation: "layout-replace",
      layoutReplaceServerAlgorithm: "rigid-product-scene-replace/v2",
      sourceNodeId: params.request.sourceNodeId || null,
      referenceWidth: referenceGeometry?.width,
      referenceHeight: referenceGeometry?.height,
      referenceAspectRatio: referenceGeometry?.aspectRatio,
      sourceProductAspectRatio: sourceProfile?.silhouetteAspectRatio,
      requestedModelRef: params.request.model || null,
      selectedProviderType: selectedModel.provider.type || null,
    },
  });
  const firstImage = generated.images[0];
  if (!firstImage?.originalUrl) {
    throw new Error("排版替换未返回结果图片");
  }

  const normalizedImage = await normalizeReplacementResultImage({
    generatedImageUrl: firstImage.originalUrl,
    generatedThumbnailUrl: firstImage.thumbnailUrl,
    referenceImageUrl,
    referenceGeometry,
    uploadPrefix: "canvas-layout-replace",
  });

  const result: CanvasGenerateImageResult = {
    images: [
      {
        originalUrl: normalizedImage.originalUrl,
        thumbnailUrl:
          normalizedImage.thumbnailUrl || normalizedImage.originalUrl,
      },
    ],
    model: generated.model,
    modelRef: generated.modelRef,
    prompt: generated.prompt,
  };

  await saveLayoutReplaceResultToGallery({
    request: params.request,
    result,
    userId: params.userId,
  }).catch((error) => {
    console.error("[Canvas Layout Replace] Failed to save result to gallery:", {
      error: error instanceof Error ? error.message : String(error),
      batchId: params.request.batchId,
    });
  });

  return result;
}