#!/usr/bin/env python3
"""Deploy/release agent: runs inside the commissioning live image when the
kernel cmdline carries rollcall.mode=deploy rollcall.job=<token>.

deploy : stream <os>/<version>/disk.raw.zst onto the target disk (sha256
         verified), seed cloud-init NoCloud files into the written root fs,
         add a UEFI boot entry when booted via UEFI, report, reboot.
release: wipe the target disk (quick: signatures + first/last 64 MiB), report,
         power off.

Everything the agent needs (image URL, checksum, disk, seed files) comes from
GET /api/v1/jobs/self; it never decides policy itself.
"""
import glob
import json
import os
import re
import subprocess
import sys
import threading
import time
import urllib.error
import urllib.request

LOG = open("/dev/console", "w", buffering=1) if os.path.exists("/dev/console") else sys.stdout
SERVER = ""
TOKEN = ""
EFI_GUID = "c12a7328-f81e-11d2-ba4b-00a0c93ec93b"
ROOT_FS = ("ext4", "xfs", "btrfs", "ext3")


def log(msg):
    line = f"[deploy] {msg}"
    print(line, file=LOG)
    if LOG is not sys.stdout:
        print(line, flush=True)


def run(cmd, timeout=600, check=False, **kw):
    p = subprocess.run(cmd, shell=isinstance(cmd, str), capture_output=True, text=True, timeout=timeout, **kw)
    if check and p.returncode != 0:
        raise RuntimeError(f"{cmd if isinstance(cmd, str) else ' '.join(cmd)}: rc={p.returncode} {p.stderr.strip()[-400:]}")
    return p


def readf(path, default=""):
    try:
        with open(path) as f:
            return f.read().strip()
    except OSError:
        return default


def cmdline():
    opts = {}
    for tok in readf("/proc/cmdline").split():
        if tok.startswith("rollcall."):
            k, _, v = tok[len("rollcall."):].partition("=")
            opts[k] = v if v else True
    return opts


# --------------------------------------------------------------------------- server
def api(method, path, body=None, tries=20):
    data = json.dumps(body).encode() if body is not None else None
    req = urllib.request.Request(SERVER + path, data=data, method=method,
                                 headers={"Content-Type": "application/json", "X-Rollcall-Job": TOKEN})
    last = None
    for i in range(tries):
        try:
            with urllib.request.urlopen(req, timeout=30) as r:
                return json.load(r)
        except urllib.error.HTTPError as e:
            if 400 <= e.code < 500:
                raise RuntimeError(f"{path}: HTTP {e.code} {e.read().decode(errors='replace')[:200]}")
            last = e
        except (urllib.error.URLError, OSError) as e:
            last = e
        time.sleep(3)
    raise RuntimeError(f"{path}: {last}")


class Progress:
    def __init__(self):
        self.phase, self.pct, self.msg = "starting", 0, ""
        self.lock = threading.Lock()
        self.stop = threading.Event()
        self.thread = threading.Thread(target=self._loop, daemon=True)
        self.thread.start()

    def set(self, phase=None, pct=None, msg=None, now=False):
        with self.lock:
            if phase is not None:
                self.phase = phase
            if pct is not None:
                self.pct = pct
            if msg is not None:
                self.msg = msg
        if now:
            self._post()

    def _post(self):
        with self.lock:
            body = {"phase": self.phase, "pct": round(self.pct, 1), "msg": self.msg}
        try:
            api("POST", "/api/v1/jobs/self/progress", body, tries=1)
        except Exception as e:  # noqa: BLE001
            log(f"progress post failed: {e}")

    def _loop(self):
        while not self.stop.wait(3):
            self._post()

    def finish(self, state, msg="", error="", details=None):
        self.stop.set()
        body = {"state": state, "phase": state, "pct": 100 if state == "done" else self.pct,
                "msg": msg, "error": error, "details": details or {}}
        return api("POST", "/api/v1/jobs/self/progress", body)


