ai-chatbot/lib/ai/providers/openai-image.ts

import "server-only";

import sharp from "sharp";
import { uploadFile } from "@/lib/storage/service";

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

export interface OpenAIImageParams {
  aspectRatio?: string;
  imageSize?: string;
  imageCount?: number;
  size?: string;
  quality?: string;
  background?: string;
  outputFormat?: string;
}

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

const DEFAULT_IMAGE_SIZE = "2K";
const DEFAULT_TIMEOUT_MS = 300000;
const SUPPORTED_ASPECT_RATIOS = [
  "1:1",
  "1:4",
  "1:8",
  "2:3",
  "3:2",
  "3:4",
  "4:1",
  "4:3",
  "4:5",
  "5:4",
  "8:1",
  "9:16",
  "16:9",
  "21:9",
] as const;

function normalizeBaseURL(baseURL: string) {
  const normalized = String(baseURL || "").trim().replace(/\/+$/, "");
  if (!normalized) return "https://api.openai.com/v1";
  if (/\/v1$/i.test(normalized)) return normalized;
  return `${normalized}/v1`;
}

function joinEndpoint(baseURL: string, path: string) {
  const normalizedPath = path.startsWith("/") ? path : `/${path}`;
  return `${normalizeBaseURL(baseURL)}${normalizedPath}`;
}

function normalizeImageSizeLabel(imageSize?: string): "1K" | "2K" | "4K" {
  const normalized = String(imageSize || "").trim().toUpperCase();
  if (normalized === "1K" || normalized === "2K" || normalized === "4K") return normalized;
  return DEFAULT_IMAGE_SIZE;
}

function normalizeAspectRatio(value?: string) {
  const normalized = String(value || "").trim();
  if (!normalized || normalized.toLowerCase() === "auto") return "auto";
  if ((SUPPORTED_ASPECT_RATIOS as readonly string[]).includes(normalized)) return normalized;
  return undefined;
}

function parseAspectRatioValue(ratio: string) {
  const [width, height] = ratio.split(":").map((item) => Number(item));
  if (!Number.isFinite(width) || !Number.isFinite(height) || width <= 0 || height <= 0) {
    return 1;
  }
  return width / height;
}

function roundEven(value: number) {
  return Math.max(2, Math.round(value / 2) * 2);
}

function inferClosestAspectRatio(aspectValue: number) {
  if (!Number.isFinite(aspectValue) || aspectValue <= 0) return "1:1";
  let closest = "1:1";
  let minDiff = Number.POSITIVE_INFINITY;
  for (const candidate of SUPPORTED_ASPECT_RATIOS) {
    const diff = Math.abs(parseAspectRatioValue(candidate) - aspectValue);
    if (diff < minDiff) {
      minDiff = diff;
      closest = candidate;
    }
  }
  return closest;
}

async function inferAspectRatioFromImageUrl(imageUrl: string): Promise<string | null> {
  const normalized = String(imageUrl || "").trim();
  if (!normalized) return null;

  try {
    let buffer: Buffer | null = null;
    if (normalized.startsWith("data:")) {
      const match = normalized.match(/^data:([A-Za-z0-9.+/-]+);base64,(.+)$/);
      if (!match) return null;
      buffer = Buffer.from(match[2], "base64");
    } else if (/^https?:\/\//i.test(normalized)) {
      const controller = new AbortController();
      const timeoutId = setTimeout(() => controller.abort(), 20000);
      try {
        const response = await fetch(normalized, {
          signal: controller.signal,
          cache: "no-store",
        });
        if (!response.ok) return null;
        buffer = Buffer.from(await response.arrayBuffer());
      } finally {
        clearTimeout(timeoutId);
      }
    }

    if (!buffer?.length) return null;
    const metadata = await sharp(buffer).metadata();
    if (!metadata.width || !metadata.height || metadata.height <= 0) return null;
    return inferClosestAspectRatio(metadata.width / metadata.height);
  } catch (error) {
    console.warn("[OpenAI Image] Failed to infer aspect ratio:", error);
    return null;
  }
}

