"""Record this machine's place in the cluster, and its measured links, in the lab notebook.

Purpose: capture the topology the rest of Level 4 assumes - which machine this is, what
         roles it holds, which interfaces it has and at what maximum transmission unit,
         whether an RDMA device is present, where the shared model library is, and the
         throughput and round-trip figures measure-links.sh produced - as one JSON line
         so that no number in a later part has to be remembered rather than looked up.
Platform: all (spark, strix, mac, nvidia). Probes that do not apply to a machine are
          absent from the record rather than causing a failure.
Minimum memory: 8 GB
Assumes: python3.9 or newer, and an .env beside this file copied from env-example.txt.
         Run measure-links.sh first if you want the link figures folded in; without it
         the topology is still recorded and the links are simply empty.

Usage: python3 record-topology.py --labbook labbook.md
       python3 record-topology.py --print          (show the record, write nothing)
       python3 record-topology.py --labbook labbook.md --note "jumbo frames enabled"
"""
import argparse
import json
import os
import platform
import re
import shutil
import subprocess
import sys
from datetime import date, datetime, timezone
from pathlib import Path

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


def read_env(path):
    """A small KEY=VALUE reader, so the script and the shell scripts share one .env."""
    values = {}
    try:
        text = Path(path).read_text(encoding="utf-8")
    except OSError:
        return values
    for line in text.splitlines():
        line = line.strip()
        if not line or line.startswith("#") or "=" not in line:
            continue
        key, _, value = line.partition("=")
        value = value.strip()
        if len(value) >= 2 and value[0] == value[-1] and value[0] in "\"'":
            value = value[1:-1]
        values[key.strip()] = value
    return values


def run(cmd, timeout=20):
    """Run a command and return its stdout, or None if it is unavailable or fails."""
    if shutil.which(cmd[0]) is None:
        return None
    try:
        out = subprocess.run(cmd, capture_output=True, text=True, timeout=timeout, check=False)
    except (OSError, subprocess.SubprocessError):
        return None
    if out.returncode != 0:
        return None
    return out.stdout


def interfaces():
    """Every network interface with its maximum transmission unit and link state."""
    found = {}
    sysnet = Path("/sys/class/net")
    if sysnet.is_dir():
        for entry in sorted(sysnet.iterdir()):
            info = {}
            for field, name in (("mtu", "mtu"), ("operstate", "state"), ("speed", "speed_mbit")):
                try:
                    raw = (entry / field).read_text(encoding="utf-8").strip()
                except OSError:
                    continue
                if name in ("mtu", "speed_mbit"):
                    try:
                        info[name] = int(raw)
                    except ValueError:
                        continue
                else:
                    info[name] = raw
            if info:
                found[entry.name] = info
        return found

    # macOS and other BSD systems: ifconfig with no arguments lists every interface.
    text = run(["ifconfig"])
    if not text:
        return found
    for line in text.splitlines():
        match = re.match(r"^(?P<name>[A-Za-z0-9._-]+):\s.*\bmtu\s+(?P<mtu>\d+)", line)
        if match:
            found[match.group("name")] = {"mtu": int(match.group("mtu"))}
    return found


def rdma():
    """Whatever the machine will say about an RDMA path, per platform."""
    info = {}
    ib = Path("/sys/class/infiniband")
    if ib.is_dir():
        info["infiniband_devices"] = sorted(p.name for p in ib.iterdir())
    if shutil.which("ib_write_bw"):
        info["ib_write_bw_present"] = True
    if sys.platform == "darwin":
        info["rdma_ctl_present"] = shutil.which("rdma_ctl") is not None
        version = run(["sw_vers", "-productVersion"])
        if version:
            info["macos"] = version.strip()
    return info or None


def storage(models_dir):
    """Where the shared model library is on this machine, and whether it is a mount."""
    if not models_dir:
        return None
    path = Path(models_dir)
    info = {"path": str(path), "exists": path.is_dir()}
    if path.is_dir():
        try:
            info["is_mount_point"] = path.is_mount()
        except OSError:
            pass
        try:
            info["entries"] = len(list(path.iterdir()))
        except OSError:
            info["entries"] = None
    return info


