ai-chatbot/lib/style-pro/task.ts

import "server-only";

import { prepareFeaturePointsCharge } from "@/lib/points/feature-charge";
import {
  generateCanvasImage,
  type CanvasGenerateImageResult,
} from "@/lib/canvas/generate-image";
import type {
  StyleProGenerationMode,
  StyleProSelectionMap,
} from "./config";

export type StyleProTaskSelectedOption = {
  categoryId: string;
  categoryTitle: string;
  jsonKey: string;
  optionId: string;
  optionLabel: string;
  prompt: string;
};

export type StyleProTaskModelRef = {
  id: string;
  displayName: string;
};

export type StyleProTaskBilling = {
  unitCost: number;
  reason: string;
  modelId: string;
  messageId?: string | null;
  note: string;
};

export type StyleProTaskRequest = {
  taskKind: "style-pro";
  canvasId?: string | null;
  sourceImageUrl: string;
  generationMode: StyleProGenerationMode;
  normalizedSelections: StyleProSelectionMap;
  selectedOptions: StyleProTaskSelectedOption[];
  selectedModels: StyleProTaskModelRef[];
  optionLabels: string[];
  prompts: {
    defaultPrompt: string;
    realisticFirstPrompt?: string;
    realisticSecondPrompt?: string;
    negativePrompt?: string;
    aspectRatio?: string;
    imageSize?: string;
  };
  billing: StyleProTaskBilling;
  requestedStyleProModelId?: string | null;
};

export type StyleProTaskResult = {
  imageUrl: string;
  thumbnailUrl?: string;
  model?: string;
  modelRef?: string;
  generationMode: StyleProGenerationMode;
  chargedPoints: number;
  selections: StyleProSelectionMap;
  optionLabels: string[];
};

function composeStyleProPrompt({
  defaultPrompt,
  selectedOptions,
  mode = "style-pro-final",
}: {
  defaultPrompt: string;
  selectedOptions: StyleProTaskSelectedOption[];
  mode?: string;
}) {
  const groupedPrompt = selectedOptions.reduce<Record<string, string[]>>(
    (acc, item) => {
      const key = item.jsonKey || item.categoryId;
      acc[key] = acc[key] || [];
      acc[key].push(item.prompt);
      return acc;
    },
    {}
  );

  return [
    defaultPrompt,
    "请严格依据以下风格化 PRO 选择执行最终图片生成。用户不可见配置 JSON：",
    JSON.stringify(
      {
        mode,
        selections: selectedOptions,
        promptByKey: groupedPrompt,
      },
      null,
      2
    ),
  ]
    .filter(Boolean)
    .join("\n\n");
}

async function consumeStyleProTaskPoints(
  userId: string,
  billing: StyleProTaskBilling
) {
  if (!Number.isFinite(billing.unitCost) || billing.unitCost <= 0) {
    return 0;
  }

  try {
    const pointsCharge = await prepareFeaturePointsCharge({
      userId,
      unitCost: billing.unitCost,
      reason: billing.reason,
      modelId: billing.modelId,
      messageId: billing.messageId || null,
      note: billing.note,
    });

    if (!pointsCharge.ok || !pointsCharge.chargeable) {
      return 0;
    }

    const consumeResult = await pointsCharge.consume();
    if (consumeResult?.ok && !consumeResult.skipped) {
      return billing.unitCost;
    }
  } catch (error) {
    console.error("[Style Pro Task] Failed to consume points:", error);
  }

  return 0;
}

function getPrimaryModelId(request: StyleProTaskRequest) {
  return String(request.selectedModels[0]?.id || "").trim();
}

export function getStyleProTaskPrimaryModelKey(request: StyleProTaskRequest) {
  return getPrimaryModelId(request) || "__default__";
}

function normalizeCanvasImageResultImage(
  result: CanvasGenerateImageResult,
  fallbackError: string
) {
  const image = result.images[0];
  const imageUrl = String(image?.originalUrl || image?.thumbnailUrl || "").trim();
  if (!imageUrl) {
    throw new Error(fallbackError);
  }

  return {
    imageUrl,
    thumbnailUrl: String(image?.thumbnailUrl || image?.originalUrl || "").trim() || imageUrl,
    model: result.model,
    modelRef: result.modelRef,
  };
}

export async function executeStyleProTask(params: {
  request: StyleProTaskRequest;
  userId: string;
}) {
  const { request, userId } = params;
  const aspectRatio = request.prompts.aspectRatio || "auto";
  const imageSize = request.prompts.imageSize || "2K";
  const negativePrompt = request.prompts.negativePrompt || undefined;

  const result =
    request.generationMode === "realistic"
      ? await (async () => {
          const firstModelId = getPrimaryModelId(request);
          const secondModelId = String(request.selectedModels[1]?.id || "").trim();
          if (!firstModelId || !secondModelId) {
            throw new Error("请先在后台配置实景化两段隐藏模型");
          }

          const firstResult = await generateCanvasImage({
            prompt: request.prompts.realisticFirstPrompt || "",
            negativePrompt,
            model: firstModelId,
            aspectRatio,
            imageSize,
            imageCount: 1,
            referenceImage: request.sourceImageUrl,
            ownerKey: userId,
            debugLabel: "style-pro-realistic-stage1",
            debugMeta: {
              canvasId: request.canvasId || null,
              selectedOptionCount: request.selectedOptions.length,
            },
          });
          const firstPassImage = normalizeCanvasImageResultImage(
            firstResult,
            "实景化第一阶段生成失败"
          );

          return generateCanvasImage({
            prompt: composeStyleProPrompt({
              defaultPrompt: request.prompts.realisticSecondPrompt || "",
              selectedOptions: request.selectedOptions,
              mode: "style-pro-realistic-final",
            }),
            negativePrompt,
            model: secondModelId,
            aspectRatio,
            imageSize,
            imageCount: 1,
            referenceImage: firstPassImage.imageUrl,
            ownerKey: userId,
            debugLabel: "style-pro-realistic-stage2",
            debugMeta: {
              canvasId: request.canvasId || null,
              selectedOptionCount: request.selectedOptions.length,
              firstStageModelRef: firstModelId,
            },
          });
        })()
      : await generateCanvasImage({
          prompt: composeStyleProPrompt({
            defaultPrompt: request.prompts.defaultPrompt,
            selectedOptions: request.selectedOptions,
          }),
          negativePrompt,
          model: getPrimaryModelId(request),
          aspectRatio,
          imageSize,
          imageCount: 1,
          referenceImage: request.sourceImageUrl,
          ownerKey: userId,
          debugLabel: "style-pro",
          debugMeta: {
            canvasId: request.canvasId || null,
            requestedStyleProModelId: request.requestedStyleProModelId || null,
            selectedOptionCount: request.selectedOptions.length,
          },
        });

  const finalImage = normalizeCanvasImageResultImage(result, "生成结果为空");
  const chargedPoints = await consumeStyleProTaskPoints(userId, request.billing);

  return {
    imageUrl: finalImage.imageUrl,
    thumbnailUrl: finalImage.thumbnailUrl,
    model: finalImage.model,
    modelRef: finalImage.modelRef,
    generationMode: request.generationMode,
    chargedPoints,
    selections: request.normalizedSelections,
    optionLabels: request.optionLabels,
  } satisfies StyleProTaskResult;
}