#!/usr/bin/python3
"""mood-surface — Microsoft Surface support (root).

Surface Pro 4+, Book, Laptop and Studio touchscreens (Intel IPTS / THC) and pens need the
linux-surface kernel plus the iptsd daemon; Debian's stock kernel doesn't see them at all.

  status            JSON: is this a Surface, kernel, iptsd, Secure Boot / key enrolment state
  install           add the linux-surface apt repo (bundled key) and install kernel + iptsd + libwacom
  auto              what mood-surface.service runs at boot: install once on an installed Surface
  remove            1.5.2 revert: purge the linux-surface kernel/iptsd/libwacom-surface, drop the repo, put Debian's
                    libwacom + stock kernel back, update-grub. Idempotent; mood-surface.service runs it every boot
  doctor            diagnose + repair "installed but touch still dead": pending MOK, stock kernel booted,
                    touch driver not loaded, iptsd not running. Report also saved to /var/log/mood-surface-doctor.txt
Progress/result: /var/lib/mood/surface.json. Live USB sessions can't swap kernels (it lives on the ISO).
"""
import json
import os
import subprocess
import sys
import time
from pathlib import Path

STATE = Path("/var/lib/mood/surface.json")
LOG = Path("/var/log/mood-surface.log")
SOURCES = Path("/etc/apt/sources.list.d/linux-surface.sources")
KEYRING = "/usr/share/keyrings/linux-surface.gpg"
PKGS = ["linux-image-surface", "iptsd", "libwacom-surface"]
APT_ENV = {**os.environ, "DEBIAN_FRONTEND": "noninteractive", "PATH": "/usr/sbin:/usr/bin:/sbin:/bin"}
APT_OPTS = ["-y", "-o", "DPkg::Lock::Timeout=900", "-o", "Dpkg::Options::=--force-confdef", "-o", "Dpkg::Options::=--force-confold"]


def dmi(k):
    try:
        return Path("/sys/class/dmi/id", k).read_text().strip()
    except OSError:
        return ""


def is_surface():
    return "microsoft" in dmi("sys_vendor").lower() and "surface" in dmi("product_name").lower()


def is_live():
    try:
        return "boot=live" in Path("/proc/cmdline").read_text()
    except OSError:
        return False


def secure_boot():
    try:
        out = subprocess.run(["mokutil", "--sb-state"], capture_output=True, text=True, timeout=10).stdout
        return "enabled" in out.lower()
    except (OSError, subprocess.SubprocessError):
        for f in Path("/sys/firmware/efi/efivars").glob("SecureBoot-*"):
            try:
                return f.read_bytes()[-1] == 1
            except OSError:
                pass
    return False


def installed(pkg):
    r = subprocess.run(["dpkg-query", "-W", "-f=${Status}", pkg], capture_output=True, text=True)
    return "install ok installed" in r.stdout


def touch_found():
    try:
        return bool(json.loads(Path("/run/mood/input.json").read_text()).get("touch"))
    except (OSError, ValueError):
        return False


def saved():
    try:
        return json.loads(STATE.read_text())
    except (OSError, ValueError):
        return {}


def save(**kw):
    st = {**saved(), **kw, "at": int(time.time())}
    STATE.parent.mkdir(parents=True, exist_ok=True)
    tmp = STATE.with_suffix(".tmp")
    tmp.write_text(json.dumps(st))
    tmp.chmod(0o644)
    os.replace(tmp, STATE)
    return st


def status():
    rel = os.uname().release
    return {
        "surface": is_surface(), "model": dmi("product_name"), "live": is_live(),
        "kernel": rel, "surfaceKernel": "surface" in rel,
        "kernelInstalled": installed("linux-image-surface"), "iptsd": installed("iptsd"),
        "secureBoot": secure_boot(), "touch": touch_found(), **{k: v for k, v in saved().items() if k in ("phase", "error", "at", "mok")},
    }


