#!/usr/bin/env python3
"""Turn a cluster definition into Talos machine configuration, and apply it.

This is fermihdi-prepare for machines with no shell. It applies the same
host_prep block to the same effect, and almost nothing about how is the same,
because Talos splits that block in two along a line no other distribution has:

  machine configuration          the Talos API, applied while the node runs
    2 MB hugepages, per node     machine.sysfs
    kernel modules               machine.kernel.modules
    vfio-pci binding             PCIDriverRebindConfig, one document per device
    CPU manager and NUMA policy  machine.kubelet.extraConfig

  the boot image                 an Image Factory schematic, applied by upgrade
    IOMMU                        intel_iommu=on / amd_iommu=on, iommu=pt
    isolcpus, rcu_nocbs, nohz_full
    the 1 GB hugepage pool       default_hugepagesz=1G hugepagesz=1G hugepages=N

The second half is the surprise. Since Talos 1.10 a UEFI machine boots a UKI,
and the kernel command line is inside it:

    "With systemd-boot, the .machine.install.extraKernelArgs field in the
     machine configuration is ignored, as kernel arguments are embedded in the
     UKI and cannot be modified without upgrading the UKI."

So there is no file to edit and no GRUB to regenerate. The arguments go into a
schematic, the schematic is uploaded and comes back as a content-addressed ID,
the ID names an installer image, and the node is upgraded to it. That upgrade
is a reboot -- which is why --apply on Talos does something --apply on bare
metal never did, and why it does it to one machine at a time.

The kernel argument list itself is not written twice: grub_tokens() is imported
from fermihdi-prepare, so the two paths cannot drift.

Usage:
  install/fermihdi-talos-prepare [cluster.yaml]            show what would happen
  install/fermihdi-talos-prepare [cluster.yaml] --apply    make it so
  install/fermihdi-talos-prepare [cluster.yaml] --verify   check running machines

  --host NAME      one host instead of every host in the file
  --render-only    write the machine configs and stop; talk to nothing
  --no-upgrade     apply machine configuration but never upgrade, so nothing
                   reboots. The boot-time half of host_prep stays unapplied and
                   is reported as outstanding.
  --out DIR        override talos.out_dir

A bare run changes nothing. --apply is the only thing that writes, and on
Talos it writes to the machine's boot loader as well as its configuration.

Exit status:  0 done   1 error   2 applied, an upgrade or reboot is outstanding
"""
import argparse
import hashlib
import importlib.machinery
import importlib.util
import json
import os
import pathlib
import re
import shutil
import subprocess
import sys
import urllib.error
import urllib.request

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

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


def _load(name, filename):
    """Import a sibling tool that has no .py extension and is not on the path."""
    spec = importlib.util.spec_from_loader(
        name, importlib.machinery.SourceFileLoader(name, str(HERE / filename)))
    mod = importlib.util.module_from_spec(spec)
    spec.loader.exec_module(mod)
    return mod


# The kernel argument list and the sizing model both live elsewhere and are
# imported rather than copied. A second table of either is a second one to get
# wrong, and the failure would be silent: a Talos host booting with different
# arguments from its bare-metal twin.
_prep = _load("fermihdi_prepare", "fermihdi-prepare")
_vc = _load("fermihdi_validate", "validate-cluster")

DEFAULT_FACTORY = "https://factory.talos.dev"
# Where the Image Factory publishes installer images for bare metal. The path
# was `installer/` before Talos 1.10 and is `metal-installer/` now.
INSTALLER_PATH = "metal-installer"

APPLY = False
CHANGES, ERRORS, OUTSTANDING = [], [], []


# ── Reporting ────────────────────────────────────────────────────────────────
# Deliberately the same vocabulary fermihdi-prepare uses: an operator reading
# both should not have to learn two.
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 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 outstanding(reason):
    if reason not in OUTSTANDING:
        OUTSTANDING.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 Exception as exc:                      # noqa: BLE001 — reported, not raised
        fail(area, f"{msg}: {exc}")
        return None
    changed(area, msg)
    return result


