// Fetch and decrypt Ente's per-file machine-learning data ("magic" search
// data: face detections + CLIP embeddings).
//
// The data lives behind `/files/data/fetch` with `type: "mldata"`. Each entry
// comes back encrypted under the file's own key and gzipped; decrypting and
// gunzipping yields the JSON payload
// `{ face: { faces: [...] }, clip: { embedding } }`. Ente caps a request at 200
// ids, so `fetchMLData` batches for callers that want many at once while
// `fetchMLDataBatch` is the single-request unit the library submits to its
// request pool.

import { gunzipSync } from "node:zlib";

import type { ApiClient } from "./api/client.js";
import { decryptBlob, fromBase64 } from "./crypto/index.js";

// The most ids one `/files/data/fetch` request may carry.
export const MLDATA_BATCH_SIZE = 200;

// The decrypted, gunzipped per-file payload. Its concrete shape is Ente's; the
// store keeps the whole object verbatim and each consumer reads the fields it
// needs, so it stays an open record rather than a fixed interface.
export type MLData = Record<string, unknown>;

interface RawRemoteFileData {
    fileID: number;
    encryptedData: string;
    decryptionHeader: string;
    updatedAt?: number;
}

// Decrypt one entry with its file key and gunzip the JSON payload. Returns
// undefined when the key is unknown or the entry does not decrypt/parse, so one
// corrupt file never fails a whole batch.
const decodeEntry = (
    entry: RawRemoteFileData,
    key: Uint8Array | undefined,
): MLData | undefined => {
    if (!key) return undefined;
    try {
        const decrypted = decryptBlob(
            fromBase64(entry.encryptedData),
            fromBase64(entry.decryptionHeader),
            key,
        );
        const json = gunzipSync(Buffer.from(decrypted)).toString("utf-8");
        return JSON.parse(json) as MLData;
    } catch {
        return undefined;
    }
};

// Fetch ML data for up to `MLDATA_BATCH_SIZE` ids in a single request. This is
// the unit the request pools schedule; callers with more ids split them into
// batches and submit each batch to the pool.
export const fetchMLDataBatch = async (
    api: ApiClient,
    fileIDs: number[],
    fileKeys: Map<number, Uint8Array>,
): Promise<Map<number, MLData>> => {
    const { data } = await api.postJSON<{ data: RawRemoteFileData[] }>(
        "/files/data/fetch",
        { type: "mldata", fileIDs },
    );
    const result = new Map<number, MLData>();
    for (const entry of data ?? []) {
        const payload = decodeEntry(entry, fileKeys.get(entry.fileID));
        if (payload) result.set(entry.fileID, payload);
    }
    return result;
};

// Fetch ML data for arbitrarily many ids, batching at `MLDATA_BATCH_SIZE`. Used
// by the one-shot metadata backup; the library fetches through its request pool
// with `fetchMLDataBatch` instead.
export const fetchMLData = async (
    api: ApiClient,
    fileIDs: number[],
    fileKeys: Map<number, Uint8Array>,
): Promise<Map<number, MLData>> => {
    const result = new Map<number, MLData>();
    for (let i = 0; i < fileIDs.length; i += MLDATA_BATCH_SIZE) {
        const batch = fileIDs.slice(i, i + MLDATA_BATCH_SIZE);
        for (const [id, payload] of await fetchMLDataBatch(
            api,
            batch,
            fileKeys,
        )) {
            result.set(id, payload);
        }
    }
    return result;
};