#!/usr/bin/env python3
"""
stress.py — Declarative stress matrix for the streaming chain.

Where scenarios.py exercises *correctness* (every option value, exact waveform
fidelity), this module exercises *capacity*: it pushes one axis of load at a time
and records how the two server processes behave. The two components under stress:

  * UDPStreamer (datasource)  — serialises + sends UDPS packets each RT cycle.
    Stress axes: signal size (bytes/packet), signal count, subscriber count
    (UDPStreamer fans a unicast source out to up to 16 clients).
  * StreamHub (hub)           — ring storage, LTTB decimation, WS fan-out, zoom.
    Stress axes: signal size, source count (independent UDPStreamer feeds),
    WS client count, and client request (zoom) rate.

What is measured (not waveform fidelity — at these rates UDP loss is expected and
the hub decimates to PushRate anyway, so end-to-end sample loss is not the gate):

  * survival      — neither server crashed/hung for the whole run (hard gate).
  * liveness      — every WS client kept receiving monotonic, wall-clock-stamped
                    pushes under the load (hard gate).
  * cpu / peakRSS — of marte and hub, captured by proc_perf.py (scaling curves;
                    soft gate against generous ceilings).
  * zoom latency  — p50/p95 round-trip of zoom queries issued under load (soft
                    gate; the headline responsiveness metric for the hub).

Stress-case dict schema (a superset of the scenarios.py scenario dict, so the
gen_data.py / gen_cfg.py generators consume it unchanged):

    { ...all scenario keys (id, network, publishing, max_payload, ws_port,
         sources, oracle, client_checks, trig_signal, row_dt/num_rows/...),
      "shape":   "hub" | "ds_fanout",
      "stress":  {
          "axis":     str,     # which load axis this case varies (for the report)
          "level":    number,  # the axis value (x for the scaling curve)
          "clients":  int,     # parallel WS clients (hub shape)
          "hubs":     int,     # parallel subscriber hubs (ds_fanout shape)
          "reqrate":  float,   # zoom requests/sec/client (0 = one-shot checks)
          "dur":      float,   # live-load duration (s)
          "gate": { "min_frames": int, "max_marte_rss_mb": float,
                    "max_hub_rss_mb": float, "max_zoom_p95_ms": float },
      } }

The matrix keeps every datagram a single UDP fragment (payload < 64 KB): the
UDPSClient reassembly buffer caps a deliverable packet at ~64 KB, so sustained
high-rate streaming must stay below it (see scenarios.py s51). Cases therefore
sweep *count* and *rate*, not oversized single packets.
"""
import itertools

import scenarios as S

# Dedicated port ranges, disjoint from scenarios.py (ws 8101-8151, udp 44610-).
_ws = itertools.count(8300)
_udp = itertools.count(45000, 2)
_data = itertools.count(46000, 2)

# Producer geometry knobs shared across the matrix. 1 kHz producer (default),
# float32 arrays sized so one cycle's array is a single sub-64 KB datagram.
F32 = "float32"


def _f32_arr(name, elements, sampling_rate=1.0e6):
    """A float32 array signal timestamped FirstSample off the shared ns anchor.

    sampling_rate defaults to elements*producer_hz so the per-cycle window equals
    one 1 kHz producer cycle (keeps the hub ring's time axis monotonic, the same
    constraint scenarios.py enforces for First/LastSample)."""
    return S._sig(name, F32, elements, time_mode="FirstSample",
                  time_signal="Tns", sampling_rate=sampling_rate, formula="ramp")


def _source(sid, n_signals, elements, multicast=False):
    """One UDPStreamer source: a uint64 ns anchor + n_signals float32 arrays.

    sampling_rate = elements / row_dt(1 ms) = elements * 1000 keeps each array's
    1 ms window aligned to the 1 kHz cycle regardless of `elements`."""
    rate = elements * 1000.0
    sigs = [S._sig("Tns", "uint64", 1, unit="ns", formula="time_ns",
                   is_time=True)]
    sigs += [_f32_arr(f"S{i}", elements, rate) for i in range(n_signals)]
    return {
        "id": sid, "udp_port": next(_udp),
        "data_port": next(_data) if multicast else None,
        "multicast_group": S.MCAST_GROUP if multicast else None,
        "signals": sigs,
    }


def _packet_bytes(n_signals, elements):
    """Approx DATA payload bytes/packet: ns anchor + n_signals*elements*4."""
    return 8 + n_signals * elements * 4


def mk_stress(sid, axis, level, sources, shape="hub", clients=1, hubs=1,
              reqrate=0.0, dur=6.0, network="unicast", publishing="Strict",
              ratio=None, min_refresh_hz=None, gate=None):
    case = {
        "id": sid, "desc": f"stress {axis}={level}",
        "network": network, "publishing": publishing,
        "ratio": ratio, "min_refresh_hz": min_refresh_hz,
        "max_payload": S.MAX_UDP_PAYLOAD - S.UDPS_HEADER_SIZE,
        "ws_port": next(_ws), "sources": sources,
        "oracle": "analytic", "client_checks": ["live", "zoom"],
        "trig_signal": None, "known_issue": None,
        "row_dt": None, "num_rows": None, "producer_hz": None,
        "shape": shape,
        "stress": {
            "axis": axis, "level": level, "clients": clients, "hubs": hubs,
            "reqrate": reqrate, "dur": dur,
            "gate": gate or {},
        },
    }
    return case