# ── talosctl ─────────────────────────────────────────────────────────────────
def talosctl(args, node=None, talosconfig=None, insecure=False,
             check=False, timeout=1800):
    cmd = ["talosctl"]
    if node:
        cmd += ["--nodes", node]
    if talosconfig:
        cmd += ["--talosconfig", str(talosconfig)]
    cmd += list(args)
    if insecure:
        cmd.append("--insecure")
    try:
        p = subprocess.run(cmd, capture_output=True, text=True, timeout=timeout)
    except FileNotFoundError:
        raise RuntimeError("talosctl is not on PATH")
    except subprocess.TimeoutExpired:
        raise RuntimeError(f"timed out after {timeout}s: {' '.join(cmd)}")
    if check and p.returncode != 0:
        raise RuntimeError((p.stderr.strip() or p.stdout.strip() or
                            f"exit {p.returncode}").splitlines()[0])
    return p


def node_addr(host):
    """Where talosctl reaches this machine.

    Separate from `address` because the Talos API and the control plane CIDR
    need not be the same network, and on a machine being installed they very
    often are not.
    """
    return (host.get("talos") or {}).get("node") or host.get("address")


def get_resource(host, resource, talosconfig=None, insecure=False):
    """One typed resource as a list of dicts, or [] if it cannot be had."""
    try:
        p = talosctl(["get", resource, "-o", "json"], node=node_addr(host),
                     talosconfig=talosconfig, insecure=insecure, timeout=60)
    except RuntimeError:
        return []
    if p.returncode != 0:
        return []
    items, buf, depth = [], "", 0
    for ch in p.stdout:
        buf += ch
        if ch == "{":
            depth += 1
        elif ch == "}":
            depth -= 1
            if depth == 0:
                try:
                    items.append(json.loads(buf))
                except json.JSONDecodeError:
                    pass
                buf = ""
    return items


# ═════════════════════════════════════════════════════════════════════════════
# The boot image: an Image Factory schematic
# ═════════════════════════════════════════════════════════════════════════════
def iommu_flag_for(host, talosconfig=None):
    """intel_iommu=on or amd_iommu=on, which is not a guess worth making.

    fermihdi-prepare reads /proc/cpuinfo, which it can do because it runs on
    the machine. Here the machine is asked over the API instead; SMBIOS names
    the CPU manufacturer and is answered even in maintenance mode. If the
    machine cannot be reached, this refuses rather than picking one: the wrong
    flag leaves the IOMMU off, and a host with the IOMMU off looks configured
    and still cannot bind a device.
    """
    prep = host.get("host_prep") or {}
    explicit = [a for a in ((prep.get("grub") or {}).get("extra_args") or [])
                if re.match(r"^(intel|amd)_iommu=", a)]
    if explicit:
        return explicit[0]

    # Secure first, then insecure. A machine that has not been configured yet
    # is exactly the one whose schematic is being built for the first time, and
    # it answers `get cpu` perfectly well over the maintenance service.
    for insecure in (False, True):
        for p in get_resource(host, "cpu", talosconfig, insecure=insecure):
            maker = (p.get("spec", {}) or {}).get("manufacturer", "")
            if re.search(r"amd|advanced micro", maker, re.I):
                return "amd_iommu=on"
            if maker:
                return "intel_iommu=on"

    raise RuntimeError(
        f"cannot tell whether {host.get('name')} is Intel or AMD, so cannot "
        f"choose between intel_iommu=on and amd_iommu=on. Either bring the "
        f"machine up so `talosctl -n {node_addr(host)} get cpu` answers, or "
        f"put the flag in host_prep.grub.extra_args yourself")


def kernel_args_for(host, talosconfig=None):
    """The boot-time half of host_prep, as kernel command line arguments.

    grub_tokens() is fermihdi-prepare's, so the same host_prep produces the
    same arguments whether it lands in GRUB or in a UKI. Only the vendor probe
    differs, and that is patched in rather than reimplemented.
    """
    prep = host.get("host_prep") or {}
    if not (prep.get("grub") or {}).get("apply", True):
        return []
    saved = _prep.iommu_flag
    try:
        _prep.iommu_flag = lambda: iommu_flag_for(host, talosconfig)
        return _prep.grub_tokens(prep)
    finally:
        _prep.iommu_flag = saved


