#!/usr/bin/env python3
"""
stress_run.py — Orchestrator for the streaming-chain stress matrix (stress.py).

Where run_e2e.sh drives scenarios.py for *correctness*, this drives
STRESS_CASES for *capacity*. Per case it:

  1. generates the FileReader input + MARTe app cfg + 1..M StreamHub cfgs
     (reusing gen_data / gen_cfg unchanged — a stress case is a scenario superset),
  2. launches the server stack:
       * "hub"       shape: 1 MARTe UDPStreamer feed + 1 StreamHub, then N
         parallel chain-clients (stress mode) on that one hub,
       * "ds_fanout" shape: 1 MARTe feed + M independent StreamHubs all
         subscribing to the same UDPStreamer (unicast fan-out), 1 client each,
  3. drives the clients in stress mode (liveness + sustained zoom-latency),
  4. snapshots CPU / peak-RSS of marte and every hub *while alive* (proc_perf),
  5. evaluates the per-case gate (survival + liveness hard gates; RSS + zoom-p95
     soft gates against the case's ceilings),

and writes stress_results.json (one record per case, carrying the axis/level for
later scaling-curve plots). Server binaries + LD_LIBRARY_PATH come from the
caller (run_stress.sh sources env.sh); this module only orchestrates.
"""
import argparse
import copy
import json
import os
import signal
import socket
import subprocess
import sys
import time

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
import gen_cfg          # noqa: E402
import gen_data         # noqa: E402
import proc_perf        # noqa: E402
import stress as ST     # noqa: E402

# WS port probing bases. StreamHub does not set SO_REUSEADDR, so a leftover
# listener (e.g. an orphaned hub from an aborted run) keeps a fixed port wedged;
# we therefore probe for genuinely-free ports at launch rather than trusting a
# static assignment. ds_fanout hubs probe from one base, single hubs from another.
FANOUT_WS_BASE = 9000
SINGLE_WS_BASE = 9100


def _popen(cmd, log_path):
    """Launch cmd in its own session, stdout+stderr → log_path."""
    f = open(log_path, "w")
    return subprocess.Popen(cmd, stdout=f, stderr=subprocess.STDOUT,
                            start_new_session=True), f


def _free_ports(n, base):
    """Find n free TCP ports on 127.0.0.1 at/above base.

    Probed ports are closed and handed straight to the hub launch; the race
    window is negligible for this sequential localhost harness, and probing is
    what makes the suite robust to leftover wedged listeners."""
    found = []
    cand = base
    while len(found) < n and cand < base + 1000:
        s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        try:
            s.bind(("127.0.0.1", cand))
            found.append(cand)
        except OSError:
            pass
        finally:
            s.close()
        cand += 1
    if len(found) < n:
        raise RuntimeError(f"could not find {n} free ports from {base}")
    return found


def _hub_ports(case):
    st = case["stress"]
    if case["shape"] == "ds_fanout":
        return _free_ports(st["hubs"], FANOUT_WS_BASE)
    return _free_ports(1, SINGLE_WS_BASE)


def _teardown(procs_files):
    for p, _f in procs_files:
        if p.poll() is None:
            try:
                os.killpg(os.getpgid(p.pid), signal.SIGTERM)
            except (ProcessLookupError, PermissionError):
                pass
    deadline = time.time() + 5
    for p, _f in procs_files:
        try:
            p.wait(timeout=max(0.1, deadline - time.time()))
        except subprocess.TimeoutExpired:
            try:
                os.killpg(os.getpgid(p.pid), signal.SIGKILL)
            except (ProcessLookupError, PermissionError):
                pass
    for _p, f in procs_files:
        try:
            f.close()
        except OSError:
            pass


