// Inference worker.
//
// Responsibilities:
//   1. add models - two ingest paths into one library:
//        url  -> resumable, pausable, chunked download
//        file -> the user's own File, read in 4 MiB slices, never copied
//   2. sha256 each model once, so a human can compare it with the publisher's
//   3. stream the bytes into the Emscripten FS (from IndexedDB chunks, or
//      straight off the local File)
//   4. load it with the C API and stream generated tokens back
//
// Both ingest paths end in the same kind of record, so loading, hashing,
// renaming and deleting are one code path (see storage.js for the schema).

import createLLMWasmModule from "./wasm/llm-wasm.js";
import { action } from "./actions.js";
import { downloadModel } from "./download.js";
import { sha256Chunks } from "./sha256.js";
import {
    CHUNK_SIZE, getModel, putModel, updateModel, listModels, newModelId, nameFromUrl,
    putChunk, iterateChunks, clearChunks, deleteModel, findModelByUrl,
} from "./storage.js";

let module = null;
let loaded = false;
let loadedId = null;
let currentNctx = 512;

// Files the page handed over in this session, by model id. A "file" model with
// no copy in IndexedDB needs its File again after a reload - see start().
const files = new Map();

let abort = null;            // AbortController for the active download
let busy = false;
let lastRequest = null;      // { id, context_size } so Resume can continue

const post = (msg) => self.postMessage(msg);
const progress = (p) => post({ event: action.PROGRESS, ...p });
const log = (line, level = 3) => post({ event: action.LOG, level, line });
const fail = (err) => post({ event: action.ERROR, error: String((err && err.message) || err) });

// --- wasm -> JS callbacks ---------------------------------------------------
self.llmjs_on_token = (text) => post({ event: action.WRITE_RESULT, line: text });
self.llmjs_on_log = (level, text) => post({ event: action.LOG, level, line: text });

const tick = () => new Promise((r) => setTimeout(r, 0));

// storage.js shaped for download.js
const storage = {
    CHUNK_SIZE, getMeta: getModel, putMeta: putModel, updateMeta: updateModel, putChunk, clearChunks,
};

// Render a message list with llama.cpp's chat stack (GGUF template, tools and
// chat_template_kwargs), inside the wasm module. Throws with the engine's
// message; the caller decides whether that is worth showing.
function renderTemplateEx(messagesJson, tools, source, kwargs) {
    const sig = ["string", "string", "string", "string", "number", "number", "number"];
    const need = module.ccall("llmjs_chat_apply_ex", "number", sig,
        [messagesJson, tools, source, kwargs, 1, 0, 0]);
    if (need < 0) {
        throw new Error(module.UTF8ToString(module._llmjs_last_error()));
    }
    const size = need + 1;
    const ptr = module._malloc(size);
    try {
        module.ccall("llmjs_chat_apply_ex", "number", sig,
            [messagesJson, tools, source, kwargs, 1, ptr, size]);
        return module.UTF8ToString(ptr, need);
    } finally {
        module._free(ptr);
    }
}

async function ensureModule() {
    if (module) return module;
    // No locateFile override: the glue resolves llm-wasm.wasm relative to its
    // own import.meta.url (web/wasm/). Passing one here resolves against *this*
    // module's URL and 404s the .wasm into an HTML page.
    module = await createLLMWasmModule();
    module._llmjs_init();
    log(`llama.cpp ${module.UTF8ToString(module._llmjs_version())}`);
    return module;
}

// --- byte sources -----------------------------------------------------------

// 4 MiB slices of a File the user picked. The bytes stay on disk; each slice is
// read, handed to the FS (or the hash), and dropped again.
async function* fileSlices(file) {
    for (let offset = 0; offset < file.size; offset += CHUNK_SIZE) {
        const end = Math.min(offset + CHUNK_SIZE, file.size);
        yield new Uint8Array(await file.slice(offset, end).arrayBuffer());
    }
}

// Single iterator over a model's bytes, wherever they live.
async function* byteStream(model) {
    if (model.source === "url" || model.copied) {
        for await (const row of iterateChunks(model.id)) yield row.data;
        return;
    }
    const file = files.get(model.id);
    if (!file) throw new Error("the local file is not attached - choose it again");
    yield* fileSlices(file);
}

