ai-chatbot/lib/ai/providers/bananaimage.ts

import "server-only";

import sharp from "sharp";
import {
  getBananaImageModelPresetById,
  resolveBananaImageRuntimeModelId,
} from "@/lib/ai/bananaimage-models";
import { createDeferredUploadedImage } from "@/lib/storage/image-archive";
import { uploadFile } from "@/lib/storage/service";

interface BananaImageConfig {
  apiKey: string;
  baseURL: string;
}

interface BananaImageParams {
  aspectRatio?: string;
  imageSize?: string;
  imageCount?: number;
}

interface BananaImageModelConfig {
  apiFormat?: "bananaimage-gemini" | "bananaimage-gpt-image-2";
  maxReferenceImages?: number;
  supportedAspectRatios?: string[];
  supportedImageSizes?: string[];
}

interface ReferenceImage {
  base64: string;
  dataUrl: string;
  mimeType: string;
  width?: number;
  height?: number;
}

interface UploadedImage {
  thumbnail: string;
  original: string;
}

const GEMINI_POLL_INTERVAL_MS = 3000;
const GEMINI_POLL_MAX_ATTEMPTS = 600;
const GPT_POLL_INTERVAL_MS = 4000;
const GPT_POLL_MAX_ATTEMPTS = 90;
const MAX_TRANSIENT_POLL_FAILURES = 8;
const REFERENCE_DOWNLOAD_TIMEOUT_MS = 120_000;

function joinApiPath(baseURL: string, path: string) {
  const cleanBase = baseURL.replace(/\/+$/, "");
  const cleanPath = path.startsWith("/") ? path : `/${path}`;
  return `${cleanBase}${cleanPath}`;
}

function delay(ms: number) {
  return new Promise((resolve) => setTimeout(resolve, ms));
}

function normalizeModelId(modelId?: string) {
  const normalized = String(modelId || "").trim();
  if (!normalized) throw new Error("BananaImage 模型 ID 未配置");
  return normalized;
}

function normalizeAspectRatioLabel(value?: string) {
  const normalized = String(value || "").trim().toLowerCase();
  if (!normalized) return "auto";
  if (normalized === "auto") return "auto";
  const match = normalized.match(/^(\d+)\s*:\s*(\d+)$/);
  if (!match) return "auto";
  return `${Number(match[1])}:${Number(match[2])}`;
}

function parseAspectRatio(value: string) {
  const match = value.match(/^(\d+):(\d+)$/);
  if (!match) return null;
  const width = Number(match[1]);
  const height = Number(match[2]);
  if (!width || !height) return null;
  return width / height;
}

function findClosestAspectRatio(
  width: number,
  height: number,
  supportedAspectRatios: string[]
) {
  const target = width / height;
  let closest = supportedAspectRatios.includes("1:1")
    ? "1:1"
    : supportedAspectRatios[0] || "1:1";
  let score = Number.POSITIVE_INFINITY;

  for (const ratio of supportedAspectRatios) {
    const value = parseAspectRatio(ratio);
    if (!value) continue;
    const candidateScore = Math.abs(Math.log(target / value));
    if (candidateScore < score) {
      closest = ratio;
      score = candidateScore;
    }
  }

  return closest;
}

function getSupportedAspectRatios(
  modelId: string,
  modelConfig?: BananaImageModelConfig | null
) {
  const configured = modelConfig?.supportedAspectRatios;
  const preset = getBananaImageModelPresetById(modelId);
  const ratios = Array.isArray(configured)
    ? configured
    : (preset?.imageGenConfig?.supportedAspectRatios as string[] | undefined);

  const normalized = (ratios || [])
    .map((ratio) => normalizeAspectRatioLabel(ratio))
    .filter((ratio) => ratio !== "auto" && parseAspectRatio(ratio) !== null);

  return Array.from(new Set(normalized.length > 0 ? normalized : ["1:1"]));
}

function resolveAspectRatio({
  requested,
  referenceImages,
  supportedAspectRatios,
}: {
  requested?: string;
  referenceImages: ReferenceImage[];
  supportedAspectRatios: string[];
}) {
  const normalized = normalizeAspectRatioLabel(requested);
  if (normalized !== "auto" && supportedAspectRatios.includes(normalized)) {
    return normalized;
  }

  const firstReference = referenceImages[0];
  if (firstReference?.width && firstReference?.height) {
    return findClosestAspectRatio(
      firstReference.width,
      firstReference.height,
      supportedAspectRatios
    );
  }

  return supportedAspectRatios.includes("1:1")
    ? "1:1"
    : supportedAspectRatios[0] || "1:1";
}

