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

import "server-only";

import { and, eq } from "drizzle-orm";
import {
  normalizeReplacementResultImage,
} 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 { prepareLayoutReplaceInputs } from "./layout-replace-inputs";
import { runValidatedLayoutReplacement } from "./layout-replace-validation";
import { verifyCanvasLayoutReplacement } from "./layout-replace-verifier";
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 inputs = await prepareLayoutReplaceInputs(sourceImageUrl, referenceImageUrl);
  const referenceGeometry = inputs.referenceGeometry;
  const sourceProfile = await resolveCanvasLayoutReplaceSourceProfile(sourceImageUrl, inputs.source.buffer);
  const finalPrompt = buildCanvasLayoutReplacePrompt({
    referenceHeight: referenceGeometry.height,
    referenceWidth: referenceGeometry.width,
    sourceProductAspectRatio: sourceProfile?.silhouetteAspectRatio,
    supplementalPrompt: normalizeLayoutReplacePromptSupplement(params.request.supplementalPrompt),
  });

  return runValidatedLayoutReplacement({
    generate: async (attempt, feedback) => {
      const candidate = rankedModels[attempt] || selectedModel;
      return generateCanvasImage({
        prompt: feedback ? `${finalPrompt}\n\n【上一结果的具体问题】${feedback}。本次必须纠正以上问题，仍严格以图1为模版底图、图2为唯一新商品。` : finalPrompt,
        displayPrompt: "排版替换结果",
        model: candidate.model.id,
        ownerKey: params.userId,
        aspectRatio: referenceGeometry.aspectRatio,
        imageSize: params.request.imageSize || "2K",
        imageCount: 1,
        promptExtend: false,
        providerImageParams: { quality: "high" },
        referenceImages: [inputs.template.url, inputs.source.url],
        requiredReferenceImageCount: 2,
        allowFailover: false,
        retryVariants: [],
        debugLabel: "canvas-layout-replace",
        debugMeta: {
          ...(params.request.debugMeta || {}),
          operation: "layout-replace",
          layoutReplaceServerAlgorithm: "template-base-product-identity/v3",
          layoutReplaceReferenceRoles: ["template", "new-product"],
          qualityAttempt: attempt + 1,
          sourceNodeId: params.request.sourceNodeId || null,
          referenceWidth: referenceGeometry.width,
          referenceHeight: referenceGeometry.height,
          referenceAspectRatio: referenceGeometry.aspectRatio,
          productReferenceWidth: inputs.source.width,
          productReferenceHeight: inputs.source.height,
          sourceProductAspectRatio: sourceProfile?.silhouetteAspectRatio,
          requestedModelRef: params.request.model || null,
          selectedProviderType: candidate.provider.type || null,
        },
      });
    },
    normalize: async (generated) => {
      const firstImage = generated.images[0];
      const normalizedImage = await normalizeReplacementResultImage({
        generatedImageUrl: firstImage.originalUrl,
        generatedThumbnailUrl: firstImage.thumbnailUrl,
        referenceImageUrl: inputs.template.url,
        referenceGeometry,
        uploadPrefix: "canvas-layout-replace",
      });
      return { ...generated, images: [{ originalUrl: normalizedImage.originalUrl,
        thumbnailUrl: normalizedImage.thumbnailUrl || normalizedImage.originalUrl }] };
    },
    verify: result => verifyCanvasLayoutReplacement({
      sourceBuffer: inputs.sourceOriginalBuffer, templateBuffer: inputs.template.buffer,
      resultUrl: result.images[0].originalUrl,
    }),
    persist: result => 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,
        });
      }),
  });
}