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

import "server-only";

import { createLogger } from "@/lib/observability/logger";
import {
  enqueueImageArchiveTasks,
  getImageArchiveDisplayHoldMs,
} 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)));
}

function isHttpUrl(value: string) {
  return /^https?:\/\//i.test(String(value || "").trim());
}

export function isDeferredImageArchiveEnabled() {
  return process.env.AI_IMAGE_DEFERRED_ARCHIVE !== "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(isHttpUrl)
    )
  );
  if (urls.length === 0) return [];

  try {
    return await enqueueImageArchiveTasks({
      urls,
      metadata: {
        ...(params.metadata || {}),
        displayHoldMs: getImageArchiveDisplayHoldMs(),
      },
    });
  } catch (error) {
    logger.warn("enqueue temporary image archive failed", {
      urls,
      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(isHttpUrl)
    )
  );
  if (urls.length === 0) {
    throw new Error("AI 未返回可用图片地址");
  }

  await registerGeneratedImageTemporaryUrls({
    urls,
    metadata: {
      providerType: params.providerType,
      modelId: params.modelId,
      imageSize: params.imageSize,
      ...(params.metadata || {}),
    },
  });

  return urls
    .map((url) => `![thumb|||${url}|||${params.imageSize}](${url})`)
    .join("\n\n");
}

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

  await registerGeneratedImageTemporaryUrls({
    urls: [url],
    metadata: metadata || null,
  });

  return {
    thumbnail: url,
    original: url,
  };
}

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 response = await fetch(params.sourceUrl, {
        signal: AbortSignal.timeout(getDownloadTimeoutMs()),
        headers: {
          "User-Agent": "AI-Chatbot-Image-Archiver/1.0",
          Accept: "image/avif,image/webp,image/png,image/jpeg,image/*,*/*;q=0.8",
        },
      });

      if (!response.ok) {
        throw new Error(`HTTP ${response.status} ${response.statusText}`);
      }

      const contentType = response.headers.get("content-type");
      const arrayBuffer = await response.arrayBuffer();
      if (!arrayBuffer || arrayBuffer.byteLength === 0) {
        throw new Error("Downloaded image is empty");
      }

      return {
        buffer: Buffer.from(arrayBuffer),
        contentType,
        extension: getImageExtension(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 filename = `ai-generated-archive-${Date.now()}.${downloaded.extension}`;
  const uploaded = await uploadFile(filename, toArrayBuffer(downloaded.buffer));
  return {
    url: uploaded.url,
    pathname: uploaded.pathname,
    contentType: uploaded.contentType,
    bytes: downloaded.buffer.length,
  };
}