ai-chatbot/app/api/agents/detail-replacement/batches/route.ts

import { and, desc, eq, inArray, sql } from "drizzle-orm";
import { NextResponse } from "next/server";
import { auth } from "@/app/(auth)/auth";
import { db } from "@/lib/db";
import { canvasImageTask, canvasProject, message } from "@/lib/db/schema";
import { extractAgentProjectStateFromCanvasData, isAgentProjectRow } from "@/lib/agents/agent-project-compat";
import {
  createCanvasImageTask,
  enqueueCanvasImageTask,
  ensureCanvasImageTaskActive,
  getCanvasImageTaskForUser,
} from "@/lib/canvas/image-task-queue";
import {
  checkUserImageQuota,
  getChatById,
  getEnabledAiModelsWithProvider,
  getPointsSystemConfig,
  getUserById,
  saveChat,
} from "@/lib/db/queries";
import { getModelGenerationCreditCost } from "@/lib/ai/image-model-pricing";
import { resolveDirectReplacementModelRef } from "@/app/(chat)/api/chat/_lib/detail-replacement";
import {
  buildAgentSlotPrompt,
  normalizeAgentSupplementPrompt,
} from "@/app/(chat)/api/chat/_lib/agent-scene-workflow";
import {
  detailReplacementSubmissionKey,
  isDetailReplacementBatchRequest,
  isDetailReplacementUuid,
  type DetailReplacementBatchReference,
  type DetailReplacementBatchRequest,
  type DetailReplacementBatchResult,
} from "@/lib/agents/detail-replacement-batch";

export const maxDuration = 1200;

function scopedChatId(projectId: string) {
  const normalized = projectId.toLowerCase().replace(/-/g, "");
  const scoped = `${normalized.slice(0, 20)}a93e17a60001`;
  return [scoped.slice(0, 8), scoped.slice(8, 12), scoped.slice(12, 16), scoped.slice(16, 20), scoped.slice(20, 32)].join("-");
}

function taskResponse(task: Awaited<ReturnType<typeof getCanvasImageTaskForUser>>) {
  const result = task?.result as DetailReplacementBatchResult | null;
  const request = task?.request as DetailReplacementBatchRequest | null;
  return {
    id: task?.id,
    status: task?.status,
    attempts: task?.attempts,
    error: task?.error || null,
    cards: Array.isArray(result?.cards) ? result.cards : [],
    currentCardId: result?.currentCardId || null,
    requestedCardIds: request?.references?.map((item) => item.id) || [],
    createdAt: task?.createdAt || null,
    queueStatus: task?.queueStatus || null,
  };
}

async function findActiveBatchId(params: {
  userId: string;
  projectId: string;
  submissionKey: string;
}) {
  const [row] = await db.select({ id: canvasImageTask.id })
    .from(canvasImageTask)
    .where(and(
      eq(canvasImageTask.userId, params.userId),
      inArray(canvasImageTask.status, ["pending", "running"]),
      sql`${canvasImageTask.request}->>'taskKind' = 'agent-detail-replacement-batch'`,
      sql`${canvasImageTask.request}->>'projectId' = ${params.projectId}`,
      sql`${canvasImageTask.request}->>'submissionKey' = ${params.submissionKey}`,
    ))
    .orderBy(desc(canvasImageTask.createdAt)).limit(1);
  return row?.id || null;
}

function isUniqueViolation(error: unknown) {
  for (let cursor: unknown = error; cursor && typeof cursor === "object"; ) {
    if ((cursor as { code?: unknown }).code === "23505") return true;
    cursor = (cursor as { cause?: unknown }).cause;
  }
  return false;
}

export async function GET(request: Request) {
  const session = await auth();
  if (!session?.user?.id) return NextResponse.json({ error: "未登录" }, { status: 401 });
  const params = new URL(request.url).searchParams;
  const projectId = String(params.get("projectId") || "").trim();
  const taskId = String(params.get("taskId") || "").trim();
  if (!isDetailReplacementUuid(projectId)) {
    return NextResponse.json({ error: "项目编号无效" }, { status: 400 });
  }
  const [project] = await db.select({ id: canvasProject.id }).from(canvasProject).where(and(
    eq(canvasProject.id, projectId), eq(canvasProject.userId, session.user.id),
  )).limit(1);
  if (!project) return NextResponse.json({ error: "项目不存在" }, { status: 404 });

  const [candidate] = await db.select({ id: canvasImageTask.id, request: canvasImageTask.request })
    .from(canvasImageTask)
    .where(and(
      eq(canvasImageTask.userId, session.user.id),
      ...(isDetailReplacementUuid(taskId) ? [eq(canvasImageTask.id, taskId)] : []),
      sql`${canvasImageTask.request}->>'taskKind' = 'agent-detail-replacement-batch'`,
      sql`${canvasImageTask.request}->>'projectId' = ${projectId}`,
    ))
    // A batch that is still working matters more than a newer finished one, so the
    // workspace keeps showing live progress after a rerun of a single card.
    .orderBy(
      sql`case when ${canvasImageTask.status} in ('pending', 'running') then 0 else 1 end`,
      desc(canvasImageTask.createdAt),
    ).limit(1);
  if (!candidate || !isDetailReplacementBatchRequest(candidate.request)) {
    return NextResponse.json({ task: null });
  }
  let task = await getCanvasImageTaskForUser({ taskId: candidate.id, userId: session.user.id });
  if (task && (task.status === "failed" ||
    (task.status === "running" && Date.now() - task.updatedAt.getTime() > 90_000))) {
    await ensureCanvasImageTaskActive({ taskId: task.id, userId: session.user.id });
    task = await getCanvasImageTaskForUser({ taskId: candidate.id, userId: session.user.id });
  }
  return NextResponse.json({ task: taskResponse(task) });
}

