#!/usr/bin/env python3
"""Phase A — bring up the Federated Grid Fault Detection slice.

This is the *weave*: it provisions the slice, installs the stack, uploads the
``fl/`` package, publishes ``results/slice_info.json``, and exits. It does **not**
run federated learning and does **not** poll in a loop — the notebook
(``federated_learning.ipynb``, Phase B) drives the FL against the ready slice.

Site placement is **discovered at submit time**, never hardcoded:
- GPU runs query FABRIC for every site's available GPU inventory (all models),
  keep the sites that have a GPU free *and* room for the client VM, and spread
  the clients across those sites — each VM gets the GPU model its site actually
  has. Set ``GPU_MODEL`` to pin one model; leave it ``auto`` to search all.
- CPU runs (and the aggregator) auto-place (``site=None``); FABRIC picks a site.
- ``SITE_MAP`` overrides placement per ISO when you need a specific site.

Actions:
    python3 experiment.py start  SLICE_NAME       # provision + configure + READY
    python3 experiment.py status SLICE_NAME       # one-shot health check
    python3 experiment.py extend SLICE_NAME [DAYS] # renew the lease (DAYS from now)
    python3 experiment.py stop   SLICE_NAME       # delete the slice

Configuration comes from environment variables (set by weave.sh); see weave.md §6.1.
``start`` auto-extends the lease to ``LEASE_DAYS`` (default 14) from now so the
slice survives long-running experimentation; re-extend any time with ``extend``.
"""

from __future__ import annotations

import json
import os
import sys
from concurrent.futures import ThreadPoolExecutor, as_completed

from fabrictestbed_extensions.fablib.fablib import FablibManager

# --- layout constants (weave.md §2/§3.1) ------------------------------------
IMAGE = "default_ubuntu_22"
ISO_ROLES = ["ercot", "caiso", "pjm", "miso", "spp", "nyiso"]
# FABRIC per-host GPU availability field (list_hosts) -> add_component model name.
GPU_FIELD_MODEL = {
    "rtx6000_available": "GPU_RTX6000",
    "a40_available": "GPU_A40",
    "a30_available": "GPU_A30",
    "tesla_t4_available": "GPU_TeslaT4",
}
MODEL_FIELD = {v: k for k, v in GPU_FIELD_MODEL.items()}
GPU_MODELS = list(GPU_FIELD_MODEL.values())
CLIENT_GPU_SPEC = {"cores": 8, "ram": 32, "disk": 100}
CLIENT_CPU_SPEC = {"cores": 4, "ram": 16, "disk": 50}
AGG_SPEC = {"cores": 4, "ram": 16, "disk": 50}
STAR_PORT = 8080
MESH_PORT = 8090
REMOTE_PY = "$HOME/fl-venv/bin/python"      # venv python created by setup-common.sh
RESULTS_DIR = os.path.join(os.path.dirname(os.path.abspath(__file__)), "results")
FL_DIR = os.path.join(os.path.dirname(os.path.abspath(__file__)), "fl")
TOOLS_DIR = os.path.join(os.path.dirname(os.path.abspath(__file__)), "tools")


# --- config from environment (weave.sh) -------------------------------------
def _env_bool(name: str, default: bool) -> bool:
    return os.environ.get(name, str(default)).strip().lower() in ("1", "true", "yes", "y")


def _parse_site_map(raw) -> dict:
    if isinstance(raw, dict):
        return {k.lower(): v for k, v in raw.items()}
    out = {}
    for pair in str(raw).split(","):
        pair = pair.strip()
        if "=" in pair:
            k, v = pair.split("=", 1)
            out[k.strip().lower()] = v.strip()
    return out


def _parse_exclude(raw) -> set:
    if isinstance(raw, (set, list, tuple)):
        return {str(s).strip().upper() for s in raw if str(s).strip()}
    return {s.strip().upper() for s in str(raw).split(",") if s.strip()}


