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

import "server-only";

import { createOpenAI } from "@ai-sdk/openai";
import sharp from "sharp";
import { createOpenAIImage } from "@/lib/ai/providers/openai-image";
import { VIP98K_BASE_URL, normalizeVip98kBaseUrl } from "@/lib/ai/vip98k-models";
import { createDeferredUploadedImage } from "@/lib/storage/image-archive";
import { uploadFile } from "@/lib/storage/service";

interface Vip98kConfig {
  apiKey: string;
  baseURL?: string | null;
}

interface Vip98kTaskParams {
  apiKey: string;
  baseURL?: string | null;
  modelId: string;
  kind: string;
  prompt: string;
  imageUrls: string[];
  videoUrls: string[];
  videoParams?: Record<string, unknown>;
  imageParams?: Record<string, unknown>;
  modelConfig: Vip98kModelRuntimeConfig;
}

type Vip98kModelRuntimeConfig = {
  capabilities?: Record<string, unknown> | null;
  imageGenConfig?: Record<string, unknown> | null;
};

type UploadedImage = {
  thumbnail: string;
  original: string;
};

const DOWNLOAD_TIMEOUT_MS = 120_000;
const DEFAULT_POLL_INTERVAL_SECONDS = 8;
const DEFAULT_MAX_POLL_SECONDS = 600;

function toRecord(value: unknown): Record<string, unknown> | null {
  if (!value || typeof value !== "object" || Array.isArray(value)) return null;
  return value as Record<string, unknown>;
}

function normalizeOpenAIBaseUrl(baseURL?: string | null) {
  const normalized = normalizeVip98kBaseUrl(baseURL || VIP98K_BASE_URL);
  return /\/v1$/i.test(normalized) ? normalized : `${normalized}/v1`;
}

function parseRuntimeConfig(value?: unknown): Vip98kModelRuntimeConfig {
  const record = toRecord(value);
  if (!record) return {};
  if ("capabilities" in record || "imageGenConfig" in record) {
    return {
      capabilities: toRecord(record.capabilities),
      imageGenConfig: toRecord(record.imageGenConfig),
    };
  }
  return {
    capabilities: null,
    imageGenConfig: record,
  };
}

function getStringField(source: unknown, key: string) {
  const record = toRecord(source);
  return String(record?.[key] || "").trim().toLowerCase();
}

function getVip98kMetadata(config: Vip98kModelRuntimeConfig) {
  const capabilitiesVip98k = toRecord(toRecord(config.capabilities)?.vip98k);
  const imageConfigVip98k = toRecord(toRecord(config.imageGenConfig)?.vip98k);
  return capabilitiesVip98k || imageConfigVip98k || {};
}

function inferKind(modelId: string, config: Vip98kModelRuntimeConfig) {
  const capabilities = config.capabilities;
  const metadata = getVip98kMetadata(config);
  const configuredKind =
    getStringField(capabilities, "kind") ||
    getStringField(capabilities, "modelType") ||
    getStringField(capabilities, "mediaType") ||
    getStringField(metadata, "modelType");
  if (configuredKind.includes("world")) return "world";
  if (configuredKind.includes("video")) return "video";
  if (configuredKind.includes("image-tool")) return "image-tool";
  if (configuredKind.includes("image")) return "image";
  if (configuredKind.includes("audio")) return "audio";
  if (configuredKind.includes("document")) return "document-tool";
  if (configuredKind.includes("chat") || configuredKind.includes("model")) return "chat";

  const normalizedModelId = modelId.trim().toLowerCase();
  if (normalizedModelId === "world-generation") return "world";
  if (/^veo[-_.]/.test(normalizedModelId) || normalizedModelId.startsWith("veo-")) {
    return "video";
  }
  if (
    normalizedModelId.startsWith("gpt-image") ||
    normalizedModelId.startsWith("nano-banana")
  ) {
    return "image";
  }
  return "chat";
}

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

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

function extractErrorMessage(payload: unknown, fallback = "VIP98K request failed") {
  if (typeof payload === "string" && payload.trim()) return payload.trim();
  const record = toRecord(payload);
  const data = toRecord(record?.data);
  const error = toRecord(record?.error);
  const value =
    data?.message ||
    data?.errorMessage ||
    data?.error ||
    error?.message ||
    record?.message ||
    record?.error ||
    fallback;
  return typeof value === "string" && value.trim() ? value.trim() : fallback;
}