def run_case(case, bins, work, out):
    sid = case["id"]
    st = case["stress"]
    input_bin = os.path.join(work, f"sinput_{sid}.bin")
    marte_cfg = os.path.join(work, f"sm_{sid}.cfg")

    gen_data.write_input(case, input_bin)
    gen_cfg.write_marte_cfg(case, marte_cfg, input_bin)

    ports = _hub_ports(case)
    hub_pf = []        # (proc, file)
    for i, p in enumerate(ports):
        hc = copy.deepcopy(case)
        hc["ws_port"] = p
        hcfg = os.path.join(work, f"sh_{sid}_{i}.cfg")
        gen_cfg.write_hub_cfg(hc, hcfg)
        proc, f = _popen([bins["hub"], "-cfg", hcfg],
                         os.path.join(out, f"shub_{sid}_{i}.log"))
        hub_pf.append((proc, f))
    time.sleep(1.0)

    marte_proc, marte_f = _popen(
        [bins["marte"], "-l", "RealTimeLoader", "-f", marte_cfg, "-s", "Running"],
        os.path.join(out, f"smarte_{sid}.log"))
    time.sleep(1.5)

    dur, reqrate = st["dur"], st["reqrate"]
    # clean stale client outputs so a missing file means "client failed".
    for i in range(max(st["clients"], st["hubs"])):
        sp = os.path.join(work, f"stress_{sid}_c{i}.json")
        if os.path.exists(sp):
            os.remove(sp)

    client_pf = []
    if case["shape"] == "ds_fanout":
        targets = [(i, ports[i]) for i in range(len(ports))]
    else:
        targets = [(i, ports[0]) for i in range(st["clients"])]
    for cid, wsport in targets:
        cmd = [bins["client"], "-mode", "stress",
               "-hub", f"127.0.0.1:{wsport}", "-scenario", sid,
               "-clientid", str(cid), "-dur", str(dur),
               "-reqrate", str(reqrate), "-out", work,
               "-timeout", f"{int(dur) + 90}s"]
        proc, f = _popen(cmd, os.path.join(out, f"sclient_{sid}_c{cid}.log"))
        client_pf.append((proc, f))

    deadline = time.time() + dur + 60
    for proc, _f in client_pf:
        try:
            proc.wait(timeout=max(1.0, deadline - time.time()))
        except subprocess.TimeoutExpired:
            proc.kill()
    for _p, f in client_pf:
        f.close()

    # survival + perf must be read while the stack is still alive.
    survival = (marte_proc.poll() is None
                and all(h.poll() is None for h, _ in hub_pf))
    perf_marte = proc_perf.snapshot(marte_proc.pid)
    perf_hubs = [proc_perf.snapshot(h.pid) for h, _ in hub_pf]

    _teardown([(marte_proc, marte_f)] + hub_pf)

    clients = []
    for cid, _ in targets:
        jp = os.path.join(work, f"stress_{sid}_c{cid}.json")
        if os.path.exists(jp):
            with open(jp) as f:
                clients.append(json.load(f))

    return _evaluate(case, survival, perf_marte, perf_hubs, clients)


def _rss_mb(rec):
    return rec.get("peak_rss_kb", 0) / 1024.0 if rec.get("avail") else 0.0


