#!/usr/bin/env python3
"""Apply a cluster definition's host_prep to the machine it runs on.

fermihdi-configure decided what this host needs and wrote it down;
validate-cluster checked that it is coherent. This applies it. Nothing here
asks a question — every decision was already made and recorded, which is what
makes the step repeatable, reviewable and runnable from Ansible.

It touches four things and nothing else:

  hugepages  the 2 MB pool, now, per NUMA node through sysfs
  mounts     a hugetlbfs mount for each pool that is actually in use
  vfio       the PCI devices whose driver is vfio, bound to vfio-pci
  grub       IOMMU, isolcpus and the 1 GB pool — the three only boot can give

Every change is idempotent: run it twice and the second run reports OK for
everything. It NEVER reboots. 1 GB pages, IOMMU and isolcpus all need one, and
when a host goes down is the operator's decision, not this script's.

Usage:
  install/fermihdi-prepare [cluster.yaml]              show what would change
  install/fermihdi-prepare [cluster.yaml] --apply      make the changes
  install/fermihdi-prepare [cluster.yaml] --verify     check the running host

A bare run changes nothing — it is a dry run, and --apply is the only thing
that writes. --verify is the pass to run after the reboot: it compares the
running kernel against the plan and says whether this host is ready to take
containers.

  --host NAME   apply the named host block instead of matching the hostname
  --no-grub     leave /etc/default/grub alone whatever the definition says

Exit status:  0 done (or nothing to do)   1 error   2 applied, reboot required
"""
import argparse
import os
import pathlib
import re
import shlex
import shutil
import socket
import subprocess
import sys
import tempfile
import time

try:
    import yaml
except ImportError:
    sys.exit("PyYAML is required:  apt-get install -y python3-yaml")

HERE = pathlib.Path(__file__).resolve().parent

# ── Where the kernel keeps the two pools ─────────────────────────────────────
# Always the per-node path, never /sys/kernel/mm/hugepages/.../nr_hugepages.
# The host-wide file lets the kernel decide which node each page comes from,
# and a reactor can only allocate from its own node — so a host total that
# looks right can still leave one instance with nothing.
HP_NODE = "/sys/devices/system/node/node{node}/hugepages/hugepages-{kb}kB/nr_hugepages"
KB_2M, KB_1G = 2048, 1048576

# Mount points, and they are not arbitrary. StorageEngine.cpp pins SPDK's EAL
# hugedir to /dev/hugepages the moment it finds 1 GB pages, so that EAL never
# eats the 2 MB networking pool. /dev/hugepages therefore has to BE the 1 GB
# mount, and the 2 MB pool needs one of its own. scripts/host_setup.sh uses the
# same pair; changing either breaks the SN at SPDK init.
MOUNT_1G = "/dev/hugepages"
MOUNT_2M = "/dev/hugepages2M"

# Generated artefacts. Everything this script installs is regenerated from
# cluster.yaml on the next --apply, so the file is the source of truth and
# these are derived — say so in each of them.
BOOT_HUGEPAGES = "/usr/local/sbin/fermihdi-hugepages"
BOOT_VFIO = "/usr/local/sbin/fermihdi-vfio-bind"
UNIT_DIR = "/etc/systemd/system"
MODULES_CONF = "/etc/modules-load.d/fermihdi-vfio.conf"
FSTAB = "/etc/fstab"
GRUB_FILE = "/etc/default/grub"
GRUB_KEY = "GRUB_CMDLINE_LINUX_DEFAULT"

# Kernel arguments this script owns. They are stripped and rewritten on every
# apply rather than appended to, so re-running after a cluster.yaml change
# leaves one isolcpus rather than three. Both IOMMU vendors are listed so that
# moving a definition between an Intel and an AMD box cleans up the stale one.
MANAGED_ARGS = {"intel_iommu", "amd_iommu", "iommu", "default_hugepagesz",
                "hugepagesz", "hugepages", "isolcpus", "rcu_nocbs", "nohz_full"}

# Matches configure and validate-cluster: kernel, page cache and everything not
# FermiHDI (4 GB), plus process stack and heap outside the engine arena (1 GB).
OS_HEADROOM_MB = 5120

APPLY = False
ERRORS = []
CHANGES = []
REBOOT = []


# ── Reporting ────────────────────────────────────────────────────────────────
def say(area, verdict, msg):
    print(f"  {area:<10} {verdict:<9} {msg}")


def ok(area, msg):
    say(area, "OK", msg)


def changed(area, msg):
    CHANGES.append(msg)
    say(area, "CHANGED", msg)


def pending(area, msg):
    say(area, "PENDING", msg)


def skip(area, msg):
    say(area, "SKIP", msg)


def warn(area, msg):
    say(area, "warning", msg)


def fail(area, msg):
    ERRORS.append(msg)
    say(area, "ERROR", msg)


def needs_reboot(reason):
    if reason not in REBOOT:
        REBOOT.append(reason)


