#!/usr/bin/env python3
"""Validate a cluster definition against the schema, then against reality.

Schema validation catches shape errors. The semantic checks after it catch the
ones that pass schema validation and still break a deployment — overlapping
cpusets, an NVMe handed to two Storage Nodes, isolcpus that does not cover what
it should. Those are the failures that show up as a wedged node hours later
rather than as an error here.

Usage:  install/validate-cluster [cluster.yaml]
"""
import json
import pathlib
import sys

try:
    import yaml
except ImportError:
    sys.exit("PyYAML is required:  pip install pyyaml")

HERE = pathlib.Path(__file__).resolve().parent
SCHEMA = HERE / "schema" / "cluster.schema.json"

# A role's floor is the smallest core count its resource controller will accept.
# Below it, FermiController::Init() throws and the node exits at startup.
ROLE_MIN_CORES = {"scc": 6, "sn": 9, "ingester": 2, "mcp": 1, "graphql": 1,
                  "proxy": 1, "dms": 1}
# Engine roles. proxy and dms are Go services and need no hugepages at all.
NEEDS_HUGEPAGES_2M = {"scc", "sn", "ingester", "mcp", "graphql"}
# 1 GB pages back SPDK NVMe DMA and nothing else. StorageEngine.cpp pins EAL's
# hugedir to /dev/hugepages when they exist, keeping SPDK off the 2 MB
# networking pool. An SN whose devices are all kernel-path (SATA/SAS/FC/iSCSI
# through bdev_aio) sets opts.no_huge and needs NONE -- so this is a property of
# the DEVICES, not of the role. The SCC needs 2 for its EAL -m 2048.
HP1G_FIXED = {"scc": 2}
# StorageEngine's own diagnostic: "minimum 2, recommend 4 per NVMe device".
HP1G_PER_NVME = 4
DMA_KINDS = {"nvme", "nvmeof"}
# Roles for which DPDK is available at all. Every role also speaks POSIX.
DPDK_CAPABLE = {"scc", "sn", "ingester", "graphql", "mcp"}
# Engine budget in MB per instance -- what the 2 MB pool has to cover. The engine
# allocates from that pool, so these set its size rather than sitting alongside
# it. libfermihdi_transport defaults FERMIHDI_SEASTAR_MEMORY to "512M"; the SN
# is the only role that overrides it, scaling FAKE_SN_TX_MEMORY_PER_SHARD_MB
# (4096) by its transmit shard count, two by default.
ENGINE_MB = {"sn": 8192, "scc": 512, "ingester": 512, "graphql": 512, "mcp": 512}
# 2 MB pages for an engine budget: the budget itself plus 44 pages (88 MB) for
# mbuf pools and EAL overhead. A 512M budget lands on 300.
def hp2m_for(rtype: str) -> int:
    mb = ENGINE_MB.get(rtype, 0)
    return mb // 2 + 44 if mb else 0
# Kernel, page cache and everything not FermiHDI (4 GB), plus process stack and
# heap outside the engine arena (1 GB). fermihdi-configure and
# fermihdi-prepare compute against the same figure.
OS_HEADROOM_MB = 5120

errors, warnings = [], []


def err(msg):
    errors.append(msg)


def warn(msg):
    warnings.append(msg)


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 check_pci_quoting(node, path=""):
    """Catch PCI addresses YAML turned into numbers.

    YAML 1.1 reads 0000:31:00.0 as sexagesimal (1860.0), but only when every
    segment is below 60 — so bus 61 survives and bus 31 does not. The schema
    reports this as "not of type string", which does not tell anyone what to do.
    """
    if isinstance(node, dict):
        for k, v in node.items():
            check_pci_quoting(v, f"{path}.{k}")
    elif isinstance(node, list):
        for i, v in enumerate(node):
            check_pci_quoting(v, f"{path}[{i}]")
    elif isinstance(node, (int, float)) and not isinstance(node, bool):
        if path.split(".")[-1].split("[")[0] in ("pci", "vfio_bind"):
            err(f"{path}: {node!r} is a number, not a PCI address. YAML read an "
                f"unquoted address as sexagesimal — quote it, e.g. \"0000:31:00.0\".")


