ai-chatbot/scripts/storage/migrate-historical-generated-images.ts

import { PutObjectCommand, S3Client } from "@aws-sdk/client-s3";
import { createHash } from "node:crypto";
import {
  appendFileSync,
  chmodSync,
  existsSync,
  mkdirSync,
  readFileSync,
} from "node:fs";
import { dirname } from "node:path";
import postgres from "postgres";
import { replaceImageUrlInDatabase } from "@/lib/storage/image-url-replacement";

type ManifestStatus =
  | "uploaded"
  | "upload_failed"
  | "completed"
  | "replace_failed";

type ManifestEvent = {
  sourceUrl: string;
  status: ManifestStatus;
  targetUrl?: string;
  pathname?: string;
  bytes?: number;
  sha256?: string;
  changed?: number;
  error?: string;
  recordedAt: string;
};

type StorageConfigRow = {
  id: string;
  name: string;
  type: string;
  endpoint: string;
  region: string | null;
  bucket: string;
  accessKeyId: string;
  secretAccessKey: string;
  publicUrl: string;
  isEnabled: boolean;
};

const MAX_IMAGE_BYTES = 100 * 1024 * 1024;

function readArg(name: string, fallback = "") {
  const prefix = `--${name}=`;
  return (
    process.argv.slice(2).find((item) => item.startsWith(prefix))?.slice(prefix.length) ||
    fallback
  ).trim();
}

function hasFlag(name: string) {
  return process.argv.slice(2).includes(`--${name}`);
}

function parseDate(value: string, name: string) {
  const parsed = new Date(`${value}T00:00:00.000Z`);
  if (!/^\d{4}-\d{2}-\d{2}$/.test(value) || Number.isNaN(parsed.getTime())) {
    throw new Error(`Invalid --${name} date: ${value}`);
  }
  return parsed;
}

function safeError(error: unknown) {
  return String(error instanceof Error ? error.message : error || "Unknown error").slice(
    0,
    2000
  );
}

function loadManifest(path: string) {
  const latest = new Map<string, ManifestEvent>();
  if (!existsSync(path)) return latest;

  for (const line of readFileSync(path, "utf8").split("\n")) {
    if (!line.trim()) continue;
    try {
      const event = JSON.parse(line) as ManifestEvent;
      if (event.sourceUrl && event.status) latest.set(event.sourceUrl, event);
    } catch {
      throw new Error(`Invalid JSONL manifest line in ${path}`);
    }
  }
  return latest;
}

function appendManifest(path: string, event: ManifestEvent) {
  mkdirSync(dirname(path), { recursive: true });
  appendFileSync(path, `${JSON.stringify(event)}\n`, { encoding: "utf8", mode: 0o600 });
  chmodSync(path, 0o600);
}

function getExtension(contentType: string, sourceUrl: string) {
  const normalizedType = 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/avif": "avif",
    "image/svg+xml": "svg",
  };
  if (byType[normalizedType]) return byType[normalizedType];

  try {
    const extension = new URL(sourceUrl).pathname
      .split(".")
      .pop()
      ?.toLowerCase()
      .replace(/[^a-z0-9]/g, "");
    if (extension && ["jpg", "jpeg", "png", "webp", "gif", "avif", "svg"].includes(extension)) {
      return extension === "jpeg" ? "jpg" : extension;
    }
  } catch {
    // Use the safe fallback below.
  }

  return "img";
}

async function downloadImage(sourceUrl: string) {
  let lastError: unknown;
  for (let attempt = 1; attempt <= 3; attempt += 1) {
    try {
      const response = await fetch(sourceUrl, {
        headers: {
          Accept: "image/avif,image/webp,image/png,image/jpeg,image/*,*/*;q=0.8",
          "User-Agent": "AI-Chatbot-Historical-Qiniu-Migrator/1.0",
        },
        signal: AbortSignal.timeout(60_000),
      });
      if (!response.ok) {
        throw new Error(`download HTTP ${response.status}`);
      }

      const contentLength = Number(response.headers.get("content-length") || 0);
      if (contentLength > MAX_IMAGE_BYTES) {
        throw new Error(`image exceeds ${MAX_IMAGE_BYTES} bytes`);
      }

      const contentType = String(response.headers.get("content-type") || "");
      if (contentType && !contentType.toLowerCase().startsWith("image/")) {
        throw new Error(`unexpected content-type: ${contentType}`);
      }

      const bytes = Buffer.from(await response.arrayBuffer());
      if (bytes.length === 0) throw new Error("downloaded image is empty");
      if (bytes.length > MAX_IMAGE_BYTES) {
        throw new Error(`image exceeds ${MAX_IMAGE_BYTES} bytes`);
      }

      return {
        bytes,
        contentType: contentType.split(";")[0]?.trim() || "application/octet-stream",
      };
    } catch (error) {
      lastError = error;
      if (attempt < 3) {
        await new Promise((resolve) => setTimeout(resolve, attempt * 1000));
      }
    }
  }
  throw lastError;
}

