#!/usr/bin/python3
"""Moodtop OS updater (root).

  mood-update check          -> refresh apt + the cloud release notes, print JSON {count, packages, moodtop, latest}
  mood-update run [--auto]   -> install updates (the daily moodtop-update.timer uses --auto)

Updates come from two apt sources: Debian (security + point releases) and the Moodtop cloud repo
(https://zandy.zo.space/moodtop/apt, signed with /usr/share/keyrings/moodtop-archive-keyring.gpg),
which ships the whole Mood layer as the `moodtop-core` package.

Scope: installed system -> full-upgrade; live USB with persistence -> Moodtop packages only (a new
kernel would never be booted from the ISO); live USB without persistence -> nothing (lost at shutdown).

Config: /etc/mood/update.json {"auto": true, "channel": "stable"}.
Status: /var/lib/mood/update-status.json, live progress /run/mood-update/progress.json (both world-readable).
"""
import fcntl
import json
import os
import re
import subprocess
import sys
import time
import urllib.request
from pathlib import Path

CONFIG = Path("/etc/mood/update.json")
STATUS = Path("/var/lib/mood/update-status.json")
LATEST = Path("/var/lib/mood/latest.json")
RUN = Path("/run/mood-update")
PROGRESS = RUN / "progress.json"
LATEST_URL = "https://zandy.zo.space/moodtop/latest.json"
APT_OPTS = ["-o", "Dpkg::Options::=--force-confdef", "-o", "Dpkg::Options::=--force-confold",
            "-o", "DPkg::Lock::Timeout=600", "-o", "Acquire::Retries=3", "-o", "APT::Get::Show-User-Simulation-Note=false"]
ENV = dict(os.environ, DEBIAN_FRONTEND="noninteractive", LC_ALL="C.UTF-8", PATH="/usr/sbin:/usr/bin:/sbin:/bin")


def jread(p, default):
    try:
        return json.loads(Path(p).read_text())
    except (OSError, ValueError):
        return default


def jwrite(p, data):
    p = Path(p)
    p.parent.mkdir(parents=True, exist_ok=True)
    tmp = p.with_name("." + p.name + ".tmp")
    tmp.write_text(json.dumps(data))
    os.chmod(tmp, 0o644)
    os.replace(tmp, p)


def config():
    c = {"auto": True, "channel": "stable"}
    c.update(jread(CONFIG, {}))
    return c


def status_update(**kw):
    s = jread(STATUS, {})
    s.update(kw)
    jwrite(STATUS, s)
    return s


def progress(running, stage="", pct=0, label=""):
    jwrite(PROGRESS, {"running": running, "stage": stage, "pct": round(pct, 1), "label": label, "t": time.time()})


def scope():
    cmd = Path("/proc/cmdline").read_text()
    if "boot=live" not in cmd:
        return "full"
    return "moodtop" if "/run/live/persistence" in Path("/proc/mounts").read_text() else "none"


def moodtop_pkgs():
    out = subprocess.run(["dpkg-query", "-W", "-f=${Package} ${Status}\n", "moodtop-*"], capture_output=True, text=True).stdout
    return [l.split()[0] for l in out.splitlines() if l.endswith("install ok installed")]


def pkg_version(name):
    return subprocess.run(["dpkg-query", "-W", "-f=${Version}", name], capture_output=True, text=True).stdout.strip()


def candidate(name):
    out = subprocess.run(["apt-cache", "policy", name], capture_output=True, text=True, env=ENV).stdout
    m = re.search(r"Candidate:\s*(\S+)", out)
    return m[1] if m and m[1] != "(none)" else ""


def upgrade_cmd(sc, simulate=False):
    base = ["apt-get", *APT_OPTS]
    if simulate:
        base.append("-s")
    else:
        base.append("-y")
    if sc == "full":
        return base + ["full-upgrade"]
    pk = moodtop_pkgs() or ["moodtop-core"]
    return base + ["install", "--only-upgrade", *pk]


def fetch_latest():
    try:
        req = urllib.request.Request(LATEST_URL, headers={"User-Agent": "moodtop-update", "Accept": "application/json"})
        with urllib.request.urlopen(req, timeout=20) as r:
            data = json.loads(r.read(200_000))
        if isinstance(data, dict) and data.get("version"):
            jwrite(LATEST, data)
            return data
    except Exception:
        pass
    return jread(LATEST, None)


def check(sc=None):
    sc = sc or scope()
    progress(True, "check", 5, "Checking for updates")
    r = subprocess.run(["apt-get", *APT_OPTS, "update"], capture_output=True, text=True, env=ENV, timeout=900)
    latest = fetch_latest()
    err = ""
    if r.returncode != 0:
        err = (r.stderr.strip().splitlines() or ["apt update failed"])[-1][:300]
    pkgs = []
    if sc != "none":
        out = subprocess.run(upgrade_cmd(sc, simulate=True), capture_output=True, text=True, env=ENV).stdout
        pkgs = [l.split()[1] for l in out.splitlines() if l.startswith("Inst ")]
    res = {"count": len(pkgs), "packages": pkgs[:300], "scope": sc,
           "moodtop": {"installed": pkg_version("moodtop-core"), "candidate": candidate("moodtop-core")},
           "latest": latest, "error": err}
    status_update(lastCheck=int(time.time()), pending=len(pkgs), lastError=err)
    progress(False)
    return res


