ai-chatbot/lib/storage/image-archive.ts

import { fetchImageBytes, cacheImageBytes } from "@/lib/storage/image-byte-cache";
import "server-only";

import { randomUUID } from "node:crypto";
import { isGeneratedImageSource } from "./image-archive-urls";
import { scheduleImageArchiveWorker } from "./image-archive-worker";
import { createLogger } from "@/lib/observability/logger";
import {
  enqueueImageArchiveTasks,
  getImageArchiveDisplayHoldMs,
  imageArchiveSourceKey,
} from "@/lib/storage/image-archive-db";
import { uploadFile } from "@/lib/storage/service";

const logger = createLogger("image-archive");

const DEFAULT_DOWNLOAD_TIMEOUT_MS = 60_000;

type DeferredImageArchiveMetadata = Record<string, unknown> & {
  providerType?: string;
  modelId?: string;
  imageSize?: string;
};

function getDownloadTimeoutMs() {
  const parsed = Number(process.env.IMAGE_ARCHIVE_DOWNLOAD_TIMEOUT_MS);
  if (!Number.isFinite(parsed)) return DEFAULT_DOWNLOAD_TIMEOUT_MS;
  return Math.max(5_000, Math.min(5 * 60_000, Math.round(parsed)));
}

export function isDeferredImageArchiveEnabled() {
  return true;
}

export function isImmediateImageArchiveEnabled() {
  return false;
}

export async function registerGeneratedImageTemporaryUrls(params: {
  urls: string[];
  metadata?: DeferredImageArchiveMetadata | null;
}) {
  if (!isDeferredImageArchiveEnabled()) return [];
  const urls = Array.from(
    new Set(
      (params.urls || []).map((url) => String(url || "").trim()).filter(isGeneratedImageSource),
    ),
  );
  if (urls.length === 0) return [];

  try {
    const tasks = await enqueueImageArchiveTasks({
      urls,
      metadata: {
        ...(params.metadata || {}),
        displayHoldMs: getImageArchiveDisplayHoldMs(),
      },
    });
    scheduleImageArchiveWorker();
    return tasks;
  } catch (error) {
    logger.warn("enqueue temporary image archive failed", {
      imageCount: urls.length,
      error: error instanceof Error ? error.message : String(error),
    });
    return [];
  }
}

export async function formatDeferredGeneratedImageMarkdown(params: {
  urls: string[];
  imageSize: string;
  providerType?: string;
  modelId?: string;
  metadata?: DeferredImageArchiveMetadata | null;
}) {
  const urls = Array.from(
    new Set(
      (params.urls || []).map((url) => String(url || "").trim()).filter(isGeneratedImageSource),
    ),
  );
  if (urls.length === 0) {
    throw new Error("AI 未返回可用图片地址");
  }

  const metadata = {
    providerType: params.providerType,
    modelId: params.modelId,
    imageSize: params.imageSize,
    ...(params.metadata || {}),
  };
  const persisted = await persistGeneratedImageUrls({ urls, metadata });

  return persisted
    .map(
      ({ originalUrl, thumbnailUrl }) =>
        `![thumb|||${originalUrl}|||${params.imageSize}](${thumbnailUrl})`,
    )
    .join("\n\n");
}

export async function createDeferredUploadedImage(
  sourceUrl: string,
  metadata?: DeferredImageArchiveMetadata | null,
) {
  const url = String(sourceUrl || "").trim();
  if (!isDeferredImageArchiveEnabled() || !isGeneratedImageSource(url)) return null;

  const [persisted] = await persistGeneratedImageUrls({
    urls: [url],
    metadata: metadata || null,
  });

  return {
    thumbnail: persisted?.thumbnailUrl || url,
    original: persisted?.originalUrl || url,
  };
}

async function persistGeneratedImageUrls(params: {
  urls: string[];
  metadata?: DeferredImageArchiveMetadata | null;
}) {
  // Only the durable queue insert is awaited. Downloading, thumbnailing and bucket
  // uploads run after delivery and never extend the provider response time.
  await registerGeneratedImageTemporaryUrls(params);
  return params.urls.map((url) => {
    const preview = imageArchiveSourceKey(url);
    return { originalUrl: preview, thumbnailUrl: preview };
  });
}

function getImageExtension(contentType: string | null, sourceUrl: string) {
  const normalized = String(contentType || "")
    .split(";")[0]
    .trim()
    .toLowerCase();
  const byType: Record<string, string> = {
    "image/jpeg": "jpg",
    "image/jpg": "jpg",
    "image/png": "png",
    "image/webp": "webp",
    "image/gif": "gif",
    "image/svg+xml": "svg",
  };
  if (byType[normalized]) return byType[normalized];

  try {
    const pathname = new URL(sourceUrl).pathname.toLowerCase();
    const ext =
      pathname
        .split(".")
        .pop()
        ?.replace(/[^a-z0-9]+/g, "") || "";
    if (["jpg", "jpeg", "png", "webp", "gif", "svg"].includes(ext)) {
      return ext === "jpeg" ? "jpg" : ext;
    }
  } catch {
    // keep fallback
  }

  return "png";
}

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

export async function downloadTemporaryImage(params: { sourceUrl: string; retries?: number }) {
  const retries = Math.max(0, Math.min(5, params.retries ?? 2));

  for (let attempt = 0; attempt <= retries; attempt += 1) {
    try {
      const data = await fetchImageBytes(params.sourceUrl, {
        timeoutMs: getDownloadTimeoutMs(),
        headers: { Accept: "image/*,*/*;q=0.8" },
      });
      return { ...data, extension: getImageExtension(data.contentType, params.sourceUrl) };
    } catch (error) {
      if (attempt >= retries) throw error;
      await new Promise((resolve) => setTimeout(resolve, 1000 * (attempt + 1)));
    }
  }

  throw new Error("Image download failed");
}

export async function archiveTemporaryImage(sourceUrl: string) {
  const downloaded = await downloadTemporaryImage({ sourceUrl });
  const objectSuffix = `${Date.now()}-${randomUUID()}`;
  const filename = `ai-generated-archive-${objectSuffix}.${downloaded.extension}`;
  const uploaded = await uploadFile(filename, toArrayBuffer(downloaded.buffer));
  await cacheImageBytes(`source::${uploaded.url}`, downloaded);
  return {
    url: uploaded.url,
    thumbnailUrl: uploaded.url,
    pathname: uploaded.pathname,
    contentType: uploaded.contentType,
    bytes: downloaded.buffer.length,
  };
}