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 { encryptBlob, toBase64 } from "./crypto/index.js";
import { downloadFile } from "./download/index.js";
import type { EnteFile } from "./model/types.js";
import { mkdtempSync, rmSync } from "node:fs";
import { join } from "node:path";
import { tmpdir } from "node:os";

const THUMB_MAX_DIMENSION = 720;
const THUMB_JPEG_QUALITY = 50;

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

export interface ThumbnailFixResult {
    fileID: number;
    title: string;
    collection: string;
    success: boolean;
    error?: string;
}

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

export const listMissingThumbnails = async (
    client: Client,
    onProgress?: ProgressCallback,
): Promise<MissingThumbnailInfo[]> => {
    const log = onProgress ?? (() => {});
    const missing: MissingThumbnailInfo[] = [];
    const seen = new Set<number>();

    const collections = await client.listCollections();
    for (const col of collections) {
        log(`[${col.name}] Checking thumbnails...`);
        const files = await client.listFiles(col.id, col.key);
        for (const file of files) {
            if (seen.has(file.id)) continue;
            seen.add(file.id);
            try {
                const api = client.getApiClient();
                const stream = await api.getThumbnailStream(file.id);
                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: file.id,
                        title: file.metadata.title,
                        collection: col.name,
                        reason: "empty thumbnail (0 bytes)",
                    });
                }
            } catch {
                missing.push({
                    fileID: file.id,
                    title: file.metadata.title,
                    collection: col.name,
                    reason: "thumbnail fetch failed",
                });
            }
        }
    }
    return missing;
};

// Bilinear resize of 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);
};

export const fixMissingThumbnails = async (
    client: Client,
    fileIDs: number[],
    onProgress?: ProgressCallback,
): Promise<ThumbnailFixResult[]> => {
    const log = onProgress ?? (() => {});
    const results: ThumbnailFixResult[] = [];
    const api = client.getApiClient();

    const collections = await client.listCollections();
    const fileMap = new Map<
        number,
        { file: EnteFile; collectionName: string }
    >();

    for (const col of collections) {
        const files = await client.listFiles(col.id, col.key);
        for (const file of files) {
            if (fileIDs.includes(file.id) && !fileMap.has(file.id)) {
                fileMap.set(file.id, {
                    file,
                    collectionName: col.name,
                });
            }
        }
    }

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

        const { file, collectionName } = entry;
        const tmpDir = mkdtempSync(join(tmpdir(), "quak-thumb-"));

        try {
            log(
                `[${collectionName}] Downloading ${file.metadata.title} for thumbnail generation...`,
            );
            const origPath = join(tmpDir, "original");
            await downloadFile(api, file, origPath);

            log(
                `[${collectionName}] Generating thumbnail for ${file.metadata.title}...`,
            );
            const fileBytes = readFileSync(origPath);
            const thumbJpeg = generateThumbnail(new Uint8Array(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 ${file.metadata.title}`,
            );
            results.push({
                fileID,
                title: file.metadata.title,
                collection: collectionName,
                success: true,
            });
        } catch (err) {
            log(
                `[${collectionName}] FAILED ${file.metadata.title}: ${err instanceof Error ? err.message : err}`,
            );
            results.push({
                fileID,
                title: file.metadata.title,
                collection: collectionName,
                success: false,
                error: err instanceof Error ? err.message : String(err),
            });
        } finally {
            rmSync(tmpDir, { recursive: true, force: true });
        }
    }

    return results;
};