// 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 callers that want many at once split them into batches of
// `MLDATA_BATCH_SIZE` and call `fetchMLDataBatch` once per batch.

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. Throws why
// when the entry does not decrypt, gunzip or parse, or parses to something
// other than a JSON object, such as `null`, a number or an array.
const decodeEntry = (entry: RawRemoteFileData, key: Uint8Array): MLData => {
    const decrypted = decryptBlob(
        fromBase64(entry.encryptedData),
        fromBase64(entry.decryptionHeader),
        key,
    );
    const json = gunzipSync(Buffer.from(decrypted)).toString("utf-8");
    const payload: unknown = JSON.parse(json);
    if (
        typeof payload !== "object" ||
        payload === null ||
        Array.isArray(payload)
    ) {
        throw new Error("not a JSON object");
    }
    return payload as MLData;
};

// 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. Returns the payload of each file
// whose entry decodes and, in `errors`, why not for each file whose entry does
// not, so one corrupt file never fails a whole batch. An entry for a file
// whose key is not known, which quak did not ask for, is skipped.
export const fetchMLDataBatch = async (
    api: ApiClient,
    fileIDs: number[],
    fileKeys: Map<number, Uint8Array>,
): Promise<{ payloads: Map<number, MLData>; errors: Map<number, unknown> }> => {
    const { data } = await api.postJSON<{ data: RawRemoteFileData[] }>(
        "/files/data/fetch",
        { type: "mldata", fileIDs },
    );
    const payloads = new Map<number, MLData>();
    const errors = new Map<number, unknown>();
    for (const entry of data ?? []) {
        const key = fileKeys.get(entry.fileID);
        if (!key) continue;
        try {
            payloads.set(entry.fileID, decodeEntry(entry, key));
        } catch (err) {
            errors.set(entry.fileID, err);
        }
    }
    return { payloads, errors };
};