// The on-disk cache of Ente's per-file machine-learning data and the CLIP
// index derived from it (issue #49).
//
// Under `<cacheDirectory>/mldata/` this keeps:
//
//   - `<fileID>.json` — one decrypted, gunzipped payload per file, written by
//     rename. Its presence means it is complete: a torn write never leaves a
//     half-file, so the set of these files is the source of truth for what is
//     cached. The full payload (face boxes, landmarks, embeddings) is read back
//     from here on demand and never held in RAM.
//
//   - `clip.f32` + `clip.json` — the derived index the content search runs on.
//     `clip.json` lists the indexed fileIDs in order plus the embedding length;
//     `clip.f32` is those CLIP embeddings packed as one `Float32Array`, so the
//     index loads in a single read with no per-vector parse. The index is
//     rebuilt from the payloads whenever it is missing or structurally
//     disagrees with the files present, and appended to as new payloads arrive.
//
//   - `fetched.json` — a small map of fileID to the `updationTime` it was
//     fetched at. This is best-effort bookkeeping for refetch decisions (a file
//     whose `updationTime` later advances is refetched); the payloads, not this
//     file, remain the record of what is cached, so losing it only forgoes
//     update-driven refetch until the next fetch rewrites it.
//
// In RAM this holds only the id list and the packed `Float32Array`.

import { mkdir, readFile, readdir } from "node:fs/promises";
import { join } from "node:path";

import { writeAtomic } from "../download/index.js";
import type { MLData } from "../mldata-fetch.js";

const CLIP_VECTORS = "clip.f32";
const CLIP_INDEX = "clip.json";
const FETCHED = "fetched.json";
// A payload file is named for its fileID alone; the derived files above are
// not, so this pattern picks out payloads and nothing else.
const PAYLOAD_RE = /^(\d+)\.json$/;
const BYTES_PER_FLOAT = 4;

// The on-disk form of `clip.json`.
interface ClipIndexFile {
    fileIDs: number[];
    embeddingLength: number;
}

// A file the model knows about, for deciding what to fetch.
export interface MLDataFile {
    id: number;
    updationTime: number;
}

// The RAM index the search reads: `fileIDs[i]` owns the `embeddingLength`
// floats of `embeddings` starting at `i * embeddingLength`.
export interface MLIndex {
    fileIDs: number[];
    embeddingLength: number;
    embeddings: Float32Array;
}

// Pull the CLIP embedding out of a payload, or undefined when it is absent or
// misshapen. Kept strict so a bad payload is skipped rather than corrupting the
// packed index.
const clipEmbedding = (payload: MLData): number[] | undefined => {
    const clip = payload.clip;
    if (typeof clip !== "object" || clip === null) return undefined;
    const embedding = (clip as { embedding?: unknown }).embedding;
    if (!Array.isArray(embedding)) return undefined;
    if (embedding.some((v) => typeof v !== "number" || !Number.isFinite(v)))
        return undefined;
    return embedding as number[];
};

export class MLDataStore {
    readonly dir: string;

    // fileIDs whose payload JSON is present on disk (present means complete).
    private readonly present = new Set<number>();
    // fileID -> updationTime it was fetched at.
    private readonly fetched = new Map<number, number>();

    // The packed index and where each id sits in it.
    private ids: number[] = [];
    private embeddingLength = 0;
    private embeddings = new Float32Array(0);
    private readonly pos = new Map<number, number>();

    private constructor(dir: string) {
        this.dir = dir;
    }

    // Open (creating the directory) and load the id list and packed index into
    // RAM, rebuilding the index from the payloads when it is missing or does
    // not match the files present.
    static async open(dir: string): Promise<MLDataStore> {
        const store = new MLDataStore(dir);
        await mkdir(dir, { recursive: true });
        await store.loadPresent();
        await store.loadFetched();
        if (!(await store.tryLoadIndex())) await store.rebuildIndex();
        return store;
    }

    // The fileIDs among `files` that must be fetched: every file with no
    // payload yet (first run, then new files), plus any whose `updationTime`
    // has advanced past the one its cached payload was fetched at. Returned
    // sorted and unique.
    neededFor(files: MLDataFile[]): number[] {
        const latest = new Map<number, number>();
        for (const f of files) {
            const seen = latest.get(f.id);
            if (seen === undefined || f.updationTime > seen)
                latest.set(f.id, f.updationTime);
        }
        const needed: number[] = [];
        for (const [id, updationTime] of latest) {
            if (!this.present.has(id)) {
                needed.push(id);
                continue;
            }
            const at = this.fetched.get(id);
            if (at !== undefined && updationTime > at) needed.push(id);
        }
        return needed.sort((a, b) => a - b);
    }