async function hashModel(model) {
    const total = model.total || 0;
    const sha256 = await sha256Chunks(byteStream(model), (done) =>
        progress({ phase: "hash", loaded: done, total }));
    return updateModel(model.id, { sha256 });
}

// --- ingest ----------------------------------------------------------------

async function addUrlModel({ url, name, expectedSha256 }) {
    let model = await findModelByUrl(url);
    if (model) {
        // same URL added again: treat it as "open", but let the name/hash be corrected
        const patch = {};
        if (name) patch.name = name;
        if (expectedSha256) patch.expectedSha256 = expectedSha256;
        if (Object.keys(patch).length) model = await updateModel(model.id, patch);
    } else {
        model = await putModel({
            id: newModelId(),
            name: (name || "").trim() || nameFromUrl(url),
            source: "url",
            url,
            fileName: null,
            total: 0,
            received: 0,
            complete: false,
            copied: true,               // url models always live in the browser
            sha256: null,
            expectedSha256: (expectedSha256 || "").trim() || null,
            created: Date.now(),
        });
    }
    post({ event: action.MODEL_ADDED, model });
    return model.id;
}

async function addFileModel({ file, name, keepCopy, expectedSha256 }) {
    if (!file) throw new Error("no file given");

    const head = new Uint8Array(await file.slice(0, 4).arrayBuffer());
    const magic = String.fromCharCode(...head);
    if (magic !== "GGUF") {
        throw new Error(`${file.name} is not a GGUF file (starts with "${magic}")`);
    }

    const model = await putModel({
        id: newModelId(),
        name: (name || "").trim() || file.name.replace(/\.gguf$/i, ""),
        source: "file",
        url: null,
        fileName: file.name,
        total: file.size,
        received: keepCopy ? 0 : file.size,
        complete: !keepCopy,
        copied: false,                  // filled in by copyFileIntoStorage below
        sha256: null,
        expectedSha256: (expectedSha256 || "").trim() || null,
        created: Date.now(),
    });

    files.set(model.id, file);
    post({ event: action.MODEL_ADDED, model });

    if (keepCopy) return (await copyFileIntoStorage(model, file)).id;
    return model.id;
}

// opt-in: put the bytes in IndexedDB so the model no longer needs the file
async function copyFileIntoStorage(model, file) {
    await clearChunks(model.id);
    let index = 0;
    let received = 0;
    for await (const bytes of fileSlices(file)) {
        const copy = bytes.buffer.slice(bytes.byteOffset, bytes.byteOffset + bytes.byteLength);
        await putChunk(model.id, index++, copy);
        received += bytes.length;
        progress({ phase: "copy", loaded: received, total: file.size });
    }
    return updateModel(model.id, { copied: true, received, total: file.size, complete: true });
}

// --- loading ---------------------------------------------------------------

async function loadIntoFS(model) {
    try { module.FS.mkdir("/models"); } catch (_) { /* already exists */ }
    try { module.FS.unlink("/models/model.gguf"); } catch (_) { /* not there */ }

    const stream = module.FS.open("/models/model.gguf", "w+");
    let written = 0;
    try {
        for await (const chunk of byteStream(model)) {
            const bytes = chunk instanceof Uint8Array ? chunk : new Uint8Array(chunk);
            module.FS.write(stream, bytes, 0, bytes.length, written);
            written += bytes.length;
            progress({ phase: "load", loaded: written, total: model.total || written });
            await tick();
        }
    } finally {
        module.FS.close(stream);
    }
    return written;
}