def notify(title, body):
    for r in Path("/run/user").glob("*"):
        try:
            uid = int(r.name)
        except ValueError:
            continue
        if uid < 1000:
            continue
        try:
            user = subprocess.run(["id", "-nu", str(uid)], capture_output=True, text=True).stdout.strip()
            subprocess.run(["runuser", "-u", user, "--", "env", f"XDG_RUNTIME_DIR={r}", f"DBUS_SESSION_BUS_ADDRESS=unix:path={r}/bus",
                            "notify-send", "-a", "Moodtop", "-i", "input-tablet", title, body], capture_output=True, timeout=10)
        except (OSError, subprocess.SubprocessError):
            pass


def run(cmd, log):
    log.write(f"$ {' '.join(cmd)}\n")
    log.flush()
    r = subprocess.run(cmd, stdout=log, stderr=subprocess.STDOUT, env=APT_ENV, timeout=3600)
    if r.returncode:
        raise RuntimeError(f"{cmd[0]} {cmd[1] if len(cmd) > 1 else ''} failed ({r.returncode}) — see {LOG}")


def install():
    if not is_surface():
        raise SystemExit("not a Microsoft Surface")
    if is_live():
        save(phase="live")
        raise SystemExit("live USB: the kernel is on the ISO — install Moodtop OS to get Surface touch")
    save(phase="installing", error="")
    try:
        with open(LOG, "a") as log:
            log.write(f"\n=== {time.ctime()} {dmi('product_name')} ===\n")
            SOURCES.write_text("Types: deb\nURIs: https://pkg.surfacelinux.com/debian\nSuites: release\nComponents: main\n"
                               f"Signed-By: {KEYRING}\nArchitectures: amd64\n")
            run(["apt-get", "-o", "DPkg::Lock::Timeout=900", "update"], log)
            pkgs = list(PKGS)
            sb = secure_boot()
            if sb:
                pkgs.append("linux-surface-secureboot-mok")   # postinst queues the key with mokutil (password "surface")
            run(["apt-get", "install", *APT_OPTS, *pkgs], log)
            subprocess.run(["systemctl", "daemon-reload"], capture_output=True)
            subprocess.run(["udevadm", "trigger", "--subsystem-match=hidraw"], capture_output=True)
        st = save(phase="reboot", error="", mok=sb)
        body = "Restart to turn on the touchscreen and pen."
        if sb:
            body += " On the blue screen after restart choose Enroll MOK → Continue → Yes, type the password: surface, then Reboot."
        notify("Surface touchscreen ready", body)
        return st
    except Exception as e:  # noqa: BLE001
        save(phase="failed", error=str(e)[:300])
        raise


def auto():
    if not is_surface() or is_live():
        return
    s = saved()
    if s.get("phase") in ("reboot", "installing") and installed("linux-image-surface") and installed("iptsd"):
        if "surface" in os.uname().release:
            save(phase="done")
        return
    if s.get("phase") == "failed" and time.time() - s.get("at", 0) < 6 * 3600:
        return
    if touch_found() and not s:   # touchscreen already works on the stock kernel (e.g. Surface Go)
        return
    # wait for the network (apt needs it); give up quietly, the next boot retries
    for _ in range(60):
        if subprocess.run(["getent", "hosts", "pkg.surfacelinux.com"], capture_output=True).returncode == 0:
            break
        time.sleep(5)
    else:
        return
    try:
        install()
    except Exception as e:  # noqa: BLE001
        print(f"mood-surface: {e}", file=sys.stderr)


def surface_pkgs():
    out = subprocess.run(["dpkg-query", "-W", "-f=${db:Status-Abbrev}|${Package}|${Version}|${Provides}\\n"],
                         capture_output=True, text=True).stdout
    pkgs = []
    for line in out.splitlines():
        st, name, ver, prov = (line.split("|") + ["", "", "", ""])[:4]
        if st.startswith("i") or st.startswith("r"):
            if "surface" in name or "surface" in ver or name == "iptsd":
                pkgs.append((name, ver, prov))
    return pkgs


