#!/usr/bin/env python3
"""Render Helm values for install/helm/fermihdi-talos from a cluster definition.

fermihdi-k8s says of the general-purpose chart that devices, DPDK and the
cpuset split do not carry across, "because a Pod does not get an IOMMU group".
On Talos they do carry across, and this is the difference: a machine
configuration binds the device to vfio-pci, a device plugin advertises it, and
a Pod requests it by name like any other resource.

So this is fermihdi-k8s plus the four things that follow from that:

  device pools     the plugin's config.json, one pool per role and device kind,
                   selecting exactly the PCI addresses cluster.yaml assigns
  device claims    fermihdi.io/<pool> on the Pods that were given those devices
  the 1 GB pool    hugepages-1Gi alongside hugepages-2Mi, from the same
                   per-device model validate-cluster sizes hosts with
  placement        nodeSelector on the role labels fermihdi-talos-prepare sets

Everything else -- images, sizing, plane CIDRs, the DMS, the secrets file --
is fermihdi-k8s's, imported rather than reimplemented. What that tool will not
do, this one does not undo: anything the file says that the chart cannot
express is still reported rather than dropped.

Usage:
    install/fermihdi-talos-k8s [cluster.yaml] [-o values.yaml]
    install/fermihdi-talos-k8s cluster.yaml --secrets values-secrets.yaml
    install/fermihdi-talos-k8s cluster.yaml --manifests | kubectl apply -f -

Then:
    helm install fermihdi install/helm/fermihdi-talos -f values.yaml
"""
import argparse
import importlib.machinery
import importlib.util
import json
import pathlib
import re
import sys

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

HERE = pathlib.Path(__file__).resolve().parent
CHART = HERE / "helm" / "fermihdi-talos"


def _load(name, filename):
    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


_k8s = _load("fermihdi_k8s", "fermihdi-k8s")
_vc = _load("fermihdi_validate", "validate-cluster")

# A multi-line value is written as a block scalar rather than an escaped
# one-liner. The device plugin's config.json is the reason: as \n-escaped JSON
# inside a quoted YAML string it is unreadable, and it is the one part of this
# file an operator is most likely to want to check by eye against `talosctl get
# pcidevices`.
yaml.SafeDumper.add_representer(
    str,
    lambda d, v: d.represent_scalar("tag:yaml.org,2002:str", v,
                                    style="|" if "\n" in v else None))

DEFAULT_PREFIX = "fermihdi.io"
# Where each role's values live in the chart, for the roles that can hold a
# device. proxy, dms, mcp and graphql never do; a filter node holds a NIC at most.
DEVICE_ROLES = {"scc": ("sc", "scc"), "sn": ("sc", "sn"), "bloomrf": ("sc", "bloomrf"),
                "ingester": ("edge", "ingester")}

notes = _k8s.notes          # one list, so both tools' notes end up in the header

# fermihdi-k8s says of the general-purpose chart that these four things have
# nowhere to go. In this chart they do, and carrying its notes across unchanged
# would state the opposite of what the values file below actually contains.
# Matched on a distinctive fragment rather than the whole sentence, so a
# reworded note there does not silently start surviving here.
#
# Each is replaced rather than merely dropped: the thing it warned about is now
# a thing the operator has to have set up, and that is worth saying.
RESOLVED = [
    ("declares dpdk, which the chart does not deploy",
     "{role} declares dpdk and this chart honours it: the container is not given "
     "FERMIHDI_FORCE_POSIX, so the engine will look for the NIC it was assigned. "
     "That NIC has to be on vfio-pci and advertised by the device plugin, or the "
     "data path falls back at run time rather than at render time."),
    ("Nothing here passes an NVMe through", None),
    ("vfio bindings and IOMMU groups have no expression in the chart",
     "vfio devices are passed through as device plugin resources. They exist only "
     "on the machines whose rebinds have actually taken — check with "
     "`fermihdi-talos-prepare cluster.yaml --verify` before deploying, because a "
     "Pod claiming a pool with no members stays Pending with no other symptom."),
    ("The chart asks for 2 MB pages only",
     "1 GB pages are requested alongside the 2 MB pool, and mounted at "
     "/dev/hugepages so SPDK's EAL finds them there. They are reserved on the "
     "kernel command line, which means the machine has to have booted its "
     "schematic image first."),
    ("container_cpuset is a cgroup split",
     "container_cpuset became the kubelet's reservedSystemCPUs, and the poll cores "
     "are held by the static CPU manager giving whole cores to these Guaranteed "
     "Pods. fermihdi-talos-prepare writes both into the machine configuration."),
]


