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

import "server-only";

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

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

interface RuxaImageParams {
  aspectRatio?: string;
  imageSize?: string;
  imageCount?: number;
  outputFormat?: "png" | "jpg" | "jpeg" | "webp";
}

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

const DEFAULT_SIZE = "1:1";
const DEFAULT_RESOLUTION = "1k";
const DEFAULT_OUTPUT_FORMAT = "png";
const POLL_INTERVAL_MS = 2000;
const POLL_MAX_ATTEMPTS = 90;

function normalizeUrlForCompare(url: string) {
  try {
    const parsed = new URL(url);
    parsed.search = "";
    parsed.hash = "";
    return parsed.toString();
  } catch {
    return url.split("?")[0]?.split("#")[0] || url;
  }
}

function filterSourceImages(urls: string[], sourceUrls: string[]) {
  if (sourceUrls.length === 0) return urls;
  const sourceSet = new Set(sourceUrls.map((url) => normalizeUrlForCompare(url)));
  return urls.filter((url) => !sourceSet.has(normalizeUrlForCompare(url)));
}

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

function normalizeAspectRatio(aspectRatio?: string) {
  if (!aspectRatio) return DEFAULT_SIZE;
  const normalized = aspectRatio.trim().toLowerCase();
  if (normalized === "auto") return DEFAULT_SIZE;
  const ratio = aspectRatio.trim();
  if (["1:1", "16:9", "9:16", "4:3", "3:4", "2:3", "3:2"].includes(ratio)) {
    return ratio;
  }
  return DEFAULT_SIZE;
}

function normalizeResolution(imageSize?: string) {
  if (!imageSize) return DEFAULT_RESOLUTION;
  const normalized = imageSize.trim().toLowerCase();
  if (["1k", "2k", "4k"].includes(normalized)) return normalized;
  return DEFAULT_RESOLUTION;
}

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

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

function extractStatus(payload: any): string | undefined {
  const raw =
    payload?.status ||
    payload?.taskStatus ||
    payload?.task_status ||
    payload?.state ||
    payload?.data?.status ||
    payload?.data?.taskStatus ||
    payload?.data?.task_status ||
    payload?.data?.state;
  if (typeof raw !== "string") return undefined;
  return raw.toLowerCase();
}

function extractErrorMessage(payload: any): string {
  let parsedResultJson: any = null;
  if (typeof payload?.resultJson === "string" && payload.resultJson.trim()) {
    try {
      parsedResultJson = JSON.parse(payload.resultJson);
    } catch {
      parsedResultJson = null;
    }
  }

  return (
    payload?.failMsg ||
    payload?.data?.failMsg ||
    payload?.message ||
    payload?.error ||
    payload?.errorMessage ||
    parsedResultJson?.failMsg ||
    parsedResultJson?.error ||
    payload?.resultJson ||
    payload?.data?.message ||
    payload?.data?.error ||
    "RUXA 任务执行失败"
  );
}

function compactPayloadSnippet(payload: any): string {
  try {
    const json = JSON.stringify(payload);
    return json.length > 320 ? `${json.slice(0, 320)}...` : json;
  } catch {
    return String(payload || "");
  }
}

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 = (url?: string) => {
    if (!url) return;
    const normalized = String(url).trim();
    if (normalized.startsWith("data:")) {
      throw new Error("RUXA 不支持 Base64 图片传输，请先上传图片后再编辑");
    }
    if (!normalized.startsWith("http")) 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 };
}