def on_battery_low():
    ps = list(Path("/sys/class/power_supply").glob("*"))
    ac = any((p / "online").exists() and (p / "online").read_text().strip() == "1" for p in ps
             if (p / "type").exists() and (p / "type").read_text().strip() in ("Mains", "USB"))
    if ac:
        return False
    for p in ps:
        if (p / "type").exists() and (p / "type").read_text().strip() == "Battery" and (p / "capacity").exists():
            try:
                return int((p / "capacity").read_text()) < 30
            except ValueError:
                pass
    return False


def sessions():
    """(user, uid) of logged-in graphical users, for notifications."""
    out = []
    for d in Path("/run/user").glob("*"):
        try:
            uid = int(d.name)
        except ValueError:
            continue
        if uid >= 1000 and (d / "bus").exists():
            import pwd
            try:
                out.append((pwd.getpwuid(uid).pw_name, uid))
            except KeyError:
                pass
    return out


def notify(title, body, reload_shell=False):
    for user, uid in sessions():
        env = ["env", f"XDG_RUNTIME_DIR=/run/user/{uid}", f"DBUS_SESSION_BUS_ADDRESS=unix:path=/run/user/{uid}/bus"]
        subprocess.run(["runuser", "-u", user, "--", *env, "notify-send", "-a", "Moodtop", "-i", "mood-logo", title, body],
                       capture_output=True, timeout=15)
        if reload_shell:
            subprocess.run(["runuser", "-u", user, "--", *env, "/usr/bin/mood-shell", "--reload"], capture_output=True, timeout=15)


def run(auto):
    cfg = config()
    sc = scope()
    if auto and not cfg.get("auto", True):
        return {"skipped": "automatic updates are off"}
    if sc == "none":
        return {"skipped": "live USB without persistence"}
    if auto and on_battery_low():
        return {"skipped": "battery low"}
    RUN.mkdir(mode=0o755, exist_ok=True)
    lock = open(RUN / "lock", "w")
    try:
        fcntl.flock(lock, fcntl.LOCK_EX | fcntl.LOCK_NB)
    except BlockingIOError:
        return {"skipped": "already running"}
    before = pkg_version("moodtop-core")
    try:
        res = check(sc)
        if res["count"] == 0:
            status_update(lastError=res["error"])
            return {"updated": 0}
        progress(True, "download", 10, f"Downloading {res['count']} update{'s' if res['count'] != 1 else ''}")
        r_fd, w_fd = os.pipe()
        cmd = upgrade_cmd(sc)
        cmd[1:1] = ["-o", f"APT::Status-Fd={w_fd}"]
        p = subprocess.Popen(cmd, stdout=subprocess.PIPE, stderr=subprocess.STDOUT, env=ENV, pass_fds=(w_fd,))
        os.close(w_fd)
        with os.fdopen(r_fd) as st:
            for line in st:
                f = line.rstrip("\n").split(":", 3)
                if len(f) < 4:
                    continue
                try:
                    pct = float(f[2])
                except ValueError:
                    continue
                if f[0] == "dlstatus":
                    progress(True, "download", 10 + pct * 0.4, f[3][:120] or "Downloading")
                elif f[0] == "pmstatus":
                    progress(True, "install", 50 + pct * 0.5, f[3][:120] or "Installing")
        log = p.stdout.read().decode(errors="replace")
        p.wait()
        Path("/var/log/mood-update.log").write_text(log[-200_000:])
        if p.returncode != 0:
            subprocess.run(["dpkg", "--configure", "-a"], capture_output=True, env=ENV)
            err = (log.strip().splitlines() or ["update failed"])[-1][:300]
            status_update(lastError=err)
            raise SystemExit(err)
        after = pkg_version("moodtop-core")
        reboot = Path("/run/reboot-required").exists() or any(re.search(r"^Unpacking (linux-image|systemd |libc6)", l) for l in log.splitlines())
        status_update(lastUpdate=int(time.time()), pending=0, lastError="", rebootNeeded=reboot or jread(STATUS, {}).get("rebootNeeded", False))
        if after and after != before:
            ver = after.split("+")[0]
            notes = (jread(LATEST, {}) or {}).get("notes") or []
            notify(f"Moodtop updated to {ver} ✨", ("; ".join(notes[:3]) if notes else "New features and fixes are ready.") + " Reopen apps to see what's new.", reload_shell=True)
        elif reboot:
            notify("Updates installed", "Restart when convenient to finish updating.")
        return {"updated": res["count"], "reboot": reboot}
    finally:
        progress(False)
        lock.close()


def main():
    if os.getuid() != 0:
        sys.exit("mood-update must run as root")
    if len(sys.argv) < 2 or sys.argv[1] not in ("check", "run"):
        sys.exit(__doc__)
    if sys.argv[1] == "check":
        print(json.dumps(check()))
    else:
        print(json.dumps(run("--auto" in sys.argv)))


main()