def check_dms_files(cfg, base):
    """The files the dms: block names, checked where this is being run.

    Paths, not contents: they are read by whatever deploys the cluster, which is
    not necessarily this machine. So a file that is not here is a warning -- it
    may be sitting on the control machine, exactly as intended -- while a pair
    that cannot work whatever machine it is on is an error.
    """
    dms = cfg.get("dms") or {}
    named = {k: dms[k] for k in ("license_file", "ica_cert_file", "ica_key_file",
                                 "schema_file") if dms.get(k)}
    if not named:
        return

    runs_dms = any(role.get("type") == "dms"
                   for host in cfg.get("hosts", [])
                   for role in host.get("roles") or [])
    if not runs_dms:
        warn(f"dms: names {', '.join(sorted(named))} but no host declares a dms role, "
             f"so nothing in this file installs them")

    # The certificate signs nothing without the key behind it, and the key
    # identifies nothing without the certificate. Either alone is a deployment
    # that comes up and then refuses every CSR.
    if bool(dms.get("ica_cert_file")) != bool(dms.get("ica_key_file")):
        missing = "ica_key_file" if dms.get("ica_cert_file") else "ica_cert_file"
        err(f"dms: names one half of the DMS-CA's identity — {missing} is missing. "
            f"The certificate is what the CA presents and the key is what it signs "
            f"with; it needs both.")

    for key, value in named.items():
        path = pathlib.Path(value)
        if not path.is_absolute():
            path = base / path
        if not path.is_file():
            warn(f"dms.{key}: {value} is not readable from here. Relative paths resolve "
                 f"against the cluster file; an absolute one has to exist on whichever "
                 f"machine runs the deploy.")
            continue
        # DMS-Core exits at startup on a license payload that is not valid JSON,
        # and the schema is served to the OPA engine as-is.
        if key in ("license_file", "schema_file"):
            try:
                json.loads(path.read_text())
            except Exception as exc:
                warn(f"dms.{key}: {value} is not valid JSON ({exc}). DMS-Core refuses to "
                     f"start on a license payload it cannot parse.")
        if key == "ica_cert_file" and "BEGIN CERTIFICATE" not in path.read_text():
            warn(f"dms.{key}: {value} holds no PEM certificate block")
        if key == "ica_key_file" and "PRIVATE KEY" not in path.read_text():
            warn(f"dms.{key}: {value} holds no PEM private key block")


