quak/src/thumbnails.ts

import { createHash } from "node:crypto";
import { readFileSync } from "node:fs";
import * as jpeg from "jpeg-js";
import type { Client } from "./client.js";
import type { Library } from "./library/index.js";
import { ApiError } from "./api/client.js";
import { encryptBlob, toBase64 } from "./crypto/index.js";
import type { EnteFile } from "./model/types.js";

const THUMB_MAX_DIMENSION = 720;
const THUMB_JPEG_QUALITY = 50;

export interface MissingThumbnailInfo {
    fileID: number;
    title: string;
    collection: string;
    reason: string;
}

// Three outcomes, not two. "fixed": a thumbnail was generated and uploaded.
// "failed": something went wrong (download, encode, upload) and the file still
// has no thumbnail. "skipped": the file is a format this helper cannot
// regenerate — a video, or an image that is not a baseline JPEG. Skipped is a
// deliberate, expected outcome, not an error (issue #17): the repair path is
// JPEG-only because `jpeg-js` is, and a PNG or HEIC is left for a format-aware
// tool rather than reported as a failure.
export type ThumbnailFixStatus = "fixed" | "skipped" | "failed";

export interface ThumbnailFixResult {
    fileID: number;
    title: string;
    collection: string;
    status: ThumbnailFixStatus;
    // Why the file was skipped or failed; unset when it was fixed.
    reason?: string;
}

export type ProgressCallback = (message: string) => void;

// Enumerate every file the library knows about, newest album first, each file
// once, and report those whose server-side thumbnail is missing. "Missing" is
// only two answers: an empty body, or a 404. Any other error reaching this
// point has already exhausted its retries — a failing server, a dropped
// connection, a deadline — and says nothing about whether the thumbnail
// exists, so it is logged and the file is left unreported. That distinction is
// what stops `fix-missing-thumbnails` from regenerating and uploading over
// thumbnails that were fine all along while the CDN was briefly returning 500s.
export const listMissingThumbnails = async (
    lib: Library,
    client: Client,
    onProgress?: ProgressCallback,
): Promise<MissingThumbnailInfo[]> => {
    const log = onProgress ?? (() => {});
    const api = client.getApiClient();
    const missing: MissingThumbnailInfo[] = [];
    const seen = new Set<number>();

    for (const album of lib.albums.list()) {
        log(`[${album.name}] Checking thumbnails...`);
        for (const photo of album.photos.list()) {
            if (seen.has(photo.fileID)) continue;
            seen.add(photo.fileID);
            try {
                const stream = await api.getThumbnailStream(photo.fileID);
                const reader = stream.getReader();
                let totalBytes = 0;
                for (;;) {
                    const { done, value } = await reader.read();
                    if (value) totalBytes += value.length;
                    if (done) break;
                }
                if (totalBytes === 0) {
                    missing.push({
                        fileID: photo.fileID,
                        title: photo.title,
                        collection: album.name,
                        reason: "empty thumbnail (0 bytes)",
                    });
                }
            } catch (err) {
                if (err instanceof ApiError && err.status === 404) {
                    missing.push({
                        fileID: photo.fileID,
                        title: photo.title,
                        collection: album.name,
                        reason: "thumbnail not found (HTTP 404)",
                    });
                } else {
                    log(
                        `[${album.name}] Could not check ${photo.title}: ${err instanceof Error ? err.message : String(err)} (not reported as missing)`,
                    );
                }
            }
        }
    }
    return missing;
};

// Bilinear resize of an RGBA pixel buffer.
const resizeRGBA = (
    src: Uint8Array,
    srcW: number,
    srcH: number,
    dstW: number,
    dstH: number,
): Uint8Array => {
    const dst = new Uint8Array(dstW * dstH * 4);
    const xRatio = srcW / dstW;
    const yRatio = srcH / dstH;
    for (let y = 0; y < dstH; y++) {
        const srcY = y * yRatio;
        const y0 = Math.floor(srcY);
        const y1 = Math.min(y0 + 1, srcH - 1);
        const fy = srcY - y0;
        for (let x = 0; x < dstW; x++) {
            const srcX = x * xRatio;
            const x0 = Math.floor(srcX);
            const x1 = Math.min(x0 + 1, srcW - 1);
            const fx = srcX - x0;
            const i00 = (y0 * srcW + x0) * 4;
            const i10 = (y0 * srcW + x1) * 4;
            const i01 = (y1 * srcW + x0) * 4;
            const i11 = (y1 * srcW + x1) * 4;
            const di = (y * dstW + x) * 4;
            for (let c = 0; c < 4; c++) {
                dst[di + c] = Math.round(
                    src[i00 + c]! * (1 - fx) * (1 - fy) +
                        src[i10 + c]! * fx * (1 - fy) +
                        src[i01 + c]! * (1 - fx) * fy +
                        src[i11 + c]! * fx * fy,
                );
            }
        }
    }
    return dst;
};

