// 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.

pragma Singleton

import Quickshell
import Quickshell.Io
import QtQuick
import "../.."

// The wallpaper side. Setting one is `wallp`'s job: it knows about the two
// outputs, the saved state and the accent run, and none of that is worth
// reimplementing here.
Singleton {
    id: root

    readonly property string home: Quickshell.env("HOME")
    readonly property string dir: `${home}/Pictures/wallpapers`
    // The distribution's wallpapers: loose images, and Plasma packages that
    // each carry one picture at several resolutions under contents/images/,
    // some with a dark version under contents/images_dark/.
    readonly property string sysDir: "/usr/share/wallpapers"

    // Which screen a pick applies to. wallp calls them H and V.
    property string target: "H"

    // Staged picks, empty until something is chosen. Selecting does not set
    // anything: both screens can be composed and applied together, which is
    // also a single wallp call rather than two.
    property string pendingH: ""
    property string pendingV: ""
    readonly property bool dirty: pendingH !== "" || pendingV !== ""

    function pick(path) {
        if (target === "V") pendingV = path;
        else pendingH = path;
    }

    function clearPending() { pendingH = ""; pendingV = ""; }

    // What a screen should show: the staged pick if there is one, else what
    // is actually set.
    function shown(which) {
        if (which === "V") return pendingV !== "" ? pendingV : currentV;
        return pendingH !== "" ? pendingH : currentH;
    }

    property list<string> files: []
    property bool scanning: true
    property bool applying: false
    signal applied(bool ok, string message)
    // Sent before wallp runs, while current still names the old picture, so
    // WallpaperSpread can cover the swap with it.
    signal swapping(string which, string from, string to)
    // Every run's end, Apply's or the engine's. applied is Apply's alone: it
    // reports to the save bar and closes the window.
    signal finished(bool ok)
    // A successful Apply, with what it set ("" for a screen it left).
    signal manualApplied(string h, string v)

    // Sorted by most recent first: the wallpaper wanted now is usually one
    // added recently. -maxdepth keeps a stray git checkout from being walked.
    function refresh() {
        scanning = true;
        scanProc.running = false;
        scanProc.running = true;
    }

    Process {
        id: scanProc
        running: true
        // img DIR "DEPTH" [ACTIONS...]: the image files under DIR. DEPTH is
        // left unquoted on purpose, so "-maxdepth 2" splits into two words
        // and "" vanishes.
        command: ["sh", "-c",
            `img() { d=$1; o=$2; shift 2; find "$d" $o -type f ` +
            `\\( -iname '*.png' -o -iname '*.jpg' -o -iname '*.jpeg' -o -iname '*.webp' \\) "$@" 2>/dev/null; }; ` +
            `img ${JSON.stringify(root.dir)} "-maxdepth 2" -not -path '*/.*' -printf '%T@ %p\\n' | sort -rn | cut -d' ' -f2-; ` +
            `img ${JSON.stringify(root.sysDir)} "-maxdepth 1" | sort; ` +
            `img ${JSON.stringify(root.sysDir)} "" -path '*/contents/images*/*'`]
        stdout: StdioCollector {
            onStreamFinished: {
                root.files = text.trim().split("\n").filter(s => s.length);
                root.scanning = false;
                sizeProc.running = false;
                sizeProc.command = ["sh", "-c", 'printf "%s\\0" "$@" | xargs -0 -r file -N -0', "sh"].concat(root.files);
                sizeProc.running = true;
            }
        }
    }

    // Each file's pixel size, by path, from `file`, which reads only the
    // header: about half a second for a thousand. -0 ends the name with a
    // NUL, since a name can hold the ": " that would otherwise end it. The
    // size is the last WxH in the description, because a JPEG's comes after
    // its "density 72x72". A file `file` gives no size for has no entry and
    // gets no quality colour.
    property var sizes: ({})
    Process {
        id: sizeProc
        stdout: StdioCollector {
            onStreamFinished: {
                const out = {};
                for (const line of text.split("\n")) {
                    const nul = line.indexOf("\0");
                    if (nul < 0) continue;
                    const re = /(\d+) ?x ?(\d+)/g;
                    let m, last = null;
                    while ((m = re.exec(line.slice(nul + 1))) !== null) last = m;
                    if (last) out[line.slice(0, nul)] = { w: +last[1], h: +last[2] };
                }
                root.sizes = out;
            }
        }
    }

    // The screens wallp sets, as the mock draws them: DP-1 wide, DP-3 turned
    // upright.
    // ponytail: fixed, like the mock; read Quickshell.screens if a monitor
    // changes.
    readonly property var panels: ({ H: { w: 2560, h: 1080 }, V: { w: 1080, h: 1920 } })

    // How far swaybg's `-m fill` has to enlarge an image to cover a screen,
    // 0 when the size is unknown. At or under 1 it is downscaled and sharp;
    // to 1.33 the softening is slight; past that it shows pixels.
    function upscale(path, which) {
        const s = sizes[path];
        const p = panels[which];
        return s && s.w > 0 && s.h > 0 ? Math.max(p.w / s.w, p.h / s.h) : 0;
    }
    function qualityColor(path, which) {
        const u = upscale(path, which);
        return u === 0 ? "transparent" : u <= 1 ? Theme.green : u <= 1.33 ? Theme.yellow : Theme.red;
    }
    // Why an image that fits will still look bad, or "" when it will not.
    function qualityNote(path, which) {
        const u = upscale(path, which);
        if (u <= 1) return "";
        const s = sizes[path], p = panels[which];
        return `${s.w}×${s.h} is enlarged ${u.toFixed(2)}× to cover the ${p.w}×${p.h} screen: `
             + (u <= 1.33 ? "each image pixel is spread over more than one screen pixel, so it will look slightly soft."
                          : "each image pixel becomes a visible block, so it will look pixelated.");
    }

    // Real-ESRGAN through stable-diffusion.cpp's sd-cli, for an image too
    // small for its screen. The file is replaced in place, so wallp's state
    // and the engine's lists, which name it by path, still find it. sd-cli
    // writes to a dot-named sibling first, which the scan skips, and only a
    // successful run is moved over the original, so a failure leaves it
    // untouched. The name keeps its extension because sd-cli writes the
    // format the extension names. Only files under `dir`: the system ones
    // are not ours to rewrite. Success is no [E] line as well as exit 0: a
    // model that fails to load logs "[E] new_upscaler_ctx failed", saves
    // the input unscaled and exits 0 (measured).
    //
    // Which Real-ESRGAN model is decided first, as its own step, so the page
    // can show the choice before anything is rewritten. Imajev-4B makes it,
    // a small Jev-style vision decision model, run once through llama-mtmd-cli on a
    // 512 px thumbnail: the answer is held to one token, A or B, by a
    // grammar, so greedy decoding returns the option with the highest logit,
    // which is the decision. About 1.7 s. Anime goes to the anime model,
    // which keeps line art clean; everything else, and any failure of the
    // decider, goes to the general one, safe on any image. Measured on 22
    // wallpapers it got every anime and photo right and called 4 of 6
    // painted 3D renders anime, which then come out smoother than painted;
    // the 2B, and a three-way prompt adding "painted or render", did worse.
    // ponytail: argmax only, so Imajev's calibration temperature is not
    // needed. One 4x pass, which covers both screens unless the image is
    // under 640 px wide or 480 px tall; --upscale-repeats if that matters.
    readonly property string upscalers: "/data/LLM-models/SD/upscalers"
    readonly property string decider: "/data/LLM-models/SD/imajev/imajev-4b"
    property string upscaling: ""
    property string upscaleError: ""
    function canUpscale(path, which) {
        return path.startsWith(dir + "/") && upscale(path, which) > 1;
    }
    function sibling(path) { return path.slice(0, path.lastIndexOf("/") + 1) + "." + basename(path); }

    // The decision for one image: decidedFor is the path it was made for, so
    // staging another image leaves it standing for the first and unread.
    property string deciding: ""
    property string decidedFor: ""
    property string decidedKind: ""   // "anime" or "photo"
    function decide(path) {
        if (deciding !== "" || upscaling !== "" || path === "") return;
        deciding = path;
        upscaleError = "";
        decideProc.command = ["sh", "-c",
            'magick "$1" -resize 512x512 "$2" 2>/dev/null; '
            + 'llama-mtmd-cli -m "$3-Q4_K_M.gguf" --mmproj "$3-mmproj-f16.gguf" --image "$2" '
            + '-ngl 99 --temp 0 -n 1 --grammar "$5" -p "$4" 2>/dev/null; rm -f "$2"',
            "sh", path, sibling(path) + ".jpg", decider,
            "Is this image A) anime, cartoon or illustration art, or B) a photograph or realistic render? Answer A or B.",
            'root ::= "A" | "B"'];
        decideProc.running = true;
    }
    Process {
        id: decideProc
        stdout: StdioCollector { id: decideOut }
        onExited: {
            root.decidedKind = decideOut.text.trim().endsWith("A") ? "anime" : "photo";
            root.decidedFor = root.deciding;
            root.deciding = "";
        }
    }

    function upscaleFile(path) {
        if (upscaling !== "" || path === "" || path !== decidedFor) return;
        upscaling = path;
        upscaleError = "";
        upscaleProc.command = ["sh", "-c",
            'log=$(sd-cli -M upscale --upscale-model "$1" -i "$2" -o "$3" 2>&1); s=$?; '
            + 'err=$(printf "%s\\n" "$log" | grep -F "[E]"); '
            + 'if [ $s -eq 0 ] && [ -z "$err" ]; then mv -f "$3" "$2"; '
            + 'else rm -f "$3"; echo "${err:-sd-cli exited $s}" >&2; exit 1; fi',
            "sh", `${upscalers}/RealESRGAN_x4plus${decidedKind === "anime" ? "_anime_6B" : ""}.pth`,
            path, sibling(path)];
        upscaleProc.running = true;
    }
    Process {
        id: upscaleProc
        stderr: StdioCollector { id: upscaleErr }
        onExited: code => {
            if (code !== 0) root.upscaleError = `upscaling ${root.basename(root.upscaling)} failed: ${upscaleErr.text.trim()}`;
            root.decidedFor = "";
            root.upscaling = "";
            // The new size clears the quality note and the red border.
            root.refresh();
        }
    }

    function basename(path) { return path.slice(path.lastIndexOf("/") + 1); }

    // The files by the subfolder they sit in, folders by name, each keeping
    // the newest-first order of `files`. Images loose in `dir` itself form a
    // group named after it, which sorts first. A folder with no images at
    // the scanned depth has no group. The system ones follow: its loose
    // images as "system", and the Plasma packages as "Plasma", one image
    // each: the dark version where the package has one, else the light, and
    // of those the largest by the WxH its file is named after. `key` is what
    // `expanded` is indexed by, so a local folder named "system" stays apart.
    readonly property var groups: {
        const local = {}, system = [], plasma = {};
        for (const f of files) {
            if (f.startsWith(dir + "/")) {
                const rel = f.slice(dir.length + 1);
                const key = rel.includes("/") ? rel.slice(0, rel.indexOf("/")) : "";
                if (!local[key]) local[key] = [];
                local[key].push(f);
            } else if (f.includes("/contents/images")) {
                const pkg = f.slice(sysDir.length + 1).split("/")[0];
                if (!plasma[pkg] || rank(f) > rank(plasma[pkg])) plasma[pkg] = f;
            } else {
                system.push(f);
            }
        }
        const byName = (a, b) => a.localeCompare(b, undefined, { sensitivity: "base" });
        return Object.keys(local).sort(byName)
            .map(k => ({ key: "local/" + k, name: k || basename(dir), files: local[k] }))
            .concat([
                { key: "system", name: "system", files: system },
                { key: "plasma", name: "Plasma", files: Object.keys(plasma).sort(byName).map(k => plasma[k]) },
            ])
            .filter(g => g.files.length > 0);
    }

    // Dark outranks any light, then size decides.
    function rank(path) {
        const m = /(\d+)x(\d+)\.\w+$/.exec(path);
        return (path.includes("/contents/images_dark/") ? 1e12 : 0) + (m ? m[1] * m[2] : 0);
    }

    // What a thumbnail is called: a Plasma package's file is named after its
    // resolution, so it goes by the package instead.
    function label(path) {
        const i = path.indexOf("/contents/images");
        return i < 0 ? basename(path) : basename(path.slice(0, i));
    }

    // Which groups are open, by key. Here rather than in the panel, which is
    // destroyed on close, so a folder left open is open next time. Reassigned
    // whole, because bindings do not see a key change inside a var.
    property var expanded: ({})
    function toggle(key) {
        const e = Object.assign({}, expanded);
        e[key] = !e[key];
        expanded = e;
    }

    // What is on each screen now. wallp writes these on every set, so
    // watching them keeps the mock honest even when set from a terminal.
    property string currentH: ""
    property string currentV: ""

    FileView {
        path: `${root.home}/.config/wallp/wall_h`
        watchChanges: true
        onFileChanged: reload()
        onLoaded: root.currentH = text().trim()
    }

    FileView {
        path: `${root.home}/.config/wallp/wall_v`
        watchChanges: true
        onFileChanged: reload()
        onLoaded: root.currentV = text().trim()
    }

    // wallp does the work, including running udt-accent, so the accent follows
    // the new wallpaper exactly as it does from the command line. Both screens
    // go in one call: wallp accepts H= and V= together.
    function apply() {
        if (dirty) run(pendingH, pendingV, "manual");
    }

    // What started the run in flight: "manual" (Apply), "tick" (the engine's
    // timer), "home" (the engine switched off) or "keep" (the pair up now,
    // run in full). A tick is the light run,
    // UDT_ROTATION, which udt-accent reads; the others are full.
    property string kind: ""

    // false when nothing was started: a run in flight, or nothing to set.
    function run(h, v, kind) {
        if (applying || (h === "" && v === "")) return false;
        applying = true;
        root.kind = kind;
        let covers = 0;
        if (h !== "" && currentH !== "" && h !== currentH) { covers++; swapping("H", currentH, h); }
        if (v !== "" && currentV !== "" && v !== currentV) { covers++; swapping("V", currentV, v); }
        const args = [];
        if (h !== "") args.push(`H=${JSON.stringify(h)}`);
        if (v !== "") args.push(`V=${JSON.stringify(v)}`);
        const env = kind === "tick" ? "UDT_ROTATION=1 " : "";
        applyProc.command = ["sh", "-c", `${env}wallp --set ${args.join(" ")} >/dev/null 2>&1`];
        uncovered = covers;
        if (covers === 0) applyProc.running = true;
        else coverWait.restart();
        return true;
    }

    // Covers still to come up before wallp may swap swaybg under them. The
    // cover decodes its pictures asynchronously, and starting wallp with
    // swapping showed the new picture through the still empty cover for a
    // moment, then the old one over it, then the burn. WallpaperSpread
    // calls covered() once it is drawing; a cover that never comes (a
    // screen missing, a picture that fails to load) is waited out.
    property int uncovered: 0
    function covered() {
        if (uncovered === 0) return;
        if (--uncovered === 0) {
            coverWait.stop();
            applyProc.running = true;
        }
    }
    Timer {
        id: coverWait
        interval: 1500
        onTriggered: if (root.uncovered > 0) {
            root.uncovered = 0;
            applyProc.running = true;
        }
    }

    Process {
        id: applyProc
        stderr: StdioCollector { id: applyErr }
        onExited: code => {
            const ok = code === 0;
            root.applying = false;
            if (root.kind === "manual") {
                if (ok) {
                    root.manualApplied(root.pendingH, root.pendingV);
                    root.clearPending();
                }
                root.applied(ok, ok ? "" : (applyErr.text.trim() || `wallp exited ${code}`));
            } else if (!ok) {
                console.warn(`Wallpapers: the engine's ${root.kind} set failed, wallp exited ${code}`);
            }
            root.finished(ok);
        }
    }
}