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

import type { AtlasCloudRuntimeConfig } from "@/lib/ai/atlascloud-models";

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

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

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];
}

export function createAtlasCloudRequestBody(params: {
  modelId: string;
  prompt: string;
  referenceImages: string[];
  referenceVideos?: 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") ||
    runtime.requiredFields.includes("model")
  ) {
    body.model = params.modelId;
  }
  if (
    runtime.propertyNames.includes("prompt") ||
    runtime.requiredFields.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) {
    body[runtime.resolutionKey] =
      matchEnumValue(runtime.enums[runtime.resolutionKey], imageSize.toLowerCase()) ||
      matchEnumValue(runtime.enums[runtime.resolutionKey], imageSize) ||
      imageSize.toLowerCase();
  } 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" as const),
            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 模型不支持参考图，请选择图片编辑模型");
  }

  const referenceVideos = params.referenceVideos || [];
  if (runtime.videoClipInput) {
    const values = referenceVideos.slice(0, runtime.videoClipInput.maxItems);
    if (values.length > 0) {
      body[runtime.videoClipInput.key] = values.map((url) => ({
        ...runtime.videoClipInput?.itemDefaults,
        url,
      }));
    } else if (runtime.videoClipInput.required) {
      throw new Error("该 Atlas Cloud 模型至少需要 1 个参考视频");
    }
  } else if (referenceVideos.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 };
}