NianAIGC/lib/jimeng/capabilities.ts

import type { EnabledImageCapability, ImageCapability } from "@/lib/types";

export type JimengCapabilityConfig = {
  id: ImageCapability;
  label: string;
  reqKey: string;
  enabled: boolean;
  route: "generate" | "inpaint" | "upscale";
  description: string;
  docsUrl: string;
};

export function getJimengCapabilities(): Record<ImageCapability, JimengCapabilityConfig> {
  return {
    "image.generate": {
      id: "image.generate",
      label: "图片生成 4.6",
      reqKey: process.env.JIMENG_IMAGE_GENERATE_46_REQ_KEY || "jimeng_seedream46_cvtob",
      enabled: true,
      route: "generate",
      description: "面向营销主图、人像写真、平面设计和风格化素材的通用生成能力。",
      docsUrl: "https://www.volcengine.com/docs/85621/2275082"
    },
    "image.inpaint": {
      id: "image.inpaint",
      label: "交互编辑 inpainting",
      reqKey: process.env.JIMENG_IMAGE_INPAINT_REQ_KEY || "jimeng_image2image_dream_inpaint",
      enabled: true,
      route: "inpaint",
      description: "局部重绘与消除笔能力，使用原图和黑白 mask 生成编辑结果。",
      docsUrl: "https://www.volcengine.com/docs/85621/1976207"
    },
    "image.upscale": {
      id: "image.upscale",
      label: "智能超清",
      reqKey: process.env.JIMENG_IMAGE_UPSCALE_REQ_KEY || "jimeng_i2i_seed3_tilesr_cvtob",
      enabled: true,
      route: "upscale",
      description: "对 AIGC 图和低质量图像做 4K/8K 超清与细节增强。",
      docsUrl: "https://www.volcengine.com/docs/85621/2164806"
    }
  };
}

export function getEnabledImageCapability(id: string): JimengCapabilityConfig {
  const capabilities = getJimengCapabilities();
  const capability = capabilities[id as ImageCapability];
  if (!capability || !capability.enabled) {
    throw new Error(`Unsupported or disabled image capability: ${id}`);
  }
  return capability;
}

export function getVisibleImageCapabilities(): JimengCapabilityConfig[] {
  return Object.values(getJimengCapabilities()).filter(
    (capability): capability is JimengCapabilityConfig & { id: EnabledImageCapability } =>
      capability.enabled
  );
}

export function isRetryableVisualCode(code?: number | string): boolean {
  return code === 50429 || code === 50430 || code === "50429" || code === "50430";
}

export function buildJimengPayload(
  capability: EnabledImageCapability,
  reqKey: string,
  input: Record<string, unknown>
): Record<string, unknown> {
  if (capability === "image.generate") {
    const payload: Record<string, unknown> = {
      req_key: reqKey,
      prompt: String(input.prompt || "").trim()
    };
    if (!payload.prompt) throw new Error("Prompt is required for image generation.");
    const imageUrls = asStringArray(input.imageUrls);
    if (imageUrls.length) payload.image_urls = imageUrls;
    assignOptional(payload, input, ["scale", "width", "height", "min_ratio", "max_ratio", "force_single"]);
    return payload;
  }

  if (capability === "image.inpaint") {
    const imageUrls = asStringArray(input.imageUrls);
    if (imageUrls.length !== 2) {
      throw new Error("Inpainting requires exactly two public image URLs: original image and mask.");
    }
    return {
      req_key: reqKey,
      image_urls: imageUrls,
      prompt: String(input.prompt || "删除").trim() || "删除",
      ...(typeof input.seed === "number" ? { seed: input.seed } : {})
    };
  }

  const imageUrls = asStringArray(input.imageUrls);
  if (imageUrls.length !== 1) {
    throw new Error("Upscale requires exactly one public image URL.");
  }
  const resolution = input.resolution === "8k" ? "8k" : "4k";
  return {
    req_key: reqKey,
    image_urls: imageUrls,
    resolution,
    ...(typeof input.scale === "number" ? { scale: input.scale } : {})
  };
}

export function buildJimengQueryPayload(reqKey: string, taskId: string): Record<string, unknown> {
  return {
    req_key: reqKey,
    task_id: taskId,
    req_json: JSON.stringify({
      return_url: true,
      logo_info: {
        add_logo: false,
        position: 0,
        language: 0,
        opacity: 1
      }
    })
  };
}

function asStringArray(value: unknown): string[] {
  if (!Array.isArray(value)) return [];
  return value.map(String).map((item) => item.trim()).filter(Boolean);
}

function assignOptional(target: Record<string, unknown>, input: Record<string, unknown>, keys: string[]) {
  for (const key of keys) {
    if (input[key] !== undefined && input[key] !== null && input[key] !== "") {
      target[key] = input[key];
    }
  }
}