function collectImageUrls(payload: any): string[] {
  const urls = new Set<string>();

  const isImageLikeKey = (key: string) =>
    /image|img|url|output|result|artifact|file/i.test(key);

  const isLikelyImageUrl = (value: string) =>
    /^https?:\/\//i.test(value) &&
    (/\.(png|jpg|jpeg|webp|gif)(\?|$)/i.test(value) || value.includes("image"));

  const walk = (value: any, parentKey = "") => {
    if (typeof value === "string") {
      if (isLikelyImageUrl(value) || (isImageLikeKey(parentKey) && /^https?:\/\//i.test(value))) {
        urls.add(value);
        return;
      }

      const trimmed = value.trim();
      if (
        trimmed.length > 1 &&
        trimmed.length < 200000 &&
        (trimmed.startsWith("{") || trimmed.startsWith("["))
      ) {
        try {
          const parsed = JSON.parse(trimmed);
          walk(parsed, parentKey);
        } catch {
          // ignore non-JSON strings
        }
      }
      return;
    }

    if (Array.isArray(value)) {
      for (const item of value) {
        walk(item, parentKey);
      }
      return;
    }

    if (value && typeof value === "object") {
      for (const [key, nestedValue] of Object.entries(value)) {
        walk(nestedValue, key);
      }
    }
  };

  walk(payload);
  return Array.from(urls);
}

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

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

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

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

async function downloadAndUploadImage(imageUrl: string): Promise<UploadedImage> {
  const response = await fetch(imageUrl);
  if (!response.ok) {
    throw new Error(`Failed to download image: ${response.status}`);
  }

  const arrayBuffer = await response.arrayBuffer();
  const buffer = Buffer.from(arrayBuffer);
  const thumbnailBuffer = await createThumbnail(buffer);

  const originalExt = detectImageExtension(imageUrl, response.headers.get("content-type"));
  const originalFilename = generateFilename("ruxa-image-orig", originalExt);
  const originalArrayBuffer = buffer.buffer.slice(
    buffer.byteOffset,
    buffer.byteOffset + buffer.byteLength
  ) as ArrayBuffer;
  const originalUploaded = await uploadFile(originalFilename, originalArrayBuffer);

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

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

async function waitForTaskResult(baseURL: string, apiKey: string, taskId: string, sourceUrls: string[]) {
  const endpointCandidates = [
    (id: string) => joinApiPath(baseURL, `/api/v1/tasks/query/${id}`),
    (id: string) => joinApiPath(baseURL, `/api/v1/tasks/${id}`),
    (id: string) => `${joinApiPath(baseURL, "/api/v1/tasks/query")}?taskId=${encodeURIComponent(id)}`,
  ];
  let endpointIndex = 0;

  for (let attempt = 1; attempt <= POLL_MAX_ATTEMPTS; attempt++) {
    try {
      const endpoint = endpointCandidates[endpointIndex]!(taskId);
      const response = await fetch(endpoint, {
        headers: {
          Authorization: `Bearer ${apiKey}`,
          "Content-Type": "application/json",
        },
      });

      if (!response.ok) {
        if (response.status === 404 && endpointIndex < endpointCandidates.length - 1) {
          endpointIndex++;
        }
        await new Promise((resolve) => setTimeout(resolve, POLL_INTERVAL_MS));
        continue;
      }

      const payload = await response.json();
      const imageUrls = filterSourceImages(collectImageUrls(payload), sourceUrls);
      const status = extractStatus(payload);

      if (imageUrls.length > 0) {
        return imageUrls;
      }

      if (status && ["failed", "fail", "error", "cancelled", "canceled"].includes(status)) {
        throw new Error(extractErrorMessage(payload));
      }

      if (status && ["success", "succeeded", "completed", "done", "finished"].includes(status)) {
        throw new Error("任务已完成但未返回图片地址");
      }
    } catch (error) {
      if (error instanceof Error && error.message !== "任务已完成但未返回图片地址") {
        throw error;
      }
    }

    await new Promise((resolve) => setTimeout(resolve, POLL_INTERVAL_MS));
  }

  throw new Error("RUXA 任务超时，请稍后重试");
}

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

export function createRuxaImage({ apiKey, baseURL }: RuxaImageConfig) {
  return {
    provider: "ruxa-image",
    chat: (modelId: string) => ({
      provider: "ruxa-image",
      doGenerate: async (options: {
        messages: Array<{ role: string; content: any }>;
        imageParams?: RuxaImageParams;
        overrideModelId?: string;
      }) => {
        const { prompt, imageUrls } = extractPromptAndImages(options.messages);
        const aspectRatio = normalizeAspectRatio(options.imageParams?.aspectRatio);
        const resolution = normalizeResolution(options.imageParams?.imageSize);
        const outputFormat = normalizeOutputFormat(options.imageParams?.outputFormat);
        const activeModelId = options.overrideModelId || modelId || "gpt-image-1.5";

        const createEndpoint = joinApiPath(baseURL, "/api/v1/tasks/create");
        const callbackUrl = process.env.RUXA_CALLBACK_URL?.trim();
        const input: Record<string, unknown> = {
          prompt,
          size: aspectRatio,
          aspect_ratio: aspectRatio,
          resolution,
          output_format: outputFormat,
        };
        if (imageUrls[0]) {
          input.image = imageUrls[0];
        }
        if (imageUrls.length > 1) {
          input.image_input = imageUrls;
        }
        const requestBody = {
          input,
          model: activeModelId,
          ...(callbackUrl ? { callBackUrl: callbackUrl } : {}),
        };
        const createResponse = await fetch(createEndpoint, {
          method: "POST",
          headers: {
            Authorization: `Bearer ${apiKey}`,
            "Content-Type": "application/json",
          },
          body: JSON.stringify(requestBody),
        });
        const createRawText = await createResponse.text();
        let createPayload: any = null;
        try {
          createPayload = JSON.parse(createRawText);
        } catch {
          createPayload = createRawText;
        }
        if (!createResponse.ok) {
          throw new Error(
            `${extractErrorMessage(createPayload)} (HTTP ${createResponse.status})`
          );
        }

        const createCode =
          typeof createPayload?.code === "number"
            ? createPayload.code
            : Number.isFinite(Number(createPayload?.code))
              ? Number(createPayload.code)
              : null;
        if (createCode !== null && createCode !== 200) {
          throw new Error(extractErrorMessage(createPayload));
        }

        let outputUrls = filterSourceImages(collectImageUrls(createPayload), imageUrls);
        if (outputUrls.length === 0) {
          const taskId = extractTaskId(createPayload);
          if (!taskId) {
            throw new Error(
              `RUXA 创建任务成功，但未返回 taskId，响应片段: ${compactPayloadSnippet(createPayload)}`
            );
          }
          outputUrls = await waitForTaskResult(baseURL, apiKey, taskId, imageUrls);
        }

        const uploadedImages: UploadedImage[] = [];
        for (const outputUrl of outputUrls.slice(0, options.imageParams?.imageCount || 1)) {
          uploadedImages.push(await downloadAndUploadImage(outputUrl));
        }

        if (uploadedImages.length === 0) {
          throw new Error("RUXA 未返回可用图片");
        }

        const text = formatMarkdown(uploadedImages, resolution);
        return {
          text,
          finishReason: "stop" as const,
          usage: {
            promptTokens: 0,
            completionTokens: text.length,
          },
        };
      },
      doStream: async function* (options: {
        messages: Array<{ role: string; content: any }>;
        imageParams?: RuxaImageParams;
        overrideModelId?: string;
      }) {
        const result = await this.doGenerate(options);
        yield {
          type: "text-delta",
          textDelta: result.text,
        };
      },
    }),
  };
}