def debian_has(name):
    return bool(subprocess.run(["apt-cache", "madison", name], capture_output=True, text=True).stdout.strip())


def remove():
    pkgs = surface_pkgs()
    if not pkgs and not SOURCES.exists():
        save(phase="removed", error="")
        return
    save(phase="removing", error="")
    with open(LOG, "a") as log:
        log.write(f"\n=== {time.ctime()} remove surface kernel ===\n")
        SOURCES.unlink(missing_ok=True)
        subprocess.run(["systemctl", "disable", "--now", "iptsd.service"], capture_output=True)
        if not pkgs:
            save(phase="removed", error="")
            return
        run(["apt-get", "-o", "DPkg::Lock::Timeout=900", "update"], log)
        purge, keep = [], ["linux-image-amd64"]
        for name, ver, prov in pkgs:
            if "surface" not in name and name != "iptsd" and debian_has(name):
                keep.append(f"{name}/trixie")          # e.g. libwacom9 rebuilt by linux-surface -> Debian's again
            else:
                purge.append(name)
                for p in prov.split(","):
                    p = p.strip().split(" ")[0].split(":")[0]
                    if p and "surface" not in p and debian_has(p):
                        keep.append(p)
        plan = keep + [f"{n}-" for n in purge]
        sim = subprocess.run(["apt-get", "-s", "install", "--allow-downgrades", *plan], capture_output=True, text=True, env=APT_ENV)
        log.write(sim.stdout + sim.stderr)
        extra = [l.split()[1] for l in sim.stdout.splitlines() if l.startswith("Remv ") and l.split()[1].split(":")[0] not in purge]
        if sim.returncode or extra:
            raise RuntimeError(f"refusing: apt would also remove {extra}" if extra else "apt could not plan the removal")
        run(["apt-get", "install", "--allow-downgrades", *APT_OPTS, *plan], log)
        left = [n for n in purge if subprocess.run(["dpkg-query", "-W", "-f=${db:Status-Abbrev}", n],
                                                   capture_output=True, text=True).stdout.startswith("rc")]
        if left:
            run(["apt-get", "purge", *APT_OPTS, *left], log)
        if not list(Path("/boot").glob("vmlinuz-*")):
            raise RuntimeError("no kernel left in /boot")
        run(["update-grub"], log)
    save(phase="removed", error="")
    if "surface" in os.uname().release:
        notify("Restart Moodtop", "The Surface kernel was removed. Restart to go back to the normal kernel (brightness keys work again).")


def sh(cmd, timeout=20):
    try:
        r = subprocess.run(cmd, shell=isinstance(cmd, str), capture_output=True, text=True, timeout=timeout)
        return (r.stdout + r.stderr).strip()
    except (OSError, subprocess.SubprocessError) as e:
        return str(e)


