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

import "server-only";

import sharp from "sharp";
import {
  ATLASCLOUD_BASE_URL,
  type AtlasCloudRuntimeConfig,
} from "@/lib/ai/atlascloud-models";
import { createDeferredUploadedImage } from "@/lib/storage/image-archive";
import { uploadFile } from "@/lib/storage/service";

interface AtlasCloudConfig {
  apiKey: string;
  baseURL?: string;
}

interface AtlasCloudImageParams {
  aspectRatio?: string;
  imageSize?: string;
  imageCount?: number;
  negativePrompt?: string;
  promptExtend?: boolean;
  quality?: string;
  seed?: number;
  watermark?: boolean;
}

interface AtlasCloudModelConfig {
  maxReferenceImages?: number;
  supportedAspectRatios?: string[];
  supportedImageSizes?: string[];
  atlasCloud?: AtlasCloudRuntimeConfig;
}

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

const POLL_INTERVAL_MS = 3000;
const POLL_TIMEOUT_MS = 20 * 60 * 1000;
const DOWNLOAD_TIMEOUT_MS = 120_000;
const TRANSIENT_FAILURE_LIMIT = 8;

function normalizeBaseUrl(baseURL?: string) {
  return String(baseURL || ATLASCLOUD_BASE_URL).replace(/\/+$/, "");
}

function joinApiPath(baseURL: string | undefined, path: string) {
  const base = normalizeBaseUrl(baseURL);
  const cleanPath = path.startsWith("/") ? path : `/${path}`;
  if (base.endsWith("/api/v1") && cleanPath.startsWith("/api/v1/")) {
    return `${base}${cleanPath.slice("/api/v1".length)}`;
  }
  return `${base}${cleanPath}`;
}

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

function parseJsonResponse(response: Response) {
  return response.text().then((text) => {
    if (!text.trim()) return null;
    try {
      return JSON.parse(text);
    } catch {
      return text;
    }
  });
}

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

function unwrapPrediction(payload: any) {
  return payload?.data && typeof payload.data === "object" ? payload.data : payload;
}

function extractTaskId(payload: any) {
  const prediction = unwrapPrediction(payload);
  return String(
    prediction?.id || prediction?.request_id || prediction?.task_id || ""
  ).trim();
}

