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 { saveGeneratedImages } from "@/lib/ai/runtime/generated-image-persistence";
import type {
  StyleProGenerationMode,
  StyleProCustomInputMap,
  StyleProSelectionMap,
} from "./config";

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

export type StyleProTaskCustomInput = {
  scope: "global" | "category";
  fieldId: string;
  fieldLabel: string;
  fieldType: string;
  value: string;
  prompt: string;
  categoryId?: string;
  categoryTitle?: 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;
  customInputs?: StyleProCustomInputMap;
  selectedOptions: StyleProTaskSelectedOption[];
  selectedCustomInputs?: StyleProTaskCustomInput[];
  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;
  gallerySavedImageCount?: number;
  gallerySaveError?: string;
  generationMode: StyleProGenerationMode;
  chargedPoints: number;
  selections: StyleProSelectionMap;
  customInputs?: StyleProCustomInputMap;
  optionLabels: string[];
};

function composeStyleProPrompt({
  defaultPrompt,
  selectedOptions,
  selectedCustomInputs = [],
  mode = "style-pro-final",
}: {
  defaultPrompt: string;
  selectedOptions: StyleProTaskSelectedOption[];
  selectedCustomInputs?: StyleProTaskCustomInput[];
  mode?: string;
}) {
  const uniqueStrings = (values: string[]) =>
    Array.from(
      new Set(values.map((value) => String(value || "").trim()).filter(Boolean))
    );
  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;
    },
    {}
  );
  const customPrompt = selectedCustomInputs.reduce<Record<string, string[]>>(
    (acc, item) => {
      const key =
        item.scope === "category" && item.categoryTitle
          ? item.categoryTitle
          : "全局自定义";
      acc[key] = acc[key] || [];
      acc[key].push(item.prompt || `${item.fieldLabel}：${item.value}`);
      return acc;
    },
    {}
  );
  const promptByKey = Object.fromEntries(
    Object.entries(groupedPrompt).map(([key, prompts]) => [
      key,
      uniqueStrings(prompts),
    ])
  );
  const customPromptByScope = Object.fromEntries(
    Object.entries(customPrompt).map(([key, prompts]) => [
      key,
      uniqueStrings(prompts),
    ])
  );
  const selectionSummary = selectedOptions.map((item) => ({
    categoryId: item.categoryId,
    categoryTitle: item.categoryTitle,
    jsonKey: item.jsonKey,
    optionId: item.optionId,
    optionLabel: item.optionLabel,
  }));
  const customInputSummary = selectedCustomInputs.map((item) => ({
    scope: item.scope,
    categoryId: item.categoryId,
    categoryTitle: item.categoryTitle,
    fieldId: item.fieldId,
    fieldLabel: item.fieldLabel,
    fieldType: item.fieldType,
    value: item.value,
  }));

  return [
    defaultPrompt,
    "请严格依据以下风格化 PRO 选择执行最终图片生成。用户不可见配置 JSON：",
    JSON.stringify(
      {
        mode,
        selections: selectionSummary,
        promptByKey,
        customInputs: customInputSummary,
        customPromptByScope,
      },
      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,
  };
}

async function saveStyleProFinalImageToGallery(params: {
  userId: string;
  request: StyleProTaskRequest;
  imageUrl: string;
  thumbnailUrl: string;
  modelRef?: string;
  prompt: string;
}) {
  const imageSize = params.request.prompts.imageSize || "2K";
  const imageMarkdown = `![thumb|||${params.imageUrl}|||${imageSize}](${params.thumbnailUrl})`;
  const persistResult = await saveGeneratedImages(
    imageMarkdown,
    params.userId,
    null,
    params.request.billing.messageId || null,
    params.prompt,
    params.modelRef || params.request.billing.modelId,
    {
      sourceImageUrls: [params.request.sourceImageUrl],
      relationType: "edited_from",
      sourceType:
        params.request.generationMode === "realistic"
          ? "style_pro_realistic"
          : "style_pro",
      projectId: params.request.canvasId || null,
      userPrompt: params.request.optionLabels.join(" / "),
      optimizedPrompt: params.prompt,
      requestedImageCount: 1,
      aspectRatio: params.request.prompts.aspectRatio || "auto",
      imageSize,
      skipBilling: true,
    }
  );
  return persistResult.savedImageCount;
}

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

  let finalPrompt = "";
  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,
            onRequestDispatched: params.onRequestDispatched,
            allowFailover: false,
            debugLabel: "style-pro-realistic-stage1",
            debugMeta: {
              canvasId: request.canvasId || null,
              selectedOptionCount: request.selectedOptions.length,
              configuredModelRef: firstModelId,
            },
          });
          const firstPassImage = normalizeCanvasImageResultImage(
            firstResult,
            "实景化第一阶段生成失败"
          );

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

  const finalImage = normalizeCanvasImageResultImage(result, "生成结果为空");
  let gallerySavedImageCount = 0;
  let gallerySaveError: string | undefined;
  try {
    gallerySavedImageCount = await saveStyleProFinalImageToGallery({
      userId,
      request,
      imageUrl: finalImage.imageUrl,
      thumbnailUrl: finalImage.thumbnailUrl,
      modelRef: finalImage.modelRef,
      prompt: finalPrompt || result.prompt,
    });
    if (gallerySavedImageCount <= 0) {
      gallerySaveError = "最终图未写入个人相册";
    }
  } catch (error) {
    gallerySaveError =
      error instanceof Error ? error.message : "保存到个人相册失败";
    console.error("[Style Pro Task] Failed to save final image to gallery:", {
      error: gallerySaveError,
      generationMode: request.generationMode,
      modelRef: finalImage.modelRef,
      canvasId: request.canvasId || null,
    });
  }
  const chargedPoints = await consumeStyleProTaskPoints(userId, request.billing);

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