def schematic_for(host, cfg, talosconfig=None):
    talos = cfg.get("talos") or {}
    exts = list(talos.get("extensions") or []) + \
        list((host.get("talos") or {}).get("extensions") or [])
    doc = {"customization": {}}
    args = kernel_args_for(host, talosconfig)
    if args:
        doc["customization"]["extraKernelArgs"] = args
    if exts:
        doc["customization"]["systemExtensions"] = {
            "officialExtensions": sorted(dict.fromkeys(exts))}
    return doc


def schematic_text(doc):
    return yaml.safe_dump(doc, sort_keys=False, default_flow_style=False)


def upload_schematic(text, factory_url, cache_dir):
    """POST a schematic and get its ID back.

    The ID is a hash of the content, so the same schematic always resolves to
    the same ID and uploading it twice is free. That is also why the answer is
    cached on disk: a re-run of a dry run should not need the network to tell
    the operator which image a machine would move to.
    """
    digest = hashlib.sha256(text.encode()).hexdigest()
    cache = cache_dir / "schematics" / f"{digest}.id"
    if cache.exists():
        return cache.read_text().strip(), True

    req = urllib.request.Request(
        factory_url.rstrip("/") + "/schematics",
        data=text.encode(), method="POST",
        headers={"Content-Type": "application/yaml"})
    try:
        with urllib.request.urlopen(req, timeout=60) as resp:
            body = json.loads(resp.read().decode())
    except urllib.error.URLError as exc:
        raise RuntimeError(
            f"could not reach the Image Factory at {factory_url}: {exc}. "
            f"An air-gapped install needs talos.factory_url pointed at a "
            f"self-hosted one") from exc
    sid = body.get("id")
    if not sid:
        raise RuntimeError(f"the Image Factory returned no schematic id: {body}")
    cache.parent.mkdir(parents=True, exist_ok=True)
    cache.write_text(sid + "\n")
    return sid, False


def installer_image(cfg, schematic_id):
    talos = cfg.get("talos") or {}
    host = (talos.get("factory_url") or DEFAULT_FACTORY)
    host = re.sub(r"^https?://", "", host).rstrip("/")
    return f"{host}/{INSTALLER_PATH}/{schematic_id}:{talos['version']}"


# ═════════════════════════════════════════════════════════════════════════════
# The running half: a machine configuration patch
# ═════════════════════════════════════════════════════════════════════════════
def hugepage_sysfs(prep):
    """2 MB pages, per NUMA node, as machine.sysfs entries.

    Always the per-node path. The host-wide file lets the kernel choose which
    node each page comes from and a reactor can only allocate from its own, so
    a total that looks right can still leave one instance with nothing.

    On bare metal this needs a script and a systemd unit, because a value
    written to sysfs does not survive a reboot. Here it needs neither: Talos
    reapplies machine.sysfs on every boot, and the whole
    fermihdi-hugepages.service arrangement simply has no counterpart.
    """
    out = {}
    nodes = prep.get("numa") or []
    if nodes:
        for n in nodes:
            want = n.get("hugepages_2m", 0) or 0
            if want:
                out[f"devices/system/node/node{n['node']}/hugepages/"
                    f"hugepages-2048kB/nr_hugepages"] = str(want)
    elif prep.get("hugepages_2m"):
        # No per-node breakdown, so node 0 is the only defensible target: it is
        # where a single-socket machine puts everything, and on a two-socket
        # machine guessing a split would be worse than saying so.
        out["devices/system/node/node0/hugepages/hugepages-2048kB/"
            "nr_hugepages"] = str(prep["hugepages_2m"])
    return out


def kernel_modules(prep):
    """Modules Talos loads at boot, in the order VFIO needs them.

    vfio_iommu_type1 is here because of a failure that costs an afternoon to
    find. SPDK's setup script tries to modprobe it from inside the container
    and cannot -- "could not insert 'vfio_iommu_type1': Operation not
    permitted" -- and what follows is "EAL: No supported IOMMU extensions
    found". A container has no business loading kernel modules; the machine
    configuration loads it instead and SPDK finds it already there.
    """
    if not prep.get("vfio_bind"):
        return []
    return [{"name": "vfio_pci"}, {"name": "vfio_iommu_type1"}]


