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

import "server-only";

import { randomUUID } from "node:crypto";
import sharp from "sharp";
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 function isImmediateImageArchiveEnabled() {
  return process.env.AI_IMAGE_IMMEDIATE_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 未返回可用图片地址");
  }

  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() || !isHttpUrl(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;
}) {
  if (!isImmediateImageArchiveEnabled()) {
    await registerGeneratedImageTemporaryUrls(params);
    return params.urls.map((url) => ({
      originalUrl: url,
      thumbnailUrl: url,
    }));
  }

  return await Promise.all(
    params.urls.map(async (url) => {
      try {
        const archived = await archiveTemporaryImage(url);
        logger.info("generated image archived before response", {
          providerType: params.metadata?.providerType,
          modelId: params.metadata?.modelId,
          bytes: archived.bytes,
        });
        return {
          originalUrl: archived.url,
          thumbnailUrl: archived.thumbnailUrl || archived.url,
        };
      } catch (error) {
        logger.error("immediate generated image archive failed; queued for retry", {
          providerType: params.metadata?.providerType,
          modelId: params.metadata?.modelId,
          error: error instanceof Error ? error.message : String(error),
        });
        await registerGeneratedImageTemporaryUrls({
          urls: [url],
          metadata: params.metadata || null,
        });
        return {
          originalUrl: url,
          thumbnailUrl: 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 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,
  };
}