def do(area, msg, fn):
    """Make a change, or in a dry run describe the one that would be made."""
    if not APPLY:
        say(area, "WOULD", msg)
        return None
    try:
        result = fn()
    except OSError as exc:
        fail(area, f"{msg}: {exc}")
        return None
    changed(area, msg)
    return result


# ── Small readers ────────────────────────────────────────────────────────────
def read_int(path, default=None):
    try:
        return int(pathlib.Path(path).read_text().strip())
    except (OSError, ValueError):
        return default


def write_text(path, text, mode=None):
    p = pathlib.Path(path)
    p.parent.mkdir(parents=True, exist_ok=True)
    p.write_text(text)
    if mode is not None:
        p.chmod(mode)


def run(cmd, **kw):
    return subprocess.run(cmd, capture_output=True, text=True, **kw)


def parse_cpuset(spec):
    """'2-7,18' -> {2,3,4,5,6,7,18}"""
    out = set()
    for part in str(spec).split(","):
        part = part.strip()
        if not part:
            continue
        if "-" in part:
            a, b = part.split("-", 1)
            out.update(range(int(a), int(b) + 1))
        else:
            out.add(int(part))
    return out


def numa_nodes():
    base = pathlib.Path("/sys/devices/system/node")
    nodes = sorted(int(p.name[4:]) for p in base.glob("node[0-9]*")) if base.is_dir() else []
    return nodes


def mem_total_mb():
    for line in pathlib.Path("/proc/meminfo").read_text().splitlines():
        if line.startswith("MemTotal:"):
            return int(line.split()[1]) // 1024
    return 0


def kernel_args():
    return pathlib.Path("/proc/cmdline").read_text().split()


def iommu_active():
    """IOMMU as the running kernel sees it, not as GRUB asks for it.

    /sys/kernel/iommu_groups is populated only once the IOMMU is actually up,
    which is the condition vfio-pci binding depends on — a host with
    intel_iommu=on in GRUB but no reboot yet looks configured and still cannot
    bind a thing.
    """
    groups = pathlib.Path("/sys/kernel/iommu_groups")
    return groups.is_dir() and any(groups.iterdir())


def iommu_flag():
    txt = pathlib.Path("/proc/cpuinfo").read_text()
    return "amd_iommu=on" if "AuthenticAMD" in txt else "intel_iommu=on"


# ── Host selection ───────────────────────────────────────────────────────────
def pick_host(cfg, name):
    hosts = cfg.get("hosts") or []
    if not hosts:
        sys.exit("cluster file defines no hosts")
    if name:
        for h in hosts:
            if h.get("name") == name:
                return h
        sys.exit(f"no host named {name!r} — the file has: "
                 + ", ".join(h.get("name", "?") for h in hosts))

    me = {socket.gethostname(), socket.gethostname().split(".")[0]}
    for h in hosts:
        hn = h.get("name", "")
        if hn in me or hn.split(".")[0] in me:
            return h
    if len(hosts) == 1:
        warn("host", f"this machine is {socket.gethostname()}, the definition names "
                     f"{hosts[0].get('name')!r} — applying it anyway, it is the only host")
        return hosts[0]
    sys.exit(f"this machine ({socket.gethostname()}) matches no host in the file. "
             f"Pass --host with one of: " + ", ".join(h.get("name", "?") for h in hosts))


# ── Hugepages ────────────────────────────────────────────────────────────────
def node_targets(total, per_node, key, nodes):
    """Per-node page targets for one pool.

    host_prep.numa says what each node must end up holding; hugepages_2m /
    hugepages_1g are the host totals, and the operator can raise those at the
    configure prompt. Spreading the surplus evenly is the honest reading of a
    raised total — reserving only the per-node sum would quietly hand back
    fewer pages than were asked for.
    """
    targets = {n: 0 for n in nodes}
    for entry in per_node or []:
        n = entry.get("node")
        if n in targets:
            targets[n] = entry.get(key, 0) or 0
        elif n is not None:
            warn("hugepages", f"host_prep.numa has an entry for node {n}, which this "
                              f"machine does not have — its pages go nowhere")
    surplus = total - sum(targets.values())
    if surplus > 0 and targets:
        order = sorted(targets)
        share, extra = divmod(surplus, len(order))
        for i, n in enumerate(order):
            targets[n] += share + (1 if i < extra else 0)
    return targets


def set_pool(kb, targets, label):
    """Raise a hugepage pool to its targets. Never lowers it.

    Lowering is the dangerous direction: pages already handed to a running
    reactor cannot be returned, so the write silently does less than it says
    and the operator is left with a number that does not match sysfs. A host
    that has more than it needs is not a problem worth creating one over.
    """
    for node in sorted(targets):
        want = targets[node]
        path = HP_NODE.format(node=node, kb=kb)
        have = read_int(path)
        if have is None:
            if want:
                fail("hugepages", f"{path} does not exist — cannot reserve {label} on node {node}")
            continue
        if want == 0:
            continue
        if have >= want:
            ok("hugepages", f"node {node} already has {have} x {label}")
            continue
        do("hugepages", f"node {node}: {have} -> {want} x {label}",
           lambda p=path, w=want: write_text(p, str(w)))
        if APPLY:
            got = read_int(path, 0)
            if got < want:
                # Runtime allocation is best-effort: it depends on finding that
                # much physically contiguous free memory, which a host that has
                # been up for a while may not have.
                warn("hugepages", f"node {node} got {got} of {want} x {label} — memory is "
                                  f"too fragmented to allocate the rest at run time")
                needs_reboot(f"{label} pool on node {node} is short "
                             f"({got}/{want}); a reboot reserves it from clean memory")