def cpu_manager(host, prep):
    """Kubelet CPU manager settings, or nothing, with a reason either way.

    isolcpus takes cores off the scheduler; it does not take them out of the
    kubelet's allocatable set. Without the static policy the kubelet will hand
    an isolated core to whatever asks first, and the HD instance that was
    pinned there shares it. The static policy gives exclusive cores to
    Guaranteed pods and keeps the rest for everything else -- which is the
    split cluster.yaml already describes, with container_cpuset naming the
    housekeeping cores and cpuset naming the poll cores.

    So reservedSystemCPUs is the union of the container_cpusets, and it is not
    invented when they are absent: a wrong reservation is worse than none.
    """
    reserved = set()
    have_numa = True
    instances = 0
    for role in host.get("roles") or []:
        for inst in role.get("instances") or []:
            instances += 1
            if inst.get("container_cpuset"):
                reserved |= _vc.parse_cpuset(inst["container_cpuset"])
            if inst.get("numa_node") is None:
                have_numa = False

    if not instances or not prep.get("isolcpus"):
        return None, "no isolated cores to protect"
    if not reserved:
        return None, ("no container_cpuset on any instance, so there is nothing "
                      "to reserve for the kubelet and system daemons — set one "
                      "and the static CPU manager can be turned on")

    cfg = {
        "cpuManagerPolicy": "static",
        "cpuManagerReconcilePeriod": "10s",
        "reservedSystemCPUs": cpuset_str(reserved),
    }
    if have_numa:
        # Every instance names the node its cores came from, so the kubelet can
        # be told to refuse a placement that would straddle one. Without that
        # on every instance the policy would reject pods it should admit.
        cfg["topologyManagerPolicy"] = "single-numa-node"
        cfg["topologyManagerScope"] = "pod"
    return cfg, None


def cpuset_str(ids):
    ids = sorted(ids)
    if not ids:
        return ""
    runs, start, prev = [], ids[0], ids[0]
    for i in ids[1:]:
        if i == prev + 1:
            prev = i
            continue
        runs.append((start, prev))
        start = prev = i
    runs.append((start, prev))
    return ",".join(str(a) if a == b else f"{a}-{b}" for a, b in runs)


def node_labels(host):
    """Labels the chart selects on, so a Pod lands where its devices are."""
    labels = {"fermihdi.io/host": host["name"]}
    for role in host.get("roles") or []:
        labels[f"fermihdi.io/role.{role['type']}"] = "true"
    prep = host.get("host_prep") or {}
    if prep.get("vfio_bind"):
        labels["fermihdi.io/vfio"] = "true"
    if prep.get("hugepages_1g"):
        labels["fermihdi.io/hugepages-1g"] = "true"
    return labels


def machine_patch(host, cfg, schematic_id, talosconfig=None):
    """The multi-document strategic merge patch for one host.

    Multi-document because PCIDriverRebindConfig is its own kind, not a field
    under machine. Talos merges each document with the one of the same kind and
    name, and appends the ones it has never seen.
    """
    prep = host.get("host_prep") or {}
    ht = host.get("talos") or {}
    talos = cfg.get("talos") or {}

    machine = {
        "install": {
            "disk": ht.get("install_disk") or talos.get("install_disk"),
            "image": installer_image(cfg, schematic_id),
        },
        "nodeLabels": node_labels(host),
    }

    sysfs = hugepage_sysfs(prep)
    if sysfs:
        machine["sysfs"] = sysfs

    mods = kernel_modules(prep)
    if mods:
        machine["kernel"] = {"modules": mods}

    kubelet, _ = cpu_manager(host, prep)
    if kubelet:
        machine["kubelet"] = {"extraConfig": kubelet}

    docs = [{"machine": machine}]

    # One document per device. Talos does the unbind, the driver_override and
    # the rebind itself; a udev rule cannot, because the kernel has already
    # bound nvme or mlx5_core by the time udev sees the device.
    for pci in prep.get("vfio_bind") or []:
        docs.append({
            "apiVersion": "v1alpha1",
            "kind": "PCIDriverRebindConfig",
            "name": Quoted(pci),
            "targetDriver": "vfio-pci",
        })
    return docs


class Quoted(str):
    """A string that must survive as a string, whatever the reader thinks."""


yaml.SafeDumper.add_representer(
    Quoted, lambda d, v: d.represent_scalar("tag:yaml.org,2002:str", str(v), style="'"))


