// llm.js-style wrapper around latest llama.cpp compiled to WASM.
//
// The public API is the one from rahuldshetty/llm.js:
//   new LLM(type, url, init_callback, write_result_callback, on_complete_callback)
//   app.load_worker(); app.run({ prompt, ... });
//
// Additions (all optional):
//   on_error_callback(err), on_progress_callback({ phase, loaded, total, ... }),
//   on_log_callback(line), on_event(message)
//   app.add_url_model({ url, name, expected_sha256 })
//   app.add_file_model({ file, name, keep_copy })
//   app.pause() / app.resume() / app.delete_model(id) / app.list_models()
//
// Deviations from the original: models live in a local library (uuid record per
// model, see storage.js), so `url` became a target object { id, url, name }. A
// plain url string is still accepted and treated as `{ url }`.
//
// The two ingest paths (url, local file) produce the same record and are loaded
// by the same code path; only where the bytes come from differs.

import { action } from "./actions.js";

export const MODELS = {
    GGUF_CPU: "GGUF_CPU",
};

class LLM {
    constructor(
        type,
        target,
        init_callback,
        write_result_callback,
        on_complete_callback,
        on_error_callback,
        on_progress_callback,
        on_log_callback,
    ) {
        this.type = type;
        this.model = typeof target === "string" ? { url: target } : (target || {});
        this.init_callback = init_callback;
        this.write_result_callback = write_result_callback;
        this.on_complete_callback = on_complete_callback;
        this.on_error_callback = on_error_callback;
        this.on_progress_callback = on_progress_callback;
        this.on_log_callback = on_log_callback;
        this.on_event = null; // any message, e.g. PAUSED / DELETED / MODELS

        this.worker = null;
        this.model_loaded = false;
        this.context_size = 2048;
        this._pending = null;
    }

    get url() {
        return this.model.url;
    }

    // autoload: post LOAD for this.model right away (the original behaviour);
    // pass false to only start the worker and drive loads yourself
    load_worker(options = {}) {
        switch (this.type) {
            case MODELS.GGUF_CPU:
                this.worker = new Worker(new URL("./worker.js", import.meta.url), { type: "module" });
                break;
            default:
                throw new Error(`unknown model type: ${this.type}`);
        }

        this.context_size = options.context_size || 2048;

        this.worker.onmessage = (event) => {
            const data = event.data;
            if (this.on_event) this.on_event(data);

            switch (data.event) {
                case action.INITIALIZED:
                    this.model_loaded = true;
                    if (this.init_callback) this.init_callback(data);
                    this._flushPending();
                    break;

                case action.WRITE_RESULT:
                    if (this.write_result_callback) this.write_result_callback(data.line);
                    break;

                case action.RUN_COMPLETED:
                    if (this.on_complete_callback) this.on_complete_callback(data);
                    break;

                case action.PROGRESS:
                    if (this.on_progress_callback) this.on_progress_callback(data);
                    break;

                case action.LOG:
                    if (this.on_log_callback) this.on_log_callback(data.line);
                    break;

                case action.ERROR:
                    if (this.on_error_callback) this.on_error_callback(data.error);
                    else console.error("llm.js:", data.error);
                    break;
            }
        };

        const autoload = options.autoload !== false;
        if (autoload && (this.model.id || this.model.url)) {
            this.worker.postMessage({
                event: action.LOAD,
                id: this.model.id || null,
                url: this.model.url || null,
                context_size: this.context_size,
            });
        }
    }

    post(message) {
        if (this.worker) this.worker.postMessage(message);
    }

    // download a GGUF and add it to the library (then load it)
    add_url_model({ url, name = "", expected_sha256 = "", context_size = 0 } = {}) {
        this.post({
            event: action.ADD_URL,
            url,
            name,
            expectedSha256: expected_sha256,
            context_size: context_size || this.context_size,
        });
    }

    // register a File the user already has; keep_copy decides whether the bytes
    // are duplicated into IndexedDB (default: no, read them in place)
    add_file_model({ file, name = "", keep_copy = false, expected_sha256 = "", context_size = 0 } = {}) {
        this.post({
            event: action.ADD_FILE,
            file,
            name,
            keepCopy: keep_copy,
            expectedSha256: expected_sha256,
            context_size: context_size || this.context_size,
        });
    }

    // load a model from the library by id
    load(id, file = null) {
        this.post({
            event: action.LOAD,
            id,
            file,
            context_size: this.context_size,
        });
    }

    // free the engine and load the same model again
    reload(id) {
        this.post({ event: action.RELOAD_MODEL, id: id || this.model.id, context_size: this.context_size });
    }

    // render a prompt through a chat template without generating; the result
    // arrives as a TEMPLATE_PREVIEW event
    preview_template({ prompt, template_source = "", messages = null,
                       tools = "[]", template_kwargs = "{}" } = {}) {
        this.post({
            event: action.PREVIEW_TEMPLATE,
            prompt,
            template_source,
            messages,
            tools,
            template_kwargs,
        });
    }

    // stop a running generation; the text streamed so far is kept and the run
    // completes with interrupted: true
    stop() {
        this.post({ event: action.STOP });
    }

    // resume an interrupted download / reload the same model
    resume() {
        this.post({ event: action.RESUME });
    }

    pause() {
        this.post({ event: action.PAUSE });
    }

    // drop a model (or the current one) from the library
    delete_model(id) {
        const target = id || this.model.id;
        if (target) this.post({ event: action.DELETE_MODEL, id: target });
    }

    hash_model(id) {
        this.post({ event: action.HASH_MODEL, id: id || this.model.id });
    }

    list_models() {
        this.post({ event: action.LIST_MODELS });
    }

    // mode "chat" runs a full message list through the chat template;
    // mode "completion" sends the prompt text as-is
    run({
        mode = "completion",
        prompt = "",
        messages = null,
        stop = [],
        chat = false,
        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 = "",
        json_schema = "",
        regex = "",
    } = {}) {
        const msg = {
            event: action.RUN_MAIN,
            mode,
            prompt,
            messages,
            stop,
            chat,
            seed,
            max_token_len,
            top_k,
            top_p,
            temp,
            repeat_last_n,
            repeat_penalty,
            context_size,
            chat_template,
            chat_template_source,
            tools,
            template_kwargs,
            grammar,
            json_schema,
            regex,
        };

        // allowed before the model finished loading, like the original docs suggest
        if (!this.model_loaded) {
            this._pending = msg;
            return;
        }
        this.post(msg);
    }

    _flushPending() {
        if (!this._pending) return;
        const msg = this._pending;
        this._pending = null;
        this.post(msg);
    }

    terminate() {
        if (this.worker) this.worker.terminate();
        this.worker = null;
        this.model_loaded = false;
    }
}

export { LLM };