def check(cfg):
    for host in cfg.get("hosts", []):
        hn = host.get("name", "<unnamed>")
        prep = host.get("host_prep", {}) or {}
        iso = parse_cpuset(prep["isolcpus"]) if prep.get("isolcpus") else set()

        seen_cores = {}       # core  -> instance that claimed it
        seen_pci = {}         # pci   -> instance that claimed it
        seen_group = {}       # iommu -> instance that claimed it
        destructive = []      # (instance, device) pairs that will be overwritten
        declared_vfio = set()
        wants_2m = False
        need_1g = 0   # 1 GB pages follow the devices, not the roles
        need_2m = 0
        # Same figures split by NUMA node. A reactor allocates from its own
        # node, so a host total that is large enough can still starve an
        # instance sitting on the node that holds none of the pages.
        node_2m, node_1g = {}, {}

        for role in host.get("roles", []):
            rtype = role["type"]
            count = role["count"]
            insts = role.get("instances", []) or []

            need_1g += HP1G_FIXED.get(rtype, 0) * count
            if rtype in NEEDS_HUGEPAGES_2M:
                wants_2m = True
                need_2m += hp2m_for(rtype) * count
            if role.get("dpdk") and rtype not in DPDK_CAPABLE:
                err(f"{hn}/{rtype}: dpdk is true, but {rtype} has no DPDK data path")

            if insts and len(insts) != count:
                err(f"{hn}/{rtype}: count is {count} but {len(insts)} instances are defined")
            if not insts:
                warn(f"{hn}/{rtype}: no instances defined — placement has not been computed yet")

            if rtype == "sn":
                for inst in insts:
                    if not inst.get("storage"):
                        err(f"{inst['name']}: an SN with no storage device will exit at startup")

            for inst in insts:
                iname = inst["name"]
                cores = parse_cpuset(inst["cpuset"])

                floor = ROLE_MIN_CORES.get(rtype, 1)
                if len(cores) < floor:
                    err(f"{iname}: cpuset has {len(cores)} core(s), {rtype} needs at least {floor}")

                for c in cores:
                    if c in seen_cores:
                        err(f"{hn}: core {c} is claimed by both {seen_cores[c]} and {iname}")
                    seen_cores[c] = iname

                # A device on one socket driven by cores on another means every
                # I/O crosses the interconnect.
                inode = inst.get("numa_node")
                if inode is not None:
                    for dev in list(inst.get("storage", [])) + ([inst["nic"]] if inst.get("nic") else []):
                        dnode = dev.get("numa_node")
                        if dnode is not None and dnode >= 0 and dnode != inode:
                            warn(f"{iname}: cores are on NUMA node {inode} but "
                                 f"{dev.get('pci') or dev.get('path')} is on node {dnode} — "
                                 f"every I/O will cross the interconnect")

                cc = inst.get("container_cpuset")
                if cc:
                    overlap = parse_cpuset(cc) & cores
                    if overlap:
                        err(f"{iname}: container_cpuset overlaps cpuset on core(s) "
                            f"{sorted(overlap)}. The runtime shim inherits that cgroup and "
                            f"will preempt poll-mode threads.")

                if iso and not cores <= iso:
                    warn(f"{iname}: cpuset {inst['cpuset']} is not fully inside "
                         f"isolcpus {prep['isolcpus']} — those cores stay on the scheduler")

                devices = list(inst.get("storage", []))
                if inst.get("nic"):
                    devices.append(inst["nic"])
                if inst.get("nic") and not role.get("dpdk"):
                    warn(f"{iname}: a NIC is assigned but the role does not set dpdk: true — "
                         f"it will use POSIX networking and the NIC will sit unused")

                nd = inst.get("numa_node")
                if nd is not None and nd >= 0:
                    node_2m[nd] = node_2m.get(nd, 0) + hp2m_for(rtype)
                    node_1g[nd] = node_1g.get(nd, 0) + HP1G_FIXED.get(rtype, 0)

                for dev in devices:
                    ident = dev.get("pci") or dev.get("path")

                    # SPDK's DMA pool scales per NVMe device. bdev_aio
                    # devices (sata/sas/fc/iscsi) never DMA and need none.
                    if dev.get("kind") in DMA_KINDS and dev in inst.get("storage", []):
                        need_1g += HP1G_PER_NVME
                        if nd is not None and nd >= 0:
                            node_1g[nd] = node_1g.get(nd, 0) + HP1G_PER_NVME

                    # A local NVMe controller has to go to SPDK through
                    # vfio-pci. Leaving it on the kernel driver means the SN
                    # cannot attach it at all.
                    if dev.get("kind") == "nvme" and dev.get("driver") != "vfio":
                        err(f"{iname}: {ident} is an NVMe controller with driver "
                            f"{dev.get('driver', 'kernel')!r}. NVMe must be bound to "
                            f"vfio-pci — switch it, or use a kernel-path kind "
                            f"(sata/sas/fc/iscsi/nvmeof) if that is what it really is.")

                    # Assigning a device to a Storage Node destroys what is on it.
                    if rtype == "sn" and dev in inst.get("storage", []):
                        destructive.append((iname, ident))
                    if ident in seen_pci:
                        err(f"{hn}: {ident} is assigned to both {seen_pci[ident]} and {iname}")
                    seen_pci[ident] = iname

                    if dev.get("driver") == "vfio":
                        if not dev.get("pci"):
                            err(f"{iname}: a device with driver vfio needs a pci address")
                        else:
                            declared_vfio.add(dev["pci"])

                        # VFIO grants a whole IOMMU group to one process. Several
                        # devices from one group may go to a single instance — a
                        # RAID-0 set often does — but splitting a group across two
                        # instances makes the second attach fail, and it fails at
                        # SPDK attach time looking like a driver fault.
                        grp = dev.get("iommu_group")
                        if grp is not None:
                            owner = seen_group.get(grp)
                            if owner and owner != iname:
                                err(f"{hn}: IOMMU group {grp} is split between {owner} and "
                                    f"{iname} ({dev['pci']}). VFIO grants a group to a single "
                                    f"process — one of them will fail to attach.")
                            seen_group[grp] = iname

        # Assigning storage to a Storage Node is not reversible. Say so plainly,
        # every run, rather than burying it in documentation.
        for iname, ident in destructive:
            warn(f"DATA LOSS: {iname} will consume {ident} entirely. The Storage Node "
                 f"writes its own layout over the whole device; anything on it is lost.")

        # Host preparation must cover exactly the devices that asked for VFIO.
        bound = set(prep.get("vfio_bind", []))
        for pci in declared_vfio - bound:
            err(f"{hn}: {pci} has driver vfio but is not in host_prep.vfio_bind")
        for pci in bound - declared_vfio:
            warn(f"{hn}: {pci} is bound to vfio-pci but no instance uses it — it will be "
                 f"taken from the kernel for nothing")

        if bound and not prep.get("iommu"):
            err(f"{hn}: VFIO devices are listed but host_prep.iommu is false — "
                f"binding will fail without IOMMU enabled in GRUB")

        have_1g = prep.get("hugepages_1g", 0) or 0
        have_2m = prep.get("hugepages_2m", 0) or 0

        # 1 GB pages. There is no fallback for a device that genuinely DMAs:
        # StorageEngine logs the shortfall and lets SPDK fail, and the SN then
        # halts on its zero-devices-attached check.
        if need_1g and not have_1g:
            err(f"{hn}: needs {need_1g} x 1 GB hugepages but allocates none. SPDK has no "
                f"fallback for NVMe DMA — it will fail to init and the SN will halt "
                f"once no devices attach.")
        elif need_1g > have_1g > 0:
            err(f"{hn}: allocates {have_1g} x 1 GB hugepages but needs {need_1g} "
                f"(4 per NVMe device for SPDK DMA, 2 for an SCC's EAL)")
        elif have_1g and not need_1g:
            warn(f"{hn}: allocates {have_1g} x 1 GB hugepages but nothing here needs them — "
                 f"no NVMe device is configured and bdev_aio does not DMA")

        # Per-node reservations. host_prep.numa is what each node must end up
        # with; the totals above are only what GRUB and sysctl are handed.
        declared = {n["node"]: n for n in (prep.get("numa") or []) if "node" in n}
        nodes_used = sorted(set(node_2m) | set(node_1g))
        if declared:
            for nd in nodes_used:
                d = declared.get(nd)
                if d is None:
                    err(f"{hn}: instances are pinned to NUMA node {nd} but host_prep.numa "
                        f"has no entry for it — nothing will be reserved there")
                    continue
                w2, w1 = node_2m.get(nd, 0), node_1g.get(nd, 0)
                if w2 > (d.get("hugepages_2m", 0) or 0):
                    err(f"{hn}: NUMA node {nd} needs {w2} x 2 MB hugepages but only "
                        f"{d.get('hugepages_2m', 0)} are reserved there")
                if w1 > (d.get("hugepages_1g", 0) or 0):
                    err(f"{hn}: NUMA node {nd} needs {w1} x 1 GB hugepages but only "
                        f"{d.get('hugepages_1g', 0)} are reserved there")
                nm = d.get("memory_mb")
                if nm:
                    want = w2 * 2 + w1 * 1024
                    usable = nm - OS_HEADROOM_MB // max(len(declared), 1)
                    if want > usable:
                        err(f"{hn}: NUMA node {nd} needs {want} MB of hugepages but has only "
                            f"{usable} MB usable locally ({nm} MB). An instance pinned there "
                            f"cannot allocate from another node.")
            # GRUB takes one number for 1 GB pages and spreads it evenly, so the
            # host total has to cover the hungriest node on every node.
            if declared and any(node_1g.values()):
                peak = max(node_1g.values())
                if have_1g and have_1g < peak * len(declared):
                    err(f"{hn}: hugepages_1g is {have_1g}, but GRUB spreads 1 GB pages evenly "
                        f"across {len(declared)} nodes and the hungriest needs {peak} — "
                        f"reserve {peak * len(declared)} so every node gets {peak}")
        elif len(nodes_used) > 1:
            warn(f"{hn}: instances span NUMA nodes {nodes_used} but host_prep.numa is absent, "
                 f"so per-node reservations cannot be checked — a node may end up with none")

        # 2 MB pages. The HD engine sizes its arena from this pool at startup.
        if need_2m > have_2m > 0:
            err(f"{hn}: allocates {have_2m} x 2 MB hugepages but the engine budgets need "
                f"{need_2m} ({need_2m * 2} MB) — reactors will fail to reserve their arena")
        # Memory. A hugepage pool larger than the machine either fails to
        # allocate or leaves the OS short enough to OOM-kill a reactor during
        # startup, when every instance claims its budget at once.
        hp_mb = (prep.get("hugepages_2m", 0) or 0) * 2 + (prep.get("hugepages_1g", 0) or 0) * 1024
        total_mb = prep.get("memory_total_mb")
        # What the instances ask the engine for, and what the pool must hold to
        # give it to them. need_hp_mb is want_mb plus the per-instance mbuf and
        # EAL overhead plus the 1 GB pool, so it is the figure to check the
        # machine against: a host too small for it is too small for want_mb too,
        # and saying so twice helps nobody.
        want_mb = sum(ENGINE_MB.get(r["type"], 0) * r["count"] for r in host.get("roles", []))
        need_hp_mb = need_2m * 2 + need_1g * 1024

        if total_mb:
            usable = total_mb - OS_HEADROOM_MB
            if hp_mb > usable:
                err(f"{hn}: hugepages total {hp_mb} MB but only {usable} MB is usable "
                    f"({total_mb} MB RAM less {OS_HEADROOM_MB} MB for the OS)")
            elif hp_mb and want_mb > hp_mb:
                # Only worth saying when the pool fits the machine. A pool that
                # does not fit is the error to act on first.
                warn(f"{hn}: instances budget {want_mb} MB but the hugepage pool is only "
                     f"{hp_mb} MB — the HD engine allocates from that pool and will fail short")
            if need_hp_mb > usable:
                err(f"{hn}: this host needs {need_hp_mb} MB of hugepages for what it runs, "
                    f"more than the {usable} MB it can spare "
                    f"({total_mb} MB RAM less {OS_HEADROOM_MB} MB for the OS, stack and heap)")
        elif hp_mb:
            warn(f"{hn}: no memory_total_mb recorded, so the {hp_mb} MB hugepage pool "
                 f"cannot be checked against the machine")

        if wants_2m and not prep.get("hugepages_2m"):
            err(f"{hn}: runs a DPDK role but allocates no 2 MB hugepages — EAL will "
                f"not initialise")
        if bound and not prep.get("grub", {}).get("apply", True):
            warn(f"{hn}: binds devices to vfio-pci but grub.apply is false — IOMMU must "
                 f"already be enabled in the running kernel")