def doctor():
    out, fixes = [], []
    say = out.append
    rel = os.uname().release
    say(f"model: {dmi('sys_vendor')} {dmi('product_name')} ({dmi('product_sku')})")
    say(f"kernel: {rel}   surface kernel installed: {installed('linux-image-surface')}   iptsd: {installed('iptsd')}")
    kernels = sorted(p.name for p in Path("/boot").glob("vmlinuz-*"))
    say("boot kernels: " + ", ".join(kernels))
    sb = secure_boot()
    pending = sh(["mokutil", "--list-new"]) if sb else ""
    say(f"secure boot: {sb}" + ("   MOK key WAITING to be enrolled" if sb and "Subject" in pending else ""))
    if "surface" not in rel:
        if not installed("linux-image-surface"):
            say("PROBLEM: Surface kernel not installed -> running install now")
            try:
                install(); fixes.append("installed linux-surface kernel + iptsd (restart needed)")
            except BaseException as e:  # noqa: BLE001
                say(f"install failed: {e}")
        elif sb and "Subject" in pending:
            say("PROBLEM: Secure Boot blocks the Surface kernel until its key is enrolled. Restart, and on the blue "
                "'MOK management' screen choose Enroll MOK -> Continue -> Yes, password: surface -> Reboot.")
        else:
            entry = next((l.split("'")[1] for l in Path("/boot/grub/grub.cfg").read_text(errors="ignore").splitlines()
                          if l.strip().startswith("menuentry ") and "surface" in l and "recovery" not in l), None) \
                if Path("/boot/grub/grub.cfg").exists() else None
            say(f"PROBLEM: booted the stock kernel although the Surface kernel is installed (grub entry: {entry})")
            sh(["update-grub"], 120)
            fixes.append("regenerated GRUB so the newest (Surface) kernel boots first - restart")
            if sb:
                say("Secure Boot is on: if the Surface kernel refuses to start, run: sudo mokutil --import "
                    "/usr/share/linux-surface-secureboot/surface.cer (password surface) and restart -> Enroll MOK")
    else:
        mods = sh("lsmod | grep -E '^(ithc|ipts|surface_hid|surface_aggregator)' | cut -d' ' -f1 | tr '\\n' ' '")
        say(f"touch modules loaded: {mods or 'none'}")
        for m in ("ithc", "ipts"):
            if m not in mods.split():
                if subprocess.run(["modprobe", m], capture_output=True).returncode == 0:
                    fixes.append(f"loaded kernel module {m}")
        subprocess.run(["udevadm", "trigger", "--action=add", "--subsystem-match=hidraw"], capture_output=True)
        subprocess.run(["udevadm", "settle", "--timeout=10"], capture_output=True)
        found = []
        for d in sorted(Path("/dev").glob("hidraw*")):
            name = ""
            try:
                name = next((l.split("=", 1)[1] for l in Path(f"/sys/class/hidraw/{d.name}/device/uevent").read_text().splitlines()
                             if l.startswith("HID_NAME=")), "")
            except OSError:
                pass
            ok = subprocess.run(["iptsd-check-device", "--quiet", str(d)], capture_output=True).returncode == 0
            say(f"  {d} {name!r} touch-device={ok}")
            if ok:
                found.append(d)
        if not found:
            say("PROBLEM: no touch controller found on the Surface kernel. Kernel messages:")
            say(sh("dmesg | grep -iE 'ithc|ipts|thc|hid-multitouch|surface' | tail -25"))
        for d in found:
            unit = "iptsd@" + sh(["systemd-escape", "--path", str(d)]) + ".service"
            if sh(["systemctl", "is-active", unit]) != "active":
                sh(["systemctl", "restart", unit])
                fixes.append(f"started {unit}")
            say(f"  {unit}: {sh(['systemctl', 'is-active', unit])}")
            say(sh(["journalctl", "-b", "-u", unit, "-n", "12", "--no-pager", "-o", "cat"]))
        time.sleep(2)
    say("touch devices seen by Moodtop: " + str(touch_found()))
    say(sh("libinput list-devices 2>/dev/null | grep -E '^(Device|Capabilities)' | paste - - | grep -iE 'touch|tablet|ipts|pen'") or "libinput: no touch/pen devices")
    say("backlight: " + sh("for d in /sys/class/backlight/*; do echo \"$(basename $d) type=$(cat $d/type) $(cat $d/brightness)/$(cat $d/max_brightness)\"; done"))
    say("FIXED: " + ("; ".join(fixes) if fixes else "nothing automatic"))
    text = "\n".join(out) + "\n"
    Path("/var/log/mood-surface-doctor.txt").write_text(text)
    save(doctor=fixes, doctorAt=int(time.time()))
    print(text)


if __name__ == "__main__":
    cmd = sys.argv[1] if len(sys.argv) > 1 else "status"
    if cmd == "status":
        print(json.dumps(status(), indent=1))
    elif cmd == "install":
        print(json.dumps(install()))
    elif cmd in ("auto", "remove"):
        try:
            remove()
        except Exception as e:  # noqa: BLE001
            save(phase="failed", error=str(e)[:300])
            sys.exit(f"mood-surface: {e}")
    elif cmd == "doctor":
        doctor()
    else:
        sys.exit(__doc__)