# --------------------------------------------------------------------------- disks
def lsblk(dev=None):
    cmd = ["lsblk", "-J", "-b", "-o", "NAME,PATH,TYPE,SIZE,SERIAL,RM,ROTA,TRAN,FSTYPE,PARTTYPE,PARTN,PKNAME,MOUNTPOINT"]
    if dev:
        cmd.append(dev)
    p = run(cmd)
    if p.returncode != 0:
        return []
    out = []

    def walk(items):
        for d in items:
            out.append(d)
            walk(d.get("children") or [])
    walk(json.loads(p.stdout).get("blockdevices", []))
    return out


def find_disk(spec):
    want, serial, need = spec.get("disk") or "", spec.get("disk_serial") or "", int(spec.get("raw_size") or 0)
    disks = [d for d in lsblk() if d.get("type") == "disk" and not d["name"].startswith(("loop", "ram", "zram", "sr"))]
    if want:
        for d in disks:
            if d["path"] == want or d["name"] == want.replace("/dev/", ""):
                if serial and d.get("serial") and d["serial"] != serial:
                    # device names move around; trust the serial over the name
                    for e in disks:
                        if e.get("serial") == serial:
                            log(f"{want} has serial {d['serial']}, expected {serial}; using {e['path']} instead")
                            return e
                    raise RuntimeError(f"{want} has serial {d['serial']}, expected {serial}; refusing to write")
                return d
        if serial:
            for d in disks:
                if d.get("serial") == serial:
                    log(f"{want} not found; serial {serial} is {d['path']}")
                    return d
        raise RuntimeError(f"target disk {want} not present")
    cands = [d for d in disks if not d.get("rm") and d.get("tran") != "usb" and int(d["size"]) >= need]
    if not cands:
        raise RuntimeError(f"no disk >= {need} bytes")
    cands.sort(key=lambda d: (int(d["size"]), d["name"]))
    return cands[0]


def release_disk(path):
    """Make sure nothing holds the disk: swap, mounts, md, lvm."""
    for d in lsblk(path):
        if d.get("mountpoint"):
            run(["umount", "-f", d["path"]])
    run("swapoff -a")
    run(["vgchange", "-an"])
    run(["mdadm", "--stop", "--scan"])
    for d in lsblk(path):
        if d.get("type") == "part":
            run(["wipefs", "-a", d["path"]])
    run(["wipefs", "-a", path])
    run(["blockdev", "--rereadpt", path])
    run(["udevadm", "settle"])


# --------------------------------------------------------------------------- deploy
def write_image(spec, disk, prog):
    url, want, size = spec["image_url"], spec["sha256"], int(spec.get("size") or 0)
    sumf, pvf = "/run/rollcall.sum", "/run/rollcall.pv"
    for f in (sumf, pvf):
        if os.path.exists(f):
            os.remove(f)
    pipeline = (f"set -o pipefail; curl -fsSL --retry 5 --retry-connrefused '{url}' "
                f"| tee >(sha256sum > {sumf}) "
                f"| pv -n -f -i 2 -s {size} 2>{pvf} "
                f"| zstd -dc -T0 "
                f"| dd of='{disk}' bs=4M iflag=fullblock oflag=direct conv=fsync status=none")
    log(f"writing {url} -> {disk}")
    proc = subprocess.Popen(["bash", "-c", pipeline], stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True)
    pct = 0
    while proc.poll() is None:
        time.sleep(2)
        try:
            with open(pvf) as f:
                lines = f.read().split()
            if lines:
                pct = int(lines[-1])
        except (OSError, ValueError):
            pass
        prog.set("writing", min(pct, 99), f"{pct}% of {size // (1 << 20)} MiB written to {disk}")
    if proc.returncode != 0:
        raise RuntimeError(f"image write failed rc={proc.returncode}: {proc.stderr.read().strip()[-400:]}")
    got = readf(sumf).split()[0] if readf(sumf) else ""
    if got != want:
        raise RuntimeError(f"sha256 mismatch: got {got} want {want}")
    run(["sync"])
    run(["blockdev", "--rereadpt", disk])
    run(["partprobe", disk])
    run(["udevadm", "settle"])
    time.sleep(2)