    // Store a batch of fetched payloads: write one file per id, fold their CLIP
    // embeddings into the packed index (in place for a refetch, appended for a
    // new file), and persist the derived files. Returns how many payloads were
    // stored and how many ids the index now holds.
    async storeFetched(
        payloads: Map<number, MLData>,
        updation: Map<number, number>,
    ): Promise<{ stored: number; indexed: number }> {
        if (payloads.size === 0) return { stored: 0, indexed: this.ids.length };

        for (const [id, payload] of payloads) {
            await this.writePayload(id, payload);
            this.present.add(id);
            const at = updation.get(id);
            if (at !== undefined) this.fetched.set(id, at);
        }

        const updates: { at: number; vector: number[] }[] = [];
        const appends: { id: number; vector: number[] }[] = [];
        for (const [id, payload] of payloads) {
            const vector = clipEmbedding(payload);
            if (!vector) continue;
            if (this.embeddingLength === 0 && this.ids.length === 0)
                this.embeddingLength = vector.length;
            // The index is fixed-width; a vector of another length (never seen
            // from Ente's CLIP model) is stored but left out of the index.
            if (vector.length !== this.embeddingLength) continue;
            const at = this.pos.get(id);
            if (at !== undefined) updates.push({ at, vector });
            else appends.push({ id, vector });
        }

        for (const { at, vector } of updates)
            this.embeddings.set(vector, at * this.embeddingLength);

        if (appends.length > 0) {
            const length = this.embeddingLength;
            const grown = new Float32Array(
                this.embeddings.length + appends.length * length,
            );
            grown.set(this.embeddings);
            let offset = this.embeddings.length;
            for (const { id, vector } of appends) {
                grown.set(vector, offset);
                this.pos.set(id, this.ids.length);
                this.ids.push(id);
                offset += length;
            }
            this.embeddings = grown;
        }

        await this.persistIndex();
        await this.persistFetched();
        return { stored: payloads.size, indexed: this.ids.length };
    }

    // The packed index the search runs on. The id list is copied so callers
    // cannot disturb the store's own order; the embeddings are the live buffer.
    getIndex(): MLIndex {
        return {
            fileIDs: [...this.ids],
            embeddingLength: this.embeddingLength,
            embeddings: this.embeddings,
        };
    }

    // The full payload for a file, read from disk, or undefined when it is not
    // cached or does not parse.
    async readPayload(fileID: number): Promise<MLData | undefined> {
        if (!this.present.has(fileID)) return undefined;
        let raw: string;
        try {
            raw = await readFile(this.payloadPath(fileID), "utf-8");
        } catch {
            return undefined;
        }
        try {
            return JSON.parse(raw) as MLData;
        } catch {
            return undefined;
        }
    }

    stats(): { stored: number; indexed: number } {
        return { stored: this.present.size, indexed: this.ids.length };
    }

    private payloadPath(id: number): string {
        return join(this.dir, `${id}.json`);
    }

    private async writePayload(id: number, payload: MLData): Promise<void> {
        await writeAtomic(
            this.payloadPath(id),
            new TextEncoder().encode(JSON.stringify(payload)),
        );
    }

    private async loadPresent(): Promise<void> {
        let names: string[];
        try {
            names = await readdir(this.dir);
        } catch {
            return;
        }
        for (const name of names) {
            const match = PAYLOAD_RE.exec(name);
            if (match) this.present.add(Number(match[1]));
        }
    }

    private async loadFetched(): Promise<void> {
        let raw: string;
        try {
            raw = await readFile(join(this.dir, FETCHED), "utf-8");
        } catch {
            return;
        }
        try {
            const parsed = JSON.parse(raw) as Record<string, unknown>;
            for (const [key, value] of Object.entries(parsed)) {
                const id = Number(key);
                if (
                    Number.isInteger(id) &&
                    typeof value === "number" &&
                    this.present.has(id)
                )
                    this.fetched.set(id, value);
            }
        } catch {
            // Corrupt bookkeeping degrades refetch decisions, never fails open.
        }
    }