function unwrapData(payload: unknown) {
  const record = toRecord(payload);
  return record?.data && typeof record.data === "object" ? record.data : payload;
}

function extractTaskId(payload: unknown) {
  const data = unwrapData(payload);
  const record = toRecord(data);
  return String(
    record?.taskId ||
      record?.task_id ||
      record?.id ||
      record?.requestId ||
      record?.request_id ||
      "",
  ).trim();
}

function extractResultPath(payload: unknown) {
  const data = unwrapData(payload);
  const record = toRecord(data);
  return String(record?.resultPath || record?.result_path || "").trim();
}

function extractStatus(payload: unknown) {
  const data = unwrapData(payload);
  const record = toRecord(data);
  return String(record?.status || record?.state || record?.taskStatus || "").trim().toLowerCase();
}

function extractDone(payload: unknown) {
  const data = unwrapData(payload);
  const record = toRecord(data);
  if (typeof record?.done === "boolean") return record.done;
  if (typeof record?.isDone === "boolean") return record.isDone;
  if (typeof record?.finished === "boolean") return record.finished;
  return false;
}

function collectHttpUrls(value: unknown, urls = new Set<string>()) {
  if (!value) return urls;
  if (typeof value === "string") {
    const rawUrlMatches = value.match(/https?:\/\/[^\s"'<>()[\]{}]+/g) || [];
    for (const match of rawUrlMatches) {
      urls.add(match.replace(/[),.;]+$/g, ""));
    }
    if (/^data:image\/[a-zA-Z0-9.+-]+;base64,/i.test(value)) urls.add(value);
    return urls;
  }
  if (Array.isArray(value)) {
    value.forEach((item) => collectHttpUrls(item, urls));
    return urls;
  }
  if (typeof value !== "object") return urls;
  Object.values(value as Record<string, unknown>).forEach((nested) =>
    collectHttpUrls(nested, urls),
  );
  return urls;
}

function collectOutputUrls(payload: unknown) {
  const data = unwrapData(payload);
  const candidateRecord = toRecord(data);
  const preferred = [
    candidateRecord?.result,
    candidateRecord?.results,
    candidateRecord?.output,
    candidateRecord?.outputs,
    candidateRecord?.url,
    candidateRecord?.urls,
    candidateRecord?.imageUrl,
    candidateRecord?.image_url,
    candidateRecord?.videoUrl,
    candidateRecord?.video_url,
    candidateRecord?.assetUrl,
    candidateRecord?.asset_url,
    candidateRecord?.files,
    candidateRecord?.assets,
    data,
  ];
  const urls = new Set<string>();
  preferred.forEach((candidate) => collectHttpUrls(candidate, urls));
  return Array.from(urls);
}

function parseDataUrl(value: string) {
  const match = value.match(/^data:([^;,]+);base64,(.+)$/s);
  if (!match) throw new Error("VIP98K returned invalid Base64 media");
  return {
    mimeType: match[1] || "image/png",
    buffer: Buffer.from(match[2] || "", "base64"),
  };
}

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";
  if (/mp4/i.test(mimeType)) return "mp4";
  if (/webm/i.test(mimeType)) return "webm";
  if (/quicktime/i.test(mimeType)) return "mov";
  return "png";
}

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