def hugepages(prep):
    want_2m = prep.get("hugepages_2m", 0) or 0
    want_1g = prep.get("hugepages_1g", 0) or 0
    per_node = prep.get("numa") or []
    nodes = numa_nodes() or [0]

    if not want_2m and not want_1g:
        skip("hugepages", "this host runs nothing that uses hugepages")
        return

    # Sanity against the machine in front of us rather than the memory_total_mb
    # recorded at configure time — hardware gets changed between the two.
    real_mb = mem_total_mb()
    want_mb = want_2m * 2 + want_1g * 1024
    usable = real_mb - OS_HEADROOM_MB
    if real_mb and want_mb > usable:
        fail("hugepages", f"the plan reserves {want_mb} MB of hugepages but this machine has "
                          f"{real_mb} MB, of which {usable} MB is spare. Reserving it would "
                          f"leave the OS short enough to OOM-kill a reactor during startup.")
        return  # and no boot unit either: do not persist a plan just refused

    if want_2m:
        # Written straight to sysfs, never through vm.nr_hugepages: that sysctl
        # applies to the DEFAULT page size, and this script sets the default to
        # 1G in GRUB whenever the 1 GB pool is in use. vm.nr_hugepages=8580 on
        # such a host asks for 8.5 TB.
        set_pool(KB_2M, node_targets(want_2m, per_node, "hugepages_2m", nodes), "2 MB")

    if want_1g:
        # Attempted at run time even though GRUB is the real answer: on a host
        # that has just booted it sometimes works, and it costs nothing when it
        # does not. StorageEngine's own error message tells operators to try
        # exactly this. Compared per node rather than against the host sum,
        # because 14 pages that all landed on node 1 satisfy a total of 14 and
        # still leave the instance on node 0 with nothing to DMA into.
        targets = node_targets(want_1g, per_node, "hugepages_1g", nodes)
        set_pool(KB_1G, targets, "1 GB")
        short = [n for n in sorted(targets)
                 if (read_int(HP_NODE.format(node=n, kb=KB_1G), 0) or 0) < targets[n]]
        if short:
            needs_reboot(f"the 1 GB pool is short on node(s) {short} — only boot-time "
                         f"reservation is reliable for 1 GB pages")

    boot_hugepages_unit(prep)


def boot_hugepages_unit(prep):
    """A oneshot that re-applies the 2 MB pool at boot.

    2 MB pages are deliberately not put in GRUB — the schema says so — which
    means nothing reserves them after a reboot unless something does it on the
    way up. sysctl cannot: it has no per-node form, and vm.nr_hugepages means
    the default size, which is 1G here.
    """
    want_2m = prep.get("hugepages_2m", 0) or 0
    if not want_2m:
        return
    nodes = numa_nodes() or [0]
    targets = node_targets(want_2m, prep.get("numa"), "hugepages_2m", nodes)
    lines = [f'set_hp "{HP_NODE.format(node=n, kb=KB_2M)}" {targets[n]}'
             for n in sorted(targets) if targets[n]]
    script = (
        "#!/bin/sh\n"
        "# Generated by fermihdi-prepare. Do not edit: change cluster.yaml and\n"
        "# re-run 'fermihdi-prepare --apply', which rewrites this file.\n"
        "#\n"
        "# Reserves the 2 MB hugepage pool for the HD engine and the DPDK mbuf pools.\n"
        "# The 1 GB pool comes from the kernel command line instead.\n"
        "set -e\n"
        "set_hp() {\n"
        "    [ -f \"$1\" ] || return 0\n"
        "    cur=$(cat \"$1\")\n"
        "    [ \"$cur\" -ge \"$2\" ] && return 0\n"
        "    echo \"$2\" > \"$1\" || echo \"fermihdi: could not reserve $2 in $1\" >&2\n"
        "}\n"
        + "\n".join(lines) + "\n"
    )
    unit = (
        "[Unit]\n"
        "Description=Reserve FermiHDI 2 MB hugepages\n"
        "After=systemd-modules-load.service\n"
        "Before=docker.service containerd.service\n"
        "\n"
        "[Service]\n"
        "Type=oneshot\n"
        "RemainAfterExit=yes\n"
        f"ExecStart={BOOT_HUGEPAGES}\n"
        "\n"
        "[Install]\n"
        "WantedBy=multi-user.target\n"
    )
    install_unit("hugepages", "fermihdi-hugepages", BOOT_HUGEPAGES, script, unit)