function resolveOpenAISize(aspectRatio: string, imageSize: "1K" | "2K" | "4K") {
  const baseMap: Record<typeof imageSize, number> = {
    "1K": 1024,
    "2K": 1536,
    "4K": 2048,
  };

  const safeRatio = parseAspectRatioValue(aspectRatio);
  const base = baseMap[imageSize] || 1536;

  if (safeRatio >= 1) {
    const width = base;
    const height = roundEven(base / safeRatio);
    return `${width}x${height}`;
  }

  const height = base;
  const width = roundEven(base * safeRatio);
  return `${width}x${height}`;
}

function normalizeQuality(value?: string) {
  const normalized = String(value || "").trim().toLowerCase();
  if (normalized === "low" || normalized === "medium" || normalized === "high") {
    return normalized;
  }
  return undefined;
}

function normalizeBackground(value?: string) {
  const normalized = String(value || "").trim().toLowerCase();
  if (normalized === "transparent" || normalized === "opaque") return normalized;
  return undefined;
}

function normalizeOutputFormat(value?: string) {
  const normalized = String(value || "").trim().toLowerCase();
  if (normalized === "jpeg" || normalized === "jpg") return "jpeg";
  if (normalized === "webp") return "webp";
  if (normalized === "png") return "png";
  return undefined;
}

function detectImageExtension(imageUrl: string, contentType?: string | null) {
  const type = String(contentType || "").toLowerCase();
  if (type.includes("jpeg") || type.includes("jpg")) return "jpg";
  if (type.includes("webp")) return "webp";
  if (type.includes("png")) return "png";
  if (type.includes("gif")) return "gif";

  const cleanUrl = imageUrl.split("?")[0]?.toLowerCase() || "";
  if (cleanUrl.endsWith(".jpg") || cleanUrl.endsWith(".jpeg")) return "jpg";
  if (cleanUrl.endsWith(".webp")) return "webp";
  if (cleanUrl.endsWith(".gif")) return "gif";
  if (cleanUrl.endsWith(".png")) return "png";
  return "png";
}

async function createThumbnail(buffer: Buffer) {
  try {
    return await sharp(buffer)
      .resize(400, 400, { fit: "inside", withoutEnlargement: true })
      .jpeg({ quality: 80 })
      .toBuffer();
  } catch (error) {
    console.error("[OpenAI Image] Failed to create thumbnail:", error);
    return buffer;
  }
}

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

async function uploadGeneratedBuffer(
  buffer: Buffer,
  extension: string,
  mimeType: string
): Promise<UploadedImage> {
  const thumbnailBuffer = await createThumbnail(buffer);

  const originalFilename = generateFilename("openai-orig", extension);
  const originalArrayBuffer = buffer.buffer.slice(
    buffer.byteOffset,
    buffer.byteOffset + buffer.byteLength
  ) as ArrayBuffer;
  const originalUploaded = await uploadFile(originalFilename, originalArrayBuffer);

  const thumbFilename = generateFilename("openai-thumb", "jpg");
  const thumbArrayBuffer = thumbnailBuffer.buffer.slice(
    thumbnailBuffer.byteOffset,
    thumbnailBuffer.byteOffset + thumbnailBuffer.byteLength
  ) as ArrayBuffer;
  const thumbUploaded = await uploadFile(thumbFilename, thumbArrayBuffer);

  console.log("[OpenAI Image] Uploaded generated image:", {
    original: originalUploaded.url,
    thumbnail: thumbUploaded.url,
    mimeType,
  });

  return {
    thumbnail: thumbUploaded.url,
    original: originalUploaded.url,
  };
}

async function downloadAndUploadImage(imageUrl: string): Promise<UploadedImage> {
  const response = await fetch(imageUrl, { cache: "no-store" });
  if (!response.ok) {
    throw new Error(`Failed to download generated image: ${response.status}`);
  }
  const buffer = Buffer.from(await response.arrayBuffer());
  const extension = detectImageExtension(imageUrl, response.headers.get("content-type"));
  const mimeType = response.headers.get("content-type") || `image/${extension === "jpg" ? "jpeg" : extension}`;
  return uploadGeneratedBuffer(buffer, extension, mimeType);
}

