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

import "server-only";

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";
import { createFalImage } from "@/lib/ai/providers/fal-image";
import { filterReachableRemoteUrls } from "@/lib/net/url-access";

export interface CanvasGenerateImageRequest {
  prompt: string;
  negativePrompt?: string;
  model?: string;
  aspectRatio?: string;
  imageCount?: number;
  referenceImage?: string;
  referenceImages?: string[];
}

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

export interface CanvasGenerateImageResult {
  images: ParsedImage[];
  model: string;
  prompt: 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 generateCanvasImage(
  request: CanvasGenerateImageRequest
): Promise<CanvasGenerateImageResult> {
  const prompt = String(request.prompt || "").trim();
  if (!prompt) {
    throw new Error("请输入提示词");
  }

  const negativePrompt =
    typeof request.negativePrompt === "string" ? request.negativePrompt : undefined;
  const model = typeof request.model === "string" ? request.model : undefined;
  const aspectRatio =
    typeof request.aspectRatio === "string" && request.aspectRatio.trim()
      ? request.aspectRatio.trim()
      : "1:1";
  const imageCount =
    typeof request.imageCount === "number" && Number.isFinite(request.imageCount)
      ? Math.max(1, Math.min(8, Math.floor(request.imageCount)))
      : 1;
  const referenceImage =
    typeof request.referenceImage === "string" ? request.referenceImage : undefined;
  const referenceImages = Array.isArray(request.referenceImages)
    ? request.referenceImages
    : [];

  const allModels = await getEnabledAiModelsWithProvider();
  const imageModels = allModels.filter((m) => m.model.supportsImageGen);

  if (!imageModels || imageModels.length === 0) {
    throw new Error("暂无可用的图片生成模型");
  }

  let selectedModel = model
    ? imageModels.find((m) => m.model.modelId === model || m.model.id === model)
    : imageModels.find((m) => m.model.isDefault) || imageModels[0];

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

  const provider = await getAiProviderById({ id: selectedModel.model.providerId });
  if (!provider) {
    throw new Error("模型提供商配置错误");
  }

  const normalizedReferenceImages = Array.from(
    new Set(
      [
        ...referenceImages,
        ...(referenceImage ? [referenceImage] : []),
      ]
        .map((item) => String(item || "").trim())
        .filter((url) => /^https?:\/\//i.test(url))
    )
  ).slice(0, 12);
  const { kept: reachableReferenceImages, dropped: unreachableReferenceImages } =
    await filterReachableRemoteUrls(normalizedReferenceImages, { timeoutMs: 3000 });
  if (unreachableReferenceImages.length > 0) {
    console.warn("[Canvas Generate Image] Dropping unreachable reference images:", {
      count: unreachableReferenceImages.length,
      urls: unreachableReferenceImages.slice(0, 6),
    });
  }
  if (normalizedReferenceImages.length > 0 && reachableReferenceImages.length === 0) {
    throw new Error("参考图片链接不可访问，请重新上传后再试");
  }

  const messages = [
    {
      role: "user" as const,
      content:
        reachableReferenceImages.length > 0
          ? [
              { type: "text" as const, text: prompt },
              ...reachableReferenceImages.map((url) => ({
                type: "image_url" as const,
                image_url: { url },
              })),
            ]
          : prompt,
    },
  ];

  const imageParams = {
    aspectRatio,
    imageSize: "2K",
    imageCount,
    negativePrompt,
  };

  let imageMarkdown = "";

  if (provider.type === "gemini") {
    const geminiProvider = createGeminiImage({
      apiKey: provider.apiKey,
      baseURL: provider.baseUrl,
    });
    const modelInstance = geminiProvider.chat(selectedModel.model.modelId);
    const result = await modelInstance.doGenerate({ messages, imageParams });
    imageMarkdown = result.text || "";
  } else if (provider.type === "alibaba-image") {
    const alibabaProvider = createAlibabaImage({
      apiKey: provider.apiKey,
      baseURL: provider.baseUrl,
    });
    const modelInstance = alibabaProvider.chat(selectedModel.model.modelId);
    const result = await modelInstance.doGenerate({ messages, imageParams });
    imageMarkdown = result.text || "";
  } else if (provider.type === "ruxa-image") {
    const ruxaProvider = createRuxaImage({
      apiKey: provider.apiKey,
      baseURL: provider.baseUrl,
    });
    const modelInstance = ruxaProvider.chat(selectedModel.model.modelId);
    const result = await modelInstance.doGenerate({ messages, imageParams });
    imageMarkdown = result.text || "";
  } else if (provider.type === "fal-image") {
    const falProvider = createFalImage({
      apiKey: provider.apiKey,
      baseURL: provider.baseUrl,
    });
    const modelInstance = falProvider.chat(selectedModel.model.modelId);
    const result = await modelInstance.doGenerate({ messages, imageParams });
    imageMarkdown = result.text || "";
  } else {
    throw new Error("不支持的图片生成模型类型");
  }

  const parsedImages = parseImagesFromMarkdown(imageMarkdown);
  if (parsedImages.length === 0) {
    throw new Error("图片生成失败");
  }

  return {
    images: parsedImages,
    model: selectedModel.model.displayName,
    prompt,
  };
}