def dump_docs(docs):
    return "---\n".join(
        yaml.dump(d, sort_keys=False, default_flow_style=False,
                  Dumper=yaml.SafeDumper) for d in docs)


# ═════════════════════════════════════════════════════════════════════════════
# Render
# ═════════════════════════════════════════════════════════════════════════════
def out_dir_for(cfg, path, override):
    d = override or (cfg.get("talos") or {}).get("out_dir") or "talos"
    d = pathlib.Path(d)
    return d if d.is_absolute() else (path.parent / d).resolve()


def render(cfg, hosts, outdir, talosconfig, factory_url, offline):
    """Write the schematic, the patch and the machine config for every host.

    Returns {host name: {...}} so the apply phase does not recompute any of it.
    """
    outdir.mkdir(parents=True, exist_ok=True)
    talos = cfg["talos"]
    plan = {}

    secrets = outdir / "secrets.yaml"
    if not secrets.exists():
        do("secrets", f"generate {secrets}",
           lambda: talosctl(["gen", "secrets", "-o", str(secrets)],
                            check=True, timeout=120))
        if not APPLY:
            # Everything below needs the bundle. Without it the machine
            # configs cannot be generated, so say so once and carry on
            # describing what would happen rather than half-failing.
            say("secrets", "WOULD", "every machine config below is generated from it")
    else:
        ok("secrets", f"{secrets} exists — reused, so node identities are stable")

    for host in hosts:
        name = host["name"]
        ht = host.get("talos") or {}
        try:
            sdoc = schematic_for(host, cfg, talosconfig)
        except RuntimeError as exc:
            fail("schematic", f"{name}: {exc}")
            continue
        stext = schematic_text(sdoc)
        (outdir / f"{name}.schematic.yaml").write_text(stext)

        args = (sdoc.get("customization") or {}).get("extraKernelArgs") or []
        if offline:
            sid, cached = "<not-uploaded>", False
        else:
            try:
                sid, cached = upload_schematic(stext, factory_url, outdir)
            except RuntimeError as exc:
                fail("schematic", f"{name}: {exc}")
                continue
        plan[name] = {"host": host, "schematic": sid, "kernel_args": args}

        shown = sid if offline else sid[:16] + "…"
        source = "not uploaded" if offline else ("cached" if cached else "uploaded")
        if args:
            ok("schematic", f"{name}: {shown} ({source}) — "
                            f"{len(args)} kernel argument(s): {' '.join(args)}")
        else:
            ok("schematic", f"{name}: {shown} ({source}) — stock kernel arguments")

        docs = machine_patch(host, cfg, sid, talosconfig)
        patch = outdir / f"{name}.patch.yaml"
        patch.write_text(dump_docs(docs))
        ok("patch", f"{name}: {patch.name} — "
                    f"{len(docs[0]['machine'].get('sysfs', {}))} sysfs, "
                    f"{len(docs) - 1} device rebind(s)")

        _, why = cpu_manager(host, host.get("host_prep") or {})
        if why:
            warn("kubelet", f"{name}: no static CPU manager — {why}")

        mtype = ht.get("machine", "worker")
        conf = outdir / f"{name}.yaml"
        gen = ["gen", "config", cfg.get("cluster_name", "fermihdi"), talos["endpoint"],
               "--with-secrets", str(secrets), "-t", mtype, "-o", str(conf),
               "--install-disk", ht.get("install_disk") or talos.get("install_disk"),
               "--config-patch", f"@{patch}",
               "--with-docs=false", "--with-examples=false", "--force"]
        if not offline:
            gen += ["--install-image", installer_image(cfg, sid)]
        if talos.get("kubernetes_version"):
            gen += ["--kubernetes-version", talos["kubernetes_version"]]

        plan[name]["config"] = conf
        do("config", f"{name}: generate {conf.name} ({mtype})",
           lambda g=gen: talosctl(g, check=True, timeout=120))

    return plan


