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

import "server-only";

import { and, asc, eq, inArray, lt, lte, sql } from "drizzle-orm";
import { db } from "@/lib/db";
import { imageArchiveTask, type ImageArchiveTask } from "@/lib/db/schema";

export type ImageArchiveTaskRecord = ImageArchiveTask;

const DEFAULT_DISPLAY_HOLD_HOURS = 1;
const MAX_IMAGE_ARCHIVE_ATTEMPTS = 12;

function getDisplayHoldHours() {
  const raw = Number(
    process.env.IMAGE_ARCHIVE_DISPLAY_HOLD_HOURS || DEFAULT_DISPLAY_HOLD_HOURS
  );
  if (!Number.isFinite(raw)) return DEFAULT_DISPLAY_HOLD_HOURS;
  return Math.max(0, Math.min(24 * 30, raw));
}

export function getImageArchiveDisplayHoldMs() {
  return getDisplayHoldHours() * 60 * 60 * 1000;
}

function isArchivableHttpUrl(value: string) {
  if (!/^https?:\/\//i.test(value)) return false;
  if (/^(data|blob):/i.test(value)) return false;
  return true;
}

export async function enqueueImageArchiveTask(params: {
  sourceUrl: string;
  archiveAfterAt?: Date;
  replaceAfterAt?: Date;
  metadata?: Record<string, unknown> | null;
}) {
  const sourceUrl = String(params.sourceUrl || "").trim();
  if (!isArchivableHttpUrl(sourceUrl)) return null;

  const now = new Date();
  const archiveAfterAt = params.archiveAfterAt || now;
  const replaceAfterAt =
    params.replaceAfterAt ||
    new Date(now.getTime() + getImageArchiveDisplayHoldMs());

  const [task] = await db
    .insert(imageArchiveTask)
    .values({
      sourceUrl,
      status: "pending",
      archiveAfterAt,
      replaceAfterAt,
      metadata: params.metadata || null,
      updatedAt: now,
    })
    .onConflictDoUpdate({
      target: imageArchiveTask.sourceUrl,
      set: {
        replaceAfterAt,
        metadata: params.metadata || undefined,
        updatedAt: now,
      },
    })
    .returning();

  return task || null;
}

export async function enqueueImageArchiveTasks(params: {
  urls: string[];
  metadata?: Record<string, unknown> | null;
}) {
  const urls = Array.from(
    new Set(
      (params.urls || [])
        .map((url) => String(url || "").trim())
        .filter(isArchivableHttpUrl)
    )
  );

  const tasks: ImageArchiveTaskRecord[] = [];
  for (const sourceUrl of urls) {
    const task = await enqueueImageArchiveTask({
      sourceUrl,
      metadata: params.metadata || null,
    });
    if (task) tasks.push(task);
  }
  return tasks;
}

export async function getImageArchiveTasksToStore(params: { limit: number }) {
  const now = new Date();
  return await db
    .select()
    .from(imageArchiveTask)
    .where(
      and(
        inArray(imageArchiveTask.status, ["pending", "failed"]),
        lte(imageArchiveTask.archiveAfterAt, now),
        lt(imageArchiveTask.attemptCount, MAX_IMAGE_ARCHIVE_ATTEMPTS)
      )
    )
    .orderBy(asc(imageArchiveTask.createdAt), asc(imageArchiveTask.id))
    .limit(Math.max(1, Math.min(50, params.limit)));
}

export async function markImageArchiveTaskArchiving(id: string) {
  await db
    .update(imageArchiveTask)
    .set({
      status: "archiving",
      attemptCount: sql`${imageArchiveTask.attemptCount} + 1`,
      lastError: null,
      updatedAt: new Date(),
    })
    .where(eq(imageArchiveTask.id, id));
}

export async function markImageArchiveTaskStored(params: {
  id: string;
  targetUrl: string;
  targetStorageConfigId?: string | null;
}) {
  await db
    .update(imageArchiveTask)
    .set({
      status: "stored",
      targetUrl: params.targetUrl,
      targetStorageConfigId: params.targetStorageConfigId || null,
      archivedAt: new Date(),
      lastError: null,
      updatedAt: new Date(),
    })
    .where(eq(imageArchiveTask.id, params.id));
}

export async function markImageArchiveTaskFailed(params: {
  id: string;
  error: unknown;
}) {
  const message = String(
    params.error instanceof Error ? params.error.message : params.error || "Unknown error"
  ).slice(0, 2000);
  const isPermanentSourceError = /\bHTTP (?:404|410)\b/i.test(message);

  await db
    .update(imageArchiveTask)
    .set({
      status: isPermanentSourceError
        ? "abandoned"
        : sql`case
            when ${imageArchiveTask.attemptCount} >= ${MAX_IMAGE_ARCHIVE_ATTEMPTS}
              then 'abandoned'
            else 'failed'
          end`,
      lastError: message,
      archiveAfterAt: new Date(Date.now() + 10 * 60 * 1000),
      updatedAt: new Date(),
    })
    .where(eq(imageArchiveTask.id, params.id));
}

export async function getImageArchiveTasksToReplace(params: { limit: number }) {
  const now = new Date();
  return await db
    .select()
    .from(imageArchiveTask)
    .where(
      and(
        eq(imageArchiveTask.status, "stored"),
        lte(imageArchiveTask.replaceAfterAt, now)
      )
    )
    .orderBy(asc(imageArchiveTask.createdAt), asc(imageArchiveTask.id))
    .limit(Math.max(1, Math.min(50, params.limit)));
}

export async function markImageArchiveTaskReplacing(id: string) {
  await db
    .update(imageArchiveTask)
    .set({
      status: "replacing",
      updatedAt: new Date(),
    })
    .where(eq(imageArchiveTask.id, id));
}

export async function markImageArchiveTaskCompleted(params: {
  id: string;
  replacedCount?: number;
}) {
  await db
    .update(imageArchiveTask)
    .set({
      status: "completed",
      replacedAt: new Date(),
      lastError: null,
      metadata:
        params.replacedCount === undefined
          ? undefined
          : sql`coalesce(${imageArchiveTask.metadata}, '{}'::jsonb) || ${JSON.stringify({
              replacedCount: params.replacedCount,
            })}::jsonb`,
      updatedAt: new Date(),
    })
    .where(eq(imageArchiveTask.id, params.id));
}

export async function markImageArchiveTaskReplaceFailed(params: {
  id: string;
  error: unknown;
}) {
  const message = String(
    params.error instanceof Error ? params.error.message : params.error || "Unknown error"
  ).slice(0, 2000);

  await db
    .update(imageArchiveTask)
    .set({
      status: "stored",
      lastError: message,
      replaceAfterAt: new Date(Date.now() + 10 * 60 * 1000),
      updatedAt: new Date(),
    })
    .where(eq(imageArchiveTask.id, params.id));
}

export async function getImageArchiveTaskStats() {
  const grouped = await db
    .select({
      status: imageArchiveTask.status,
      count: sql<number>`count(*)::int`,
    })
    .from(imageArchiveTask)
    .groupBy(imageArchiveTask.status);

  const pendingStore = grouped
    .filter((item) => item.status === "pending" || item.status === "failed")
    .reduce((sum, item) => sum + (item.count || 0), 0);
  const pendingReplace = grouped
    .filter((item) => item.status === "stored")
    .reduce((sum, item) => sum + (item.count || 0), 0);

  return {
    grouped,
    pendingStore,
    pendingReplace,
  };
}