function extractOutputUrls(payload: any) {
  const prediction = unwrapPrediction(payload);
  const candidates = [
    prediction?.outputs,
    prediction?.output,
    prediction?.images,
    prediction?.urls?.outputs,
  ];
  const urls: string[] = [];
  for (const candidate of candidates) {
    const values = Array.isArray(candidate) ? candidate : candidate ? [candidate] : [];
    for (const value of values) {
      const url =
        typeof value === "string"
          ? value
          : String(value?.url || value?.image_url || value?.src || "");
      if (/^(?:https?:\/\/|data:image\/)/i.test(url) && !urls.includes(url)) {
        urls.push(url);
      }
    }
  }
  return urls;
}

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

  const pushText = (value: unknown) => {
    if (typeof value === "string" && value.trim()) texts.push(value.trim());
  };
  const pushImage = (value: unknown) => {
    const url = String(
      typeof value === "string"
        ? value
        : (value as any)?.url ||
            (value as any)?.imageUrl ||
            (value as any)?.fileUrl ||
            (value as any)?.src ||
            ""
    ).trim();
    if (
      url &&
      (/^https?:\/\//i.test(url) || url.startsWith("data:image/")) &&
      !imageUrls.includes(url)
    ) {
      imageUrls.push(url);
    }
  };
  const visitPart = (part: any) => {
    if (!part) return;
    if (typeof part === "string") return pushText(part);
    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}`
      );
    }
  };

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

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

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

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

function detectExtension(mimeType: string) {
  if (/jpe?g/i.test(mimeType)) return "jpg";
  if (/webp/i.test(mimeType)) return "webp";
  if (/gif/i.test(mimeType)) return "gif";
  return "png";
}

async function uploadReferenceImage(
  baseURL: string | undefined,
  apiKey: string,
  imageUrl: string
) {
  if (/^https:\/\//i.test(imageUrl)) return imageUrl;
  const source = await loadImageSource(imageUrl);
  const form = new FormData();
  form.append(
    "file",
    new Blob([source.buffer], { type: source.mimeType }),
    `reference-${Date.now()}.${detectExtension(source.mimeType)}`
  );
  const response = await fetch(joinApiPath(baseURL, "/api/v1/model/uploadMedia"), {
    method: "POST",
    headers: { Authorization: `Bearer ${apiKey}` },
    body: form,
    signal: AbortSignal.timeout(DOWNLOAD_TIMEOUT_MS),
  });
  const payload = await parseJsonResponse(response);
  if (!response.ok) {
    throw new Error(
      `${extractErrorMessage(payload, "Atlas Cloud 参考图上传失败")}（HTTP ${response.status}）`
    );
  }
  const data = payload?.data && typeof payload.data === "object" ? payload.data : payload;
  const url = String(
    data?.url || data?.file_url || data?.fileUrl || data?.media_url || data?.mediaUrl || ""
  ).trim();
  if (!url) throw new Error("Atlas Cloud 参考图上传成功但未返回文件地址");
  return url;
}

function normalizeAspectRatio(
  requested: string | undefined,
  supported: string[] | undefined
) {
  const values = (supported || []).map(String);
  const normalized = String(requested || "auto").trim().toLowerCase();
  const exact = values.find((value) => value.toLowerCase() === normalized);
  if (exact) return exact;
  return values.find((value) => value === "auto") || values.find((value) => value === "1:1") || values[0] || "1:1";
}

function normalizeImageSize(requested: string | undefined, supported?: string[]) {
  const values = (supported || []).map((value) => String(value).trim().toUpperCase());
  const normalized = String(requested || "2K").trim().toUpperCase();
  if (values.includes(normalized)) return normalized;
  return values.includes("2K") ? "2K" : values[0] || normalized;
}

function matchEnumValue(values: Array<string | number> | undefined, value: unknown) {
  if (!values?.length) return value;
  const normalized = String(value || "").trim().toLowerCase();
  return values.find((item) => String(item).trim().toLowerCase() === normalized);
}

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

function pickClosestSizeValue(
  map: Record<string, Record<string, string>>,
  imageSize: string,
  aspectRatio: string
) {
  const tiers = [imageSize, "2K", "1K", "4K"];
  const values = tiers.map((tier) => map[tier]).find(Boolean);
  if (!values) return undefined;
  if (values[aspectRatio]) return values[aspectRatio];
  const target = parseRatio(aspectRatio);
  return Object.entries(values).sort(
    ([a], [b]) =>
      Math.abs(Math.log(target / parseRatio(a))) -
      Math.abs(Math.log(target / parseRatio(b)))
  )[0]?.[1];
}

function createRequestBody(params: {
  modelId: string;
  prompt: string;
  referenceImages: string[];
  imageParams?: AtlasCloudImageParams;
  modelConfig: AtlasCloudModelConfig;
}) {
  const runtime = params.modelConfig.atlasCloud;
  if (!runtime) throw new Error("Atlas Cloud 模型运行参数缺失，请重新同步模型");
  const body: Record<string, unknown> = {
    ...runtime.defaults,
  };
  if (runtime.propertyNames.includes("model")) body.model = params.modelId;
  if (runtime.propertyNames.includes("prompt")) body.prompt = params.prompt;
  const aspectRatio = normalizeAspectRatio(
    params.imageParams?.aspectRatio,
    params.modelConfig.supportedAspectRatios
  );
  const imageSize = normalizeImageSize(
    params.imageParams?.imageSize,
    params.modelConfig.supportedImageSizes
  );

  if (runtime.aspectRatioKey) {
    body[runtime.aspectRatioKey] =
      runtime.aspectRatioValueMap?.[aspectRatio] ||
      matchEnumValue(runtime.enums[runtime.aspectRatioKey], aspectRatio) || aspectRatio;
  }
  if (runtime.resolutionKey) {
    const resolutionValue =
      matchEnumValue(runtime.enums[runtime.resolutionKey], imageSize.toLowerCase()) ||
      matchEnumValue(runtime.enums[runtime.resolutionKey], imageSize) ||
      imageSize.toLowerCase();
    body[runtime.resolutionKey] = resolutionValue;
  } else {
    const sizeValue = pickClosestSizeValue(
      runtime.sizeValueMap,
      imageSize,
      aspectRatio === "auto" ? "1:1" : aspectRatio
    );
    if (sizeValue) body.size = sizeValue;
  }

  const imageInputs = runtime.imageInputs?.length
    ? runtime.imageInputs
    : runtime.imageInputKey
      ? [
          {
            key: runtime.imageInputKey,
            mode: runtime.imageInputMode || "single",
            maxItems: Math.max(1, runtime.maxReferenceImages || 1),
            required: runtime.requiredFields.includes(runtime.imageInputKey),
          },
        ]
      : [];
  let referenceCursor = 0;
  for (const input of imageInputs) {
    const values = params.referenceImages.slice(
      referenceCursor,
      referenceCursor + input.maxItems
    );
    if (values.length > 0) {
      body[input.key] = input.mode === "single" ? values[0] : values;
      referenceCursor += values.length;
    } else if (input.required) {
      throw new Error(
        `该 Atlas Cloud 模型至少需要 ${referenceCursor + 1} 张参考图`
      );
    }
  }
  if (imageInputs.length === 0 && params.referenceImages.length > 0) {
    throw new Error("当前 Atlas Cloud 模型不支持参考图，请选择图片编辑模型");
  }

  if (runtime.outputCountKey) {
    const requestedCount = Math.max(1, params.imageParams?.imageCount || 1);
    body[runtime.outputCountKey] = Math.min(runtime.maxImages, requestedCount);
  }
  const optionalParams: Array<[string, unknown]> = [
    ["negative_prompt", params.imageParams?.negativePrompt],
    ["prompt_extend", params.imageParams?.promptExtend],
    ["quality", params.imageParams?.quality],
    ["seed", params.imageParams?.seed],
    ["watermark", params.imageParams?.watermark],
  ];
  for (const [key, value] of optionalParams) {
    if (value !== undefined && runtime.propertyNames.includes(key)) {
      body[key] = matchEnumValue(runtime.enums[key], value) ?? value;
    }
  }
  if (runtime.propertyNames.includes("enable_sync_mode")) {
    body.enable_sync_mode = false;
  }
  if (runtime.propertyNames.includes("enable_base64_output")) {
    body.enable_base64_output = false;
  }
  const unsupportedRequiredFields = runtime.requiredFields.filter(
    (key) => body[key] === undefined
  );
  if (unsupportedRequiredFields.length > 0) {
    throw new Error(
      `当前界面暂不支持该 Atlas Cloud 模型的必填参数：${unsupportedRequiredFields.join(", ")}`
    );
  }
  return { body, imageSize };
}

async function submitTask(params: {
  apiKey: string;
  baseURL?: string;
  endpoint: string;
  body: Record<string, unknown>;
}) {
  const response = await fetch(joinApiPath(params.baseURL, params.endpoint), {
    method: "POST",
    headers: {
      Authorization: `Bearer ${params.apiKey}`,
      "Content-Type": "application/json",
    },
    body: JSON.stringify(params.body),
    signal: AbortSignal.timeout(120_000),
  });
  const payload = await parseJsonResponse(response);
  if (!response.ok) {
    throw new Error(
      `${extractErrorMessage(payload)}（HTTP ${response.status}）`
    );
  }
  const immediateOutputs = extractOutputUrls(payload);
  if (immediateOutputs.length > 0) return { outputs: immediateOutputs, taskId: "" };
  const taskId = extractTaskId(payload);
  if (!taskId) throw new Error("Atlas Cloud 未返回任务 ID");
  return { outputs: [] as string[], taskId };
}

async function waitForTask(params: {
  apiKey: string;
  baseURL?: string;
  endpoint: string;
  taskId: string;
}) {
  const startedAt = Date.now();
  let transientFailures = 0;
  while (Date.now() - startedAt < POLL_TIMEOUT_MS) {
    await delay(POLL_INTERVAL_MS);
    try {
      const endpoint = params.endpoint.replace(
        /\{(?:request_id|id|task_id)\}/g,
        encodeURIComponent(params.taskId)
      );
      const response = await fetch(joinApiPath(params.baseURL, endpoint), {
        cache: "no-store",
        headers: { Authorization: `Bearer ${params.apiKey}` },
        signal: AbortSignal.timeout(60_000),
      });
      const payload = await parseJsonResponse(response);
      if (!response.ok) {
        if (response.status >= 500 || response.status === 429) {
          throw new Error(`HTTP ${response.status}`);
        }
        throw new Error(extractErrorMessage(payload));
      }
      transientFailures = 0;
      const prediction = unwrapPrediction(payload);
      const status = String(prediction?.status || "").toLowerCase();
      const outputs = extractOutputUrls(payload);
      if (status === "completed" || status === "succeeded" || status === "success") {
        if (outputs.length === 0) {
          throw new Error("Atlas Cloud 任务完成但未返回图片");
        }
        return outputs;
      }
      if (["failed", "error", "timeout", "canceled", "cancelled"].includes(status)) {
        throw new Error(extractErrorMessage(payload, "Atlas Cloud 图片生成失败"));
      }
    } catch (error) {
      const message = error instanceof Error ? error.message : String(error);
      if (!/^HTTP (?:429|5\d\d)$/.test(message) && !/fetch|network|abort/i.test(message)) {
        throw error;
      }
      transientFailures += 1;
      if (transientFailures >= TRANSIENT_FAILURE_LIMIT) throw error;
    }
  }
  throw new Error("Atlas Cloud 图片生成超时，任务可能仍在处理，请稍后重试");
}

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

async function uploadImageBuffer(buffer: Buffer, mimeType: string) {
  const extension = detectExtension(mimeType);
  const thumbnailBuffer = await sharp(buffer, { failOn: "none" })
    .resize(400, 400, { fit: "inside", withoutEnlargement: true })
    .jpeg({ quality: 82 })
    .toBuffer()
    .catch(() => buffer);
  const suffix = `${Date.now()}-${Math.random().toString(36).slice(2, 8)}`;
  const [original, thumbnail] = await Promise.all([
    uploadFile(`atlascloud-orig-${suffix}.${extension}`, toArrayBuffer(buffer)),
    uploadFile(`atlascloud-thumb-${suffix}.jpg`, toArrayBuffer(thumbnailBuffer)),
  ]);
  return { original: original.url, thumbnail: thumbnail.url };
}

async function archiveOutputImage(imageUrl: string): Promise<UploadedImage> {
  if (imageUrl.startsWith("data:")) {
    const parsed = parseDataUrl(imageUrl);
    return uploadImageBuffer(parsed.buffer, parsed.mimeType);
  }
  const deferred = await createDeferredUploadedImage(imageUrl, {
    providerType: "atlascloud",
  });
  if (deferred) return deferred;
  const response = await fetch(imageUrl, {
    cache: "no-store",
    redirect: "follow",
    signal: AbortSignal.timeout(DOWNLOAD_TIMEOUT_MS),
    headers: { "User-Agent": "yedir-ai-atlascloud/1.0" },
  });
  if (!response.ok) {
    throw new Error(`Atlas Cloud 结果图片下载失败（HTTP ${response.status}）`);
  }
  return uploadImageBuffer(
    Buffer.from(await response.arrayBuffer()),
    response.headers.get("content-type") || "image/png"
  );
}

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

export function createAtlasCloud({ apiKey, baseURL }: AtlasCloudConfig) {
  return {
    provider: "atlascloud",
    chat: (modelId: string, defaultModelConfig?: AtlasCloudModelConfig | null) => ({
      provider: "atlascloud",
      doGenerate: async (options: {
        messages: Array<{ role: string; content?: any; parts?: any[] }>;
        imageParams?: AtlasCloudImageParams;
        imageGenConfig?: AtlasCloudModelConfig | null;
        overrideModelId?: string;
      }) => {
        const runtimeModelId = String(options.overrideModelId || modelId || "").trim();
        if (!runtimeModelId) throw new Error("Atlas Cloud 模型 ID 未配置");
        const modelConfig = options.imageGenConfig || defaultModelConfig || {};
        if (!modelConfig.atlasCloud) {
          throw new Error("Atlas Cloud 模型能力配置缺失，请在系统设置中重新同步模型");
        }
        const { prompt, imageUrls } = extractPromptAndImages(options.messages);
        const maxReferenceImages = Math.max(
          0,
          Math.min(14, modelConfig.maxReferenceImages ?? modelConfig.atlasCloud.maxReferenceImages)
        );
        const referenceImages = await Promise.all(
          imageUrls
            .slice(0, maxReferenceImages)
            .map((url) => uploadReferenceImage(baseURL, apiKey, url))
        );
        const { body, imageSize } = createRequestBody({
          modelId: runtimeModelId,
          prompt,
          referenceImages,
          imageParams: options.imageParams,
          modelConfig,
        });
        const submitted = await submitTask({
          apiKey,
          baseURL,
          endpoint: modelConfig.atlasCloud.submitEndpoint,
          body,
        });
        const outputUrls = submitted.outputs.length
          ? submitted.outputs
          : await waitForTask({
              apiKey,
              baseURL,
              endpoint: modelConfig.atlasCloud.pollEndpoint,
              taskId: submitted.taskId,
            });
        const requestedCount = Math.max(1, options.imageParams?.imageCount || 1);
        const archived = await Promise.all(
          outputUrls.slice(0, requestedCount).map(archiveOutputImage)
        );
        const text = formatMarkdown(archived, 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?: AtlasCloudImageParams;
        imageGenConfig?: AtlasCloudModelConfig | null;
        overrideModelId?: string;
      }) {
        const result = await this.doGenerate(options);
        yield { type: "text-delta", textDelta: result.text };
      },
    }),
  };
}