// Copyright (C) 2026 Danilo M. <danix@danix.xyz>
//
// This program is free software; you can redistribute it and/or modify
// it under the terms of the GNU General Public License version 2 as
// published by the Free Software Foundation.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
// GNU General Public License for more details.

import Quickshell
import Quickshell.Io
import QtQuick
import "../.."
import "Config.js" as C

// desktop-assistant's daemon settings: its config.toml, the tools it loads,
// and a restart, since it reads both only at startup. The module holds the
// data and the writes; AssistantPage is the view. See README.md.
Module {
    id: am
    name: "assistant"
    label: "Assistant"
    // nf-fa-comments, present in Inconsolata Nerd Font.
    icon: "\uf086"
    settings: true
    group: "Assistant"
    keywords: ["llm", "model", "llama", "server", "prompt", "temperature", "thinking",
               "rounds", "tools", "environment", "daemon", "restart", "cloud", "provider"]

    readonly property string home: Quickshell.env("HOME")
    readonly property string daemonDir: home + "/Programming/GIT/desktop-assistant"
    readonly property string cfgDir: (Quickshell.env("XDG_CONFIG_HOME") || home + "/.config")
                                     + "/desktop-assistant"
    readonly property string cfgPath: cfgDir + "/config.toml"
    readonly property string toolsDir: cfgDir + "/tools"
    // The system interpreter, as for calctl: python3 on PATH is a pyenv shim.
    readonly property string python: "/usr/bin/python3"

    readonly property var keys: ["llm", "system", "tool_rounds", "temperature", "thinking"]
    readonly property var labels: ({ llm: "Server", system: "System prompt", tool_rounds: "Tool rounds",
                                     temperature: "Temperature", thinking: "Thinking" })

    // config.toml as last read or written; "" when there is none. The page
    // edits draft, which save() writes; base is the file the draft began
    // from, so a save over a file changed meanwhile is refused.
    property string text: ""
    property string base: ""
    property string draft: ""
    readonly property var values: C.scalars(am.draft, am.keys)
    readonly property var envList: C.env(am.draft)
    // The [env] keys, assigned only when the set changes: a model bound to the
    // draft would rebuild every row on each keystroke and drop the focus of
    // the one being typed in. Rows read their values from envValue().
    property var envKeys: []
    onDraftChanged: {
        const k = C.env(am.draft).map(e => e.key);
        if (JSON.stringify(k) !== JSON.stringify(am.envKeys)) am.envKeys = k;
        const n = (C.clouds(am.draft) ?? []).length;
        if (n !== am.cloudCount) am.cloudCount = n;
    }
    function envValue(key) { return (am.envList.find(e => e.key === key) ?? { value: "" }).value; }
    // The [[cloud]] tables in the draft, null when one cannot be read (the
    // page then leaves them alone), and the saved ones, which the canvas and
    // Prefer offer. cloudCount follows the draft's count only, so rows are
    // not rebuilt on each keystroke (as envKeys).
    readonly property var clouds: C.clouds(am.draft)
    readonly property var savedClouds: C.clouds(am.text) ?? []
    property int cloudCount: 0
    readonly property string cloudError: am.clouds ? C.cloudError(am.clouds) : ""
    // name -> the provider's ids, or why they could not be fetched.
    property var fetched: ({})
    property string fetching: ""
    // Tool switches flipped but not saved: name -> wanted on.
    property var toolsWanted: ({})
    // Fields holding a value that cannot be written; Save waits for none.
    property var invalid: ({})
    readonly property bool valid: Object.keys(am.invalid).length === 0 && am.cloudError === ""
    dirty: am.draft !== am.base || Object.keys(am.toolsWanted).length > 0
    // The daemon's built-in values, from assistant.py --defaults.
    property var defaults: ({})
    // [{name, on, description}]
    property var tools: []
    // A write since the last restart, which the daemon has not read yet.
    property bool pending: false

    // The router's models, [{id, status, inputs}], status as it reports it:
    // "loaded", "loading", "unloaded" and so on; inputs its
    // input_modalities, null when not reported. Empty when it is down.
    property var models: []
    // The id being loaded, "" when none is.
    property string loading: ""
    // Whether the assistant's model takes audio, for the canvas's mic: the
    // one loading, else the daemon's pick, the first not unloaded. Unknown
    // counts as yes, and the daemon's own error says otherwise.
    readonly property bool hears: {
        const m = am.models.find(x => am.loading !== "" ? x.id === am.loading : x.status !== "unloaded");
        return !m?.inputs || m.inputs.includes("audio");
    }
    property string modelError: ""
    // The router answers /models and /models/load on the server's root, and
    // the OpenAI API under /v1, which is what the daemon's llm points at.
    readonly property string llm: am.values.llm ?? am.defaults.llm ?? "http://localhost:8181/v1"
    readonly property string routerRoot: am.llm.replace(/\/v1\/?$/, "")
    property bool restarting: false
    property string status: ""
    property string error: ""

    searchItems: am.keys.map(k => ({ text: am.labels[k], key: "key:" + k }))
        .concat(am.envList.map(e => ({ text: e.key + " = " + e.value, key: "env:" + e.key })))
        .concat(am.tools.map(t => ({ text: "tool: " + t.name, key: "tool:" + t.name })))
        .concat(am.models.map(m => ({ text: "model: " + m.id, key: "model:" + m.id })))

    function refresh() {
        fv.reload();
        am.listTools();
        am.listModels();
        if (Object.keys(am.defaults).length === 0)
            am.run([am.python, am.daemonDir + "/assistant.py", "--defaults"], null, (code, out) => {
                try { if (code === 0) am.defaults = JSON.parse(out); } catch (e) {}
            });
    }

    property FileView fv: FileView {
        path: am.cfgPath
        watchChanges: true
        printErrors: false
        onFileChanged: reload()
        // A clean draft follows the file; a dirty one is kept for save().
        onLoaded: {
            am.text = text();
            if (am.draft === am.base) { am.base = am.text; am.draft = am.text; }
        }
        onLoadFailed: if (am.draft === am.base) { am.text = ""; am.base = ""; am.draft = ""; }
    }

    // One Process per call, so calls never share one (a running Process
    // ignores a second start, AGENTS.md). input, when not null, is written
    // to stdin and stdin closed. cb(code, stdout, stderr) runs once the exit
    // and both streams are in, which arrive in no fixed order.
    function run(cmd, input, cb) {
        op.createObject(am, { command: cmd, input: input, cb: cb, stdinEnabled: input !== null }).running = true;
    }
    property Component op: Component {
        Process {
            id: proc
            property var input: null
            property var cb: null
            property int left: 3
            property int code: -1
            property string out: ""
            property string err: ""
            function step() {
                if (--left > 0) return;
                try { if (proc.cb) proc.cb(proc.code, proc.out, proc.err); }
                finally { proc.destroy(); }
            }
            onStarted: if (proc.input !== null) { proc.write(proc.input); proc.stdinEnabled = false; }
            onExited: code => { proc.code = code; proc.step(); }
            stdout: StdioCollector { onStreamFinished: { proc.out = text; proc.step(); } }
            stderr: StdioCollector { onStreamFinished: { proc.err = text; proc.step(); } }
        }
    }

    // Validates the whole file with tomllib before it replaces the old one,
    // atomically, so a bad edit can never leave a config the daemon refuses
    // to start with.
    readonly property string writer: "import os, sys, tomllib\n"
        + "t = sys.stdin.read()\n"
        + "try:\n    tomllib.loads(t)\n"
        + "except tomllib.TOMLDecodeError as e:\n    sys.exit(f'not valid TOML: {e}')\n"
        + "p = sys.argv[1]\nos.makedirs(os.path.dirname(p), exist_ok=True)\n"
        + "with open(p + '.tmp', 'w') as f:\n    f.write(t)\n"
        + "os.replace(p + '.tmp', p)\n"

    // The config first, then the tool renames, each only when there is
    // something to do; any failure stops there and says why.
    function save(done) {
        if (am.text !== am.base) {
            am.error = "config.toml changed on disk since these edits began. Discard to load it, then redo them.";
            return;
        }
        am.error = "";
        const t = am.draft;
        const changed = t !== am.text || Object.keys(am.toolsWanted).length > 0;
        const tools = () => am.applyTools(Object.keys(am.toolsWanted), () => {
            if (changed) am.pending = true;
            am.listTools();
            if (done) done();
        });
        if (t === am.text) { tools(); return; }
        am.run([am.python, "-c", am.writer, am.cfgPath], t, (code, out, err) => {
            if (code !== 0) { am.error = err.trim() || "config.toml could not be written"; return; }
            am.text = t;
            am.base = t;
            tools();
        });
    }
    function discard() {
        am.error = "";
        am.invalid = ({});
        am.toolsWanted = ({});
        am.base = am.text;
        am.draft = am.text;
    }
    function markInvalid(id, bad) {
        if (!!am.invalid[id] === bad) return;
        const next = Object.assign({}, am.invalid);
        if (bad) next[id] = true; else delete next[id];
        am.invalid = next;
    }
    // toml is the value written as TOML, or null to fall back to the default.
    function setKey(key, toml) { am.draft = C.setScalar(am.draft, key, toml); }
    function setEnv(key, value) { am.draft = C.setEnv(am.draft, key, value); }
    function quote(s) { return C.quote(s); }

    // patch is {name}, {url}, {key_cmd} or {models}; i === cloudCount adds one.
    function editCloud(i, patch) {
        const list = (am.clouds ?? []).map(c => Object.assign({}, c));
        if (i === list.length) list.push({ name: "", url: "", key_cmd: "", models: [] });
        Object.assign(list[i], patch);
        am.draft = C.setClouds(am.draft, list);
    }
    function removeCloud(i) { am.draft = C.setClouds(am.draft, am.clouds.filter((c, j) => j !== i)); }

    // The daemon reads the saved config, so Fetch saves the draft first.
    // The key command runs in that process; the key never reaches here.
    function fetchModels(name) {
        am.save(() => {
            am.fetching = name;
            am.run([am.python, am.daemonDir + "/assistant.py", "--cloud-models", name], null,
                   (code, out, err) => {
                am.fetching = "";
                let got;
                try { got = code === 0 ? JSON.parse(out) : (err.trim().split("\n").pop() || `exit ${code}`); }
                catch (e) { got = "the assistant sent an unreadable list"; }
                am.fetched = Object.assign({}, am.fetched, { [name]: got });
            });
        });
    }

    // Every tools/*.toml is loaded; *.toml.off is the same file switched off.
    readonly property string lister: "import glob, json, os, sys, tomllib\n"
        + "out = []\n"
        + "for p in glob.glob(sys.argv[1] + '/*.toml') + glob.glob(sys.argv[1] + '/*.toml.off'):\n"
        + "    on = p.endswith('.toml')\n"
        + "    name = os.path.basename(p).removesuffix('.off').removesuffix('.toml')\n"
        + "    try:\n        d = tomllib.load(open(p, 'rb')).get('description', '')\n"
        + "    except Exception as e:\n        d = f'unreadable: {e}'\n"
        + "    out.append({'name': name, 'on': on, 'description': d})\n"
        + "print(json.dumps(sorted(out, key=lambda t: t['name'])))\n"

    function listTools() {
        am.run([am.python, "-c", am.lister, am.toolsDir], null, (code, out) => {
            try { if (code === 0) am.tools = JSON.parse(out); } catch (e) {}
        });
    }

    // Whether a tool is on in the draft: its wanted state, else the file's.
    function toolOn(t) { return am.toolsWanted[t.name] ?? t.on; }
    function toggleTool(t) {
        const next = Object.assign({}, am.toolsWanted);
        if (!am.toolOn(t) === t.on) delete next[t.name]; else next[t.name] = !am.toolOn(t);
        am.toolsWanted = next;
    }
    // One rename at a time, then done(). -n never overwrites a file that
    // already holds the other name.
    function applyTools(names, done) {
        if (names.length === 0) { done(); return; }
        const name = names[0];
        const on = am.toolsDir + "/" + name + ".toml";
        const off = on + ".off";
        const want = am.toolsWanted[name];
        am.run(["mv", "-n", want ? off : on, want ? on : off], null, (code, out, err) => {
            if (code !== 0) { am.error = err.trim() || "could not switch " + name; am.listTools(); return; }
            const next = Object.assign({}, am.toolsWanted);
            delete next[name];
            am.toolsWanted = next;
            am.applyTools(names.slice(1), done);
        });
    }

    function listModels(done) {
        am.run(["curl", "-sf", "-m", "3", am.llm + "/models"], null, (code, out) => {
            try {
                am.models = code === 0 ? JSON.parse(out).data.map(m => ({
                    id: m.id, status: m.status?.value ?? "",
                    inputs: m.architecture?.input_modalities ?? null })) : [];
            } catch (e) { am.models = []; }
            if (done) done();
        });
    }

    // The router keeps one model resident, so this swaps out whatever is
    // loaded, for every client. It does not touch the preferred model: a
    // load from the page is for now, and the canvas brings the preferred one
    // back when it opens. The list is polled until the router reports the
    // model loaded, since the request may return before the load finishes.
    function loadModel(id) {
        am.modelError = "";
        am.loading = id;
        am.run(["curl", "-s", "-m", "600", "-X", "POST", "-H", "Content-Type: application/json",
                "-d", JSON.stringify({ model: id }), am.routerRoot + "/models/load"], null, (code, out) => {
            let err = code === 0 ? "" : "the router did not answer";
            try {
                const e = JSON.parse(out).error;
                // Already running is the outcome asked for, not a failure.
                if (e && e.message !== "model is already running") err = e.message;
            } catch (x) {}
            if (err) { am.modelError = `${id}: ${err}`; am.loading = ""; am.listModels(); }
            else { poll.left = 300; poll.start(); }
        });
    }
    // Called when the canvas opens and when its chat's model changes, with
    // the chat's model, else the preferred one: loaded if it is a router
    // model neither loaded nor loading. A cloud id is not in the router's
    // list, so nothing is swapped for it. One model and one client at a time
    // is how this desktop runs, so this is the moment to bring the model
    // back. The list is read every time, so hears is current.
    function ensureModel(id) {
        if (am.loading !== "") return;
        am.listModels(() => {
            const m = am.models.find(x => x.id === id);
            if (m && m.status !== "loaded" && m.status !== "loading") am.loadModel(id);
        });
    }

    property Timer poll: Timer {
        property int left: 0
        interval: 1000
        repeat: true
        onTriggered: am.listModels(() => {
            const m = am.models.find(x => x.id === am.loading);
            if (m && m.status === "loaded") { am.loading = ""; am.poll.stop(); }
            else if (--am.poll.left <= 0) {
                am.modelError = `${am.loading} did not finish loading`;
                am.loading = "";
                am.poll.stop();
            }
        })
    }

    function restart() {
        am.restarting = true;
        am.run([am.daemonDir + "/assistant.sh", "restart"], null, (code, out) => {
            am.restarting = false;
            // assistant.sh ends on its status: the process, the socket, and
            // the log's tail, which is where a refused config says why.
            am.status = code === 0 ? "Restarted." : out.trim();
            if (code === 0) am.pending = false;
        });
    }

    page: Component {
        Page {
            title: "Assistant"
            actions: Component {
                SaveBar { dirty: am.dirty; valid: am.valid; onSave: am.save(); onDiscard: am.discard() }
            }
            AssistantPage { width: parent.width; asst: am }
        }
    }
}