def install_unit(area, name, script_path, script, unit):
    """Write a generated boot script and its unit, and enable it. Idempotent."""
    unit_path = f"{UNIT_DIR}/{name}.service"

    def unchanged(path, want):
        p = pathlib.Path(path)
        return p.exists() and p.read_text() == want

    if not shutil.which("systemctl"):
        warn(area, f"no systemd here, so {script_path} is written but nothing will run it "
                   f"at boot — wire it into whatever init this host uses")
        if not unchanged(script_path, script):
            do(area, f"write {script_path}", lambda: write_text(script_path, script, 0o755))
        return

    enabled = run(["systemctl", "is-enabled", f"{name}.service"]).stdout.strip() == "enabled"
    if unchanged(script_path, script) and unchanged(unit_path, unit) and enabled:
        ok(area, f"{name}.service is installed and enabled")
        return

    def install():
        write_text(script_path, script, 0o755)
        write_text(unit_path, unit)
        run(["systemctl", "daemon-reload"])
        r = run(["systemctl", "enable", f"{name}.service"])
        if r.returncode != 0:
            raise OSError(r.stderr.strip() or "systemctl enable failed")

    do(area, f"install and enable {name}.service", install)


# ── hugetlbfs mounts ─────────────────────────────────────────────────────────
def pagesize_kb(text):
    """'2M' / '1024M' / '1G' -> kB. hugetlbfs reports whichever the kernel likes."""
    m = re.match(r"^(\d+)([KMG])?$", text or "")
    if not m:
        return None
    n = int(m.group(1))
    return {"K": n, "M": n * 1024, "G": n * 1024 * 1024, None: n // 1024}[m.group(2)]


def hugetlbfs_mounts():
    out = {}
    for line in pathlib.Path("/proc/mounts").read_text().splitlines():
        f = line.split()
        if len(f) > 3 and f[2] == "hugetlbfs":
            opts = dict(o.split("=", 1) for o in f[3].split(",") if "=" in o)
            out[f[1]] = pagesize_kb(opts.get("pagesize", ""))
    return out


def mounts(prep):
    want_2m = prep.get("hugepages_2m", 0) or 0
    want_1g = prep.get("hugepages_1g", 0) or 0
    if not want_2m and not want_1g:
        return
    mounted = hugetlbfs_mounts()

    if want_2m:
        have = mounted.get(MOUNT_2M)
        if have == KB_2M:
            ok("mounts", f"{MOUNT_2M} is mounted (pagesize=2M)")
        elif have is not None:
            fail("mounts", f"{MOUNT_2M} is mounted with a {have} kB page size, not 2 MB — "
                           f"unmount it and re-run")
        else:
            # No size= cap: the pool size is set in sysfs, and capping the mount
            # as well just adds a second number that has to be kept in step.
            do("mounts", f"mount hugetlbfs at {MOUNT_2M} (pagesize=2M)",
               lambda: mount_hugetlbfs(MOUNT_2M, "2M"))
        fstab_entry(MOUNT_2M, "2M")

    if want_1g:
        have = mounted.get(MOUNT_1G)
        if have == KB_1G:
            ok("mounts", f"{MOUNT_1G} is mounted (pagesize=1G) — SPDK's EAL hugedir")
        elif have is None:
            do("mounts", f"mount hugetlbfs at {MOUNT_1G} (pagesize=1G)",
               lambda: mount_hugetlbfs(MOUNT_1G, "1G"))
        elif (prep.get("grub") or {}).get("apply", True):
            # systemd mounts /dev/hugepages at the default page size, so this
            # corrects itself once default_hugepagesz=1G is on the command line.
            # Remounting it now would mean unmounting a filesystem something may
            # already hold open, to gain nothing before the reboot: there are no
            # 1 GB pages to put in it yet either.
            pending("mounts", f"{MOUNT_1G} has a {have} kB page size; default_hugepagesz=1G "
                              f"makes it 1 GB at the next boot")
            needs_reboot(f"{MOUNT_1G} must be the 1 GB mount — SPDK pins EAL to it")
        else:
            fail("mounts", f"{MOUNT_1G} has a {have} kB page size and must be 1 GB "
                           f"(StorageEngine pins SPDK's EAL hugedir to it), but grub.apply "
                           f"is false so default_hugepagesz cannot be set. Enable it, or "
                           f"mount {MOUNT_1G} with pagesize=1G by hand.")


def mount_hugetlbfs(where, pagesize):
    pathlib.Path(where).mkdir(parents=True, exist_ok=True)
    r = run(["mount", "-t", "hugetlbfs", "-o", f"pagesize={pagesize}", "hugetlbfs", where])
    if r.returncode != 0:
        raise OSError(r.stderr.strip() or "mount failed")


def fstab_entry(where, pagesize):
    line = f"hugetlbfs {where} hugetlbfs pagesize={pagesize} 0 0"
    fstab = pathlib.Path(FSTAB)
    current = fstab.read_text() if fstab.exists() else ""
    if any(l.split()[1:2] == [where] for l in current.splitlines()
           if l.strip() and not l.startswith("#")):
        ok("mounts", f"{FSTAB} already mounts {where} at boot")
        return
    do("mounts", f"add {where} to {FSTAB}",
       lambda: fstab.write_text(current.rstrip("\n") + "\n"
                                + "# FermiHDI 2 MB hugetlbfs — added by fermihdi-prepare\n"
                                + line + "\n"))


# ── VFIO ─────────────────────────────────────────────────────────────────────
def pci_dir(addr):
    return pathlib.Path(f"/sys/bus/pci/devices/{addr}")


def current_driver(addr):
    link = pci_dir(addr) / "driver"
    return os.path.basename(os.readlink(link)) if link.is_symlink() else None


def iommu_group_of(addr):
    link = pci_dir(addr) / "iommu_group"
    return os.path.basename(os.readlink(link)) if link.is_symlink() else None


def block_devices_of(addr):
    """Block devices sitting behind this PCI function, e.g. nvme0n1 for its controller."""
    found = []
    for blk in pathlib.Path("/sys/block").glob("*"):
        try:
            real = str((blk / "device").resolve())
        except OSError:
            continue
        if f"/{addr}/" in real + "/":
            found.append(blk.name)
    return found


def in_use_paths(names):
    """Anything mounted or swapped-on that lives on one of these devices."""
    used = []
    sources = []
    for line in pathlib.Path("/proc/mounts").read_text().splitlines():
        f = line.split()
        if len(f) > 1:
            sources.append((f[0], f"mounted at {f[1]}"))
    if pathlib.Path("/proc/swaps").exists():
        for line in pathlib.Path("/proc/swaps").read_text().splitlines()[1:]:
            f = line.split()
            if f:
                sources.append((f[0], "in use as swap"))
    for src, why in sources:
        if not src.startswith("/dev/"):
            continue
        base = os.path.basename(src)
        for name in names:
            # The device itself or one of its partitions. Plain startswith would
            # read /dev/sdaa as a partition of /dev/sda and refuse to bind a
            # device nothing is using.
            if base == name or re.match(rf"^{re.escape(name)}p?\d+$", base):
                used.append(f"{src} is {why}")
    return used


def net_ifaces_of(addr):
    net = pci_dir(addr) / "net"
    return sorted(p.name for p in net.iterdir()) if net.is_dir() else []


def default_route_iface():
    for line in pathlib.Path("/proc/net/route").read_text().splitlines()[1:]:
        f = line.split()
        if len(f) > 2 and f[1] == "00000000":
            return f[0]
    return None


def ssh_local_ip():
    """The address this session came in on, when there is one.

    Unbinding the NIC that carries the SSH session is a mistake nobody gets to
    correct remotely, and the default route is not always the one in play — a
    management network on a second NIC is exactly the case that bites.
    """
    conn = os.environ.get("SSH_CONNECTION", "").split()
    return conn[2] if len(conn) >= 3 else None


def ifaces_holding(ip):
    if not ip or not shutil.which("ip"):
        return set()
    r = run(["ip", "-o", "addr", "show"])
    out = set()
    for line in r.stdout.splitlines():
        f = line.split()
        if len(f) > 3 and f[3].split("/")[0] == ip:
            out.add(f[1])
    return out


def refusals(addr):
    """Reasons this device must not be taken from the kernel right now.

    None of these are overridable. Every one of them means something is using
    the device this second, and vfio-pci binding is not a request the kernel
    can decline politely on their behalf.
    """
    reasons = []
    blocks = block_devices_of(addr)
    if blocks:
        reasons += in_use_paths(blocks)
    for ifn in net_ifaces_of(addr):
        if ifn == default_route_iface():
            reasons.append(f"{ifn} carries the default route — binding it cuts this host off")
        if ifn in ifaces_holding(ssh_local_ip()):
            reasons.append(f"{ifn} holds the address this SSH session arrived on")
    return reasons


def bind_one(addr):
    """driver_override, not new_id.

    new_id binds every device with the same vendor:device id, which on a box
    whose boot drive is the same model as its data drives takes the boot drive
    too. driver_override names one function and nothing else.
    """
    dev = pci_dir(addr)
    cur = current_driver(addr)
    write_text(dev / "driver_override", "vfio-pci")
    if cur:
        write_text(f"/sys/bus/pci/drivers/{cur}/unbind", addr)
    write_text("/sys/bus/pci/drivers_probe", addr)
    time.sleep(0.5)
    now = current_driver(addr)
    if now != "vfio-pci":
        raise OSError(f"driver is still {now or 'none'!r} after the bind")


def vfio(prep):
    binds = prep.get("vfio_bind") or []
    if not binds:
        skip("vfio", "nothing on this host is bound to vfio-pci")
        return

    if not prep.get("iommu"):
        fail("vfio", "host_prep lists devices to bind but iommu is false — vfio-pci "
                     "cannot claim a device without an IOMMU")
        return

    if APPLY and not run(["modprobe", "vfio-pci"]).returncode == 0:
        warn("vfio", "modprobe vfio-pci failed — the module may be built in")
    modules = ("# Generated by fermihdi-prepare — load vfio-pci before the\n"
               "# bind unit runs.\nvfio-pci\n")
    if pathlib.Path(MODULES_CONF).exists() and pathlib.Path(MODULES_CONF).read_text() == modules:
        ok("vfio", f"{MODULES_CONF} already loads vfio-pci at boot")
    else:
        do("vfio", f"list vfio-pci in {MODULES_CONF} so it loads at boot",
           lambda: write_text(MODULES_CONF, modules))

    live = iommu_active()
    if not live:
        # Not an error when GRUB is about to fix it: this is the ordinary first
        # run on a host that has never had IOMMU turned on.
        pending("vfio", "the IOMMU is not up in the running kernel, so nothing can be bound "
                        "yet — the devices below are bound at the next boot")
        needs_reboot("the IOMMU has to be enabled before vfio-pci can claim a device")

    for addr in binds:
        if not pci_dir(addr).is_dir():
            fail("vfio", f"{addr} is not present on this machine")
            continue
        cur = current_driver(addr)
        if cur == "vfio-pci":
            grp = iommu_group_of(addr)
            ok("vfio", f"{addr} is on vfio-pci (IOMMU group {grp}, /dev/vfio/{grp})")
            continue
        stop = refusals(addr)
        if stop:
            for reason in stop:
                fail("vfio", f"{addr}: {reason}")
            fail("vfio", f"{addr} was NOT bound. Release it first — this script will not "
                         f"take a device out from under something that is using it.")
            continue
        if not live:
            pending("vfio", f"{addr}: {cur or 'no driver'} -> vfio-pci, after the reboot")
            continue
        do("vfio", f"{addr}: {cur or 'no driver'} -> vfio-pci", lambda a=addr: bind_one(a))

    boot_vfio_unit(binds)


def boot_vfio_unit(binds):
    """vfio-pci binding does not survive a reboot; driver_override lives in sysfs."""
    script = (
        "#!/bin/sh\n"
        "# Generated by fermihdi-prepare. Do not edit: change cluster.yaml and\n"
        "# re-run 'fermihdi-prepare --apply', which rewrites this file.\n"
        "#\n"
        "# Hands the data-plane devices to vfio-pci before anything that wants\n"
        "# them starts. driver_override names one PCI function, unlike new_id,\n"
        "# which would claim every device sharing the same vendor:device id.\n"
        "modprobe vfio-pci 2>/dev/null || true\n"
        "for addr in " + " ".join(binds) + "; do\n"
        "    dev=/sys/bus/pci/devices/$addr\n"
        "    if [ ! -d \"$dev\" ]; then\n"
        "        echo \"fermihdi: $addr is not present\" >&2\n"
        "        continue\n"
        "    fi\n"
        "    cur=\"\"\n"
        "    [ -L \"$dev/driver\" ] && cur=$(basename \"$(readlink \"$dev/driver\")\")\n"
        "    [ \"$cur\" = vfio-pci ] && continue\n"
        "    echo vfio-pci > \"$dev/driver_override\"\n"
        "    [ -n \"$cur\" ] && echo \"$addr\" > \"/sys/bus/pci/drivers/$cur/unbind\"\n"
        "    echo \"$addr\" > /sys/bus/pci/drivers_probe\n"
        "done\n"
    )
    unit = (
        "[Unit]\n"
        "Description=Bind FermiHDI data-plane devices to vfio-pci\n"
        "After=systemd-modules-load.service\n"
        "Before=docker.service containerd.service\n"
        "\n"
        "[Service]\n"
        "Type=oneshot\n"
        "RemainAfterExit=yes\n"
        f"ExecStart={BOOT_VFIO}\n"
        "\n"
        "[Install]\n"
        "WantedBy=multi-user.target\n"
    )
    install_unit("vfio", "fermihdi-vfio-bind", BOOT_VFIO, script, unit)


# ── GRUB ─────────────────────────────────────────────────────────────────────
def grub_tokens(prep):
    t = []
    if prep.get("iommu"):
        t += [iommu_flag(), "iommu=pt"]
    hp1g = prep.get("hugepages_1g", 0) or 0
    if hp1g:
        # default_hugepagesz matters as much as the count. systemd mounts
        # /dev/hugepages at whatever the default size is, and that mount is the
        # one StorageEngine pins SPDK's EAL to; leave the default at 2 MB and
        # the SN's hugedir points straight at the networking pool.
        t += ["default_hugepagesz=1G", "hugepagesz=1G", f"hugepages={hp1g}"]
    iso = prep.get("isolcpus")
    if iso:
        # rcu_nocbs and nohz_full alongside isolcpus: taking a core off the
        # scheduler still leaves it servicing RCU callbacks and the tick, and
        # both of those are exactly the microsecond stalls a poll-mode thread
        # cannot absorb.
        t += [f"isolcpus={iso}", f"rcu_nocbs={iso}", f"nohz_full={iso}"]
    t += list((prep.get("grub") or {}).get("extra_args") or [])
    return t


def rewrite_cmdline(existing, want):
    keys = MANAGED_ARGS | {t.split("=", 1)[0] for t in want}
    kept = [t for t in shlex.split(existing) if t.split("=", 1)[0] not in keys]
    return " ".join(kept + want)


def grub(prep, no_grub):
    want = grub_tokens(prep)
    if not want:
        skip("grub", "nothing here needs a kernel argument")
        return

    running = kernel_args()
    missing = [t for t in want if t not in running]
    apply_grub = (prep.get("grub") or {}).get("apply", True)

    if no_grub or not apply_grub:
        why = "--no-grub was given" if no_grub else "grub.apply is false"
        if missing:
            fail("grub", f"the running kernel is missing {' '.join(missing)} and {why}. "
                         f"Add them by hand, or set grub.apply: true and re-run.")
        else:
            ok("grub", f"the running kernel has every argument this host needs ({why}, "
                       f"so {GRUB_FILE} is left alone)")
        return

    path = pathlib.Path(GRUB_FILE)
    if not path.exists():
        fail("grub", f"{GRUB_FILE} does not exist — this host does not use GRUB, so "
                     f"{' '.join(want)} has to reach the kernel some other way")
        return

    # The file is checked even when the running kernel is already correct. A
    # host that boots right today because someone typed the arguments in once,
    # or whose /etc/default/grub was replaced by an image update, loses them
    # silently at the next reboot — and that reboot is usually unattended.
    lines = path.read_text().splitlines()
    current = ""
    index = None
    for i, line in enumerate(lines):
        m = re.match(rf"^\s*{GRUB_KEY}\s*=\s*(.*?)\s*$", line)
        if m:
            index = i
            current = m.group(1)
            if len(current) > 1 and current[0] in "\"'" and current[-1] == current[0]:
                current = current[1:-1]
    new = rewrite_cmdline(current, want)

    if new == current:
        ok("grub", f"{GRUB_FILE} already asks for every argument this host needs")
    else:
        say("grub", "DIFF", f'-{GRUB_KEY}="{current}"')
        say("grub", "DIFF", f'+{GRUB_KEY}="{new}"')

        def edit():
            backup = f"{GRUB_FILE}.fermihdi-{time.strftime('%Y%m%d%H%M%S')}.bak"
            shutil.copy2(GRUB_FILE, backup)
            out = list(lines)
            if index is None:
                out.append(f'{GRUB_KEY}="{new}"')
            else:
                out[index] = f'{GRUB_KEY}="{new}"'
            path.write_text("\n".join(out) + "\n")
            try:
                regenerate_grub_config()
            except OSError:
                shutil.copy2(backup, GRUB_FILE)
                raise
            return backup

        backup = do("grub", f"rewrite {GRUB_KEY} and regenerate the boot config", edit)
        if backup:
            ok("grub", f"the previous file is kept at {backup}")

    if missing:
        needs_reboot(f"the kernel is running without {' '.join(missing)}")
    else:
        ok("grub", "the running kernel already has every argument this host needs")


def regenerate_grub_config():
    """Generate to a scratch file first, and only install a config that built.

    A grub.cfg that fails to generate but gets installed anyway is an
    unbootable machine, and the failure mode is a remote host that never comes
    back. Generate, check, then install; restore the backup on any failure.
    """
    mkconfig = shutil.which("grub-mkconfig") or shutil.which("grub2-mkconfig")
    if not mkconfig:
        raise OSError("neither grub-mkconfig nor grub2-mkconfig is installed")
    with tempfile.NamedTemporaryFile(suffix=".cfg") as tmp:
        r = run([mkconfig, "-o", tmp.name])
        if r.returncode != 0:
            raise OSError("grub-mkconfig failed on the new file: "
                          + (r.stderr.strip().splitlines() or ["no output"])[-1])
        if os.path.getsize(tmp.name) < 512:
            raise OSError("grub-mkconfig produced an implausibly small config")
    target = "/boot/grub/grub.cfg" if pathlib.Path("/boot/grub").is_dir() \
        else "/boot/grub2/grub.cfg"
    updater = shutil.which("update-grub")
    r = run([updater]) if updater else run([mkconfig, "-o", target])
    if r.returncode != 0:
        raise OSError("installing the new boot config failed: " + r.stderr.strip())


# ── Verify ───────────────────────────────────────────────────────────────────
def verify(host, prep):
    """Compare the running kernel with the plan. This is the post-reboot pass."""
    want = grub_tokens(prep)
    running = kernel_args()
    for token in want:
        if token in running:
            ok("cmdline", token)
        else:
            fail("cmdline", f"{token} is not on the running kernel command line")

    if prep.get("iommu"):
        if iommu_active():
            ok("iommu", "IOMMU groups are present")
        else:
            fail("iommu", "no IOMMU groups — vfio-pci cannot claim anything")

    nodes = numa_nodes() or [0]
    for kb, key, label in ((KB_2M, "hugepages_2m", "2 MB"), (KB_1G, "hugepages_1g", "1 GB")):
        want_total = prep.get(key, 0) or 0
        if not want_total:
            continue
        targets = node_targets(want_total, prep.get("numa"), key, nodes)
        for node in sorted(targets):
            if not targets[node]:
                continue
            have = read_int(HP_NODE.format(node=node, kb=kb), 0) or 0
            free = read_int(HP_NODE.format(node=node, kb=kb).replace(
                "nr_hugepages", "free_hugepages"), have) or 0
            if have >= targets[node]:
                ok("hugepages", f"node {node}: {have} x {label} ({free} free)")
            else:
                fail("hugepages", f"node {node} has {have} x {label}, needs {targets[node]}")

    mounted = hugetlbfs_mounts()
    if (prep.get("hugepages_2m", 0) or 0) and mounted.get(MOUNT_2M) != KB_2M:
        fail("mounts", f"{MOUNT_2M} is not mounted with a 2 MB page size — EAL aborts with "
                       f"\"no mounted hugetlbfs found for that size\"")
    elif prep.get("hugepages_2m"):
        ok("mounts", f"{MOUNT_2M} (pagesize=2M)")
    if (prep.get("hugepages_1g", 0) or 0) and mounted.get(MOUNT_1G) != KB_1G:
        fail("mounts", f"{MOUNT_1G} is not the 1 GB mount — SPDK pins its EAL hugedir there "
                       f"and will find the wrong pool")
    elif prep.get("hugepages_1g"):
        ok("mounts", f"{MOUNT_1G} (pagesize=1G)")

    for addr in prep.get("vfio_bind") or []:
        drv = current_driver(addr)
        grp = iommu_group_of(addr)
        if drv != "vfio-pci":
            fail("vfio", f"{addr} is on {drv or 'no driver'}, not vfio-pci")
        elif grp and not pathlib.Path(f"/dev/vfio/{grp}").exists():
            fail("vfio", f"{addr} is bound but /dev/vfio/{grp} does not exist")
        else:
            ok("vfio", f"{addr} on vfio-pci (/dev/vfio/{grp})")

    # The instances are not this script's business except here: cores that were
    # promised to a poll-mode thread and left on the scheduler are a stall that
    # shows up as latency nobody can source months later.
    iso = parse_cpuset(prep["isolcpus"]) if prep.get("isolcpus") else set()
    if iso:
        for role in host.get("roles", []):
            for inst in role.get("instances", []) or []:
                cores = parse_cpuset(inst["cpuset"])
                if not cores <= iso:
                    warn("isolcpus", f"{inst['name']} runs on {inst['cpuset']}, which is not "
                                     f"fully isolated — {sorted(cores - iso)} stay on the "
                                     f"scheduler")


# ── Entry point ──────────────────────────────────────────────────────────────
def main():
    global APPLY
    ap = argparse.ArgumentParser(
        prog="fermihdi-prepare", add_help=True,
        description="Apply a cluster definition's host_prep to this machine.")
    ap.add_argument("cluster", nargs="?", default=str(HERE / "cluster.yaml"),
                    help="cluster definition (default: install/cluster.yaml)")
    ap.add_argument("--host", metavar="NAME",
                    help="host block to use instead of matching this machine's hostname")
    ap.add_argument("--apply", action="store_true", help="make the changes")
    ap.add_argument("--verify", action="store_true",
                    help="check the running host against the plan and change nothing")
    ap.add_argument("--no-grub", action="store_true",
                    help="leave /etc/default/grub alone whatever the definition says")
    args = ap.parse_args()

    path = pathlib.Path(args.cluster)
    if not path.exists():
        sys.exit(f"no such file: {path}")
    cfg = yaml.safe_load(path.read_text())
    host = pick_host(cfg, args.host)
    prep = host.get("host_prep") or {}

    APPLY = args.apply and not args.verify
    if APPLY and os.geteuid() != 0:
        sys.exit("--apply changes hugepages, device drivers and GRUB: run it as root.")

    mode = "verify" if args.verify else ("apply" if APPLY else "dry run — nothing is changed")
    print()
    print(f"  host       {host.get('name')}   cluster {cfg.get('cluster_name', '-')}")
    print(f"  source     {path}")
    print(f"  mode       {mode}")
    print()

    if not prep:
        print("  nothing to do — this host has no host_prep block.")
        print()
        return 0

    if args.verify:
        verify(host, prep)
    else:
        hugepages(prep)
        mounts(prep)
        vfio(prep)
        grub(prep, args.no_grub)

    print()
    if ERRORS:
        print(f"  {len(ERRORS)} error(s) — this host is not ready.")
        print()
        return 1
    if args.verify:
        print("  This host matches its definition and is ready to take containers.")
        print()
        return 0
    if not APPLY:
        print("  Nothing was changed. Re-run with --apply to make it so.")
        print()
        return 0

    print(f"  {len(CHANGES)} change(s) applied.")
    if REBOOT:
        print()
        print("  A REBOOT is required before this host can run FermiHDI:")
        for reason in REBOOT:
            print(f"    - {reason}")
        print()
        print("  Nothing here reboots for you. When the host is back:")
        print(f"    {sys.argv[0]} {path} --verify")
        print()
        return 2
    print()
    print("  No reboot needed. Confirm with:")
    print(f"    {sys.argv[0]} {path} --verify")
    print()
    return 0


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