    // Load the packed index if it is present and agrees with the payloads in
    // both directions: every id it names must still be present, its vector file
    // must be exactly the size the id count and embedding length imply, and no
    // embedding-bearing payload on disk may be missing from it. Returns whether
    // it loaded.
    private async tryLoadIndex(): Promise<boolean> {
        let metaRaw: string;
        try {
            metaRaw = await readFile(join(this.dir, CLIP_INDEX), "utf-8");
        } catch {
            return false;
        }
        let meta: ClipIndexFile;
        try {
            meta = JSON.parse(metaRaw) as ClipIndexFile;
        } catch {
            return false;
        }
        if (
            !Array.isArray(meta.fileIDs) ||
            typeof meta.embeddingLength !== "number"
        )
            return false;
        if (meta.fileIDs.some((id) => !this.present.has(id))) return false;

        // The reverse must hold too. A payload carrying an embedding but absent
        // from the index means the index is stale — realistically the process
        // died after storeFetched renamed the payloads into place but before it
        // rewrote clip.json/clip.f32. Loading such an index as "consistent"
        // would drop those embeddings for good (neededFor sees the payloads
        // present and never refetches), so treat it as a disagreement and
        // rebuild. Only present ids the index omits are read; a payload
        // legitimately without an embedding stays out and forces no rebuild.
        const indexed = new Set(meta.fileIDs);
        for (const id of this.present) {
            if (indexed.has(id)) continue;
            const payload = await this.readPayload(id);
            if (payload && clipEmbedding(payload)) return false;
        }

        let bytes: Buffer;
        try {
            bytes = await readFile(join(this.dir, CLIP_VECTORS));
        } catch {
            return false;
        }
        const expected =
            meta.fileIDs.length * meta.embeddingLength * BYTES_PER_FLOAT;
        if (bytes.byteLength !== expected) return false;

        // One read, no parse: copy into an aligned buffer and view it as
        // floats. The copy is needed because a Buffer from the pool can start
        // at an offset a Float32Array cannot be laid over.
        const aligned = new Uint8Array(bytes.byteLength);
        aligned.set(bytes);
        this.embeddings = new Float32Array(aligned.buffer);
        this.embeddingLength = meta.embeddingLength;
        this.ids = [...meta.fileIDs];
        this.pos.clear();
        this.ids.forEach((id, i) => this.pos.set(id, i));
        return true;
    }

    // Rebuild the packed index by reading every payload present, then persist
    // it. Payloads without a CLIP embedding (or of an unexpected length) are
    // simply not indexed.
    private async rebuildIndex(): Promise<void> {
        this.ids = [];
        this.pos.clear();
        this.embeddingLength = 0;
        const vectors: number[][] = [];
        for (const id of [...this.present].sort((a, b) => a - b)) {
            const payload = await this.readPayload(id);
            if (!payload) continue;
            const vector = clipEmbedding(payload);
            if (!vector) continue;
            if (this.embeddingLength === 0)
                this.embeddingLength = vector.length;
            if (vector.length !== this.embeddingLength) continue;
            this.pos.set(id, this.ids.length);
            this.ids.push(id);
            vectors.push(vector);
        }
        const length = this.embeddingLength;
        const packed = new Float32Array(this.ids.length * length);
        vectors.forEach((vector, i) => packed.set(vector, i * length));
        this.embeddings = packed;
        await this.persistIndex();
    }

    private async persistIndex(): Promise<void> {
        const meta: ClipIndexFile = {
            fileIDs: this.ids,
            embeddingLength: this.embeddingLength,
        };
        await writeAtomic(
            join(this.dir, CLIP_INDEX),
            new TextEncoder().encode(JSON.stringify(meta)),
        );
        await writeAtomic(
            join(this.dir, CLIP_VECTORS),
            new Uint8Array(
                this.embeddings.buffer,
                this.embeddings.byteOffset,
                this.embeddings.byteLength,
            ),
        );
    }

    private async persistFetched(): Promise<void> {
        const record: Record<string, number> = {};
        for (const [id, at] of this.fetched) record[id] = at;
        await writeAtomic(
            join(this.dir, FETCHED),
            new TextEncoder().encode(JSON.stringify(record)),
        );
    }
}