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

import "server-only";

import { randomUUID } from "node:crypto";
import sharp from "sharp";
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 }) {
  if (params.sourceUrl.startsWith("data:")) {
    const match = params.sourceUrl.match(/^data:(image\/[a-z0-9.+-]+);base64,([a-z0-9+/=]+)$/i);
    if (!match) throw new Error("Invalid inline generated image");
    return {
      buffer: Buffer.from(match[2], "base64"),
      contentType: match[1],
      extension: getImageExtension(match[1], ""),
    };
  }
  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 objectSuffix = `${Date.now()}-${randomUUID()}`;
  const filename = `ai-generated-archive-${objectSuffix}.${downloaded.extension}`;
  const uploaded = await uploadFile(filename, toArrayBuffer(downloaded.buffer));
  const thumbnailBuffer = await sharp(downloaded.buffer, { failOn: "none" })
    .resize(640, 640, { fit: "inside", withoutEnlargement: true })
    .webp({ quality: 82 })
    .toBuffer()
    .catch(() => null);
  const thumbnailUploaded = thumbnailBuffer
    ? await uploadFile(
        `ai-generated-archive-thumb-${objectSuffix}.webp`,
        toArrayBuffer(thumbnailBuffer),
      ).catch((error) => {
        logger.warn("generated image thumbnail upload failed", {
          error: error instanceof Error ? error.message : String(error),
        });
        return null;
      })
    : null;
  return {
    url: uploaded.url,
    thumbnailUrl: thumbnailUploaded?.url || uploaded.url,
    pathname: uploaded.pathname,
    contentType: uploaded.contentType,
    bytes: downloaded.buffer.length,
  };
}