ai-chatbot/app/api/canvas/generate/image/route.ts

import { NextResponse } from "next/server";
import { auth } from "@/app/(auth)/auth";
import { getEnabledAiModelsWithProvider, getAiProviderById } from "@/lib/db/queries";
import { createGeminiImage } from "@/lib/ai/providers/gemini-image";
import { createAlibabaImage } from "@/lib/ai/providers/alibaba-image";
import { createRuxaImage } from "@/lib/ai/providers/ruxa-image";

export const maxDuration = 120;

interface GenerateImageRequest {
  prompt: string;
  negativePrompt?: string;
  model?: string; // 模型 ID
  aspectRatio?: string; // 宽高比 "1:1", "16:9" 等
  imageCount?: number; // 生成数量
  referenceImage?: string; // 参考图片（用于图生图）
}

interface ParsedImage {
  thumbnailUrl: string;
  originalUrl: string;
}

function parseImagesFromMarkdown(markdown: string): ParsedImage[] {
  const imageRegex = /!\[([^\]]*)\]\(([^)]+)\)/g;
  const matches = [...markdown.matchAll(imageRegex)];
  const parsed: ParsedImage[] = [];

  for (const match of matches) {
    const altText = match[1] || "";
    const thumbnailUrl = match[2];
    let originalUrl = thumbnailUrl;

    if (altText.includes("thumb|||")) {
      const content = altText.substring(altText.indexOf("thumb|||") + 8);
      const parts = content.split("|||");
      originalUrl = parts[0] || thumbnailUrl;
    } else if (altText.includes("thumb$$")) {
      const content = altText.substring(altText.indexOf("thumb$$") + 7);
      const parts = content.split("$$");
      originalUrl = parts[0] || thumbnailUrl;
    } else if (altText.includes("thumb|")) {
      const content = altText.substring(altText.indexOf("thumb|") + 6);
      const parts = content.split("|");
      originalUrl = parts[0] || thumbnailUrl;
    }

    parsed.push({ thumbnailUrl, originalUrl });
  }

  return parsed;
}

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: GenerateImageRequest = await request.json();
    const { prompt, negativePrompt, model, aspectRatio = "1:1", imageCount = 1, referenceImage } = body;

    if (!prompt) {
      return NextResponse.json(
        { success: false, error: "请输入提示词" },
        { status: 400 }
      );
    }

    // 获取可用的图片生成模型（筛选 supportsImageGen 的模型）
    const allModels = await getEnabledAiModelsWithProvider();
    const imageModels = allModels.filter(m => m.model.supportsImageGen);
    
    if (!imageModels || imageModels.length === 0) {
      return NextResponse.json(
        { success: false, error: "暂无可用的图片生成模型" },
        { status: 400 }
      );
    }

    // 选择模型（如果指定了 model，使用指定的；否则使用默认的）
    let selectedModel = model 
      ? imageModels.find(m => m.model.modelId === model)
      : imageModels.find(m => m.model.isDefault) || imageModels[0];

    if (!selectedModel) {
      selectedModel = imageModels[0];
    }

    const provider = await getAiProviderById({ id: selectedModel.model.providerId });
    if (!provider) {
      return NextResponse.json(
        { success: false, error: "模型提供商配置错误" },
        { status: 500 }
      );
    }

    // 构建消息
    const messages = [{
      role: "user" as const,
      content: referenceImage
        ? [
            { type: "text" as const, text: prompt },
            { type: "image_url" as const, image_url: { url: referenceImage } },
          ]
        : prompt,
    }];

    // 图片参数
    const imageParams = {
      aspectRatio,
      imageSize: "2K",
      imageCount,
      negativePrompt,
    };

    let imageMarkdown = "";

    // 根据提供商类型调用不同的生成器
    if (provider.type === "gemini") {
      const geminiProvider = createGeminiImage({
        apiKey: provider.apiKey,
        baseURL: provider.baseUrl,
      });

      const model = geminiProvider.chat(selectedModel.model.modelId);
      const result = await model.doGenerate({
        messages,
        imageParams,
      });
      imageMarkdown = result.text || "";
    } else if (provider.type === "alibaba-image") {
      const alibabaProvider = createAlibabaImage({
        apiKey: provider.apiKey,
        baseURL: provider.baseUrl,
      });

      const model = alibabaProvider.chat(selectedModel.model.modelId);
      const result = await model.doGenerate({
        messages,
        imageParams,
      });
      imageMarkdown = result.text || "";
    } else if (provider.type === "ruxa-image") {
      const ruxaProvider = createRuxaImage({
        apiKey: provider.apiKey,
        baseURL: provider.baseUrl,
      });

      const model = ruxaProvider.chat(selectedModel.model.modelId);
      const result = await model.doGenerate({
        messages,
        imageParams,
      });
      imageMarkdown = result.text || "";
    } else {
      return NextResponse.json(
        { success: false, error: "不支持的图片生成模型类型" },
        { status: 400 }
      );
    }

    const parsedImages = parseImagesFromMarkdown(imageMarkdown);

    if (parsedImages.length > 0) {
      return NextResponse.json({
        success: true,
        data: {
          images: parsedImages,
          model: selectedModel.model.displayName,
          prompt,
        }
      });
    } else {
      return NextResponse.json(
        { success: false, error: "图片生成失败" },
        { status: 500 }
      );
    }

  } catch (error) {
    console.error("[Canvas Generate Image] Error:", error);
    return NextResponse.json(
      { success: false, error: error instanceof Error ? error.message : "生成失败" },
      { status: 500 }
    );
  }
}