def resolve_inherited_notes(cfg):
    """Drop what fermihdi-k8s could not do and this chart can, saying what replaced it."""
    kept = []
    for n in notes:
        for fragment, replacement in RESOLVED:
            if fragment in n:
                if replacement:
                    role = n.split(" ", 1)[0].rstrip(":")
                    kept.append(replacement.replace("{role}", role))
                break
        else:
            kept.append(n)
    notes[:] = list(dict.fromkeys(kept))


def note(msg):
    if msg not in notes:
        notes.append(msg)


def pool_name(rtype, kind):
    """A device plugin resource name: no dots, no dashes, no surprises."""
    return re.sub(r"[^a-z0-9_]", "_", f"{rtype}_{kind}".lower())


def vfio_devices(cfg):
    """Every vfio device in the file, grouped into pools.

    Grouped by role and kind rather than by host, because a resource pool is
    cluster-wide: the plugin runs on every machine that has any, each node
    advertises the ones it has, and the scheduler puts the Pod where the count
    it asked for exists. That is also what makes an SN portable between two
    machines with the same devices.
    """
    pools = {}                      # pool -> {pci addresses}
    per_instance = {}               # (rtype, instance name) -> {pool: count}
    for host in cfg.get("hosts") or []:
        for role in host.get("roles") or []:
            rtype = role["type"]
            if rtype not in DEVICE_ROLES:
                continue
            for inst in role.get("instances") or []:
                devs = list(inst.get("storage") or [])
                if inst.get("nic"):
                    devs.append(inst["nic"])
                for dev in devs:
                    if dev.get("driver") != "vfio" or not dev.get("pci"):
                        continue
                    pool = pool_name(rtype, dev.get("kind") or "device")
                    pools.setdefault(pool, set()).add(dev["pci"])
                    counts = per_instance.setdefault((rtype, inst["name"]), {})
                    counts[pool] = counts.get(pool, 0) + 1
    return pools, per_instance


def device_plugin_config(pools, prefix):
    """The plugin's config.json.

    deviceType is "accelerator" for everything, including the NIC. It is the
    generic PCI class -- the one whose selectors are just vendor, device,
    driver and address -- and it returns /dev/vfio/<group>, which is what both
    SPDK and DPDK open. "netDevice" would be the usual choice for a NIC, and it
    goes looking for a netlink interface that a device on vfio-pci no longer
    has.

    The driver selector is not redundant beside an explicit address list. It is
    what stops the plugin advertising a device whose rebind silently failed:
    without it a Pod could be scheduled onto a machine where the NVMe is still
    on the kernel's nvme driver, and SPDK would find nothing to attach.
    """
    return {
        "resourceList": [
            {
                "resourceName": pool,
                "resourcePrefix": prefix,
                "deviceType": "accelerator",
                "selectors": [{
                    "drivers": ["vfio-pci"],
                    "pciAddresses": sorted(addrs),
                }],
            }
            for pool, addrs in sorted(pools.items())
        ]
    }