def find_parts(disk):
    parts = [d for d in lsblk(disk) if d.get("type") == "part"]
    root = esp = None
    for p in sorted(parts, key=lambda d: -int(d["size"])):
        if p.get("fstype") in ROOT_FS and root is None:
            root = p
        if (p.get("parttype") or "").lower() == EFI_GUID and p.get("fstype") == "vfat" and esp is None:
            esp = p
    return root, esp, parts


def mount(dev, where, opts=None):
    os.makedirs(where, exist_ok=True)
    cmd = ["mount"] + (["-o", opts] if opts else []) + [dev, where]
    run(cmd, check=True)


def seed(root, files):
    mnt = "/mnt/rollcall-root"
    # xfs root written from an image can carry a duplicate UUID with nothing else; nouuid is harmless
    try:
        mount(root["path"], mnt)
    except RuntimeError:
        mount(root["path"], mnt, "nouuid")
    try:
        if not os.path.isdir(f"{mnt}/etc"):
            raise RuntimeError(f"{root['path']} has no /etc; not a root filesystem")
        for f in files:
            dst = os.path.join(mnt, f["path"].lstrip("/"))
            os.makedirs(os.path.dirname(dst), exist_ok=True)
            with open(dst, "w") as fh:
                fh.write(f["content"])
            os.chmod(dst, int(f.get("mode", "0644"), 8))
            log(f"wrote {f['path']}")
        # cloud-init's stale state would stop it from re-running on a reused image
        run(f"rm -rf {mnt}/var/lib/cloud/instance {mnt}/var/lib/cloud/instances/* {mnt}/var/lib/cloud/data/* 2>/dev/null")
        sel = readf(f"{mnt}/etc/selinux/config")
        if re.search(r"^SELINUX=(enforcing|permissive)", sel, re.M):
            open(f"{mnt}/.autorelabel", "w").close()
            log("SELinux enabled in image: scheduled autorelabel for first boot")
        os_release = readf(f"{mnt}/etc/os-release")
        m = re.search(r'^PRETTY_NAME="?([^"\n]*)', os_release, re.M)
        run(["sync"])
        return m.group(1) if m else ""
    finally:
        run(["umount", mnt])


def efi_entry(disk, esp, label):
    if not os.path.isdir("/sys/firmware/efi") or esp is None:
        return None
    mnt = "/mnt/rollcall-esp"
    try:
        mount(esp["path"], mnt, "ro")
    except RuntimeError as e:
        return f"esp mount failed: {e}"
    try:
        loader = None
        for name in ("shimx64.efi", "grubx64.efi"):
            for d in ("debian", "ubuntu", "rocky", "almalinux", "centos", "redhat", "fedora"):
                if os.path.exists(f"{mnt}/EFI/{d}/{name}"):
                    loader = f"\\EFI\\{d}\\{name}"
                    break
            if loader:
                break
        if not loader:
            hits = glob.glob(f"{mnt}/EFI/*/shimx64.efi") + glob.glob(f"{mnt}/EFI/*/grubx64.efi") + glob.glob(f"{mnt}/EFI/BOOT/BOOTX64.EFI")
            if hits:
                loader = "\\" + os.path.relpath(hits[0], mnt).replace("/", "\\")
        if not loader:
            return "no EFI loader found on ESP"
    finally:
        run(["umount", mnt])
    # drop stale entries with our label, then create a new one first in BootOrder
    cur = run(["efibootmgr"]).stdout
    for m in re.finditer(r"^Boot([0-9A-Fa-f]{4})\*?\s+(.*)$", cur, re.M):
        if m.group(2).strip().startswith(label):
            run(["efibootmgr", "-q", "-b", m.group(1), "-B"])
    p = run(["efibootmgr", "-c", "-d", disk, "-p", str(esp.get("partn")), "-L", label, "-l", loader])
    if p.returncode != 0:
        return f"efibootmgr failed: {p.stderr.strip()[-200:]}"
    return f"{label} -> {loader} (part {esp.get('partn')})"