def make_config(num_clients: int = 4, use_gpu: bool = True, gpu_model: str = "auto",
                agg_gpu: bool = False, site_map="", cuda_tag: str = "cu121",
                exclude_sites="", lease_days: int = 14) -> dict:
    """Build a provisioning config directly (used by the notebook / callers).

    ``gpu_model``: "auto" searches all GPU models; else pin one (GPU_RTX6000 |
    GPU_A40 | GPU_A30 | GPU_TeslaT4). ``site_map`` / ``exclude_sites`` accept a
    string ("ercot=TACC,agg=RENC" / "EDUKY,EDC") or a dict/list. ``lease_days``
    is how far out ``start`` renews the lease (FABRIC caps this, typically 14).
    """
    return {
        "num_clients": max(1, min(int(num_clients), len(ISO_ROLES))),
        "use_gpu": bool(use_gpu),
        "gpu_model": str(gpu_model).strip(),
        "agg_gpu": bool(agg_gpu),
        "site_map": _parse_site_map(site_map),
        "cuda_tag": cuda_tag,
        "exclude_sites": _parse_exclude(exclude_sites),
        "lease_days": max(1, int(lease_days)),
    }


def load_config() -> dict:
    """Build the provisioning config from environment variables (weave.sh path)."""
    return make_config(
        num_clients=int(os.environ.get("NUM_CLIENTS", "4")),
        use_gpu=_env_bool("USE_GPU", True),
        gpu_model=os.environ.get("GPU_MODEL", "auto"),
        agg_gpu=_env_bool("AGG_GPU", False),
        site_map=os.environ.get("SITE_MAP", ""),
        cuda_tag=os.environ.get("CUDA_TAG", "cu121"),
        exclude_sites=os.environ.get("EXCLUDE_SITES", ""),
        lease_days=int(os.environ.get("LEASE_DAYS", "14")),
    )


# --- dynamic host discovery & placement (host-level, like pegasus-fabric) ----
def _gpu_fields(cfg: dict):
    """Which per-host GPU availability fields to consider."""
    m = cfg["gpu_model"].lower()
    if m in ("", "auto", "any", "all"):
        return list(GPU_FIELD_MODEL.keys())
    if cfg["gpu_model"] in MODEL_FIELD:
        return [MODEL_FIELD[cfg["gpu_model"]]]
    return list(GPU_FIELD_MODEL.keys())


def _host_site_map(res) -> dict:
    """host_name -> site_name across the whole testbed."""
    mapping = {}
    for site in res.get_site_names():
        try:
            for hname in res.get_hosts_by_site(site):
                mapping[hname] = site
        except Exception:
            pass
    return mapping


def _best_gpu_on_host(row: dict, fields) -> tuple:
    """(field, available) for the most-available GPU model *on this host*."""
    best_field, best_avail = None, 0
    for f in fields:
        a = int(row.get(f, 0) or 0)
        if a > best_avail:
            best_field, best_avail = f, a
    return best_field, best_avail