async function uploadBuffer(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(`vip98k-orig-${suffix}.${extension}`, toArrayBuffer(buffer)),
    uploadFile(`vip98k-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 uploadBuffer(parsed.buffer, parsed.mimeType);
  }
  const deferred = await createDeferredUploadedImage(imageUrl, {
    providerType: "vip98k",
  });
  if (deferred) return deferred;
  const response = await fetch(imageUrl, {
    cache: "no-store",
    redirect: "follow",
    headers: { "User-Agent": "yedir-ai-vip98k/1.0" },
    signal: AbortSignal.timeout(DOWNLOAD_TIMEOUT_MS),
  });
  if (!response.ok) {
    throw new Error(`VIP98K result image download failed (HTTP ${response.status})`);
  }
  return uploadBuffer(
    Buffer.from(await response.arrayBuffer()),
    response.headers.get("content-type") || "image/png",
  );
}

async function ensurePublicMediaUrl(mediaUrl: string) {
  const normalized = String(mediaUrl || "").trim();
  if (/^https?:\/\//i.test(normalized)) return normalized;
  if (!normalized.startsWith("data:")) return normalized;
  const parsed = parseDataUrl(normalized);
  const extension = detectExtension(parsed.mimeType);
  const suffix = `${Date.now()}-${Math.random().toString(36).slice(2, 8)}`;
  const uploaded = await uploadFile(
    `vip98k-ref-${suffix}.${extension}`,
    toArrayBuffer(parsed.buffer),
  );
  return uploaded.url;
}

function extractPromptAndMedia(messages: Array<{ role: string; content?: any; parts?: any[] }>) {
  const message = messages.filter((item) => item.role === "user").pop();
  if (!message) throw new Error("No user message found");

  const texts: string[] = [];
  const imageUrls: string[] = [];
  const videoUrls: string[] = [];

  const pushText = (value: unknown) => {
    if (typeof value === "string" && value.trim()) texts.push(value.trim());
  };

  const pushMedia = (value: unknown, mediaType?: 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?:\/\/|data:)/i.test(url)) return;
    const normalizedMediaType = String(mediaType || "").toLowerCase();
    const isVideo =
      normalizedMediaType.startsWith("video/") ||
      url.startsWith("data:video/") ||
      /\.(?:mp4|mov|webm|m4v)(?:[?#].*)?$/i.test(url);
    const target = isVideo ? videoUrls : imageUrls;
    if (!target.includes(url)) target.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") pushMedia(part.image_url, "image/jpeg");
    if (part.type === "image") pushMedia(part.image || part.url, "image/jpeg");
    if (part.type === "file") {
      pushMedia(
        part.url || part.file || part.data,
        part.mediaType || part.contentType || part.mimeType,
      );
    }
    if (part.inlineData?.data) {
      pushMedia(
        `data:${part.inlineData.mimeType || "image/png"};base64,${part.inlineData.data}`,
        part.inlineData.mimeType,
      );
    }
  };

  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() || "Generate media";
  return { prompt, imageUrls, videoUrls };
}

function getConfiguredPollSeconds(config: Vip98kModelRuntimeConfig, key: string, fallback: number) {
  const metadata = getVip98kMetadata(config);
  const value = Number(metadata[key]);
  if (!Number.isFinite(value) || value <= 0) return fallback;
  return Math.min(Math.max(value, 1), 3600);
}

async function submitTask(params: {
  apiKey: string;
  baseURL?: string | null;
  body: Record<string, unknown>;
}) {
  const baseURL = normalizeVip98kBaseUrl(params.baseURL);
  const response = await fetch(`${baseURL}/api/external/task/generate`, {
    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})`);
  }
  return {
    payload,
    taskId: extractTaskId(payload),
    resultPath: extractResultPath(payload),
    outputs: collectOutputUrls(payload),
  };
}

async function waitForTask(params: {
  apiKey: string;
  baseURL?: string | null;
  taskId: string;
  resultPath?: string;
  modelConfig: Vip98kModelRuntimeConfig;
}) {
  const baseURL = normalizeVip98kBaseUrl(params.baseURL);
  const intervalMs =
    getConfiguredPollSeconds(
      params.modelConfig,
      "pollIntervalSeconds",
      DEFAULT_POLL_INTERVAL_SECONDS,
    ) * 1000;
  const timeoutMs =
    getConfiguredPollSeconds(params.modelConfig, "maxPollSeconds", DEFAULT_MAX_POLL_SECONDS) *
    1000;
  const startedAt = Date.now();

  while (Date.now() - startedAt < timeoutMs) {
    await delay(intervalMs);
    const resultUrl = params.resultPath
      ? params.resultPath.startsWith("http")
        ? params.resultPath
        : `${baseURL}${params.resultPath.startsWith("/") ? "" : "/"}${params.resultPath}`
      : `${baseURL}/api/external/task/result?taskId=${encodeURIComponent(params.taskId)}`;
    const response = await fetch(resultUrl, {
      cache: "no-store",
      headers: { Authorization: `Bearer ${params.apiKey}` },
      signal: AbortSignal.timeout(60_000),
    });
    const payload = await parseJsonResponse(response);
    if (!response.ok) {
      throw new Error(`${extractErrorMessage(payload)} (HTTP ${response.status})`);
    }

    const status = extractStatus(payload);
    const done = extractDone(payload);
    const outputs = collectOutputUrls(payload);
    if (
      ["succeeded", "success", "completed", "done", "finished"].includes(status) ||
      (done && outputs.length > 0)
    ) {
      return { payload, outputs };
    }
    if (["failed", "error", "timeout", "canceled", "cancelled"].includes(status)) {
      throw new Error(extractErrorMessage(payload, "VIP98K task failed"));
    }
  }

  throw new Error("VIP98K task timed out; it may still be running");
}