# ═════════════════════════════════════════════════════════════════════════════
# Apply
# ═════════════════════════════════════════════════════════════════════════════
def in_maintenance(host, talosconfig):
    """Has this machine been configured yet?

    A machine in maintenance mode serves the API with a self-signed
    certificate and no client authentication, so a secure call fails and an
    insecure one works. That is the test, and it is also the reason apply-config
    needs --insecure exactly once in a machine's life.
    """
    if talosctl(["get", "machineconfig"], node=node_addr(host),
                talosconfig=talosconfig, timeout=30).returncode == 0:
        return False
    return talosctl(["version"], node=node_addr(host),
                    insecure=True, timeout=30).returncode == 0


def apply_config(name, host, conf, talosconfig):
    """Give a machine its configuration. Returns True if this was a fresh install.

    The probe is read-only and runs in a dry run too, because whether a machine
    is in maintenance mode changes what the rest of the plan says: a first
    install writes the schematic image to disk itself, so the machine boots
    with the right kernel arguments and never needs the upgrade that an
    already-running machine does.
    """
    node = node_addr(host)
    maint = in_maintenance(host, talosconfig)

    if maint:
        # First contact. The machine has no certificates yet, so this is the
        # one call that legitimately skips verification, and it installs Talos
        # to the disk and reboots into it.
        do("apply", f"{name}: apply {conf.name} to {node} (maintenance mode — "
                    f"this installs to disk and reboots)",
           lambda: talosctl(["apply-config", "-f", str(conf), "-i"],
                            node=node, check=True))
        if APPLY:
            outstanding(f"{name} is installing Talos and will reboot into it")
        return True

    ok("apply", f"{name}: {node} is already configured — updating in place")
    do("apply", f"{name}: apply {conf.name} to {node}",
       lambda: talosctl(["apply-config", "-f", str(conf)],
                        node=node, talosconfig=talosconfig, check=True))
    return False


def bootstrap(cfg, plan, talosconfig, outdir):
    """Start etcd, once, on the one node the file nominates."""
    picked = [(n, p) for n, p in plan.items()
              if (p["host"].get("talos") or {}).get("bootstrap")]
    if not picked:
        skip("bootstrap", "no host sets talos.bootstrap")
        return
    name, p = picked[0]
    node = node_addr(p["host"])

    if APPLY and talosctl(["get", "etcdmembers"], node=node,
                          talosconfig=talosconfig, timeout=60).returncode == 0:
        ok("bootstrap", f"{name}: etcd is already up — bootstrapping is once per cluster")
    else:
        do("bootstrap", f"{name}: bootstrap etcd on {node}",
           lambda: talosctl(["bootstrap"], node=node,
                            talosconfig=talosconfig, check=True, timeout=600))

    kubeconfig = outdir / "kubeconfig"
    do("kubeconfig", f"write {kubeconfig}",
       lambda: talosctl(["kubeconfig", str(kubeconfig), "--force"],
                        node=node, talosconfig=talosconfig, check=True, timeout=300))


def running_image(host, talosconfig):
    """The installer image the machine last booted, or ''.

    This is what says whether an upgrade is needed. Comparing schematic IDs
    rather than versions matters: two nodes on the same Talos release can have
    different kernel arguments, and only the ID shows it.
    """
    for u in get_resource(host, "upgradestatus", talosconfig):
        img = (u.get("spec", {}) or {}).get("image", "")
        if img:
            return img
    for m in get_resource(host, "machineconfig", talosconfig):
        spec = m.get("spec", {}) or {}
        if isinstance(spec, dict):
            img = ((spec.get("machine") or {}).get("install") or {}).get("image", "")
            if img:
                return img
    return ""