function normalizeImageSize(
  requested?: string,
  modelConfig?: BananaImageModelConfig | null
) {
  const supported = (modelConfig?.supportedImageSizes || ["1K", "2K", "4K"])
    .map((size) => String(size).trim().toUpperCase())
    .filter((size) => ["1K", "2K", "4K"].includes(size));
  const normalized = String(requested || "2K").trim().toUpperCase();
  if (supported.includes(normalized)) return normalized;
  return supported.includes("2K") ? "2K" : supported[0] || "2K";
}

function resolveApiFormat(
  modelId: string,
  modelConfig?: BananaImageModelConfig | null
) {
  if (modelConfig?.apiFormat) return modelConfig.apiFormat;
  const preset = getBananaImageModelPresetById(modelId);
  const configured = preset?.imageGenConfig?.apiFormat;
  if (
    configured === "bananaimage-gemini" ||
    configured === "bananaimage-gpt-image-2"
  ) {
    return configured;
  }
  return modelId === "gpt-image-2"
    ? "bananaimage-gpt-image-2"
    : "bananaimage-gemini";
}

function extractPromptAndImages(
  messages: Array<{ role: string; content?: any; parts?: any[] }>
) {
  const lastUserMessage = messages.filter((message) => message.role === "user").pop();
  if (!lastUserMessage) throw new Error("未找到用户消息");

  const texts: string[] = [];
  const imageUrls: string[] = [];
  const seenImages = new Set<string>();

  const pushText = (value: unknown) => {
    if (typeof value === "string" && value.trim()) texts.push(value.trim());
  };
  const pushImage = (value: unknown) => {
    let url = "";
    if (typeof value === "string") url = value;
    else if (value instanceof URL) url = value.toString();
    else if (value && typeof value === "object") {
      url = String(
        (value as any).url ||
          (value as any).imageUrl ||
          (value as any).fileUrl ||
          (value as any).src ||
          ""
      );
    }
    url = url.trim();
    if (!url || seenImages.has(url)) return;
    if (!url.startsWith("data:image/") && !/^https?:\/\//i.test(url)) return;
    seenImages.add(url);
    imageUrls.push(url);
  };

  const visitPart = (part: any) => {
    if (!part) return;
    if (typeof part === "string") {
      pushText(part);
      return;
    }
    if (part.type === "text" || typeof part.text === "string") {
      pushText(part.text);
    }
    if (part.type === "image_url") pushImage(part.image_url);
    if (part.type === "image") pushImage(part.image || part.url);
    if (part.type === "file") pushImage(part.url || part.file || part.data);
    if (part.inlineData?.data) {
      pushImage(
        `data:${part.inlineData.mimeType || "image/png"};base64,${part.inlineData.data}`
      );
    }
  };

  const content = lastUserMessage.content;
  if (typeof content === "string") pushText(content);
  else if (Array.isArray(content)) content.forEach(visitPart);
  if (Array.isArray(lastUserMessage.parts)) lastUserMessage.parts.forEach(visitPart);

  const prompt = texts.join("\n").trim();
  if (!prompt) throw new Error("请输入图片生成或编辑要求");
  return { prompt, imageUrls };
}

function parseDataUrl(dataUrl: string) {
  const match = dataUrl.match(/^data:([^;,]+);base64,(.+)$/s);
  if (!match) throw new Error("参考图 Base64 格式无效");
  const mimeType = match[1] || "image/png";
  const base64 = match[2] || "";
  return { mimeType, base64, buffer: Buffer.from(base64, "base64") };
}

async function loadReferenceImage(imageUrl: string): Promise<ReferenceImage> {
  let mimeType = "image/png";
  let base64 = "";
  let buffer: Buffer;

  if (imageUrl.startsWith("data:")) {
    const parsed = parseDataUrl(imageUrl);
    mimeType = parsed.mimeType;
    base64 = parsed.base64;
    buffer = parsed.buffer;
  } else {
    const response = await fetch(imageUrl, {
      cache: "no-store",
      redirect: "follow",
      signal: AbortSignal.timeout(REFERENCE_DOWNLOAD_TIMEOUT_MS),
      headers: {
        Accept: "image/avif,image/webp,image/png,image/jpeg,image/*,*/*;q=0.8",
        "User-Agent": "yedir-ai-bananaimage/1.0",
      },
    });
    if (!response.ok) {
      throw new Error(`BananaImage 参考图下载失败（HTTP ${response.status}）`);
    }
    mimeType = response.headers.get("content-type")?.split(";")[0] || "image/png";
    buffer = Buffer.from(await response.arrayBuffer());
    base64 = buffer.toString("base64");
  }

  const metadata = await sharp(buffer, { failOn: "none" }).metadata().catch(() => null);
  return {
    base64,
    dataUrl: `data:${mimeType};base64,${base64}`,
    mimeType,
    width: metadata?.width,
    height: metadata?.height,
  };
}

function detectImageExtension(mimeType?: string | null) {
  const normalized = String(mimeType || "").toLowerCase();
  if (normalized.includes("jpeg") || normalized.includes("jpg")) return "jpg";
  if (normalized.includes("webp")) return "webp";
  if (normalized.includes("gif")) return "gif";
  return "png";
}

function generateFilename(prefix: string, extension: string) {
  return `${prefix}-${Date.now()}-${Math.random().toString(36).slice(2, 8)}.${extension}`;
}

function toArrayBuffer(buffer: Buffer) {
  return buffer.buffer.slice(
    buffer.byteOffset,
    buffer.byteOffset + buffer.byteLength
  ) as ArrayBuffer;
}

async function uploadImageBuffer(
  buffer: Buffer,
  mimeType?: string | null
): Promise<UploadedImage> {
  const extension = detectImageExtension(mimeType);
  const thumbnailBuffer = await sharp(buffer, { failOn: "none" })
    .resize(400, 400, { fit: "inside", withoutEnlargement: true })
    .jpeg({ quality: 82 })
    .toBuffer()
    .catch(() => buffer);

  const original = await uploadFile(
    generateFilename("bananaimage-orig", extension),
    toArrayBuffer(buffer)
  );
  const thumbnail = await uploadFile(
    generateFilename("bananaimage-thumb", "jpg"),
    toArrayBuffer(thumbnailBuffer)
  );
  return { original: original.url, thumbnail: thumbnail.url };
}

async function archiveOutputImage(imageUrl: string): Promise<UploadedImage> {
  const immediatePreview = await createDeferredUploadedImage(imageUrl, { providerType: "bananaimage" });
  if (immediatePreview) return immediatePreview;
  if (imageUrl.startsWith("data:")) {
    const parsed = parseDataUrl(imageUrl);
    return uploadImageBuffer(parsed.buffer, parsed.mimeType);
  }


  const response = await fetch(imageUrl, {
    cache: "no-store",
    redirect: "follow",
    signal: AbortSignal.timeout(REFERENCE_DOWNLOAD_TIMEOUT_MS),
  });
  if (!response.ok) {
    throw new Error(`BananaImage 结果图片下载失败（HTTP ${response.status}）`);
  }
  return uploadImageBuffer(
    Buffer.from(await response.arrayBuffer()),
    response.headers.get("content-type")
  );
}

async function parseJsonResponse(response: Response) {
  const rawText = await response.text();
  if (!rawText.trim()) return null;
  try {
    return JSON.parse(rawText);
  } catch {
    return rawText;
  }
}

function extractErrorMessage(payload: any, fallback = "BananaImage 请求失败") {
  if (typeof payload === "string" && payload.trim()) return payload.trim();
  return (
    payload?.error?.message ||
    payload?.error ||
    payload?.message ||
    payload?.detail ||
    fallback
  );
}

function extractTaskId(payload: any) {
  const value =
    payload?.task_id ||
    payload?.taskId ||
    payload?.id ||
    payload?.data?.task_id ||
    payload?.data?.taskId ||
    payload?.data?.id;
  return typeof value === "string" && value.trim() ? value.trim() : "";
}

function extractGeminiImages(payload: any) {
  const parts = payload?.candidates?.[0]?.content?.parts;
  if (!Array.isArray(parts)) return [];
  const images: string[] = [];
  for (const part of parts) {
    if (part?.inlineData?.data) {
      images.push(
        `data:${part.inlineData.mimeType || "image/png"};base64,${part.inlineData.data}`
      );
    } else if (part?.fileData?.fileUri) {
      images.push(String(part.fileData.fileUri));
    } else if (/^https?:\/\/\S+$/i.test(String(part?.text || "").trim())) {
      images.push(String(part.text).trim());
    }
  }
  return Array.from(new Set(images));
}

async function submitGeminiTask({
  baseURL,
  apiKey,
  model,
  prompt,
  referenceImages,
  aspectRatio,
  imageSize,
}: {
  baseURL: string;
  apiKey: string;
  model: string;
  prompt: string;
  referenceImages: ReferenceImage[];
  aspectRatio: string;
  imageSize: string;
}) {
  const response = await fetch(
    joinApiPath(baseURL, "/api/gemini/v1beta/asyncGenerateContent"),
    {
      method: "POST",
      headers: {
        Authorization: `Bearer ${apiKey}`,
        "Content-Type": "application/json",
        "X-Client-Source": "web",
      },
      body: JSON.stringify({
        contents: [
          {
            role: "user",
            parts: [
              { text: prompt },
              ...referenceImages.map((image) => ({
                inlineData: { mimeType: image.mimeType, data: image.base64 },
              })),
            ],
          },
        ],
        generationConfig: {
          responseModalities: ["IMAGE"],
          imageConfig: { aspectRatio, imageSize },
        },
        enableFallback: true,
        model,
      }),
    }
  );
  const payload = await parseJsonResponse(response);
  if (!response.ok) {
    throw new Error(`${extractErrorMessage(payload)}（HTTP ${response.status}）`);
  }
  const taskId = extractTaskId(payload);
  if (!taskId) throw new Error("BananaImage 未返回 task_id");
  return taskId;
}

async function waitForGeminiTask(baseURL: string, apiKey: string, taskId: string) {
  let transientFailures = 0;
  for (let attempt = 0; attempt < GEMINI_POLL_MAX_ATTEMPTS; attempt += 1) {
    await delay(GEMINI_POLL_INTERVAL_MS);
    try {
      const response = await fetch(
        joinApiPath(baseURL, `/api/gemini/v1beta/tasks/${encodeURIComponent(taskId)}`),
        { headers: { Authorization: `Bearer ${apiKey}` }, cache: "no-store" }
      );
      if (!response.ok) {
        if (response.status >= 500) throw new Error(`HTTP ${response.status}`);
        const payload = await parseJsonResponse(response);
        throw new Error(extractErrorMessage(payload));
      }
      const payload = await parseJsonResponse(response);
      transientFailures = 0;
      if (payload?.status === "completed") {
        const images = extractGeminiImages(payload.response);
        if (images.length === 0) throw new Error("BananaImage 任务完成但未返回图片");
        return images;
      }
      if (payload?.status === "failed") {
        throw new Error(extractErrorMessage(payload, "BananaImage 图片生成失败"));
      }
    } catch (error) {
      const message = error instanceof Error ? error.message : String(error);
      if (!/^HTTP 5\d\d$/.test(message) && !/fetch|network|abort/i.test(message)) {
        throw error;
      }
      transientFailures += 1;
      if (transientFailures >= MAX_TRANSIENT_POLL_FAILURES) throw error;
    }
  }
  throw new Error("BananaImage 生成超时，任务可能仍在处理，请稍后重试");
}

async function submitGptImageTask({
  baseURL,
  apiKey,
  prompt,
  referenceImages,
  aspectRatio,
  imageSize,
}: {
  baseURL: string;
  apiKey: string;
  prompt: string;
  referenceImages: ReferenceImage[];
  aspectRatio: string;
  imageSize: string;
}) {
  const body: Record<string, unknown> = {
    model: "gpt-image-2",
    prompt,
    size: aspectRatio.toLowerCase(),
    resolution: imageSize.toLowerCase(),
  };
  if (referenceImages.length > 0) {
    body.image_urls = referenceImages.map((image) => image.dataUrl);
  }

  const response = await fetch(joinApiPath(baseURL, "/api/v1/images/generations"), {
    method: "POST",
    headers: {
      Authorization: `Bearer ${apiKey}`,
      "Content-Type": "application/json",
      "X-Client-Source": "web",
    },
    body: JSON.stringify(body),
  });
  const payload = await parseJsonResponse(response);
  if (!response.ok) {
    throw new Error(`${extractErrorMessage(payload)}（HTTP ${response.status}）`);
  }
  const taskId = extractTaskId(payload);
  if (!taskId) throw new Error("BananaImage GPT Image 2 未返回 task_id");
  return taskId;
}

async function waitForGptImageTask(baseURL: string, apiKey: string, taskId: string) {
  let transientFailures = 0;
  for (let attempt = 0; attempt < GPT_POLL_MAX_ATTEMPTS; attempt += 1) {
    await delay(GPT_POLL_INTERVAL_MS);
    try {
      const response = await fetch(
        joinApiPath(baseURL, `/api/v1/images/tasks/${encodeURIComponent(taskId)}`),
        { headers: { Authorization: `Bearer ${apiKey}` }, cache: "no-store" }
      );
      if (!response.ok) {
        if (response.status >= 500) throw new Error(`HTTP ${response.status}`);
        const payload = await parseJsonResponse(response);
        throw new Error(extractErrorMessage(payload));
      }
      const payload = await parseJsonResponse(response);
      transientFailures = 0;
      if (payload?.status === "completed") {
        const url = String(payload?.url || "").trim();
        if (!url) throw new Error("BananaImage GPT Image 2 任务完成但未返回图片");
        return [url];
      }
      if (payload?.status === "failed") {
        throw new Error(extractErrorMessage(payload, "BananaImage GPT Image 2 生成失败"));
      }
    } catch (error) {
      const message = error instanceof Error ? error.message : String(error);
      if (!/^HTTP 5\d\d$/.test(message) && !/fetch|network|abort/i.test(message)) {
        throw error;
      }
      transientFailures += 1;
      if (transientFailures >= MAX_TRANSIENT_POLL_FAILURES) throw error;
    }
  }
  throw new Error("BananaImage GPT Image 2 生成超时，请稍后重试");
}

function formatMarkdown(images: UploadedImage[], imageSize: string) {
  return images
    .map((image) => `![thumb|||${image.original}|||${imageSize}](${image.thumbnail})`)
    .join("\n\n");
}

export function createBananaImage({ apiKey, baseURL }: BananaImageConfig) {
  return {
    provider: "bananaimage",
    chat: (modelId: string) => ({
      provider: "bananaimage",
      doGenerate: async (options: {
        messages: Array<{ role: string; content?: any; parts?: any[] }>;
        imageParams?: BananaImageParams;
        imageGenConfig?: BananaImageModelConfig | null;
        overrideModelId?: string;
      }) => {
        const requestedModelId = normalizeModelId(options.overrideModelId || modelId);
        const runtimeModelId = resolveBananaImageRuntimeModelId(requestedModelId);
        const apiFormat = resolveApiFormat(requestedModelId, options.imageGenConfig);
        const { prompt, imageUrls } = extractPromptAndImages(options.messages);
        const maxReferenceImages = Math.max(
          0,
          Math.min(14, options.imageGenConfig?.maxReferenceImages || 14)
        );
        const referenceImages = await Promise.all(
          imageUrls.slice(0, maxReferenceImages).map(loadReferenceImage)
        );
        const supportedAspectRatios = getSupportedAspectRatios(
          requestedModelId,
          options.imageGenConfig
        );
        const aspectRatio = resolveAspectRatio({
          requested: options.imageParams?.aspectRatio,
          referenceImages,
          supportedAspectRatios,
        });
        const imageSize = normalizeImageSize(
          options.imageParams?.imageSize,
          options.imageGenConfig
        );

        let outputUrls: string[];
        if (apiFormat === "bananaimage-gpt-image-2") {
          const taskId = await submitGptImageTask({
            baseURL,
            apiKey,
            prompt,
            referenceImages,
            aspectRatio,
            imageSize,
          });
          outputUrls = await waitForGptImageTask(baseURL, apiKey, taskId);
        } else {
          const taskId = await submitGeminiTask({
            baseURL,
            apiKey,
            model: runtimeModelId,
            prompt,
            referenceImages,
            aspectRatio,
            imageSize,
          });
          outputUrls = await waitForGeminiTask(baseURL, apiKey, taskId);
        }

        const uploadedImages = await Promise.all(
          outputUrls.slice(0, Math.max(1, options.imageParams?.imageCount || 1)).map(archiveOutputImage)
        );
        const text = formatMarkdown(uploadedImages, imageSize);
        return {
          text,
          finishReason: "stop" as const,
          usage: { promptTokens: 0, completionTokens: text.length },
        };
      },
      doStream: async function* (options: {
        messages: Array<{ role: string; content?: any; parts?: any[] }>;
        imageParams?: BananaImageParams;
        imageGenConfig?: BananaImageModelConfig | null;
        overrideModelId?: string;
      }) {
        const result = await this.doGenerate(options);
        yield { type: "text-delta", textDelta: result.text };
      },
    }),
  };
}