def main():
    path = pathlib.Path(sys.argv[1]) if len(sys.argv) > 1 else HERE / "cluster.yaml"
    if not path.exists():
        sys.exit(f"no such file: {path}")

    cfg = yaml.safe_load(path.read_text())

    # Before the schema, because "1860.0 is not of type string" does not tell
    # anyone that their PCI address needed quoting.
    check_pci_quoting(cfg)
    if errors:
        for e in errors:
            print(f"  ERROR    {e}")
        sys.exit(f"\n{len(errors)} error(s) — not deployable")

    try:
        import jsonschema
        jsonschema.validate(cfg, json.loads(SCHEMA.read_text()))
        print(f"  schema   OK   ({path})")
    except ImportError:
        print("  schema   SKIPPED — pip install jsonschema to enable")
    except Exception as exc:  # jsonschema.ValidationError
        first = str(exc).split("\n")[0]
        loc = getattr(exc, "json_path", "?")
        sys.exit(f"  schema   FAILED at {loc}: {first}")

    check_dms_files(cfg, path.parent.resolve())
    check(cfg)

    for w in warnings:
        print(f"  warning  {w}")
    for e in errors:
        print(f"  ERROR    {e}")

    if errors:
        sys.exit(f"\n{len(errors)} error(s), {len(warnings)} warning(s) — not deployable")
    print(f"  checks   OK   ({len(warnings)} warning(s))")


if __name__ == "__main__":
    main()