function normalizeRatio(value: unknown) {
  const normalized = String(value || "").trim();
  return normalized && normalized !== "auto" ? normalized : undefined;
}

function normalizeDuration(value: unknown) {
  const numeric = Number(value);
  if (!Number.isFinite(numeric) || numeric <= 0) return undefined;
  return Math.round(numeric);
}

function buildTaskBody(params: Vip98kTaskParams) {
  const videoParams = params.videoParams || {};
  const imageParams = params.imageParams || {};
  const promptTitle = params.prompt.split(/\s+/).slice(0, 8).join(" ").slice(0, 80);
  const body: Record<string, unknown> = {
    model: params.modelId,
    displayName: promptTitle || params.modelId,
    prompt: params.prompt,
  };

  const aspectRatio =
    normalizeRatio(videoParams.ratio) ||
    normalizeRatio(videoParams.aspectRatio) ||
    normalizeRatio(imageParams.aspectRatio);
  const duration = normalizeDuration(videoParams.duration);
  const resolution = String(videoParams.resolution || imageParams.imageSize || "").trim();
  const negativePrompt = String(videoParams.negativePrompt || imageParams.negativePrompt || "").trim();

  if (aspectRatio) {
    body.ratio = aspectRatio;
    body.aspectRatio = aspectRatio;
    body.aspect_ratio = aspectRatio;
  }
  if (duration) body.duration = duration;
  if (resolution) {
    body.resolution = resolution;
    body.imageSize = resolution;
    body.image_size = resolution;
  }
  if (negativePrompt) {
    body.negativePrompt = negativePrompt;
    body.negative_prompt = negativePrompt;
  }
  if (typeof videoParams.generateAudio === "boolean") {
    body.generateAudio = videoParams.generateAudio;
    body.generate_audio = videoParams.generateAudio;
  }

  if (params.imageUrls.length > 0) {
    body.images = params.imageUrls;
    body.imageUrls = params.imageUrls;
    body.referenceImages = params.imageUrls;
    body.reference_images = params.imageUrls;
    body.image = params.imageUrls[0];
    body.imageUrl = params.imageUrls[0];
    body.image_url = params.imageUrls[0];
    body.firstFrameImage = params.imageUrls[0];
    body.first_frame_image = params.imageUrls[0];
    body.firstFrameImageUrl = params.imageUrls[0];
    body.first_frame_image_url = params.imageUrls[0];
    if (params.imageUrls[1]) {
      body.lastFrameImage = params.imageUrls[1];
      body.last_frame_image = params.imageUrls[1];
      body.endImageUrl = params.imageUrls[1];
      body.end_image_url = params.imageUrls[1];
    }
  }

  if (params.videoUrls.length > 0) {
    body.videos = params.videoUrls;
    body.videoUrls = params.videoUrls;
    body.referenceVideos = params.videoUrls;
  }

  if (params.kind === "world") {
    body.mode = "world-generation";
  } else if (params.kind === "video") {
    body.mode = params.imageUrls.length > 0 ? "image-to-video" : "text-to-video";
  }

  return body;
}

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

function formatVideoMarkdown(videoUrl: string) {
  return `<video src="${videoUrl}" controls style="max-width: 100%; border-radius: 8px; margin-bottom: 12px;"></video>\n\n[下载视频](${videoUrl})`;
}

function formatWorldMarkdown(urls: string[]) {
  if (urls.length === 0) return "3D 世界生成成功，但未返回可访问的结果地址";
  return urls.map((url, index) => `[查看 3D 世界 ${index + 1}](${url})`).join("\n\n");
}