def do_deploy(job, prog):
    spec, files = job["spec"], job.get("files") or []
    started = time.time()
    prog.set("preparing", 0, "selecting disk", now=True)
    disk = find_disk(spec)
    size = int(disk["size"])
    if int(spec.get("raw_size") or 0) > size:
        raise RuntimeError(f"{disk['path']} ({size} bytes) smaller than image ({spec['raw_size']} bytes)")
    log(f"target {disk['path']} serial={disk.get('serial')} size={size}")
    prog.set("preparing", 1, f"wiping {disk['path']}", now=True)
    release_disk(disk["path"])
    write_image(spec, disk["path"], prog)
    prog.set("seeding", 99, "writing cloud-init seed", now=True)
    root, esp, parts = find_parts(disk["path"])
    if root is None:
        raise RuntimeError("no root filesystem found after writing image: " + ", ".join(f"{p['name']}:{p.get('fstype')}" for p in parts))
    pretty = seed(root, files)
    efi = efi_entry(disk["path"], esp, f"rollcall {spec.get('os', '')}")
    if efi:
        log(f"efi: {efi}")
    details = {"disk": disk["path"], "serial": disk.get("serial"), "root": root["path"], "root_fs": root.get("fstype"),
               "esp": esp["path"] if esp else None, "efi": efi, "os_release": pretty, "firmware": "uefi" if os.path.isdir("/sys/firmware/efi") else "bios",
               "seconds": round(time.time() - started, 1)}
    return details


def do_release(job, prog):
    spec = job["spec"]
    prog.set("preparing", 0, "selecting disk", now=True)
    disk = find_disk(spec)
    path, size = disk["path"], int(disk["size"])
    prog.set("wiping", 5, f"wiping {path}", now=True)
    release_disk(path)
    mib = 64
    run(f"dd if=/dev/zero of={path} bs=1M count={mib} conv=fsync status=none")
    run(f"dd if=/dev/zero of={path} bs=1M count={mib} seek={max(0, size // (1 << 20) - mib)} conv=fsync status=none")
    if spec.get("wipe") == "full":
        p = run(["blkdiscard", "-f", path], timeout=3600)
        if p.returncode != 0:
            prog.set("wiping", 10, "full zero fill", now=True)
            run(f"dd if=/dev/zero of={path} bs=16M oflag=direct status=none", timeout=24 * 3600)
    run(["blockdev", "--rereadpt", path])
    return {"disk": path, "serial": disk.get("serial"), "wipe": spec.get("wipe")}


def main():
    global SERVER, TOKEN
    opts = cmdline()
    SERVER = (opts.get("server") or "").rstrip("/")
    TOKEN = opts.get("job") if isinstance(opts.get("job"), str) else ""
    if not SERVER or not TOKEN:
        log("rollcall.server= and rollcall.job= required")
        return 2
    log(f"fetching job from {SERVER}")
    try:
        job = api("GET", "/api/v1/jobs/self")
    except Exception as e:  # noqa: BLE001
        log(f"cannot fetch job: {e}; staying up")
        return 1
    log(f"job {job['id']} kind={job['kind']} spec={json.dumps({k: v for k, v in job['spec'].items() if k not in ('ssh_keys', 'user_data')})}")
    prog = Progress()
    try:
        details = do_deploy(job, prog) if job["kind"] == "deploy" else do_release(job, prog)
        resp = prog.finish("done", "finished", details=details)
    except Exception as e:  # noqa: BLE001
        log(f"FAILED: {e!r}")
        try:
            resp = prog.finish("failed", str(e), error=str(e))
        except Exception as e2:  # noqa: BLE001
            log(f"could not report failure: {e2}")
            resp = {"next": "stay"}
    nxt = resp.get("next", "stay")
    log(f"server says: {nxt}")
    if opts.get("nopoweroff") or nxt == "stay":
        log("staying up for inspection (ssh root@<ip>)")
        return 0 if nxt != "stay" else 1
    time.sleep(3)
    run(["systemctl", "reboot" if nxt == "reboot" else "poweroff"])
    return 0


if __name__ == "__main__":
    sys.exit(main())