ai-chatbot/lib/canvas/generate-image-response.ts

import "server-only";
import { NextResponse } from "next/server";
import {
  createCanvasImageTask,
  enqueueCanvasImageTask,
  getCanvasImageTaskForUser,
} from "@/lib/canvas/image-task-queue";
import { generateCanvasImage, type CanvasGenerateImageRequest } from "@/lib/canvas/generate-image";
import { getPointsSystemConfig } from "@/lib/db/queries";
import { prepareFeaturePointsCharge } from "@/lib/points/feature-charge";
import { getImageToolCreditCost, type ImageToolCreditFeature } from "@/lib/points/config";

interface GenerateImageRequest extends CanvasGenerateImageRequest {
  async?: boolean;
}

function normalizeBillingFeature(
  value: unknown,
): Extract<ImageToolCreditFeature, "multi-view" | "repaint"> | null {
  if (value === "multi-view" || value === "repaint") return value;
  return null;
}

function toGenerateImagePayload(body: GenerateImageRequest): CanvasGenerateImageRequest {
  return {
    allowFailover: typeof body.allowFailover === "boolean" ? body.allowFailover : undefined,
    prompt: String(body.prompt || "").trim(),
    negativePrompt: typeof body.negativePrompt === "string" ? body.negativePrompt : undefined,
    model: typeof body.model === "string" ? body.model : undefined,
    aspectRatio:
      typeof body.aspectRatio === "string" && body.aspectRatio.trim()
        ? body.aspectRatio.trim()
        : "auto",
    imageSize:
      typeof body.imageSize === "string" && body.imageSize.trim()
        ? body.imageSize.trim()
        : undefined,
    imageCount:
      typeof body.imageCount === "number" && Number.isFinite(body.imageCount) ? body.imageCount : 1,
    promptExtend: typeof body.promptExtend === "boolean" ? body.promptExtend : undefined,
    referenceImage: typeof body.referenceImage === "string" ? body.referenceImage : undefined,
    referenceImages: Array.isArray(body.referenceImages)
      ? body.referenceImages.map((item) => String(item || "").trim()).filter(Boolean)
      : [],
  };
}

export async function executeCanvasImageResponse(body: GenerateImageRequest, userId: string) {
  try {
    const payload: CanvasGenerateImageRequest = {
      ...toGenerateImagePayload(body),
      ownerKey: userId,
    };
    if (!payload.prompt) {
      return NextResponse.json({ success: false, error: "请输入提示词" }, { status: 400 });
    }

    const billingFeature = normalizeBillingFeature(body.billingFeature);
    const pointsCharge = billingFeature
      ? await (async () => {
          const pointsConfig = await getPointsSystemConfig();
          return prepareFeaturePointsCharge({
            userId: userId,
            unitCost: getImageToolCreditCost(pointsConfig, billingFeature),
            reason: billingFeature === "multi-view" ? "image_multi_view" : "image_repaint",
            modelId: "canvas-generate-image",
            note: billingFeature === "multi-view" ? "图片多视角生成消耗积分" : "图片重绘消耗积分",
            pointsConfig,
          });
        })()
      : null;
    if (pointsCharge && !pointsCharge.ok) {
      return NextResponse.json(pointsCharge.payload, { status: pointsCharge.status });
    }

    const asyncMode = body.async === true;
    if (asyncMode) {
      const task = await createCanvasImageTask({
        userId: userId,
        request: payload as unknown as Record<string, unknown>,
      });
      await enqueueCanvasImageTask(task.id);
      const currentTask = await getCanvasImageTaskForUser({
        taskId: task.id,
        userId: userId,
      });
      if (pointsCharge?.chargeable) {
        try {
          await pointsCharge.consume({
            note:
              billingFeature === "multi-view"
                ? `图片多视角生成消耗积分（${task.id}）`
                : `图片重绘消耗积分（${task.id}）`,
          });
        } catch (pointsError) {
          console.error("[Points] Failed to consume canvas image task points:", pointsError);
        }
      }
      return NextResponse.json({
        success: true,
        taskId: task.id,
        status: currentTask?.status || task.status,
        queueStatus: currentTask?.queueStatus,
      });
    }

    const result = await generateCanvasImage(payload);
    if (pointsCharge?.chargeable) {
      try {
        await pointsCharge.consume();
      } catch (pointsError) {
        console.error("[Points] Failed to consume canvas image points:", pointsError);
      }
    }
    return NextResponse.json({
      success: true,
      data: result,
    });
  } catch (error) {
    console.error("[Canvas Generate Image] Error:", error);
    return NextResponse.json(
      { success: false, error: error instanceof Error ? error.message : "生成失败" },
      { status: 500 },
    );
  }
}