async function start({ id, file, context_size }) {
    if (busy) {
        fail(new Error("a load is already running"));
        return;
    }
    if (file) files.set(id, file);

    busy = true;
    abort = new AbortController();
    lastRequest = { id, context_size };

    try {
        await ensureModule();

        let model = await getModel(id);
        if (!model) throw new Error(`unknown model ${id}`);

        // the wasm instance can only hold one model at a time
        if (loaded && loadedId !== id) {
            module._llmjs_free();
            loaded = false;
            loadedId = null;
        }

        // 1. get the bytes into the browser (url models only; a local file is
        //    already where it needs to be)
        if (model.source === "url" && !model.complete) {
            await downloadModel(
                { id: model.id, url: model.url, total: model.total },
                { signal: abort.signal, onProgress: progress, storage },
            );
            model = await getModel(id);
        }
        if (model.source === "file" && !model.copied && !files.has(model.id)) {
            throw new Error("the local file is not attached - choose it again");
        }

        // 2. identify the bytes once
        if (!model.sha256) model = await hashModel(model);

        // 3. into the wasm FS, then load
        const bytes = await loadIntoFS(model);
        if (model.received && bytes !== model.received) {
            log(`stored model size ${bytes} != expected ${model.received}`);
        }

        currentNctx = context_size || currentNctx;
        const n = module.ccall("llmjs_load_model", "number", ["string", "number"],
                               ["/models/model.gguf", currentNctx]);
        if (n < 0) throw new Error(`llmjs_load_model failed (${n})`);

        loaded = true;
        loadedId = id;
        busy = false;
        post({
            event: action.INITIALIZED,
            model,
            bytes,
            n_ctx: n,
            chat_template_name: module.UTF8ToString(module._llmjs_chat_template_name()),
            // the GGUF's own Jinja template, rendered by the engine inside the wasm build
            chat_template: module.UTF8ToString(module._llmjs_chat_template()),
            has_chat_template: module._llmjs_has_chat_template() !== 0,
        });
        post({ event: action.MODEL_UPDATED, model });
    } catch (err) {
        busy = false;
        if (err && err.name === "AbortError") {
            post({ event: action.PAUSED, id });
        } else {
            fail(err);
        }
    }
}

// Set by a STOP message. The generation loop below checks it between tokens,
// which is only possible because generation is stepped from JS instead of one
// long blocking call: inside a single ccall this worker could not read messages.
let cancelRequested = false;
let generating = false;

// mode "chat": a message list rendered through the chat template.
// mode "completion": the prompt is the prompt, no template involved.
// Both honour the same sampling limits and stop sequences.
async function run(opts) {
    const {
        mode = "completion",
        prompt = "",
        messages = null,
        stop = [],
        seed = -1,
        max_token_len = 64,
        top_k = 40,
        top_p = 0.9,
        temp = 0.7,
        repeat_last_n = 32,
        repeat_penalty = 1.1,
        context_size = 0,
        chat_template = "",
        chat_template_source = "",
        tools = "[]",
        template_kwargs = "{}",
        grammar = "",
    } = opts;

    if (!loaded || !module) {
        fail(new Error("no model loaded"));
        return;
    }
    if (generating) {
        fail(new Error("a run is already in progress"));
        return;
    }

    const n = module.ccall("llmjs_create_context", "number", ["number"],
                           [context_size || currentNctx]);
    if (n < 0) {
        fail(new Error(`llmjs_create_context failed (${n})`));
        return;
    }

    // stop sequences are matched in C++ against the decoded text, so a stop
    // string is never streamed to the page
    module.ccall("llmjs_set_stop_sequences", null, ["string"],
                 [JSON.stringify(Array.isArray(stop) ? stop : [])]);

    // arm the engine: everything promised by the parameters goes in here
    let rc;
    if (mode === "chat") {
        const list = Array.isArray(messages) && messages.length
            ? messages
            : [{ role: "user", content: prompt }];
        rc = module.ccall(
            "llmjs_gen_begin_chat",
            "number",
            ["string", "string", "string", "string", "number", "number", "number",
             "number", "number", "number", "number", "string"],
            [JSON.stringify(list), tools || "[]", chat_template_source, template_kwargs || "{}",
             max_token_len, temp, top_k, top_p, repeat_penalty, repeat_last_n, seed, grammar],
        );
    } else {
        rc = module.ccall(
            "llmjs_gen_begin_raw",
            "number",
            ["string", "number", "number", "number", "number", "number",
             "number", "number", "string"],
            [prompt, max_token_len, temp, top_k, top_p, repeat_penalty,
             repeat_last_n, seed, grammar],
        );
    }
    if (rc < 0) {
        fail(new Error(`could not start generation (${rc})`));
        return;
    }

    generating = true;
    cancelRequested = false;
    const t0 = performance.now();
    let tokens = 0;
    try {
        // one token per turn of the event loop: a STOP message is handled
        // during the await, and the loop then exits with the text so far
        while (!cancelRequested) {
            if (module.ccall("llmjs_gen_step", "number", [], []) <= 0) break;
            await tick();
        }
        tokens = module.ccall("llmjs_gen_end", "number", [], []);
    } finally {
        generating = false;
    }

    const ms = performance.now() - t0;
    const interrupted = cancelRequested;
    cancelRequested = false;

    post({
        event: action.RUN_COMPLETED,
        mode,
        tokens,
        ms,
        interrupted,
        tokens_per_second: tokens > 0 ? (tokens * 1000) / ms : 0,
        // the final text, with any stop sequence removed: whatever the page
        // streamed may still contain the beginning of a stop string
        output: module.UTF8ToString(module._llmjs_last_output()),
        // jinja / jinja:override / builtin:<name> / raw
        template_strategy: module.UTF8ToString(module._llmjs_chat_template_strategy()),
        // { content, reasoning_content, tool_calls } when the template's parser
        // recognised the reply, null when it did not
        parsed: (() => {
            const text = module.UTF8ToString(module._llmjs_parsed_output());
            if (!text) return null;
            try { return JSON.parse(text); } catch (_) { return null; }
        })(),
    });
}