def discover_hosts(res, spec: dict, require_gpu: bool, gpu_fields, excluded) -> list:
    """Find individual HOSTS that can actually provision the VM.

    Uses per-host availability (list_hosts) so we never pick a site whose
    aggregate looks free but whose chosen host can't fit the VM or lacks the GPU.
    Returns [{site, host, model, gpu_avail, fit}] (model=None when require_gpu=False).
    """
    host_site = _host_site_map(res)
    try:
        rows = res.list_hosts(output="list", quiet=True) or []
    except Exception as exc:
        print(f"### PROGRESS: WARNING — list_hosts failed ({exc}); auto-placement only")
        return []

    cands = []
    for row in rows:
        host = row.get("name")
        site = host_site.get(host)
        if not host or not site or site in excluded:
            continue
        if str(row.get("state", "")).lower() != "active":
            continue
        cores = int(row.get("cores_available", 0) or 0)
        ram = int(row.get("ram_available", 0) or 0)
        disk = int(row.get("disk_available", 0) or 0)
        if cores < spec["cores"] or ram < spec["ram"] or disk < spec["disk"]:
            continue
        fit = min(cores // spec["cores"], ram // spec["ram"], disk // spec["disk"])
        if fit < 1:
            continue
        if require_gpu:
            field, avail = _best_gpu_on_host(row, gpu_fields)
            if not field or avail < 1:
                continue
            cands.append({"site": site, "host": host, "model": GPU_FIELD_MODEL[field],
                          "gpu_avail": avail, "fit": min(fit, avail)})
        else:
            cands.append({"site": site, "host": host, "model": None,
                          "gpu_avail": 0, "fit": fit})
    return cands


def allocate_hosts(cands: list, n: int) -> list:
    """Pick n host records, one per distinct site first (geographic spread), then fill."""
    by_site = {}
    for c in cands:
        by_site.setdefault(c["site"], []).append(c)
    for s in by_site:
        by_site[s].sort(key=lambda x: (-x["fit"], x["host"]))
    site_order = sorted(by_site, key=lambda s: (-by_site[s][0]["fit"], s))
    remaining = {(c["site"], c["host"]): c["fit"] for c in cands}

    picks = []
    while len(picks) < n:
        progressed = False
        for s in site_order:
            for c in by_site[s]:
                if remaining[(s, c["host"])] > 0:
                    picks.append(c)
                    remaining[(s, c["host"])] -= 1
                    progressed = True
                    break
            if len(picks) == n:
                break
        if not progressed:
            break
    return picks


def resolve_placement(res, cfg: dict):
    """Return (agg, [client, ...]) where each entry is (iso|None, site, host, model).

    site/host None => auto-placement; model None => CPU node. Placement is
    host-level: each client is pinned to a specific host that actually has room
    for the VM and (GPU runs) a free GPU of the model requested on that host.
    """
    n = cfg["num_clients"]
    roles = ISO_ROLES[:n]
    site_map = cfg["site_map"]
    excluded = cfg["exclude_sites"]

    spec = CLIENT_GPU_SPEC if cfg["use_gpu"] else CLIENT_CPU_SPEC
    gpu_fields = _gpu_fields(cfg)
    want = ("all models" if len(gpu_fields) > 1 else GPU_FIELD_MODEL[gpu_fields[0]]) \
        if cfg["use_gpu"] else "cpu"
    print(f"### PROGRESS: Searching hosts for a {spec['cores']}c/{spec['ram']}G/"
          f"{spec['disk']}G client ({'GPU: ' + want if cfg['use_gpu'] else 'CPU'})...")

    cands = discover_hosts(res, spec, cfg["use_gpu"], gpu_fields, excluded)
    if cands:
        kind = "GPU-capable hosts" if cfg["use_gpu"] else "capable hosts"
        print(f"### PROGRESS: {kind} found: {len(cands)} across "
              f"{len(set(c['site'] for c in cands))} sites")
        for c in sorted(cands, key=lambda x: (-x["fit"], x["site"]))[:12]:
            print(f"###   {c['site']:6s} host={c['host']:14s} "
                  f"gpu={c['model'] or '-':12s} fits={c['fit']}")
    else:
        print("### PROGRESS: WARNING — no host matched; falling back to auto-placement")

    picks = allocate_hosts(cands, n)
    if len(picks) < n:
        print(f"### PROGRESS: WARNING — only {len(picks)} host slot(s) for {n} clients; "
              f"remaining clients auto-place (may fail to schedule)")

    fallback_model = (picks[0]["model"] if (picks and picks[0]["model"]) else
                      (cfg["gpu_model"] if cfg["gpu_model"] in MODEL_FIELD else GPU_MODELS[0]))

    def _cand_on_site(site):
        for c in cands:
            if c["site"] == site:
                return c
        return None

    clients, queue = [], list(picks)
    for iso in roles:
        if iso in site_map:                          # explicit site override
            site = site_map[iso]
            c = _cand_on_site(site)
            host = c["host"] if c else None
            model = (c["model"] if c and c["model"] else fallback_model) if cfg["use_gpu"] else None
            clients.append((iso, site, host, model))
        elif queue:                                  # discovered host (pinned)
            c = queue.pop(0)
            clients.append((iso, c["site"], c["host"], c["model"]))
        else:                                        # auto-place fallback
            clients.append((iso, None, None, fallback_model if cfg["use_gpu"] else None))

    # Aggregator (CPU unless AGG_GPU; auto-placed — its site is not the experiment).
    agg_site = site_map.get("agg")
    agg = (None, agg_site, None, fallback_model if cfg["agg_gpu"] else None)
    return agg, clients


def extend_lease(slice_obj, days: int) -> bool:
    """Renew the slice lease to ``days`` from now. Best-effort: warns, never raises.

    FABRIC provisions a short default lease (often 24h); long experimentation
    needs a longer one. ``renew(days=...)`` re-submits with the new lease only
    (``post_boot_config=False``), so it does not re-run node setup. Returns True
    on success.
    """
    try:
        before = slice_obj.get_lease_end()
        slice_obj.renew(days=int(days))
        after = slice_obj.get_lease_end()
        print(f"### PROGRESS: Lease extended to ~{days} day(s) from now "
              f"(was {before}, now {after})")
        return True
    except Exception as exc:
        name = ""
        try:
            name = slice_obj.get_name()
        except Exception:
            pass
        print(f"### PROGRESS: WARNING — could not extend lease ({exc}); the slice "
              f"keeps its default lease. Retry later with: "
              f"python3 experiment.py extend {name} <days>")
        return False


def get_fabnet_ip(node):
    """FABNetv4 IP by network-name pattern (add_fabnet() returns None)."""
    site = node.get_site()
    iface = node.get_interface(network_name=f"FABNET_IPv4_{site}")
    return str(iface.get_ip_addr()) if iface else None


# --- provisioning ------------------------------------------------------------
def preview_placement(cfg: dict = None) -> dict:
    """Dry run — discover hosts/GPUs and show the planned placement, no provisioning.

    Handy from the provisioning notebook to inspect what a run *would* create
    before committing. Returns {"agg": ..., "clients": [...]}.
    """
    cfg = cfg if cfg is not None else load_config()
    fablib = FablibManager()
    res = fablib.get_resources(update=True)
    agg, clients = resolve_placement(res, cfg)
    return {"agg": agg, "clients": clients}


def start(slice_name: str, cfg: dict = None) -> None:
    cfg = cfg if cfg is not None else load_config()
    fablib = FablibManager()
    n = cfg["num_clients"]
    profile = "GPU" if cfg["use_gpu"] else "CPU"
    print(f"### PROGRESS: Bringing up '{slice_name}' — {profile} profile, "
          f"agg + {n} clients")

    res = fablib.get_resources(update=True)   # host-level availability (GPU & CPU)
    agg_plan, client_plan = resolve_placement(res, cfg)
    _, agg_site, agg_host, agg_model = agg_plan

    slice_obj = fablib.new_slice(name=slice_name)

    # Aggregator.
    agg = slice_obj.add_node(name="agg", site=agg_site, host=agg_host,
                             cores=AGG_SPEC["cores"], ram=AGG_SPEC["ram"],
                             disk=AGG_SPEC["disk"], image=IMAGE)
    if agg_model:
        agg.add_component(model=agg_model, name="gpu1")
    agg.add_fabnet()

    # Clients — each pinned to the discovered host, requesting that host's GPU model.
    spec = CLIENT_GPU_SPEC if cfg["use_gpu"] else CLIENT_CPU_SPEC
    client_names = []
    print("### PROGRESS: Client placement:")
    for iso, site, host, gpu_model in client_plan:
        name = f"client-{iso}"
        client_names.append(name)
        node = slice_obj.add_node(name=name, site=site, host=host,
                                  cores=spec["cores"], ram=spec["ram"],
                                  disk=spec["disk"], image=IMAGE)
        if gpu_model:
            node.add_component(model=gpu_model, name="gpu1")
        print(f"###   {name:16s} site={site or 'auto':6s} "
              f"host={host or 'auto':16s} gpu={gpu_model or '-'}")
        node.add_fabnet()

    print("### PROGRESS: Submitting slice (5-15 min; longer for GPU)...")
    slice_obj.submit()
    slice_obj.wait_ssh(progress=True)

    # Re-fetch — original node objects go stale after submit.
    slice_obj = fablib.get_slice(name=slice_name)
    all_nodes = slice_obj.get_nodes()

    # Upload code + tools, then run setup, in parallel across nodes.
    gpu_hint = "true" if cfg["use_gpu"] else "false"

    def provision_node(node):
        name = node.get_name()
        node.upload_directory(FL_DIR, ".")
        node.upload_directory(TOOLS_DIR, ".")
        script = "setup-agg.sh" if name == "agg" else "setup-client.sh"
        # Do NOT pass fablib's timeout= — it wraps the command in `sudo timeout`,
        # which runs setup as root with HOME=/root and puts the venv in the wrong
        # place. Run as the login user; cap runtime with our own `timeout`.
        node.execute(f"timeout 2400 bash ~/tools/{script} {gpu_hint} {cfg['cuda_tag']}", quiet=True, 
                     output_file=f"results/setup-{node.get_name()}.log")
        # execute() does NOT raise on a non-zero script exit, so verify explicitly.
        out, _ = node.execute(
            f"{REMOTE_PY} -c 'import torch, flwr' 2>&1 && echo VERIFY_OK", quiet=True)
        if "VERIFY_OK" not in out:
            raise RuntimeError(out.strip()[-300:] or "torch/flwr import failed")
        return name

    print("### PROGRESS: Uploading fl/ + tools/ and installing stack (parallel)...")
    failed = []
    with ThreadPoolExecutor(max_workers=len(all_nodes)) as pool:
        futures = {pool.submit(provision_node, nd): nd.get_name() for nd in all_nodes}
        for fut in as_completed(futures):
            nm = futures[fut]
            try:
                fut.result()
                print(f"### PROGRESS: {nm} configured + verified")
            except Exception as exc:
                failed.append(nm)
                print(f"### PROGRESS: WARNING — setup on {nm} FAILED: {exc}")

    # Discover the actual GPU model landed on each node (for slice_info).
    def _gpu_model_of(node):
        try:
            for comp in node.get_components():
                m = str(comp.get_model())
                if "GPU" in m:
                    return m
        except Exception:
            pass
        return None

    def _host_of(node):
        try:
            return node.get_host()
        except Exception:
            return None

    # Publish the Phase-A -> Phase-B hand-off.
    nodes_info = []
    for node in slice_obj.get_nodes():
        name = node.get_name()
        site = node.get_site()
        iso = name[len("client-"):] if name.startswith("client-") else None
        lat, lon = (0.0, 0.0)
        try:
            lat, lon = res.get_location_lat_long(site)
        except Exception:
            pass
        nodes_info.append({
            "name": name,
            "role": "aggregator" if name == "agg" else "client",
            "iso": iso,
            "client_index": ISO_ROLES.index(iso) if iso in ISO_ROLES else None,
            "site": site,
            "host": _host_of(node),
            "gpu_model": _gpu_model_of(node),
            "lat": lat, "long": lon,
            "fabnet_ip": get_fabnet_ip(node),
            "mgmt_ip": node.get_management_ip(),
        })

    info = {
        "slice_name": slice_name,
        "use_gpu": cfg["use_gpu"],
        "num_clients": n,
        "agg_node": "agg",
        "client_nodes": client_names,
        "star_port": STAR_PORT,
        "mesh_port": MESH_PORT,
        "remote_python": REMOTE_PY,
        "remote_code_dir": "~",          # fl/ uploaded to ~/fl
        "results_dir_remote": "~/results",
        "nodes": nodes_info,
    }
    os.makedirs(RESULTS_DIR, exist_ok=True)
    info_path = os.path.join(RESULTS_DIR, "slice_info.json")
    with open(info_path, "w") as fh:
        json.dump(info, fh, indent=2)

    print("### PROGRESS: slice_info.json written:")
    for nd in nodes_info:
        print(f"###   {nd['name']:16s} {nd['role']:11s} site={nd['site']} "
              f"host={nd['host'] or '-'} gpu={nd['gpu_model'] or '-'} "
              f"fabnet={nd['fabnet_ip']}")

    # Extend the lease so the slice survives long-running experimentation.
    print(f"### PROGRESS: Extending lease to {cfg['lease_days']} day(s)...")
    extend_lease(slice_obj, cfg["lease_days"])

    if failed:
        print(f"### PROGRESS: WARNING — stack install failed/unverified on: "
              f"{', '.join(failed)}")
        print("### PROGRESS: The slice is up but NOT ready to train. Repair without "
              "re-provisioning by re-running the setup cell (§3) in "
              "federated_learning.ipynb (it re-uploads fl/ + reinstalls the stack).")
    else:
        print(f"### PROGRESS: READY — slice '{slice_name}' is up and verified. "
              f"Open federated_learning.ipynb to run the FL.")


def status(slice_name: str) -> None:
    """One-shot health check (NOT a polling loop)."""
    fablib = FablibManager()
    try:
        slice_obj = fablib.get_slice(name=slice_name)
        state = str(slice_obj.get_state())
        if "StableOK" not in state:
            print(f"### PROGRESS: slice state is {state}")
            sys.exit(1)
        for node in slice_obj.get_nodes():
            out, _ = node.execute("echo ok", quiet=True)
            if "ok" not in out:
                raise RuntimeError(f"{node.get_name()} not responding")
        print(f"### PROGRESS: Healthy — {state}")
    except Exception as exc:
        print(f"### PROGRESS: ERROR — {exc}")
        sys.exit(1)


def extend(slice_name: str, days: int = None) -> None:
    """Renew an already-running slice's lease to ``days`` from now.

    ``days`` defaults to ``LEASE_DAYS`` (env, default 14). Does not touch the
    nodes or re-run setup — just extends the lease.
    """
    days = int(days) if days is not None else int(os.environ.get("LEASE_DAYS", "14"))
    fablib = FablibManager()
    try:
        slice_obj = fablib.get_slice(name=slice_name)
    except Exception as exc:
        print(f"### PROGRESS: ERROR — slice '{slice_name}' not found: {exc}")
        sys.exit(1)
    ok = extend_lease(slice_obj, days)
    sys.exit(0 if ok else 1)


def stop(slice_name: str) -> None:
    fablib = FablibManager()
    print(f"### PROGRESS: Deleting slice '{slice_name}'...")
    try:
        fablib.get_slice(name=slice_name).delete()
        print("### PROGRESS: Slice deleted.")
    except Exception as exc:
        print(f"### PROGRESS: Could not delete (may already be gone): {exc}")


if __name__ == "__main__":
    actions = {"start": start, "status": status, "stop": stop, "extend": extend}
    args = sys.argv[1:]
    # Accept either "ACTION SLICE_NAME [DAYS]", just "SLICE_NAME" (LoomAI passes
    # the slice name as the first positional arg), or nothing (env/default).
    if args and args[0] in actions:
        action = args[0]
        name = args[1] if len(args) > 1 else os.environ.get("SLICE_NAME", "fl-grid")
    elif args:
        action, name = "start", args[0]
    else:
        action, name = "start", os.environ.get("SLICE_NAME", "fl-grid")
    if action == "extend":
        # optional 3rd positional: number of days
        days = int(args[2]) if len(args) > 2 and args[2].isdigit() else None
        extend(name, days)
    else:
        actions[action](name)