def upgrade(name, p, cfg, talosconfig, no_upgrade):
    """Move a machine to the installer image its schematic names.

    This is where isolcpus, the IOMMU and the 1 GB pool actually arrive, and it
    is a reboot. One machine at a time, and never in parallel: an upgrade is
    the one operation here that can take a node out of the cluster.
    """
    host = p["host"]
    want = installer_image(cfg, p["schematic"])

    if not p["kernel_args"]:
        skip("upgrade", f"{name}: no kernel arguments, so the stock image is right")
        return

    if p.get("fresh"):
        # The install-config carried install.image, so the installer that just
        # wrote the disk was the schematic's own. There is nothing to upgrade
        # from -- the machine comes up on the right kernel command line.
        ok("upgrade", f"{name}: installed from {want.rsplit('/', 1)[-1][:24]}… — "
                      f"the arguments arrive with the first boot, not an upgrade")
        return

    if no_upgrade:
        outstanding(f"{name} still needs `talosctl -n {node_addr(host)} upgrade "
                    f"--image {want}` for: {' '.join(p['kernel_args'])}")
        skip("upgrade", f"{name}: --no-upgrade, so nothing reboots. "
                        f"{' '.join(p['kernel_args'])} is not in effect")
        return

    have = running_image(host, talosconfig) if APPLY else ""
    if have and have == want:
        ok("upgrade", f"{name}: already running {want.rsplit('/', 1)[-1][:24]}…")
        return

    do("upgrade", f"{name}: upgrade {node_addr(host)} to {want} — THIS REBOOTS IT",
       lambda: talosctl(["upgrade", "--image", want, "--wait"],
                        node=node_addr(host), talosconfig=talosconfig,
                        check=True, timeout=3600))


# ═════════════════════════════════════════════════════════════════════════════
# Verify
# ═════════════════════════════════════════════════════════════════════════════
def read_file(host, path, talosconfig):
    try:
        p = talosctl(["read", path], node=node_addr(host),
                     talosconfig=talosconfig, timeout=60)
    except RuntimeError:
        return None
    return p.stdout.strip() if p.returncode == 0 else None


def verify_host(host, cfg, talosconfig):
    """Compare the running machine against the definition. Changes nothing."""
    name = host["name"]
    prep = host.get("host_prep") or {}
    node = node_addr(host)
    print(f"\n  ── {name} ({node}) " + "─" * max(0, 50 - len(name) - len(node)))

    cmdline = read_file(host, "/proc/cmdline", talosconfig)
    if cmdline is None:
        fail("verify", f"{name}: cannot read /proc/cmdline — the machine is not "
                       f"reachable, or is still in maintenance mode")
        return
    running = set(cmdline.split())
    try:
        want = kernel_args_for(host, talosconfig)
    except RuntimeError as exc:
        want = []
        warn("cmdline", f"{name}: {exc}")
    missing = [a for a in want if a not in running]
    if not want:
        skip("cmdline", "this host asks for no kernel arguments")
    elif missing:
        fail("cmdline", f"{' '.join(missing)} not in the running kernel — the "
                        f"upgrade to the schematic image has not happened, or "
                        f"the machine has not rebooted into it")
    else:
        ok("cmdline", f"all {len(want)} argument(s) present")

    for path, want_n in sorted(hugepage_sysfs(prep).items()):
        got = read_file(host, "/sys/" + path, talosconfig)
        node_n = re.search(r"node(\d+)", path).group(1)
        if got is None:
            fail("hugepages", f"node {node_n}: cannot read {path}")
        elif got.strip() != want_n:
            fail("hugepages", f"node {node_n}: {got.strip()} 2 MB pages, want {want_n}"
                              + ("  (1 GB pages are reserved at boot and take the "
                                 "contiguous memory first — a late 2 MB reservation "
                                 "can simply fail)" if int(got.strip() or 0) == 0 else ""))
        else:
            ok("hugepages", f"node {node_n}: {want_n} x 2 MB")

    binds = prep.get("vfio_bind") or []
    if binds:
        status = {(s.get("spec", {}) or {}).get("pciid"):
                  (s.get("spec", {}) or {}).get("targetDriver")
                  for s in get_resource(host, "pcidriverrebindstatus", talosconfig)}
        for pci in binds:
            if status.get(pci) == "vfio-pci":
                ok("vfio", f"{pci} → vfio-pci")
            else:
                fail("vfio", f"{pci} is not rebound (Talos reports "
                             f"{status.get(pci) or 'nothing'}) — without the IOMMU "
                             f"up, the rebind cannot succeed")
        loaded = {(m.get("metadata", {}) or {}).get("id", "")
                  for m in get_resource(host, "loadedkernelmodules", talosconfig)}
        for mod in ("vfio_pci", "vfio_iommu_type1"):
            if mod in loaded or mod.replace("_", "-") in loaded:
                ok("modules", f"{mod} loaded")
            else:
                fail("modules", f"{mod} is not loaded — SPDK cannot load it from "
                                f"inside a container and will fail to initialise")