async function verifyPublicUrl(url: string) {
  let lastStatus = 0;
  for (let attempt = 1; attempt <= 4; attempt += 1) {
    const head = await fetch(url, {
      method: "HEAD",
      signal: AbortSignal.timeout(20_000),
      headers: { "User-Agent": "AI-Chatbot-Historical-Qiniu-Verifier/1.0" },
    }).catch(() => null);
    lastStatus = head?.status || 0;
    if (head?.ok) return;

    const ranged = await fetch(url, {
      signal: AbortSignal.timeout(20_000),
      headers: {
        Range: "bytes=0-0",
        "User-Agent": "AI-Chatbot-Historical-Qiniu-Verifier/1.0",
      },
    }).catch(() => null);
    lastStatus = ranged?.status || lastStatus;
    if (ranged?.ok) return;
    await new Promise((resolve) => setTimeout(resolve, attempt * 1000));
  }
  throw new Error(`uploaded object is not publicly readable (HTTP ${lastStatus || "unknown"})`);
}

function buildTargetUrl(publicUrl: string, key: string) {
  return `${publicUrl.replace(/\/+$/, "")}/${key.replace(/^\/+/, "")}`;
}

async function runPool<T>(
  values: T[],
  concurrency: number,
  worker: (value: T, index: number) => Promise<void>
) {
  let cursor = 0;
  await Promise.all(
    Array.from({ length: Math.min(concurrency, values.length) }, async () => {
      while (cursor < values.length) {
        const index = cursor;
        cursor += 1;
        await worker(values[index], index);
      }
    })
  );
}

