Veskslop-t/media-scrape/src/download.ts

/*
 * Attachment downloader for the media archive. Serial, hash-addressed, conservative.
 * Rules: docs/media-scraper.md. Third-party embeds are never targets here.
 */

import { execFile } from "node:child_process";
import { createHash } from "node:crypto";
import { writeFileSync } from "node:fs";
import { join } from "node:path";
import { promisify } from "node:util";
import type { MediaRow } from "./store.js";

const execFileAsync = promisify(execFile);

/** Read pixel dimensions with ffprobe. Null when the tool is missing or the file is not media. */
export async function probeDimensions(
    path: string,
    ffprobe = process.env.VESKSLOP_FFPROBE ?? "ffprobe"
): Promise<{ width: number; height: number } | null> {
    try {
        const { stdout } = await execFileAsync(
            ffprobe,
            ["-v", "error", "-select_streams", "v:0", "-show_entries", "stream=width,height", "-of", "json", path],
            { timeout: 15_000 }
        );
        const stream = JSON.parse(stdout)?.streams?.[0];
        if (Number.isFinite(stream?.width) && Number.isFinite(stream?.height)) {
            return { width: stream.width, height: stream.height };
        }
    } catch {
        // missing ffprobe, a timeout, or a non-visual file
    }
    return null;
}

export interface DownloadOptions {
    fetchImpl?: typeof fetch;
    /** Skip files larger than this. Default 50 MiB. */
    maxFileBytes?: number;
    /** Gap between files. Default 1000-3000 ms. */
    minGapMs?: number;
    maxGapMs?: number;
    /** Retries after the first attempt, for network and 5xx errors only. */
    maxRetries?: number;
    /** Preview size cap. Default 5 MiB. */
    previewMaxBytes?: number;
    /** Cache third-party embed thumbnails. Download mode only. */
    embedPreviews?: boolean;
    sleep?: (ms: number) => Promise<void>;
    random?: () => number;
    log?: (line: string) => void;
}

export interface DownloadOutcome {
    id: number;
    state: "downloaded" | "skipped" | "failed";
    path?: string;
    hash?: string;
    bytes?: number;
    error?: string;
}

/** 4xx that will not improve on retry: the signed URL expired or the file is gone. */
const SKIP_STATUS = new Set([400, 401, 403, 404, 410]);

const TYPE_EXTENSIONS: Record<string, string> = {
    "image/png": "png",
    "image/jpeg": "jpg",
    "image/gif": "gif",
    "image/webp": "webp",
    "image/avif": "avif",
    "video/mp4": "mp4",
    "video/webm": "webm",
    "video/quicktime": "mov",
    "video/x-matroska": "mkv",
    "audio/mpeg": "mp3",
    "audio/ogg": "ogg",
    "audio/wav": "wav",
    "audio/x-wav": "wav",
    "audio/mp4": "m4a",
    "audio/flac": "flac"
};

function extensionFor(contentType: string | null, url: string): string {
    const type = (contentType ?? "").split(";")[0].trim().toLowerCase();
    if (TYPE_EXTENSIONS[type]) return TYPE_EXTENSIONS[type];

    try {
        const match = /\.([a-z0-9]{2,5})$/i.exec(new URL(url).pathname);
        if (match) return match[1].toLowerCase();
    } catch {
        // fall through
    }
    return "bin";
}

const defaultSleep = (ms: number) => new Promise<void>(resolve => setTimeout(resolve, ms));