# ═════════════════════════════════════════════════════════════════════════════
def main():
    global APPLY
    ap = argparse.ArgumentParser(
        prog="fermihdi-talos-prepare",
        description="Render and apply Talos machine configuration from a cluster definition.")
    ap.add_argument("cluster", nargs="?", default=str(HERE / "cluster.yaml"))
    ap.add_argument("--host", metavar="NAME", help="one host instead of all of them")
    ap.add_argument("--apply", action="store_true", help="make the changes")
    ap.add_argument("--verify", action="store_true",
                    help="check the running machines and change nothing")
    ap.add_argument("--render-only", action="store_true",
                    help="write the files and talk to no machine")
    ap.add_argument("--no-upgrade", action="store_true",
                    help="apply machine configuration but never reboot")
    ap.add_argument("--out", metavar="DIR", help="override talos.out_dir")
    ap.add_argument("--talosconfig", metavar="PATH",
                    help="default: <out_dir>/talosconfig")
    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())
    if not cfg.get("talos"):
        sys.exit(f"{path} has no talos block. This tool is for hosts running Talos; "
                 f"use fermihdi-prepare for everything else.")

    hosts = cfg.get("hosts") or []
    if args.host:
        hosts = [h for h in hosts if h.get("name") == args.host]
        if not hosts:
            sys.exit(f"no host named {args.host!r} — the file has: " +
                     ", ".join(h.get("name", "?") for h in cfg.get("hosts") or []))

    APPLY = args.apply and not args.verify and not args.render_only
    outdir = out_dir_for(cfg, path, args.out)
    talosconfig = pathlib.Path(args.talosconfig) if args.talosconfig \
        else outdir / "talosconfig"
    tc = talosconfig if talosconfig.exists() else None

    if not args.render_only and not shutil.which("talosctl"):
        sys.exit("talosctl is not on PATH, and every mode but --render-only needs "
                 "it.\n  See https://docs.siderolabs.com/talos/")

    mode = ("verify" if args.verify else
            "render only" if args.render_only else
            "apply" if APPLY else "dry run — nothing is changed")
    print()
    print(f"  cluster    {cfg.get('cluster_name', '-')}   Talos {cfg['talos']['version']}")
    print(f"  source     {path}")
    print(f"  output     {outdir}")
    print(f"  hosts      {', '.join(h['name'] for h in hosts)}")
    print(f"  mode       {mode}")
    print()

    if args.verify:
        for host in hosts:
            verify_host(host, cfg, tc)
    else:
        factory = (cfg["talos"].get("factory_url") or DEFAULT_FACTORY)
        plan = render(cfg, hosts, outdir, tc, factory, args.render_only)
        if not args.render_only:
            print()
            tc = talosconfig if talosconfig.exists() else tc
            for name, p in plan.items():
                p["fresh"] = apply_config(name, p["host"], p["config"], tc)
            print()
            bootstrap(cfg, plan, tc, outdir)
            print()
            for name, p in plan.items():
                upgrade(name, p, cfg, tc, args.no_upgrade)

    print()
    if ERRORS:
        print(f"  {len(ERRORS)} error(s) — this cluster is not ready.")
        print()
        return 1
    if args.verify:
        print("  Every machine matches its definition and is ready to take Pods.")
        print()
        return 0
    if args.render_only:
        print(f"  Written to {outdir}. Nothing was uploaded and no machine was "
              f"contacted.")
        print()
        return 0
    if not APPLY:
        print("  Nothing was changed. Re-run with --apply to make it so.")
        print("  Note that --apply reboots any machine whose kernel arguments changed,")
        print("  one at a time. --no-upgrade applies the rest and reboots nothing.")
        print()
        return 0

    print(f"  {len(CHANGES)} change(s) applied.")
    if OUTSTANDING:
        print()
        print("  Outstanding:")
        for reason in OUTSTANDING:
            print(f"    - {reason}")
        print()
        print("  When the machines are back:")
        print(f"    {sys.argv[0]} {path} --verify")
        print()
        return 2
    print()
    print("  Confirm with:")
    print(f"    {sys.argv[0]} {path} --verify")
    print(f"    kubectl --kubeconfig {outdir / 'kubeconfig'} get nodes")
    print()
    return 0


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