def hugepages_1g(cfg, rtype):
    """The 1 GB pool one Pod of this role needs, as a Kubernetes quantity.

    The same model validate-cluster uses, and for the same reason: 1 GB pages
    back SPDK's NVMe DMA and nothing else, so an SN whose devices are all
    kernel-path needs none at all -- bdev_aio pread/pwrites through the block
    layer and never DMAs. The SCC's two are for EAL's -m 2048.
    """
    fixed = _vc.HP1G_FIXED.get(rtype, 0)
    per_pod = []
    for host in cfg.get("hosts") or []:
        for role in host.get("roles") or []:
            if role["type"] != rtype:
                continue
            for inst in role.get("instances") or []:
                dma = sum(1 for d in inst.get("storage") or []
                          if d.get("kind") in _vc.DMA_KINDS)
                per_pod.append(fixed + dma * _vc.HP1G_PER_NVME)
    if not per_pod:
        return None, 0
    n = max(per_pod)
    if len(set(per_pod)) > 1:
        note(f"{rtype}: instances need different numbers of 1 GB pages "
             f"({sorted(set(per_pod))}), and replicas of one workload are "
             f"identical — the largest wins, so the others reserve more than "
             f"they use")
    return (f"{n}Gi" if n else None), n


def build(cfg, path, **kw):
    values = _k8s.build(cfg, path, **kw)
    resolve_inherited_notes(cfg)
    talos = cfg.get("talos") or {}
    prefix = ((talos.get("device_plugin") or {}).get("resource_prefix")
              or DEFAULT_PREFIX)

    pools, per_instance = vfio_devices(cfg)
    plugin_wanted = (talos.get("device_plugin") or {}).get("enabled", True)

    values["talos"] = {
        "privilegedNamespaces": True,
        "devicePlugin": {
            # No devices means nothing to advertise, and a DaemonSet that would
            # come up on every machine to find nothing.
            "enabled": bool(pools) and plugin_wanted,
            "logLevel": "info",
            "resourcePrefix": prefix,
            "config": json.dumps(device_plugin_config(pools, prefix), indent=2) + "\n",
        },
    }
    if img := (talos.get("device_plugin") or {}).get("image"):
        values["talos"]["devicePlugin"]["image"] = img

    for rtype, (group, key) in DEVICE_ROLES.items():
        block = values.get(group, {}).get(key)
        if not block or not block.get("enabled"):
            continue

        # DPDK is a property of the role in cluster.yaml; the chart needs it to
        # decide whether to force POSIX on the container.
        block["dpdk"] = any(r.get("dpdk") for h in cfg.get("hosts") or []
                            for r in h.get("roles") or [] if r["type"] == rtype)

        hp1g, n1g = hugepages_1g(cfg, rtype)
        if hp1g:
            block["hugepages1Gi"] = hp1g

        claims = {}
        for (rt, _), counts in per_instance.items():
            if rt != rtype:
                continue
            for pool, count in counts.items():
                claims[f"{prefix}/{pool}"] = max(claims.get(f"{prefix}/{pool}", 0), count)
        if claims and plugin_wanted:
            block["vfioResources"] = claims
            per_pod = [sum(c.values()) for (rt, _), c in per_instance.items() if rt == rtype]
            if len(set(per_pod)) > 1:
                note(f"{rtype}: instances hold different numbers of devices "
                     f"({sorted(set(per_pod))}) — every replica claims the largest, "
                     f"so a machine holding fewer will not be scheduled onto")
        elif claims:
            note(f"{rtype}: {len(claims)} device pool(s) are declared but "
                 f"talos.device_plugin.enabled is false, so the Pods claim nothing "
                 f"and the devices sit bound to vfio-pci with no user")

        if rtype == "sn" and any(
                d.get("driver") == "vfio"
                for h in cfg.get("hosts") or [] for r in h.get("roles") or []
                if r["type"] == "sn" for i in r.get("instances") or []
                for d in i.get("storage") or []):
            # Real devices: the pre-allocated file and the PVC behind it are
            # storage nothing would read.
            block["useRawDevices"] = True

    # The machines this is deployed onto. Worth saying out loud, because a
    # values file that mentions no machine at all reads as portable and is not.
    workers = [h["name"] for h in cfg.get("hosts") or []
               if (h.get("talos") or {}).get("machine") == "worker"]
    if not workers:
        note("no host is a Talos worker, so every Pod would have to run on a "
             "controlplane — which is where etcd is")
    return values