def _evaluate(case, survival, perf_marte, perf_hubs, clients):
    g = case["stress"].get("gate", {})
    fails = []

    if not survival:
        fails.append("server process died during run")
    if not clients:
        fails.append("no client results recorded")

    min_frames = g.get("min_frames", 1)
    for c in clients:
        cid = c.get("clientId")
        if c.get("frames", 0) < min_frames:
            fails.append(f"c{cid} frames {c.get('frames')} < {min_frames}")
        if not c.get("monotonic", False):
            fails.append(f"c{cid} non-monotonic time axis")
        if not c.get("wallclock", False):
            fails.append(f"c{cid} non-wallclock timestamps")

    marte_rss = _rss_mb(perf_marte)
    if "max_marte_rss_mb" in g and marte_rss > g["max_marte_rss_mb"]:
        fails.append(f"marte peakRSS {marte_rss:.0f}MB > {g['max_marte_rss_mb']}MB")
    hub_rss_max = max((_rss_mb(h) for h in perf_hubs), default=0.0)
    if "max_hub_rss_mb" in g and hub_rss_max > g["max_hub_rss_mb"]:
        fails.append(f"hub peakRSS {hub_rss_max:.0f}MB > {g['max_hub_rss_mb']}MB")

    p95s = [c["zoomP95ms"] for c in clients if c.get("zoomCount", 0) > 0]
    zoom_p95 = max(p95s) if p95s else 0.0
    p50s = [c["zoomP50ms"] for c in clients if c.get("zoomCount", 0) > 0]
    zoom_p50 = max(p50s) if p50s else 0.0
    if "max_zoom_p95_ms" in g and zoom_p95 > g["max_zoom_p95_ms"]:
        fails.append(f"zoom p95 {zoom_p95:.0f}ms > {g['max_zoom_p95_ms']}ms")

    st = case["stress"]
    return {
        "id": case["id"], "shape": case["shape"],
        "axis": st["axis"], "level": st["level"],
        "status": "PASS" if not fails else "FAIL",
        "survival": survival,
        "clients": len(clients),
        "min_frames": min(((c.get("frames", 0)) for c in clients), default=0),
        "marte_cpu_s": perf_marte.get("cpu_s", 0.0),
        "marte_rss_mb": round(marte_rss, 1),
        "hub_cpu_s": round(sum(h.get("cpu_s", 0.0) for h in perf_hubs), 2),
        "hub_rss_mb": round(hub_rss_max, 1),
        "zoom_count": sum(c.get("zoomCount", 0) for c in clients),
        "zoom_fail": sum(c.get("zoomFail", 0) for c in clients),
        "zoom_p50_ms": round(zoom_p50, 1),
        "zoom_p95_ms": round(zoom_p95, 1),
        "fails": fails,
    }


def main():
    p = argparse.ArgumentParser(description="Run the streaming-chain stress matrix")
    p.add_argument("--marte", required=True, help="MARTeApp.ex path")
    p.add_argument("--hub", required=True, help="StreamHub.ex path")
    p.add_argument("--client", required=True, help="chain-client path")
    p.add_argument("--work", required=True, help="scratch dir")
    p.add_argument("--out", required=True, help="report/log dir")
    p.add_argument("--only", default="", help="run a single case id")
    p.add_argument("--axis", default="", help="run only cases on this axis")
    args = p.parse_args()

    os.makedirs(args.work, exist_ok=True)
    os.makedirs(args.out, exist_ok=True)
    bins = {"marte": args.marte, "hub": args.hub, "client": args.client}

    cases = ST.STRESS_CASES
    if args.only:
        cases = [c for c in cases if c["id"] == args.only]
    if args.axis:
        cases = [c for c in cases if c["stress"]["axis"] == args.axis]
    if not cases:
        print("no stress cases selected", file=sys.stderr)
        sys.exit(1)

    results = []
    for c in cases:
        errs = ST.validate_case(c)
        if errs:
            print(f"══ {c['id']}: INVALID {errs} ══")
            results.append({"id": c["id"], "axis": c["stress"]["axis"],
                            "level": c["stress"]["level"], "status": "FAIL",
                            "fails": errs})
            continue
        print(f"\n══ stress {c['id']}  ({c['shape']} {c['stress']['axis']}="
              f"{c['stress']['level']}) ══")
        rec = run_case(c, bins, args.work, args.out)
        results.append(rec)
        tag = rec["status"]
        print(f"  {tag}  frames>={rec.get('min_frames')} "
              f"marteCPU={rec.get('marte_cpu_s', 0):.1f}s "
              f"marteRSS={rec.get('marte_rss_mb', 0):.0f}MB "
              f"hubRSS={rec.get('hub_rss_mb', 0):.0f}MB "
              f"zoom p50/p95={rec.get('zoom_p50_ms', 0):.0f}/"
              f"{rec.get('zoom_p95_ms', 0):.0f}ms")
        if rec.get("fails"):
            for fmsg in rec["fails"]:
                print(f"    - {fmsg}")

    overall = bool(results) and all(r["status"] == "PASS" for r in results)
    doc = {"overall": "PASS" if overall else "FAIL", "cases": results}
    rp = os.path.join(args.out, "stress_results.json")
    with open(rp, "w") as f:
        json.dump(doc, f, indent=2)
    npass = sum(r["status"] == "PASS" for r in results)
    print(f"\nstress_results.json: {npass}/{len(results)} pass → {doc['overall']}")
    sys.exit(0 if overall else 1)


if __name__ == "__main__":
    main()