self.addEventListener("message", (e) => {
    const msg = e.data;
    switch (msg.event) {
        case action.ADD_URL:
            (async () => {
                const id = await addUrlModel(msg);
                await start({ id, context_size: msg.context_size });
            })().catch(fail);
            break;

        case action.ADD_FILE:
            (async () => {
                const id = await addFileModel(msg);
                await start({ id, context_size: msg.context_size });
            })().catch(fail);
            break;

        case action.LOAD:
            start(msg);
            break;

        case action.PAUSE:
            if (abort) abort.abort();
            break;

        case action.RESUME:
            if (lastRequest) start(lastRequest);
            else fail(new Error("nothing to resume"));
            break;

        case action.PREVIEW_TEMPLATE:
            (async () => {
                await ensureModule();
                const messages = JSON.stringify(msg.messages && msg.messages.length
                    ? msg.messages
                    : [{ role: "user", content: msg.prompt || "" }]);
                const text = renderTemplateEx(messages, msg.tools || "[]",
                    msg.template_source || "", msg.template_kwargs || "{}");
                post({ event: action.TEMPLATE_PREVIEW, text });
            })().catch(fail);
            break;

        case action.RELOAD_MODEL:
            (async () => {
                // a real reload: drop the model and the backend, then load again
                if (module) {
                    module._llmjs_free();
                    module._llmjs_init();
                }
                loaded = false;
                loadedId = null;
                await start({ id: msg.id, file: files.get(msg.id) || null, context_size: msg.context_size });
            })().catch(fail);
            break;

        case action.HASH_MODEL:
            (async () => {
                const model = await getModel(msg.id);
                if (!model) throw new Error(`unknown model ${msg.id}`);
                const hashed = await hashModel(model);
                post({ event: action.MODEL_UPDATED, model: hashed });
            })().catch(fail);
            break;

        case action.DELETE_MODEL:
            (async () => {
                if (loaded && loadedId === msg.id) {
                    module._llmjs_free();
                    loaded = false;
                    loadedId = null;
                }
                files.delete(msg.id);
                const chunks = await deleteModel(msg.id);
                post({ event: action.DELETED, id: msg.id, chunks });
            })().catch(fail);
            break;

        case action.LIST_MODELS:
            listModels()
                .then((models) => post({ event: action.MODELS, models }))
                .catch(fail);
            break;

        case action.STOP:
            if (generating) cancelRequested = true;
            break;

        case action.RUN_MAIN:
            run(msg).catch(fail);
            break;
    }
}, false);