def main():
    ap = argparse.ArgumentParser(description=__doc__.splitlines()[0])
    ap.add_argument("cluster", nargs="?", default=str(HERE / "cluster.yaml"))
    ap.add_argument("-o", "--output", help="write the values here instead of stdout")
    ap.add_argument("--secrets", metavar="PATH",
                    help="also write the DMS's keys and tokens. It contains private "
                         "keys: treat it as one.")
    ap.add_argument("--manifests", action="store_true",
                    help="render the chart with these values and print the manifests")
    ap.add_argument("--chart", default=str(CHART))
    ap.add_argument("--domain", help="ingress domain, e.g. fermihdi.example.com")
    ap.add_argument("--storage-class", help="storage class for every PVC")
    args = ap.parse_args()

    path = pathlib.Path(args.cluster).resolve()
    if not path.exists():
        sys.exit(f"no such file: {path}")
    cfg = yaml.safe_load(path.read_text())
    if not isinstance(cfg, dict) or "hosts" not in cfg:
        sys.exit(f"{path} is not a cluster definition (no top-level 'hosts')")
    if not cfg.get("talos"):
        sys.exit(f"{path} has no talos block. Use fermihdi-k8s and "
                 f"install/helm/fermihdi for a cluster that is not on Talos.")

    values = build(cfg, path, domain=args.domain, storage_class=args.storage_class)

    header = [f"# Generated by fermihdi-talos-k8s from {path}.",
              "# Re-generate rather than editing: this file is derived, and the cluster",
              "# definition is the thing that is kept.",
              "#",
              "# helm install fermihdi install/helm/fermihdi-talos -f this-file.yaml",
              "#",
              "# The machines have to have been prepared first — the device pools below",
              "# name PCI addresses that only exist as /dev/vfio entries once",
              "# fermihdi-talos-prepare has applied the rebinds and the machines have",
              "# booted the schematic image:",
              "#",
              "#   install/fermihdi-talos-prepare cluster.yaml --verify"]
    if notes:
        header += ["#", "# What did not come across:"]
        header += [f"#   - {n}" for n in notes]
    body = "\n".join(header) + "\n" + yaml.safe_dump(values, sort_keys=False, width=88)

    if args.output:
        pathlib.Path(args.output).write_text(body)
        print(f"  wrote {args.output}", file=sys.stderr)
    elif not args.manifests:
        print(body)

    if args.secrets:
        sec = _k8s.secrets_values(cfg, path)
        empty = [k for k, v in sec["secrets"]["dms"].items() if not v]
        p = pathlib.Path(args.secrets)
        p.write_text(
            "# Generated by fermihdi-talos-k8s. HOLDS PRIVATE KEYS -- do not commit it.\n"
            "# helm install ... -f values.yaml -f this-file.yaml\n"
            + yaml.safe_dump(sec, sort_keys=False, width=88))
        p.chmod(0o600)
        print(f"  wrote {args.secrets} (mode 0600)", file=sys.stderr)
        if empty:
            print(f"  empty, and the deployment needs them: {', '.join(sorted(empty))}",
                  file=sys.stderr)

    for n in notes:
        print(f"  note  {n}", file=sys.stderr)

    if args.manifests:
        import subprocess
        import tempfile
        with tempfile.NamedTemporaryFile("w", suffix=".yaml") as tmp:
            tmp.write(body)
            tmp.flush()
            cmd = ["helm", "template", cfg.get("cluster_name", "fermihdi"),
                   args.chart, "-f", tmp.name]
            if args.secrets:
                cmd += ["-f", args.secrets]
            try:
                p = subprocess.run(cmd, capture_output=True, text=True)
            except FileNotFoundError:
                sys.exit("helm is required for --manifests")
            if p.returncode != 0:
                sys.exit(p.stderr.strip())
            print(p.stdout)
    return 0


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