async function main() {
  const databaseUrl = process.env.POSTGRES_URL;
  if (!databaseUrl) throw new Error("POSTGRES_URL is required");

  const fromText = readArg("from", "2026-06-01");
  const toText = readArg("to", "2026-09-01");
  const from = parseDate(fromText, "from");
  const to = parseDate(toText, "to");
  if (to <= from) throw new Error("--to must be later than --from");

  const targetStorageId = readArg("target-storage-id");
  const manifestPath = readArg("manifest");
  const concurrency = Math.max(
    1,
    Math.min(48, Number(readArg("concurrency", "4")) || 4)
  );
  const limit = Math.max(0, Number(readArg("limit", "0")) || 0);
  const execute = hasFlag("execute");

  if (execute && (!targetStorageId || !manifestPath)) {
    throw new Error("--execute requires --target-storage-id and --manifest");
  }

  const sqlClient = postgres(databaseUrl, {
    connect_timeout: 15,
    idle_timeout: 20,
    max: Math.max(2, concurrency),
    prepare: false,
  });

  try {
    const candidateRows = await sqlClient<{ url: string }[]>`
      WITH candidate_urls AS (
        SELECT "originalUrl" AS url
        FROM "GeneratedImage"
        WHERE "createdAt" >= ${from} AND "createdAt" < ${to}
        UNION
        SELECT "thumbnailUrl" AS url
        FROM "GeneratedImage"
        WHERE "createdAt" >= ${from} AND "createdAt" < ${to}
      )
      SELECT url
      FROM candidate_urls
      WHERE url ~ '^https?://'
      ORDER BY url
    `;

    const hostCounts = new Map<string, number>();
    for (const { url } of candidateRows) {
      const host = new URL(url).hostname.toLowerCase();
      hostCounts.set(host, (hostCounts.get(host) || 0) + 1);
    }

    console.log(
      JSON.stringify({
        mode: execute ? "execute" : "plan",
        from: fromText,
        to: toText,
        distinctUrls: candidateRows.length,
        hosts: Object.fromEntries(
          [...hostCounts.entries()].sort((left, right) => right[1] - left[1])
        ),
      })
    );

    if (!execute) return;

    const storageRows = await sqlClient<StorageConfigRow[]>`
      SELECT
        id, name, type, endpoint, region, bucket,
        "accessKeyId", "secretAccessKey", "publicUrl", "isEnabled"
      FROM "StorageConfig"
      WHERE id = ${targetStorageId}
      LIMIT 1
    `;
    const storage = storageRows[0];
    if (!storage || !storage.isEnabled || storage.type !== "qiniu-kodo") {
      throw new Error("Enabled Qiniu target storage config not found");
    }
    if (!storage.endpoint || !storage.bucket || !storage.accessKeyId || !storage.secretAccessKey || !storage.publicUrl) {
      throw new Error("Qiniu target storage config is incomplete");
    }

    const targetHost = new URL(storage.publicUrl).hostname.toLowerCase();
    const manifest = loadManifest(manifestPath);
    const pendingCandidates = candidateRows
      .map((row) => row.url)
      .filter((url) => new URL(url).hostname.toLowerCase() !== targetHost)
      .filter((url) => {
        const status = manifest.get(url)?.status;
        return status !== "uploaded" && status !== "completed";
      });
    const uploadCandidates = limit > 0 ? pendingCandidates.slice(0, limit) : pendingCandidates;

    const s3 = new S3Client({
      endpoint: storage.endpoint,
      region: storage.region || "cn-east-1",
      credentials: {
        accessKeyId: storage.accessKeyId,
        secretAccessKey: storage.secretAccessKey,
      },
      forcePathStyle: true,
    });

    let uploaded = 0;
    let uploadFailed = 0;
    await runPool(uploadCandidates, concurrency, async (sourceUrl) => {
      try {
        const downloaded = await downloadImage(sourceUrl);
        const contentSha = createHash("sha256").update(downloaded.bytes).digest("hex");
        const sourceSha = createHash("sha256").update(sourceUrl).digest("hex");
        const extension = getExtension(downloaded.contentType, sourceUrl);
        const key = `historical-images/${fromText}_${toText}/${sourceSha.slice(0, 24)}-${contentSha.slice(0, 12)}.${extension}`;
        await s3.send(
          new PutObjectCommand({
            Bucket: storage.bucket,
            Key: key,
            Body: downloaded.bytes,
            ContentType: downloaded.contentType,
            CacheControl: "public, max-age=31536000, immutable",
          })
        );
        const targetUrl = buildTargetUrl(storage.publicUrl, key);
        await verifyPublicUrl(targetUrl);
        const event: ManifestEvent = {
          sourceUrl,
          targetUrl,
          pathname: key,
          status: "uploaded",
          bytes: downloaded.bytes.length,
          sha256: contentSha,
          recordedAt: new Date().toISOString(),
        };
        appendManifest(manifestPath, event);
        manifest.set(sourceUrl, event);
        uploaded += 1;
      } catch (error) {
        const event: ManifestEvent = {
          sourceUrl,
          status: "upload_failed",
          error: safeError(error),
          recordedAt: new Date().toISOString(),
        };
        appendManifest(manifestPath, event);
        manifest.set(sourceUrl, event);
        uploadFailed += 1;
      }

      const processed = uploaded + uploadFailed;
      if (processed % 25 === 0 || processed === uploadCandidates.length) {
        console.log(
          JSON.stringify({ phase: "upload", processed, total: uploadCandidates.length, uploaded, failed: uploadFailed })
        );
      }
    });

    const currentCandidateSet = new Set(candidateRows.map((row) => row.url));
    const replaceCandidates = [...manifest.values()].filter(
      (event) =>
        Boolean(event.targetUrl) &&
        (event.status === "uploaded" ||
          event.status === "replace_failed" ||
          (event.status === "completed" && currentCandidateSet.has(event.sourceUrl)))
    );

    let replaced = 0;
    let replaceFailed = 0;
    for (const event of replaceCandidates) {
      try {
        let changed = 0;
        let drained = false;
        for (let pass = 0; pass < 20; pass += 1) {
          const result = await replaceImageUrlInDatabase(
            event.sourceUrl,
            event.targetUrl as string
          );
          changed += result.changed;
          if (result.changed === 0) {
            drained = true;
            break;
          }
        }
        if (!drained) throw new Error("replacement did not drain within 20 passes");

        const completed: ManifestEvent = {
          ...event,
          status: "completed",
          changed,
          error: undefined,
          recordedAt: new Date().toISOString(),
        };
        appendManifest(manifestPath, completed);
        manifest.set(event.sourceUrl, completed);
        replaced += 1;
      } catch (error) {
        const failed: ManifestEvent = {
          ...event,
          status: "replace_failed",
          error: safeError(error),
          recordedAt: new Date().toISOString(),
        };
        appendManifest(manifestPath, failed);
        manifest.set(event.sourceUrl, failed);
        replaceFailed += 1;
      }

      const processed = replaced + replaceFailed;
      if (processed % 25 === 0 || processed === replaceCandidates.length) {
        console.log(
          JSON.stringify({ phase: "replace", processed, total: replaceCandidates.length, replaced, failed: replaceFailed })
        );
      }
    }

    const remainingRows = await sqlClient<{ count: number }[]>`
      WITH candidate_urls AS (
        SELECT "originalUrl" AS url
        FROM "GeneratedImage"
        WHERE "createdAt" >= ${from} AND "createdAt" < ${to}
        UNION
        SELECT "thumbnailUrl" AS url
        FROM "GeneratedImage"
        WHERE "createdAt" >= ${from} AND "createdAt" < ${to}
      )
      SELECT count(*)::int AS count
      FROM candidate_urls
      WHERE url ~ '^https?://' AND url NOT LIKE ${`${storage.publicUrl.replace(/\/+$/, "")}/%`}
    `;

    const finalStates = [...manifest.values()];
    console.log(
      JSON.stringify({
        phase: "complete",
        uploaded,
        uploadFailed,
        replaced,
        replaceFailed,
        remainingNonQiniuUrls: remainingRows[0]?.count || 0,
        manifest: {
          completed: finalStates.filter((event) => event.status === "completed").length,
          uploadFailed: finalStates.filter((event) => event.status === "upload_failed").length,
          replaceFailed: finalStates.filter((event) => event.status === "replace_failed").length,
        },
      })
    );
  } finally {
    await sqlClient.end({ timeout: 5 });
  }
}

main()
  .then(() => process.exit(0))
  .catch((error) => {
    console.error(safeError(error));
    process.exit(1);
  });