async function uploadBase64ToS3(base64Data: string, mimeType: string) {
  const buffer = Buffer.from(base64Data, "base64");
  const extension = detectImageExtension("", mimeType);
  return uploadGeneratedBuffer(
    buffer,
    extension,
    mimeType || `image/${extension === "jpg" ? "jpeg" : extension}`
  );
}

function extractPromptAndImages(messages: Array<{ role: string; content: any }>) {
  const lastUserMessage = messages.filter((message) => message.role === "user").pop();
  if (!lastUserMessage) {
    throw new Error("No user message found");
  }

  let prompt = "";
  const imageUrls: string[] = [];

  const pushImageUrl = (value?: string) => {
    if (!value) return;
    const normalized = String(value).trim();
    if (!normalized) return;
    if (!imageUrls.includes(normalized)) {
      imageUrls.push(normalized);
    }
  };

  if (typeof lastUserMessage.content === "string") {
    prompt = lastUserMessage.content.trim();
  } else if (Array.isArray(lastUserMessage.content)) {
    for (const part of lastUserMessage.content) {
      if (part?.type === "text" && typeof part.text === "string") {
        prompt = `${prompt} ${part.text}`.trim();
      } else if (part?.type === "image_url" && part.image_url?.url) {
        pushImageUrl(part.image_url.url);
      } else if (part?.type === "file") {
        pushImageUrl(part.url || part.data);
      } else if (part?.type === "image") {
        pushImageUrl(part.image || part.url);
      }
    }
  }

  if (!prompt) {
    throw new Error("No prompt found in user message");
  }

  return { prompt, imageUrls };
}

async function loadInputImage(source: string) {
  const normalized = String(source || "").trim();
  if (!normalized) {
    throw new Error("输入图片地址为空");
  }

  if (normalized.startsWith("data:")) {
    const match = normalized.match(/^data:([A-Za-z0-9.+/-]+);base64,(.+)$/);
    if (!match) {
      throw new Error("不支持的 Base64 图片格式");
    }
    const mimeType = match[1] || "image/png";
    const buffer = Buffer.from(match[2], "base64");
    const extension = detectImageExtension("", mimeType);
    return {
      buffer,
      mimeType,
      filename: generateFilename("openai-edit-input", extension),
    };
  }

  const response = await fetch(normalized, { cache: "no-store" });
  if (!response.ok) {
    throw new Error(`Failed to fetch input image: ${response.status}`);
  }
  const buffer = Buffer.from(await response.arrayBuffer());
  const mimeType = response.headers.get("content-type") || "image/png";
  const extension = detectImageExtension(normalized, mimeType);
  return {
    buffer,
    mimeType,
    filename: generateFilename("openai-edit-input", extension),
  };
}

function extractErrorMessage(payload: any) {
  if (!payload) return "OpenAI Images 请求失败";
  if (typeof payload === "string") return payload;
  if (payload.error?.message) return String(payload.error.message);
  if (payload.message) return String(payload.message);
  try {
    return JSON.stringify(payload);
  } catch {
    return "OpenAI Images 请求失败";
  }
}

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

