ai-chatbot/app/api/style-pro/generate/route.ts

import { NextResponse } from "next/server";
import { z } from "zod";
import { auth } from "@/app/(auth)/auth";
import { generateCanvasImage } from "@/lib/canvas/generate-image";
import { getSystemSetting, getUserById } from "@/lib/db/queries";
import {
  findStyleProImageModelByRef,
  getStyleProImageModelEntries,
} from "@/lib/style-pro/model-options";
import {
  STYLE_PRO_CONFIG_KEY,
  type StyleProCategory,
  type StyleProOption,
  normalizeStyleProConfig,
} from "@/lib/style-pro/config";

export const maxDuration = 120;
export const dynamic = "force-dynamic";

const generateRequestSchema = z.object({
  canvasId: z.string().trim().min(1).max(160).optional(),
  sourceImageUrl: z.string().trim().min(1).max(4000),
  modelId: z.string().trim().min(1).max(160).optional(),
  selections: z.record(z.array(z.string().trim().min(1).max(120)).max(20)),
});

type SelectedOption = {
  category: StyleProCategory;
  option: StyleProOption;
};

function isDisallowedImageUrl(value: string) {
  const trimmed = value.trim().toLowerCase();
  return trimmed.startsWith("data:") || trimmed.startsWith("javascript:");
}

function buildSelectionPayload(selectedOptions: SelectedOption[]) {
  return selectedOptions.map(({ category, option }) => ({
    categoryId: category.id,
    categoryTitle: category.title,
    jsonKey: option.jsonKey || category.id,
    optionId: option.id,
    optionLabel: option.label,
    prompt: option.prompt,
  }));
}

function composePrompt({
  defaultPrompt,
  selectedOptions,
}: {
  defaultPrompt: string;
  selectedOptions: SelectedOption[];
}) {
  const selectionPayload = buildSelectionPayload(selectedOptions);
  const groupedPrompt = selectionPayload.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: "style-pro-final",
        selections: selectionPayload,
        promptByKey: groupedPrompt,
      },
      null,
      2
    ),
  ]
    .filter(Boolean)
    .join("\n\n");
}

function validateSelections(
  categories: StyleProCategory[],
  rawSelections: Record<string, string[]>
) {
  const categoryById = new Map(
    categories.map((category) => [category.id, category])
  );
  const selectedOptions: SelectedOption[] = [];
  const normalizedSelections: Record<string, string[]> = {};

  for (const [categoryId, optionIds] of Object.entries(rawSelections)) {
    const category = categoryById.get(categoryId);
    if (!category) {
      return {
        ok: false as const,
        error: `无效分类：${categoryId}`,
      };
    }

    const enabledOptions = new Map(
      category.options
        .filter((option) => option.enabled)
        .map((option) => [option.id, option])
    );
    const uniqueOptionIds = Array.from(new Set(optionIds));
    if (uniqueOptionIds.length > category.maxSelections) {
      return {
        ok: false as const,
        error: `${category.title} 最多选择 ${category.maxSelections} 项`,
      };
    }

    for (const optionId of uniqueOptionIds) {
      const option = enabledOptions.get(optionId);
      if (!option) {
        return {
          ok: false as const,
          error: `无效选项：${category.title} / ${optionId}`,
        };
      }
      selectedOptions.push({ category, option });
    }
    normalizedSelections[category.id] = uniqueOptionIds;
  }

  const missingCategory = categories.find((category) => {
    if (!category.required && category.minSelections <= 0) return false;
    const selectedCount = normalizedSelections[category.id]?.length || 0;
    return (
      selectedCount <
      Math.max(category.required ? 1 : 0, category.minSelections)
    );
  });
  if (missingCategory) {
    return {
      ok: false as const,
      error: `请完成必选项：${missingCategory.title}`,
    };
  }

  return {
    ok: true as const,
    selectedOptions,
    normalizedSelections,
  };
}

export async function POST(request: Request) {
  try {
    const session = await auth();
    if (!session?.user?.id) {
      return NextResponse.json(
        { success: false, error: "未登录" },
        { status: 401 }
      );
    }

    const body = generateRequestSchema.parse(await request.json());
    if (isDisallowedImageUrl(body.sourceImageUrl)) {
      return NextResponse.json(
        { success: false, error: "请使用已上传的图片链接生成" },
        { status: 400 }
      );
    }

    const setting = await getSystemSetting({ key: STYLE_PRO_CONFIG_KEY }).catch(
      () => null
    );
    const config = normalizeStyleProConfig(setting?.value);
    if (!config.enabled) {
      return NextResponse.json(
        { success: false, error: "风格化 PRO 暂未启用" },
        { status: 403 }
      );
    }

    const enabledCategories = config.categories
      .filter((category) => category.enabled)
      .map((category) => ({
        ...category,
        options: category.options.filter((option) => option.enabled),
      }))
      .filter((category) => category.options.length > 0);
    const selectionResult = validateSelections(
      enabledCategories,
      body.selections
    );
    if (!selectionResult.ok) {
      return NextResponse.json(
        { success: false, error: selectionResult.error },
        { status: 400 }
      );
    }

    let selectedModelRef = config.modelId || undefined;
    if (body.modelId) {
      const currentUser = await getUserById({ id: session.user.id }).catch(
        () => null
      );
      const includeHidden =
        session.user.type === "admin" || Boolean(currentUser?.isHighVip);
      const imageModelEntries = await getStyleProImageModelEntries({
        includeHidden,
      });
      const selectedModel = findStyleProImageModelByRef(
        imageModelEntries,
        body.modelId
      );
      if (!selectedModel) {
        return NextResponse.json(
          { success: false, error: "所选生图模型不可用" },
          { status: 400 }
        );
      }
      selectedModelRef = selectedModel.model.id;
    }

    const prompt = composePrompt({
      defaultPrompt: config.defaultPrompt,
      selectedOptions: selectionResult.selectedOptions,
    });
    const result = await generateCanvasImage({
      prompt,
      negativePrompt: config.negativePrompt || undefined,
      model: selectedModelRef,
      aspectRatio: config.aspectRatio || "auto",
      imageSize: config.imageSize || "2K",
      imageCount: 1,
      referenceImage: body.sourceImageUrl,
      ownerKey: session.user.id,
      debugLabel: "style-pro",
      debugMeta: {
        canvasId: body.canvasId || null,
        requestedStyleProModelId: body.modelId || null,
        selectedOptionCount: selectionResult.selectedOptions.length,
      },
    });
    const image = result.images[0];
    if (!image?.originalUrl) {
      return NextResponse.json(
        { success: false, error: "生成结果为空" },
        { status: 502 }
      );
    }

    return NextResponse.json({
      success: true,
      data: {
        imageUrl: image.originalUrl,
        thumbnailUrl: image.thumbnailUrl || image.originalUrl,
        model: result.model,
        modelRef: result.modelRef,
        selections: selectionResult.normalizedSelections,
        optionLabels: selectionResult.selectedOptions.map(
          ({ category, option }) => `${category.title}：${option.label}`
        ),
      },
    });
  } catch (error) {
    console.error("[Style Pro Generate] Failed:", error);
    if (error instanceof z.ZodError) {
      return NextResponse.json(
        { success: false, error: "请求参数不完整" },
        { status: 400 }
      );
    }
    return NextResponse.json(
      {
        success: false,
        error: error instanceof Error ? error.message : "风格化 PRO 生成失败",
      },
      { status: 500 }
    );
  }
}