const generateThumbnail = (fileBytes: Uint8Array): Uint8Array => {
    const decoded = jpeg.decode(fileBytes, {
        useTArray: true,
        formatAsRGBA: true,
    });
    const { width: srcW, height: srcH } = decoded;
    const scale = Math.min(
        THUMB_MAX_DIMENSION / srcW,
        THUMB_MAX_DIMENSION / srcH,
        1,
    );
    const dstW = Math.round(srcW * scale);
    const dstH = Math.round(srcH * scale);

    let pixels: Uint8Array;
    if (scale < 1) {
        pixels = resizeRGBA(decoded.data, srcW, srcH, dstW, dstH);
    } else {
        pixels = decoded.data;
    }

    const encoded = jpeg.encode(
        { data: pixels, width: dstW, height: dstH },
        THUMB_JPEG_QUALITY,
    );
    return new Uint8Array(encoded.data);
};

// A baseline/JFIF JPEG starts with the SOI marker 0xFFD8. `jpeg-js` decodes
// only JPEG, so this signature check is what separates a file the helper can
// regenerate from one it must skip: a PNG, HEIC, or the odd non-image byte
// stream all fail this and are reported as skipped rather than crashing the
// decoder into an opaque failure (issue #17).
const isJpeg = (bytes: Uint8Array): boolean =>
    bytes.length >= 2 && bytes[0] === 0xff && bytes[1] === 0xd8;

// The reason a file cannot have a JPEG thumbnail regenerated for it from its
// metadata alone, before any bytes are fetched, or undefined when it might. A
// non-image (video, live photo) is unsupported outright; a still image still
// has to be checked against its actual bytes once downloaded.
const unsupportedByType = (file: EnteFile): string | undefined => {
    if (file.metadata.fileType !== "image") {
        return `unsupported file type: ${file.metadata.fileType} (only JPEG images can be regenerated)`;
    }
    return undefined;
};

// Regenerate and upload a thumbnail for each requested file. Originals are read
// through the library's content cache (`photo.original()`); the generated
// thumbnail is JPEG-encoded, encrypted under the file's own key, and registered
// with the server — the encrypt-and-upload path is unchanged. Each file is
// resolved to one outcome (fixed / skipped / failed) and a failure on one file
// never stops the others.
export const fixMissingThumbnails = async (
    lib: Library,
    client: Client,
    fileIDs: number[],
    onProgress?: ProgressCallback,
): Promise<ThumbnailFixResult[]> => {
    const log = onProgress ?? (() => {});
    const results: ThumbnailFixResult[] = [];
    const api = client.getApiClient();

    // Resolve each requested fileID to its file record and owning album by
    // enumerating the library, each file taken from the first album that holds
    // it. The raw `EnteFile` carries the per-file key the thumbnail is
    // encrypted under, which the projected records deliberately do not.
    const wanted = new Set(fileIDs);
    const fileMap = new Map<
        number,
        { file: EnteFile; collectionName: string }
    >();
    for (const album of lib.albums.list()) {
        for (const photo of album.photos.list()) {
            if (!wanted.has(photo.fileID) || fileMap.has(photo.fileID))
                continue;
            const file = lib.getFile(album.collectionID, photo.fileID);
            if (file) {
                fileMap.set(photo.fileID, {
                    file,
                    collectionName: album.name,
                });
            }
        }
    }

    for (const fileID of fileIDs) {
        const entry = fileMap.get(fileID);
        if (!entry) {
            results.push({
                fileID,
                title: "unknown",
                collection: "unknown",
                status: "failed",
                reason: "file not found in any collection",
            });
            continue;
        }

        const { file, collectionName } = entry;
        const title = file.metadata.title;

        const typeReason = unsupportedByType(file);
        if (typeReason) {
            log(`[${collectionName}] Skipping ${title}: ${typeReason}`);
            results.push({
                fileID,
                title,
                collection: collectionName,
                status: "skipped",
                reason: typeReason,
            });
            continue;
        }

        try {
            const photo = lib.photos.byID({ fileID });
            if (!photo) {
                throw new Error("file not present in the library cache");
            }

            log(
                `[${collectionName}] Downloading ${title} for thumbnail generation...`,
            );
            const { path } = await photo.original();
            const fileBytes = new Uint8Array(readFileSync(path));

            if (!isJpeg(fileBytes)) {
                const reason =
                    "unsupported image format (only baseline JPEG can be regenerated)";
                log(`[${collectionName}] Skipping ${title}: ${reason}`);
                results.push({
                    fileID,
                    title,
                    collection: collectionName,
                    status: "skipped",
                    reason,
                });
                continue;
            }

            log(`[${collectionName}] Generating thumbnail for ${title}...`);
            const thumbJpeg = generateThumbnail(fileBytes);

            log(
                `[${collectionName}] Encrypting and uploading thumbnail (${thumbJpeg.length} bytes)...`,
            );
            const { header, ciphertext } = encryptBlob(thumbJpeg, file.key);
            const md5 = createHash("md5").update(ciphertext).digest("base64");
            const { objectKey, url } = await api.getUploadURL(
                ciphertext.length,
                md5,
            );
            await api.putFile(url, ciphertext);
            await api.updateThumbnail(file.id, objectKey, toBase64(header));

            log(`[${collectionName}] Thumbnail uploaded for ${title}`);
            results.push({
                fileID,
                title,
                collection: collectionName,
                status: "fixed",
            });
        } catch (err) {
            log(
                `[${collectionName}] FAILED ${title}: ${err instanceof Error ? err.message : err}`,
            );
            results.push({
                fileID,
                title,
                collection: collectionName,
                status: "failed",
                reason: err instanceof Error ? err.message : String(err),
            });
        }
    }

    return results;
};