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

cluster.yaml describes machines: this host runs an SCC on cores 2-7, that one
runs two Storage Nodes with these NVMe devices bound to vfio. Kubernetes
describes workloads, and most of that detail has nowhere to go -- a Pod does
not get an IOMMU group. What does carry across is the shape of the deployment
(which roles exist and how many of each), the images, the plane CIDRs and the
sizing, and that is what this writes out.

What it will not do is pretend. Anything in the file that cannot be expressed
in the chart is reported rather than dropped silently, because a values file
that looks complete is worse than one that says what it left behind.

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

Then:
    helm install fermihdi install/helm/fermihdi -f values.yaml
"""
import argparse
import importlib.machinery
import importlib.util
import json
import pathlib
import subprocess
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"

# The sizing model lives in validate-cluster and is imported rather than copied:
# a second table of engine budgets is a second table to get wrong. It has no
# extension and is not on the path, so it is loaded by location.
_spec = importlib.util.spec_from_loader(
    "fermihdi_validate",
    importlib.machinery.SourceFileLoader("fermihdi_validate", str(HERE / "validate-cluster")),
)
_vc = importlib.util.module_from_spec(_spec)
_spec.loader.exec_module(_vc)

# Role -> where it sits in the chart's values. The chart groups by plane, the
# cluster file by host, and this is the whole of the translation between them.
ROLE_PATH = {
    "scc": ("sc", "scc"),
    "sn": ("sc", "sn"),
    "ingester": ("edge", "ingester"),
    "proxy": ("edge", "proxy"),
    "mcp": ("extra", "mcp"),
    "graphql": ("extra", "graphql"),
}
# Published image names, which are not always the role name.
IMAGE_NAMES = {"scc": "scc", "sn": "sn", "proxy": "proxy", "ingester": "ingester",
               "mcp": "hd_mcp", "graphql": "hd_graphql"}
DMS_IMAGES = {"dms-core": ("core", "dms-core"), "dms-ca": ("ca", "dms-ca"),
              "dms-opa-engine": ("opaEngine", "dms-opa"), "dms-webux": ("webux", "dms-webux")}
# Kernel, page cache and everything that is not the engine arena.
NON_ARENA_MB = 1024

notes = []


def note(msg):
    notes.append(msg)


def instances(cfg, rtype):
    out = []
    for host in cfg.get("hosts") or []:
        for role in host.get("roles") or []:
            if role.get("type") != rtype:
                continue
            for inst in role.get("instances") or []:
                out.append((host, role, inst))
    return out


def cores(inst):
    """Width of an instance's cpuset, which is what a CPU request means here.

    Requests and limits are set to the same integer so the Pod lands in the
    Guaranteed QoS class -- with the static CPU manager policy that is what
    gets a reactor whole cores instead of a share of them, which is the nearest
    Kubernetes has to the pinning cluster.yaml describes.
    """
    spec = inst.get("cpuset")
    if not spec or spec == "FIXME":
        return None
    return len(_vc.parse_cpuset(spec))


# The roles the chart asks the scheduler for hugepages for. The GraphQL and MCP
# services run the engine too, but they allocate from /dev/shm rather than from a
# hugepage pool -- so their footprint is the chart's shmSize plus its memory
# limit, not something this can work out from a cpuset.
HUGEPAGE_ROLES = {"scc", "sn", "ingester"}


def sizing(rtype, inst):
    """CPU, memory and hugepages for one instance, from the same model
    validate-cluster sizes hosts with."""
    out = {}
    n = cores(inst)
    if n:
        out["resources"] = {"requests": {"cpu": str(n)}, "limits": {"cpu": str(n)}}
    arena = _vc.ENGINE_MB.get(rtype)
    if arena and rtype in HUGEPAGE_ROLES:
        mem = f"{arena + NON_ARENA_MB}Mi"
        out.setdefault("resources", {}).setdefault("requests", {})["memory"] = mem
        out["resources"].setdefault("limits", {})["memory"] = mem
    pages = _vc.hp2m_for(rtype)
    if pages and rtype in HUGEPAGE_ROLES:
        # Pages of 2 MB, as a quantity the scheduler understands.
        out["hugepages2Mi"] = f"{pages * 2}Mi"
    return out


def image_split(repo, default_registry):
    """A full repository reference back into the registry / name the chart wants."""
    repo = repo.split(":")[0] if repo.count(":") and "/" not in repo.split(":")[-1] else repo
    if "/" in repo:
        registry, name = repo.rsplit("/", 1)
    else:
        registry, name = "", repo
    return (registry if registry and registry != default_registry else None), name


def build(cfg, cluster_path, domain=None, storage_class=None):
    registry = (cfg.get("registry") or {}).get("url", "")
    images = cfg.get("images") or {}
    tag = images.get("tag", "latest")
    # The namespace under the registry that the images live in, matching what
    # the Ansible layer assumes when cluster.yaml names no image explicitly.
    default_registry = f"{registry}/hd" if registry else ""

    values = {
        "global": {
            "imageRegistry": default_registry,
            "imageTag": str(tag),
        },
        "dms": {},
        "sc": {},
        "edge": {},
        "extra": {},
    }
    if storage_class:
        values["global"]["storageClass"] = storage_class

    nets = cfg.get("networks") or {}
    for key, field in (("control_plane", "controlPlaneCIDR"),
                       ("data_plane", "dataPlaneCIDR"),
                       ("observability", "observabilityCIDR")):
        if nets.get(key):
            values["global"][field] = nets[key]
    if not nets:
        note("no networks: block, so the nodes are left to pick their own interfaces. "
             "Fine on a single-homed node pool and wrong on any other.")

    if domain:
        values["ingress"] = {"domain": domain}

    # --- The node roles ------------------------------------------------------
    for rtype, (group, key) in ROLE_PATH.items():
        found = instances(cfg, rtype)
        block = {"enabled": bool(found), "replicas": len(found)}
        if found:
            _, _, first = found[0]
            block.update(sizing(rtype, first))
            widths = {cores(i) for _, _, i in found if cores(i)}
            if len(widths) > 1:
                note(f"{rtype}: instances differ in size ({sorted(widths)} cores) and a "
                     f"Deployment's replicas are identical -- the largest wins, so the "
                     f"smaller ones are over-provisioned. Split them into releases to "
                     f"size them apart.")
                n = max(widths)
                block["resources"]["requests"]["cpu"] = str(n)
                block["resources"]["limits"]["cpu"] = str(n)
            if repo := images.get(rtype):
                reg, name = image_split(repo, default_registry)
                block["image"] = name
                if reg:
                    block["imageRegistry"] = reg
            else:
                block["image"] = IMAGE_NAMES[rtype]
        else:
            block["replicas"] = 0
        values[group][key] = block

    # The chart names Storage Cluster members after the cluster; cluster.yaml
    # names them in full. sc-1-scc and sc-1-sn-1 both give up "sc-1".
    scc = instances(cfg, "scc")
    sn = instances(cfg, "sn")
    name = None
    if scc:
        name = scc[0][2]["name"].rsplit("-scc", 1)[0]
    elif sn:
        name = sn[0][2]["name"].rsplit("-sn", 1)[0]
    if name:
        values["sc"]["clusterName"] = name

    # --- The DMS -------------------------------------------------------------
    dms_cfg = cfg.get("dms") or {}
    runs_dms = bool(instances(cfg, "dms"))
    for chart_key, (values_key, image) in DMS_IMAGES.items():
        block = values["dms"].setdefault(values_key, {})
        block["enabled"] = runs_dms
        if repo := images.get(chart_key):
            reg, name_ = image_split(repo, default_registry)
            block["image"] = name_
            if reg:
                block["imageRegistry"] = reg
        else:
            block["image"] = image
    values["dms"].setdefault("otelCollector", {})["enabled"] = runs_dms

    if not runs_dms:
        # No dms role: the DMS is somewhere else, and dms.url says where. The
        # nodes reach it through the reference Services, so the host goes there
        # and nothing else in the chart has to know.
        url = dms_cfg.get("url", "")
        host = url.split("://", 1)[-1].split("/", 1)[0]
        host, _, port = host.partition(":")
        if host:
            values["dms"]["externalHost"] = host
            values["dms"]["otelExternalHost"] = host
            note(f"no host declares a dms role, so the DMS is treated as external and the "
                 f"nodes are pointed at {host}. The collector is assumed to be there too; "
                 f"set dms.otelExternalHost if it is not.")
            if port and port != "6986":
                note(f"dms.url names port {port}, but the reference Service the nodes use "
                     f"assumes 6986. A DMS on another port needs the chart's node "
                     f"templates changed, or a DNS name that resolves to one on 6986.")
        else:
            note("no host declares a dms role and dms.url is unset: the nodes have nowhere "
                 "to register.")

    if schema := dms_cfg.get("schema_file"):
        path = pathlib.Path(schema)
        if not path.is_absolute():
            path = cluster_path.parent / path
        if path.is_file():
            values["dms"]["core"]["datasetSchema"] = path.read_text()
        else:
            note(f"dms.schema_file: {schema} is not readable from here, so no active "
                 f"schema is installed.")

    # --- What has nowhere to go ---------------------------------------------
    devices, vfio, overflow = [], [], False
    for host in cfg.get("hosts") or []:
        for role in host.get("roles") or []:
            for inst in role.get("instances") or []:
                for dev in (inst.get("storage") or []) + ([inst["nic"]] if inst.get("nic") else []):
                    devices.append((inst["name"], dev.get("pci") or dev.get("path")))
                    if dev.get("driver") == "vfio":
                        vfio.append(inst["name"])
                if inst.get("container_cpuset"):
                    overflow = True
            if role.get("dpdk"):
                note(f"{role['type']} declares dpdk, which the chart does not deploy: every "
                     f"node runs with FERMIHDI_FORCE_POSIX=1. A DPDK data path needs SR-IOV "
                     f"or a host-network Pod with the NIC bound, neither of which this "
                     f"renders.")
    if devices:
        note(f"{len(devices)} device(s) are named for {len(set(n for n, _ in devices))} "
             f"instance(s) -- Storage Nodes in the chart write to a PVC-backed file, not to "
             f"the devices this file binds. Nothing here passes an NVMe through.")
    if vfio:
        note("vfio bindings and IOMMU groups have no expression in the chart; those hosts "
             "would need a device plugin.")
    if overflow:
        note("container_cpuset is a cgroup split that keeps the container runtime off the "
             "poll cores. Kubernetes does the equivalent with the static CPU manager policy "
             "and integer CPU limits, which is what the resources here ask for -- the node "
             "pool has to be configured for it.")
    if any((h.get("host_prep") or {}).get("hugepages_1g") for h in cfg.get("hosts") or []):
        note("1 GB hugepages are requested by this file for SPDK's NVMe DMA. The chart asks "
             "for 2 MB pages only, because nothing in it DMAs to a device.")

    return values


def secrets_values(cfg, cluster_path):
    """The material the chart's Secrets carry, read from what cluster.yaml names.

    Kept out of the main values file on purpose: this one holds private keys.
    """
    import os
    dms = cfg.get("dms") or {}

    def from_file_or_env(key, env):
        if path := dms.get(key):
            p = pathlib.Path(path)
            if not p.is_absolute():
                p = cluster_path.parent / p
            if p.is_file():
                return p.read_text()
        return os.environ.get(env, "")

    out = {
        "caCert": from_file_or_env("ica_cert_file", "FERMIHDI_DMS_CA_CERT_PEM"),
        "caKey": from_file_or_env("ica_key_file", "FERMIHDI_DMS_CA_KEY_PEM"),
        "licensePayload": from_file_or_env("license_file", "FERMIHDI_DMS_LICENSE_PAYLOAD"),
    }
    for values_key, env in (("caBundle", "FERMIHDI_CA_CERT_PEM"),
                            ("coreCert", "FERMIHDI_DMS_CORE_CERT_PEM"),
                            ("coreKey", "FERMIHDI_DMS_CORE_KEY_PEM"),
                            ("tlsCert", "FERMIHDI_DMS_WEBUX_CERT_PEM"),
                            ("tlsKey", "FERMIHDI_DMS_WEBUX_KEY_PEM"),
                            ("webuxFeCert", "FERMIHDI_DMS_WEBUX_FE_CERT_PEM"),
                            ("webuxFeKey", "FERMIHDI_DMS_WEBUX_FE_KEY_PEM"),
                            ("encryptionKey", "FERMIHDI_DMS_ENCRYPTION_KEY"),
                            ("identityKey", "FERMIHDI_DMS_IDENTITY_KEY"),
                            ("feToken", "FERMIHDI_DMS_FE_TOKEN"),
                            ("icaChainPem", "FERMIHDI_DMS_ICA_CHAIN_PEM"),
                            ("opaToken", "FERMIHDI_BOOTSTRAP_TOKEN_DMS_OPA_ENGINE")):
        out[values_key] = os.environ.get(env, "")
    import os as _os
    token = _os.environ.get(
        (cfg.get("dms") or {}).get("bootstrap_token_env", "FERMIHDI_BOOTSTRAP_TOKEN"), "")
    return {
        "secrets": {
            "create": True,
            "dms": out,
            "sc": {"caBundle": out["caBundle"], "sccToken": token, "snTokenPrefix": token},
            "edge": {"caBundle": out["caBundle"], "ingesterToken": token, "proxyToken": token},
        }
    }


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 a second values file holding the DMS's keys and "
                         "tokens, read from the files cluster.yaml names and from the "
                         "environment. 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), help="chart to render (default: %(default)s)")
    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')")

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

    header = [f"# Generated by fermihdi-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 -f this-file.yaml"]
    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 = secrets_values(cfg, path)
        empty = [k for k, v in sec["secrets"]["dms"].items() if not v]
        pathlib.Path(args.secrets).write_text(
            "# Generated by fermihdi-k8s. HOLDS PRIVATE KEYS -- do not commit it.\n"
            "# helm install ... -f values.yaml -f this-file.yaml\n"
            "#\n"
            "# In production, prefer External Secrets or Sealed Secrets over a file:\n"
            "# the chart reads every one of these from Secret objects, so anything that\n"
            "# creates them with these keys will do.\n"
            + yaml.safe_dump(sec, sort_keys=False, width=88))
        pathlib.Path(args.secrets).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 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:
                out = subprocess.run(cmd, capture_output=True, text=True, check=False)
            except FileNotFoundError:
                sys.exit("helm is not installed. Write the values with -o and render them "
                         "wherever helm is.")
            if out.returncode != 0:
                sys.exit(out.stderr.strip())
            print(out.stdout)


if __name__ == "__main__":
    main()
