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";

// The server refuses a thumbnail larger than the one it already records for the
// file (`thumbnail.size`, the encrypted size), so these encodings are tried
// from largest to smallest and the first that fits is uploaded.
const THUMB_ENCODINGS = [
    { maxDimension: 720, quality: 50 },
    { maxDimension: 720, quality: 30 },
    { maxDimension: 480, quality: 30 },
    { maxDimension: 320, quality: 20 },
    { maxDimension: 160, quality: 20 },
];

// The server accepts a new thumbnail only from the file's owner, so files other
// people own in albums shared with this account are never checked or repaired.
const NOT_OWNED_REASON =
    "owned by another account (only the owner can replace its thumbnail)";

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 server would refuse any thumbnail for the
// file or this helper cannot regenerate it — a file another account owns, a
// recorded thumbnail size nothing fits within, 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.
// Files another account owns are logged as skipped and not checked.
export const listMissingThumbnails = async (
    lib: Library,
    client: Client,
    onProgress?: ProgressCallback,
): Promise<MissingThumbnailInfo[]> => {
    const log = onProgress ?? (() => {});
    const api = client.getApiClient();
    const { userID } = client.whoami();
    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);
            const file = lib.getFile(album.collectionID, photo.fileID);
            if (file && file.ownerID !== userID) {
                log(
                    `[${album.name}] Skipping ${photo.title}: ${NOT_OWNED_REASON}`,
                );
                continue;
            }
            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 = (
    decoded: { data: Uint8Array; width: number; height: number },
    maxDimension: number,
    quality: number,
): Uint8Array => {
    const { width: srcW, height: srcH } = decoded;
    const scale = Math.min(maxDimension / srcW, maxDimension / 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 },
        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, known from
// its record alone before any bytes are fetched, or undefined when it might. A
// still image still has to be checked against its actual bytes once
// downloaded.
const reasonToSkip = (file: EnteFile, userID: number): string | undefined => {
    if (file.ownerID !== userID) {
        return NOT_OWNED_REASON;
    }
    if (file.metadata.fileType !== "image") {
        return `unsupported file type: ${file.metadata.fileType} (only JPEG images can be regenerated)`;
    }
    if (!file.thumbnail.size) {
        return `recorded thumbnail size is ${file.thumbnail.size ?? "unknown"} (the server refuses a thumbnail larger than the one it records)`;
    }
    return undefined;
};

// Encrypt the largest encoding of the decoded image whose ciphertext is no
// larger than `maxSize`, or return undefined when even the smallest is larger.
const encryptThumbnailWithin = (
    decoded: { data: Uint8Array; width: number; height: number },
    key: Uint8Array,
    maxSize: number,
): { header: Uint8Array; ciphertext: Uint8Array } | undefined => {
    for (const { maxDimension, quality } of THUMB_ENCODINGS) {
        const thumbJpeg = generateThumbnail(decoded, maxDimension, quality);
        const encrypted = encryptBlob(thumbJpeg, key);
        if (encrypted.ciphertext.length <= maxSize) return encrypted;
    }
    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();
    const { userID } = client.whoami();

    // 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 skipReason = reasonToSkip(file, userID);
        if (skipReason) {
            log(`[${collectionName}] Skipping ${title}: ${skipReason}`);
            results.push({
                fileID,
                title,
                collection: collectionName,
                status: "skipped",
                reason: skipReason,
            });
            continue;
        }
        const maxSize = file.thumbnail.size!;

        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 decoded = jpeg.decode(fileBytes, {
                useTArray: true,
                formatAsRGBA: true,
            });
            const fitting = encryptThumbnailWithin(decoded, file.key, maxSize);
            if (!fitting) {
                const reason = `no thumbnail encoding fits the recorded thumbnail size of ${maxSize} bytes`;
                log(`[${collectionName}] Skipping ${title}: ${reason}`);
                results.push({
                    fileID,
                    title,
                    collection: collectionName,
                    status: "skipped",
                    reason,
                });
                continue;
            }
            const { header, ciphertext } = fitting;

            log(
                `[${collectionName}] Uploading thumbnail (${ciphertext.length} bytes)...`,
            );
            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;
};