def link_results(results_dir):
    """Fold in whatever measure-links.sh left behind, without re-running anything."""
    directory = Path(results_dir)
    if not directory.is_dir():
        return {}
    links = {}
    for forward in sorted(directory.glob("*-forward.json")):
        peer = forward.name[: -len("-forward.json")]
        entry = {}
        for label, filename in (("forward", forward),
                                ("reverse", directory / f"{peer}-reverse.json")):
            try:
                data = json.loads(Path(filename).read_text(encoding="utf-8"))
            except (OSError, ValueError):
                continue
            sent = data.get("end", {}).get("sum_sent", {})
            if "bits_per_second" in sent:
                entry[f"{label}_bits_per_second"] = round(float(sent["bits_per_second"]))
        ping_file = directory / f"{peer}-ping.txt"
        try:
            ping_text = ping_file.read_text(encoding="utf-8")
        except OSError:
            ping_text = ""
        match = re.search(r"=\s*([\d.]+)/([\d.]+)/([\d.]+)", ping_text)
        if match:
            entry["rtt_avg_ms"] = float(match.group(2))
        if entry:
            links[peer] = entry
    return links


def build_record(env, results_dir, note):
    peers = [p for p in env.get("CLUSTER_PEERS", "").split() if p]
    roles = [r.strip() for r in env.get("MACHINE_ROLES", "").split(",") if r.strip()]
    return {
        "lab": "part-18/build-and-measure-your-cluster-network",
        "record": "topology",
        "date": date.today().isoformat(),
        "recorded_at": datetime.now(timezone.utc).replace(microsecond=0).isoformat(),
        "machine": env.get("MACHINE_NAME", ""),
        "roles": roles,
        "peers": peers,
        "os": platform.platform(),
        "arch": platform.machine(),
        "python": platform.python_version(),
        "cluster_interface": env.get("CLUSTER_IFACE", "") or None,
        "intended_mtu": int(env["CLUSTER_MTU"]) if env.get("CLUSTER_MTU", "").isdigit() else None,
        "interfaces": interfaces(),
        "rdma": rdma(),
        "model_library": storage(env.get("MODELS_DIR", "")),
        "links": link_results(results_dir),
        "note": note,
    }


def main():
    parser = argparse.ArgumentParser(description=__doc__.splitlines()[0])
    parser.add_argument("--labbook", default=None,
                        help="append one JSON line to this file (for example labbook.md)")
    parser.add_argument("--env", default=str(HERE / ".env"),
                        help="the settings file to read (default: .env beside this script)")
    parser.add_argument("--results-dir", default=str(HERE / "link-results"),
                        help="where measure-links.sh left its raw iperf3 and ping output")
    parser.add_argument("--note", default="", help="free text kept with the record")
    parser.add_argument("--print", dest="show", action="store_true",
                        help="print the record and write nothing")
    args = parser.parse_args()

    env = read_env(args.env)
    # Anything already exported in the shell wins over the file, so a one-off run can
    # override a value without editing .env.
    for key in ("MACHINE_NAME", "MACHINE_ROLES", "CLUSTER_PEERS", "CLUSTER_IFACE",
                "CLUSTER_MTU", "MODELS_DIR"):
        if os.environ.get(key):
            env[key] = os.environ[key]

    if not env.get("MACHINE_NAME"):
        print("MACHINE_NAME is not set. Fill it in in .env, or export it, so the record "
              "can be told apart from the other machines'.", file=sys.stderr)
        return 1

    record = build_record(env, args.results_dir, args.note)
    print(json.dumps(record, indent=2, sort_keys=True))

    if args.show or not args.labbook:
        if not args.labbook:
            print("\nNothing written: pass --labbook labbook.md to record this.")
        return 0

    path = Path(args.labbook)
    with path.open("a", encoding="utf-8") as handle:
        handle.write(json.dumps(record, sort_keys=True) + "\n")
    print(f"\nrecorded in {path}")
    return 0


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