#!/usr/bin/env python3
"""Ansible dynamic inventory generated from cluster.yaml.

The cluster file already says which machines exist, how to reach them, what
each one runs and where every instance is pinned. Writing that a second time as
an Ansible inventory would mean two files that have to agree, and they would
stop agreeing on the first host anyone added. So there is no inventory file
here — this reads cluster.yaml and answers Ansible's questions from it.

Groups are one per role, so a play can target `fermihdi_sn` and get exactly the
hosts running a Storage Node:

    fermihdi          every host in the file
    fermihdi_scc      hosts running one or more SCCs
    fermihdi_sn       ... Storage Nodes
    fermihdi_proxy    ... and so on, one group per role type in the file

Each host carries three variables:

    fermihdi_instances   every instance on this host, with its role type,
                         cpuset, container_cpuset, devices and NUMA node
                         flattened onto it — the deploy role loops over this
    fermihdi_host_prep   the host_prep block fermihdi-prepare applies
    fermihdi_roles       the raw roles list, for anything the flattening loses

Cluster-wide settings (registry, images, networks, dms) arrive as a single
`fermihdi_cluster` group variable rather than as a dozen separate ones. The one
thing this rewrites on the way through is the file paths in the dms: block,
which are made absolute against the cluster file so they mean the same thing
from whatever directory the deploy is run in.
group_vars/all.yml derives the individual settings from it, so the defaults and
fallbacks live in a file an operator can read and edit, and inventory data
never has to compete with group_vars for precedence.

Usage:
    ansible-playbook -i fermihdi-inventory site.yml
    FERMIHDI_CLUSTER=/path/to/cluster.yaml ansible-playbook -i ... site.yml
    ./fermihdi-inventory --list | jq .          # to see what Ansible sees
"""
import argparse
import json
import os
import pathlib
import sys

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

HERE = pathlib.Path(__file__).resolve().parent
# install/ansible/ -> install/cluster.yaml, the same default validate-cluster
# and fermihdi-prepare use.
DEFAULT_CLUSTER = HERE.parent / "cluster.yaml"


def cluster_path(explicit=None):
    if explicit:
        return pathlib.Path(explicit)
    env = os.environ.get("FERMIHDI_CLUSTER")
    if env:
        return pathlib.Path(env)
    return DEFAULT_CLUSTER


def load(path):
    if not path.exists():
        # Ansible swallows a traceback but shows stderr, so say the useful thing.
        sys.exit(f"fermihdi-inventory: no cluster file at {path}. Point at one with "
                 f"FERMIHDI_CLUSTER=/path/to/cluster.yaml, or write one with "
                 f"install/fermihdi-configure.")
    doc = yaml.safe_load(path.read_text())
    if not isinstance(doc, dict) or "hosts" not in doc:
        sys.exit(f"fermihdi-inventory: {path} is not a cluster definition "
                 f"(no top-level 'hosts').")
    return doc


def instances_of(host):
    """Every instance on a host, with its role's type and dpdk flag folded in.

    cluster.yaml nests instances under roles because that is how they are
    decided. Deploying them is a flat loop over containers, and a play that has
    to reach up to the parent role for the type ends up written twice.
    """
    out = []
    for role in host.get("roles") or []:
        rtype = role.get("type")
        for inst in role.get("instances") or []:
            merged = dict(inst)
            merged["type"] = rtype
            merged["dpdk"] = bool(role.get("dpdk", False))
            out.append(merged)
    return out


# The dms: block names files rather than holding their contents -- the license,
# the DMS-CA's identity, the dataset schema. They are read by whatever runs the
# deploy, which is rarely the directory the cluster file happens to be in, so a
# relative path is resolved against the file itself here. That is the only thing
# this script changes about what it reads: everything else is handed over as
# written and interpreted in group_vars/all.yml.
DMS_FILE_KEYS = ("license_file", "ica_cert_file", "ica_key_file", "schema_file")


def resolve_dms_files(cluster_vars, cluster_file):
    dms = cluster_vars.get("dms")
    if not isinstance(dms, dict):
        return
    for key in DMS_FILE_KEYS:
        value = dms.get(key)
        if value and not os.path.isabs(os.path.expanduser(value)):
            dms[key] = str((cluster_file.parent / value).resolve())
        elif value:
            dms[key] = os.path.expanduser(value)


def build(doc, path):
    hosts = doc.get("hosts") or []
    cluster_vars = {k: v for k, v in doc.items() if k != "hosts"}
    resolve_dms_files(cluster_vars, path)

    inventory = {
        "_meta": {"hostvars": {}},
        "all": {"children": ["fermihdi", "ungrouped"]},
        "fermihdi": {"hosts": [], "vars": {
            "fermihdi_cluster": cluster_vars,
            # The plays copy this file to each host and hand it to
            # fermihdi-prepare, so they have to be looking at the same one the
            # inventory was built from -- not whatever the default resolves to
            # in the play's working directory.
            "fermihdi_cluster_file": str(path),
        }},
    }

    for host in hosts:
        name = host.get("name")
        if not name:
            sys.exit("fermihdi-inventory: a host in the cluster file has no name")

        hostvars = {
            "ansible_host": host.get("address", name),
            "fermihdi_roles": host.get("roles") or [],
            "fermihdi_instances": instances_of(host),
            "fermihdi_host_prep": host.get("host_prep") or {},
        }
        if host.get("ssh_user"):
            hostvars["ansible_user"] = host["ssh_user"]
        if host.get("ssh_key"):
            # Expanded here rather than in a play: ~ in an Ansible variable is
            # not expanded by the ssh connection plugin, and the failure looks
            # like a missing key file.
            hostvars["ansible_ssh_private_key_file"] = os.path.expanduser(host["ssh_key"])

        inventory["_meta"]["hostvars"][name] = hostvars
        inventory["fermihdi"]["hosts"].append(name)

        for role in host.get("roles") or []:
            group = f"fermihdi_{role.get('type')}"
            entry = inventory.setdefault(group, {"hosts": []})
            if name not in entry["hosts"]:
                entry["hosts"].append(name)

    return inventory


def main():
    ap = argparse.ArgumentParser(description=__doc__.splitlines()[0])
    ap.add_argument("--list", action="store_true", help="emit the whole inventory")
    ap.add_argument("--host", metavar="NAME", help="emit one host's variables")
    ap.add_argument("--cluster", metavar="PATH", help="cluster file to read")
    args = ap.parse_args()

    path = cluster_path(args.cluster).resolve()
    doc = load(path)
    inventory = build(doc, path)

    if args.host:
        # _meta means Ansible never calls this, but the protocol requires it and
        # it is how a person checks one host by hand.
        print(json.dumps(inventory["_meta"]["hostvars"].get(args.host, {}), indent=2))
    else:
        print(json.dumps(inventory, indent=2))


if __name__ == "__main__":
    main()