# Default gate: at least a handful of frames per client, generous RSS ceilings,
# and a 1 s p95 zoom round-trip cap. Stress cases tune these per axis.
def _gate(min_frames=5, marte_rss=512.0, hub_rss=1024.0, zoom_p95=1000.0):
    return {"min_frames": min_frames, "max_marte_rss_mb": marte_rss,
            "max_hub_rss_mb": hub_rss, "max_zoom_p95_ms": zoom_p95}


# ── UDPStreamer (datasource) stress ───────────────────────────────────────────
# Single source + single hub; the load lands on the UDPStreamer serialise/send
# path and is read back through one WS client (proc_perf measures marte).

# size: one float32 array, growing element count → bigger single-datagram packet.
_DS_SIZE = [
    mk_stress(f"ds_size_{e}", "ds_signal_elements", e,
              [_source("src", 1, e)], gate=_gate())
    for e in (1000, 4000, 8000, 15000)   # 4 KB → 60 KB packets (sub-64 KB cap)
]

# count: many 1000-element float32 arrays in one source → wider packet + more
# per-cycle serialise work.
_DS_COUNT = [
    mk_stress(f"ds_count_{n}", "ds_signal_count", n,
              [_source("src", n, 1000)], gate=_gate())
    for n in (1, 4, 8, 15)   # 15*4 KB ≈ 60 KB packet (sub-64 KB cap)
]

# clients (fan-out): one source, M subscriber hubs (ds_fanout shape). The
# UDPStreamer copies every packet to each unicast client (max 16); CPU should
# scale with M. Each hub gets its own WS port + one client.
_DS_CLIENTS = [
    mk_stress(f"ds_clients_{m}", "ds_subscriber_hubs", m,
              [_source("src", 4, 2000)], shape="ds_fanout", hubs=m,
              gate=_gate(hub_rss=1024.0))
    for m in (1, 2, 4, 8)
]

# ── StreamHub (hub) stress ────────────────────────────────────────────────────
# All measured against a single hub; proc_perf measures the hub process.

# size: hub ring/decimation cost vs per-signal width.
_HUB_SIZE = [
    mk_stress(f"hub_size_{e}", "hub_signal_elements", e,
              [_source("src", 1, e)], gate=_gate())
    for e in (1000, 4000, 8000, 15000)
]

# sources: N independent UDPStreamer feeds into one hub (each its own udp_port).
_HUB_SOURCES = [
    mk_stress(f"hub_sources_{n}", "hub_source_count", n,
              [_source(f"s{i}", 2, 1000) for i in range(n)],
              gate=_gate(marte_rss=1024.0))
    for n in (1, 2, 4, 8)
]

# clients: one source, C parallel WS clients all recording live + zooming.
_HUB_CLIENTS = [
    mk_stress(f"hub_clients_{c}", "hub_ws_clients", c,
              [_source("src", 2, 2000)], clients=c, gate=_gate())
    for c in (1, 4, 8, 16)
]

# request rate: one source, a few clients each issuing zoom queries at a sustained
# rate; the headline gate is zoom p95 latency under that query load.
_HUB_REQRATE = [
    mk_stress(f"hub_reqrate_{r}", "hub_zoom_reqrate_hz", r,
              [_source("src", 2, 2000)], clients=4, reqrate=r, dur=8.0,
              gate=_gate(zoom_p95=1500.0))
    for r in (5, 20, 50)
]

STRESS_CASES = (_DS_SIZE + _DS_COUNT + _DS_CLIENTS +
                _HUB_SIZE + _HUB_SOURCES + _HUB_CLIENTS + _HUB_REQRATE)


def validate_case(c):
    """Stress-case validity = scenario validity + stress-specific bounds."""
    errs = S.validate_scenario(c)
    st = c.get("stress", {})
    if c["shape"] == "ds_fanout":
        if st["hubs"] < 1 or st["hubs"] > 16:
            errs.append(f"{c['id']}: ds_fanout hubs must be 1..16 "
                        f"(UDPStreamer max unicast clients)")
    if c["shape"] == "hub" and (st["clients"] < 1):
        errs.append(f"{c['id']}: hub clients must be >= 1")
    # Single-datagram ceiling: keep each source's packet < 64 KB so it never
    # needs reassembly (the deliverable-packet cap).
    for src in c["sources"]:
        n_data = sum(1 for s in src["signals"] if not s["is_time"])
        elem = max((s["elements"] for s in src["signals"]
                    if not s["is_time"]), default=0)
        pb = _packet_bytes(n_data, elem)
        if pb >= 65536:
            errs.append(f"{c['id']}: source {src['id']} packet {pb} B exceeds "
                        f"the 64 KB single-datagram cap")
    return errs


if __name__ == "__main__":
    import sys
    ok = True
    seen_ids, seen_ws, seen_udp = set(), set(), set()
    by_axis = {}
    for c in STRESS_CASES:
        errs = validate_case(c)
        if c["id"] in seen_ids:
            errs.append("duplicate id")
        seen_ids.add(c["id"])
        if c["ws_port"] in seen_ws:
            errs.append(f"duplicate ws_port {c['ws_port']}")
        seen_ws.add(c["ws_port"])
        for src in c["sources"]:
            if src["udp_port"] in seen_udp:
                errs.append(f"duplicate udp_port {src['udp_port']}")
            seen_udp.add(src["udp_port"])
        by_axis.setdefault(c["stress"]["axis"], []).append(c["id"])
        print(f"{c['id']:20s} {c['shape']:9s} "
              f"{'OK' if not errs else errs}")
        ok = ok and not errs
    print(f"\n{len(STRESS_CASES)} stress cases across {len(by_axis)} axes, "
          f"{'ALL VALID' if ok else 'INVALID PRESENT'}")
    sys.exit(0 if ok else 1)