export async function POST(request: Request) {
  try {
    const session = await auth();
    if (!session?.user?.id) return NextResponse.json({ error: "未登录" }, { status: 401 });
    const body = await request.json();
    const projectId = String(body.projectId || "").trim();
    const taskId = String(body.requestId || "").trim();
    if (!isDetailReplacementUuid(projectId) || !isDetailReplacementUuid(taskId)) {
      return NextResponse.json({ error: "任务编号或项目编号无效" }, { status: 400 });
    }
    const [project] = await db.select().from(canvasProject).where(and(
      eq(canvasProject.id, projectId), eq(canvasProject.userId, session.user.id),
    )).limit(1);
    if (!project || !isAgentProjectRow(project)) {
      return NextResponse.json({ error: "风格化项目不存在" }, { status: 404 });
    }

    const state = extractAgentProjectStateFromCanvasData(project.canvasData);
    const knownReferences = new Map(
      (state.workflowState.agentModeState?.detailReplacement?.references || [])
        .map((item) => [String(item.id), String(item.imageUrl)]),
    );
    const requested = Array.isArray(body.references) ? body.references : [];
    const references: DetailReplacementBatchReference[] = requested.map((raw: unknown) => {
      const item = raw && typeof raw === "object" ? raw as Record<string, unknown> : {};
      return {
        id: String(item.id || "").slice(0, 128),
        label: String(item.label || "详情参考图").slice(0, 255),
        imageUrl: String(item.imageUrl || "").trim(),
        candidateId: String(item.candidateId || "").slice(0, 128) || null,
        passthrough: item.passthrough === true,
        rerunReferenceImages: Array.isArray(item.rerunReferenceImages)
          ? item.rerunReferenceImages.slice(0, 6).map((entry: { url?: string }) => ({ url: String(entry?.url || "") }))
          : [],
      };
    });
    if (
      references.length === 0 || references.length > 40 ||
      new Set(references.map((item) => item.id)).size !== references.length ||
      references.some((item) => !item.id || knownReferences.get(item.id) !== item.imageUrl)
    ) {
      return NextResponse.json({ error: "详情参考图尚未保存或已发生变化，请刷新后重试" }, { status: 400 });
    }
    const productImageUrls = Array.isArray(body.productImageUrls)
      ? Array.from(new Set(body.productImageUrls.map((url: unknown) => String(url || "").trim())
          .filter((url: string) => /^https?:\/\//i.test(url)))).slice(0, 16) as string[]
      : [];
    if (productImageUrls.length === 0 && !references.every((item) => item.passthrough)) {
      return NextResponse.json({ error: "请先在输入图区选择至少一张产品图" }, { status: 400 });
    }
    const strategy = body.replacementStrategy === "product_only" ? "product_only" : "scene_replace";
    const prompt = String(body.prompt || "").trim().slice(0, 24_000);
    if (strategy !== "product_only" && !prompt) {
      return NextResponse.json({ error: "请先配置风格化提示词" }, { status: 400 });
    }
    const models = await getEnabledAiModelsWithProvider();
    const model = await resolveDirectReplacementModelRef({
      selectedChatModel: String(body.model || "auto"),
      finalModelId: String(body.model || "auto"),
      models,
    });
    if (!model || model === "auto" || !models.some((entry) => entry.model.id === model)) {
      return NextResponse.json({ error: "详情替换模型不可用，请检查模型配置" }, { status: 400 });
    }
    const quota = await checkUserImageQuota(session.user.id);
    if (!quota.allowed) {
      return NextResponse.json({ error: "本月图片生成配额已用完" }, { status: 429 });
    }
    const requestedImageCount = references.filter((item) => !item.passthrough).length;
    if (quota.limit != null && quota.used + requestedImageCount > quota.limit) {
      return NextResponse.json({ error: `本月剩余生图配额不足 ${requestedImageCount} 张` }, { status: 429 });
    }
    const pointsConfig = await getPointsSystemConfig();
    if (pointsConfig.enabled && pointsConfig.blockOnInsufficient !== false) {
      const selectedModel = models.find((entry) => entry.model.id === model)?.model;
      const unitCost = getModelGenerationCreditCost(
        selectedModel as Parameters<typeof getModelGenerationCreditCost>[0], "image", {
        imageSize: ["1K", "2K", "4K"].includes(body.imageSize) ? body.imageSize : "2K",
      });
      const user = await getUserById({ id: session.user.id });
      const expectedCost = unitCost * requestedImageCount;
      if (!user?.isHighVip && expectedCost > Number(user?.pointsBalance || 0)) {
        return NextResponse.json({ error: `积分不足，本批次预计消耗 ${expectedCost} 积分` }, { status: 402 });
      }
    }
    const chatId = scopedChatId(projectId);
    const existingChat = await getChatById({ id: chatId });
    if (existingChat && existingChat.userId !== session.user.id) {
      return NextResponse.json({ error: "无权访问项目对话" }, { status: 403 });
    }
    if (!existingChat) await saveChat({
      id: chatId, userId: session.user.id, title: project.title || "详情替换", visibility: "private",
    });
    const taskRequest: DetailReplacementBatchRequest = {
      taskKind: "agent-detail-replacement-batch",
      projectId,
      chatId,
      model,
      prompt: strategy === "product_only"
        ? "仅替换参考详情图中的产品主体"
        : buildAgentSlotPrompt({
            workflowPrompt: prompt,
            slotPrompt: "",
            userSupplementPrompt: normalizeAgentSupplementPrompt(
              String(body.userPrompt || ""), "详情参考图替换",
            ),
            includeSlotTitle: false,
            modeType: "detail_replacement",
          }),
      userPrompt: String(body.userPrompt || "").slice(0, 12_000),
      imageSize: ["1K", "2K", "4K"].includes(body.imageSize) ? body.imageSize : "2K",
      replacementStrategy: strategy,
      productImageUrls,
      modelImageUrl: /^https?:\/\//i.test(String(body.modelImageUrl || ""))
        ? String(body.modelImageUrl) : null,
      references,
    };
    const submissionKey = detailReplacementSubmissionKey(taskRequest);
    taskRequest.submissionKey = submissionKey;

    const duplicateId = await findActiveBatchId({
      userId: session.user.id, projectId, submissionKey,
    });
    if (duplicateId) {
      await ensureCanvasImageTaskActive({ taskId: duplicateId, userId: session.user.id });
      const duplicate = await getCanvasImageTaskForUser({ taskId: duplicateId, userId: session.user.id });
      return NextResponse.json({ task: taskResponse(duplicate) });
    }

    let task: Awaited<ReturnType<typeof createCanvasImageTask>>;
    try {
      task = await createCanvasImageTask({
        taskId, userId: session.user.id, request: taskRequest as unknown as Record<string, unknown>,
      });
    } catch (error) {
      if (!isUniqueViolation(error)) throw error;
      // Two identical submissions raced past the lookup above; the database
      // unique index kept one of them, so answer with that batch.
      const racedId = await findActiveBatchId({
        userId: session.user.id, projectId, submissionKey,
      });
      if (!racedId) throw error;
      const raced = await getCanvasImageTaskForUser({ taskId: racedId, userId: session.user.id });
      return NextResponse.json({ task: taskResponse(raced) });
    }
    const existingTask = await getCanvasImageTaskForUser({ taskId: task.id, userId: session.user.id });
    const storedRequest = existingTask?.request as DetailReplacementBatchRequest | null;
    if (
      !isDetailReplacementBatchRequest(storedRequest) ||
      storedRequest.projectId !== projectId ||
      (storedRequest.submissionKey && storedRequest.submissionKey !== submissionKey)
    ) {
      return NextResponse.json({ error: "任务编号已被其他请求使用" }, { status: 409 });
    }
    await db.insert(message).values({
      id: taskId,
      chatId,
      role: "user",
      parts: [{ type: "text", text: taskRequest.userPrompt || `批量替换 ${references.length} 张详情图` }],
      attachments: [],
      createdAt: new Date(),
    }).onConflictDoNothing({ target: message.id });
    await enqueueCanvasImageTask(task.id);
    const current = await getCanvasImageTaskForUser({ taskId, userId: session.user.id });
    return NextResponse.json({ task: taskResponse(current) });
  } catch (error) {
    console.error("[Detail Replacement Batch] Failed to start", error);
    return NextResponse.json({ error: error instanceof Error ? error.message : "提交详情替换任务失败" }, { status: 500 });
  }
}