/** Download one media row into `dir`. The file name is the SHA-256 of the bytes. */
export async function downloadOne(
    media: Pick<MediaRow, "id" | "url" | "kind">,
    dir: string,
    options: DownloadOptions = {}
): Promise<DownloadOutcome> {
    const fetchImpl = options.fetchImpl ?? fetch;
    const maxBytes = options.maxFileBytes ?? 50 * 1024 * 1024;
    const maxRetries = options.maxRetries ?? 2;
    const sleep = options.sleep ?? defaultSleep;

    let lastError = "download failed";

    for (let attempt = 0; attempt <= maxRetries; attempt++) {
        try {
            const response = await fetchImpl(media.url);

            if (!response.ok) {
                if (SKIP_STATUS.has(response.status)) {
                    return { id: media.id, state: "skipped", error: `http ${response.status}` };
                }
                lastError = `http ${response.status}`;
            } else {
                const declared = Number(response.headers.get("content-length") ?? 0);
                if (declared > maxBytes) {
                    return { id: media.id, state: "skipped", error: `too large (${declared} bytes)` };
                }

                const buffer = Buffer.from(await response.arrayBuffer());
                if (buffer.byteLength > maxBytes) {
                    return { id: media.id, state: "skipped", error: "too large" };
                }

                const hash = createHash("sha256").update(buffer).digest("hex");
                const extension = extensionFor(response.headers.get("content-type"), media.url);
                const file = `${hash}.${extension}`;
                writeFileSync(join(dir, file), buffer);

                return { id: media.id, state: "downloaded", path: join("media", file), hash, bytes: buffer.byteLength };
            }
        } catch (error) {
            lastError = error instanceof Error ? error.message : String(error);
        }

        if (attempt < maxRetries) {
            options.log?.(`download ${media.id} failed (${lastError}), retrying`);
            await sleep(2 ** attempt * 1000 + 1000);
        }
    }

    return { id: media.id, state: "failed", error: lastError };
}

const DISCORD_HOSTS = new Set(["cdn.discordapp.com", "media.discordapp.net"]);

/**
 * The preview target for a media row.
 *
 * Attachments and stickers: their Discord-hosted URL, resized to a small WebP.
 * Embeds: the stored thumbnail or image URL, used as-is (Discord-hosted ones are still resized).
 * Returns null when the row has no preview source.
 */
export function previewUrlFor(media: {
    url: string;
    kind: string;
    previewUrl?: string | null;
}): string | null {
    const candidate =
        media.previewUrl ??
        (media.kind === "image" || media.kind === "gif" ? media.url : null);
    if (!candidate) return null;

    try {
        const url = new URL(candidate);
        if (DISCORD_HOSTS.has(url.hostname)) {
            url.hostname = "media.discordapp.net";
            url.searchParams.set("width", "320");
            url.searchParams.set("height", "320");
            url.searchParams.set("format", "webp");
        }
        return url.toString();
    } catch {
        return null;
    }
}

/** Fetch a small preview into `dir`. One attempt; a failure never fails the scrape. */
export async function downloadPreview(
    media: Pick<MediaRow, "id" | "url" | "kind" | "source" | "preview_url">,
    dir: string,
    options: DownloadOptions = {}
): Promise<DownloadOutcome> {
    // Third-party thumbnails only in download mode; the full embed is never fetched.
    if (media.source === "embed" && options.embedPreviews !== true) {
        return { id: media.id, state: "skipped", error: "embed previews off" };
    }

    const target = previewUrlFor({
        url: media.url,
        kind: media.kind,
        previewUrl: media.preview_url
    });
    if (!target) return { id: media.id, state: "skipped", error: "no preview url" };

    const fetchImpl = options.fetchImpl ?? fetch;
    const maxBytes = options.previewMaxBytes ?? 5 * 1024 * 1024;

    try {
        const response = await fetchImpl(target);
        if (!response.ok) return { id: media.id, state: "skipped", error: `http ${response.status}` };

        const buffer = Buffer.from(await response.arrayBuffer());
        if (buffer.byteLength > maxBytes) {
            return { id: media.id, state: "skipped", error: "too large" };
        }

        const hash = createHash("sha256").update(buffer).digest("hex");
        const file = `${hash}.webp`;
        writeFileSync(join(dir, file), buffer);

        return { id: media.id, state: "downloaded", path: file, hash, bytes: buffer.byteLength };
    } catch (error) {
        return {
            id: media.id,
            state: "failed",
            error: error instanceof Error ? error.message : String(error)
        };
    }
}

/** The gap before the next file, in milliseconds. */
export function downloadGapMs(options: DownloadOptions = {}): number {
    const min = options.minGapMs ?? 1000;
    const max = Math.max(min, options.maxGapMs ?? 3000);
    const random = options.random ?? Math.random;
    return min + Math.floor(random() * (max - min + 1));
}