async function requestOpenAIImageGeneration(params: {
  baseURL: string;
  apiKey: string;
  model: string;
  prompt: string;
  size?: string;
  imageUrls: string[];
  quality?: string;
  background?: string;
  outputFormat?: string;
}) {
  const { baseURL, apiKey, model, prompt, size, imageUrls, quality, background, outputFormat } = params;

  const controller = new AbortController();
  const timeoutId = setTimeout(() => controller.abort(), DEFAULT_TIMEOUT_MS);

  try {
    if (imageUrls.length > 0) {
      const endpoint = joinEndpoint(baseURL, "/images/edits");
      const formData = new FormData();
      formData.append("model", model);
      formData.append("prompt", prompt);
      if (size) formData.append("size", size);
      formData.append("response_format", "url");

      if (quality) formData.append("quality", quality);
      if (background) formData.append("background", background);
      if (outputFormat) formData.append("output_format", outputFormat);

      if (imageUrls.length === 1) {
        const image = await loadInputImage(imageUrls[0]);
        formData.append("image", new Blob([image.buffer], { type: image.mimeType }), image.filename);
      } else {
        for (const source of imageUrls) {
          const image = await loadInputImage(source);
          formData.append("image[]", new Blob([image.buffer], { type: image.mimeType }), image.filename);
        }
      }

      const response = await fetch(endpoint, {
        method: "POST",
        headers: {
          Authorization: `Bearer ${apiKey}`,
        },
        body: formData,
        signal: controller.signal,
      });
      const payload = await parseJsonResponse(response);
      if (!response.ok) {
        throw new Error(`${extractErrorMessage(payload)} (HTTP ${response.status})`);
      }
      return payload;
    }

    const endpoint = joinEndpoint(baseURL, "/images/generations");
    const response = await fetch(endpoint, {
      method: "POST",
      headers: {
        Authorization: `Bearer ${apiKey}`,
        "Content-Type": "application/json",
      },
      body: JSON.stringify({
        model,
        prompt,
        ...(size ? { size } : {}),
        response_format: "url",
        ...(quality ? { quality } : {}),
        ...(background ? { background } : {}),
        ...(outputFormat ? { output_format: outputFormat } : {}),
      }),
      signal: controller.signal,
    });
    const payload = await parseJsonResponse(response);
    if (!response.ok) {
      throw new Error(`${extractErrorMessage(payload)} (HTTP ${response.status})`);
    }
    return payload;
  } finally {
    clearTimeout(timeoutId);
  }
}

async function extractUploadedImages(payload: any): Promise<UploadedImage[]> {
  const data = Array.isArray(payload?.data) ? payload.data : [];
  const uploadedImages: UploadedImage[] = [];

  for (const item of data) {
    if (item?.url && typeof item.url === "string") {
      uploadedImages.push(await downloadAndUploadImage(item.url));
      continue;
    }

    if (item?.b64_json && typeof item.b64_json === "string") {
      uploadedImages.push(await uploadBase64ToS3(item.b64_json, "image/png"));
    }
  }

  return uploadedImages;
}

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

export function createOpenAIImage({ apiKey, baseURL }: OpenAIImageConfig) {
  return {
    provider: "openai-image",
    chat: (modelId: string) => ({
      provider: "openai-image",
      doGenerate: async (options: {
        messages: Array<{ role: string; content: any }>;
        imageParams?: OpenAIImageParams;
        overrideModelId?: string;
      }) => {
        const { prompt, imageUrls } = extractPromptAndImages(options.messages);
        const activeModelId = options.overrideModelId || modelId;
        const normalizedAspectRatio = normalizeAspectRatio(options.imageParams?.aspectRatio);
        const aspectRatio =
          normalizedAspectRatio === "auto"
            ? await inferAspectRatioFromImageUrl(imageUrls[0] || "")
            : normalizedAspectRatio;
        const imageSize = normalizeImageSizeLabel(options.imageParams?.imageSize);
        const size =
          options.imageParams?.size ||
          (aspectRatio ? resolveOpenAISize(aspectRatio, imageSize) : undefined);
        const quality = normalizeQuality(options.imageParams?.quality);
        const background = normalizeBackground(options.imageParams?.background);
        const outputFormat = normalizeOutputFormat(options.imageParams?.outputFormat);

        console.log("[OpenAI Image] Submitting request", {
          model: activeModelId,
          size,
          aspectRatio: aspectRatio || "auto",
          imageCount: imageUrls.length,
          baseURL: normalizeBaseURL(baseURL),
        });

        const payload = await requestOpenAIImageGeneration({
          baseURL,
          apiKey,
          model: activeModelId,
          prompt,
          size,
          imageUrls,
          quality,
          background,
          outputFormat,
        });

        const uploadedImages = await extractUploadedImages(payload);
        if (uploadedImages.length === 0) {
          throw new Error("OpenAI Images 未返回可用图片");
        }

        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 }>;
        imageParams?: OpenAIImageParams;
        overrideModelId?: string;
      }) {
        const result = await this.doGenerate(options);
        yield {
          type: "text-delta",
          textDelta: result.text,
        };
      },
    }),
  };
}