async function runTask(params: Vip98kTaskParams) {
  const body = buildTaskBody(params);
  const submitted = await submitTask({
    apiKey: params.apiKey,
    baseURL: params.baseURL,
    body,
  });

  const outputUrls = submitted.outputs.length
    ? submitted.outputs
    : await waitForTask({
        apiKey: params.apiKey,
        baseURL: params.baseURL,
        taskId: submitted.taskId,
        resultPath: submitted.resultPath,
        modelConfig: params.modelConfig,
      }).then((result) => result.outputs);

  if (params.kind === "video") {
    const videoUrl =
      outputUrls.find((url) => /\.(mp4|webm|mov|m4v)(?:[?#].*)?$/i.test(url)) ||
      outputUrls[0];
    if (!videoUrl) throw new Error("VIP98K task completed but returned no video URL");
    return formatVideoMarkdown(videoUrl);
  }

  if (params.kind === "image" || params.kind === "image-tool") {
    if (outputUrls.length === 0) {
      throw new Error("VIP98K task completed but returned no image URL");
    }
    const imageSize = String(params.imageParams?.imageSize || "2K").trim() || "2K";
    const archived = await Promise.all(outputUrls.slice(0, 1).map(archiveOutputImage));
    return formatImageMarkdown(archived, imageSize);
  }

  if (params.kind === "world") {
    return formatWorldMarkdown(outputUrls);
  }

  return outputUrls.length ? outputUrls.join("\n") : "VIP98K task completed";
}

export function createVip98k({ apiKey, baseURL }: Vip98kConfig) {
  const rootBaseURL = normalizeVip98kBaseUrl(baseURL);
  const openAIBaseURL = normalizeOpenAIBaseUrl(rootBaseURL);
  const openAITextProvider = createOpenAI({
    apiKey,
    baseURL: openAIBaseURL,
  });
  const openAIImageProvider = createOpenAIImage({
    apiKey,
    baseURL: openAIBaseURL,
  });

  return {
    provider: "vip98k",
    chat: (modelId: string, defaultModelConfig?: unknown): any => {
      const modelConfig = parseRuntimeConfig(defaultModelConfig);
      const kind = inferKind(modelId, modelConfig);

      if (kind === "image") {
        return openAIImageProvider.chat(modelId) as any;
      }

      if (kind === "chat") {
        return openAITextProvider.chat(modelId) as any;
      }

      return {
        provider: "vip98k",
        doGenerate: async (options: {
          messages: Array<{ role: string; content?: any; parts?: any[] }>;
          imageParams?: Record<string, unknown>;
          videoParams?: Record<string, unknown>;
          imageGenConfig?: unknown;
          overrideModelId?: string;
        }) => {
          const runtimeConfig =
            options.imageGenConfig !== undefined
              ? parseRuntimeConfig({
                  capabilities: modelConfig.capabilities,
                  imageGenConfig: options.imageGenConfig,
                })
              : modelConfig;
          const activeModelId = String(options.overrideModelId || modelId || "").trim();
          if (!activeModelId) throw new Error("VIP98K model ID is not configured");
          const activeKind = inferKind(activeModelId, runtimeConfig);
          const { prompt, imageUrls, videoUrls } = extractPromptAndMedia(options.messages);
          const [publicImageUrls, publicVideoUrls] = await Promise.all([
            Promise.all(imageUrls.map(ensurePublicMediaUrl)),
            Promise.all(videoUrls.map(ensurePublicMediaUrl)),
          ]);
          const text = await runTask({
            apiKey,
            baseURL: rootBaseURL,
            modelId: activeModelId,
            kind: activeKind,
            prompt,
            imageUrls: publicImageUrls,
            videoUrls: publicVideoUrls,
            videoParams: options.videoParams,
            imageParams: options.imageParams,
            modelConfig: runtimeConfig,
          });
          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?: Record<string, unknown>;
          videoParams?: Record<string, unknown>;
          imageGenConfig?: unknown;
          overrideModelId?: string;
        }) {
          const result = await this.doGenerate(options);
          yield { type: "text-delta", textDelta: result.text };
        },
      };
    },
  };
}