#!/usr/bin/env python3
"""
m10_bench_smoke.py — M10-bench broker+router ceiling test (10k/s on bench profile)

Usage:
  docker compose --profile bench up -d
  python3 scripts/m10_bench_smoke.py

Exit code 0 = all green. Exit code 1 = assertion failed.

Step 1 — pre-flight (bench profile up, deliverd-bench consuming)
Step 2 — 10k/s soak (5 min, ramp 15s)
Step 3 — broker ceiling check (p99 ≤ 50ms, NATS queue depth < 1000)
Step 4 — teardown

NOTE: The bench profile has 2 loadgen-bench instances × 5000/s = 10k/s.
M10_PLAN.md described 10 instances × 5000/s = 50k/s; only 2 exist in
docker-compose.yml. This smoke targets 10k/s and adapts thresholds accordingly.
"""
import subprocess
import sys
import time
import urllib.request
import urllib.parse
import json

sys.path.insert(0, __file__.rsplit("/", 1)[0])
import m10_lib as lib

PROM = "http://localhost:9090"
NATS_MONITOR = "http://localhost:8222"
BENCH_DURATION_MIN = 5       # minutes
BENCH_RAMP_SEC = 15           # ramp-up seconds
CLUSTER_TARGET = 10000        # alerts/sec cluster-wide target (2 × 5000)
P99_ROUTER_THRESHOLD = 0.050  # seconds — router p99 must be under this
NATS_QUEUE_DEPTH_MAX = 1000   # NATS queue depth must stay below this
RATE_TOLERANCE = 0.10         # ±10% (loadgen ramp variance is higher at 10k/s)


def pass_(msg: str):
    print(f"  ✅ {msg}")


def warn_(msg: str):
    print(f"  ⚠️  {msg}")


def fail_(msg: str):
    print(f"  ❌ {msg}")
    sys.exit(1)


# ---------------------------------------------------------------------------
# NATS monitoring (port 8222)
# ---------------------------------------------------------------------------

def nats_varz() -> dict:
    """Fetch NATS /varz JSON. Returns {} on error."""
    try:
        with urllib.request.urlopen(f"{NATS_MONITOR}/varz", timeout=10) as r:
            return json.loads(r.read())
    except Exception:
        return {}


def nats_queue_depth() -> int:
    """Return current NATS JetStream total message count across all streams.
    
    Uses GET /jsz (JetStream info endpoint, no auth required on default NATS).
    Returns total messages stored in JetStream. High values indicate the
    consumer is not keeping up with the producer.
    """
    try:
        with urllib.request.urlopen(f"{NATS_MONITOR}/jsz", timeout=10) as r:
            d = json.loads(r.read())
        return int(d.get("messages", 0))
    except Exception:
        return -1


# ---------------------------------------------------------------------------
# Prometheus-backed assertions
# ---------------------------------------------------------------------------

def assert_router_p99_under(threshold: float, window_seconds: int = 60) -> float:
    """
    Assert p99 router recipient expansion latency is under threshold.
    Metric: ba_recipient_expansion_seconds_bucket
    Returns the actual p99, or raises AssertionError.
    """
    query = (
        f'histogram_quantile(0.99, '
        f'rate(ba_recipient_expansion_seconds_bucket[{window_seconds}s]))'
    )
    results = lib.scrape(query)
    if not results:
        raise AssertionError(
            f"router p99 query returned no data. Is routerd receiving alerts? "
            f"(query: {query})"
        )
    p99 = float(results[0]["value"][1])
    if p99 > threshold:
        raise AssertionError(
            f"router p99 {p99*1000:.1f}ms exceeds threshold {threshold*1000:.1f}ms "
            f"(window={window_seconds}s)"
        )
    return p99


def assert_bench_rate_near(target: float, tolerance: float = 0.10,
                           window_seconds: int = 30) -> float:
    """
    Assert the bench cluster-wide accept rate is within tolerance of target.
    tolerance=0.10 means ±10%.
    Returns the actual rate.
    """
    query = (
        f'rate(ba_ingestd_alerts_received_total{{result="accepted"}}[{window_seconds}s])'
    )
    results = lib.scrape(query)
    if not results:
        raise AssertionError(
            f"rate query returned no data. Is ingestd receiving alerts? (query: {query})"
        )
    actual = sum(float(r["value"][1]) for r in results)
    lower = target * (1 - tolerance)
    upper = target * (1 + tolerance)
    if not (lower <= actual <= upper):
        raise AssertionError(
            f"bench rate {actual:.1f}/s is outside ±{tolerance*100:.0f}% "
            f"tolerance of target {target}/s (range: {lower:.1f}–{upper:.1f})"
        )
    return actual


# ---------------------------------------------------------------------------
# Steps
# ---------------------------------------------------------------------------

def step1_preflight() -> None:
    """Verify all bench services are up."""
    print("Step 1 — pre-flight (bench profile)")
    targets = lib.get_prometheus_targets()
    required = ["ingestd", "routerd", "prometheus",
                "loadgen-bench-1", "loadgen-bench-2"]
    all_up = True
    for svc in required:
        status = targets.get(svc, "0")
        if status == "1":
            pass_(f"{svc} is up")
        else:
            warn_(f"{svc} is {'not scraped' if status == '0' else status}")
            all_up = False

    # Verify deliverd-bench is consuming (check NATS consumer lag)
    # A lagging consumer shows up as non-zero "pending" on the deliverd-bench
    # subject. We check that the subject has active consumers.
    nats_info = nats_varz()
    js = nats_info.get("jetstream", {})
    if js:
        pass_("NATS JetStream is available")
    else:
        warn_("NATS JetStream stats unavailable — cannot verify consumer lag")

    if not all_up:
        fail_("not all bench services are up — fix before running bench smoke")


def step2_start_loadgen() -> subprocess.Popen:
    """Start the bench loadgen cluster."""
    print("\nStep 2 — starting 2-instance bench loadgen cluster (10k/s)")
    proc = subprocess.Popen(
        ["docker", "compose", "--profile", "bench", "up", "-d"],
        stdout=subprocess.DEVNULL,
        stderr=subprocess.DEVNULL,
    )
    code = proc.wait()
    if code != 0:
        fail_(f"docker compose --profile bench up -d failed (exit {code})")
    pass_("2 loadgen-bench instances started")
    print(f"  waiting {BENCH_RAMP_SEC}s for ramp-up to complete...", flush=True)
    time.sleep(BENCH_RAMP_SEC)
    pass_(f"ramp-up complete — now targeting {CLUSTER_TARGET}/s")
    return proc


def step2_monitor_soak() -> list[dict]:
    """
    Monitor the bench soak: sample rate + router p99 + NATS queue depth every 30s.
    Fails fast on any breach.
    Returns list of samples for the log.
    """
    print(f"\n  Monitoring bench soak for {BENCH_DURATION_MIN} minutes...")
    samples = []
    start = time.time()
    deadline = start + BENCH_DURATION_MIN * 60
    sample_interval = 30  # seconds

    while time.time() < deadline:
        time.sleep(sample_interval)
        elapsed_min = int(time.time() - start) // 60
        elapsed_sec = int(time.time() - start)

        try:
            rate = assert_bench_rate_near(CLUSTER_TARGET, tolerance=RATE_TOLERANCE,
                                           window_seconds=30)
            p99 = assert_router_p99_under(P99_ROUTER_THRESHOLD, window_seconds=60)
            qd = nats_queue_depth()
            qd_ok = qd < NATS_QUEUE_DEPTH_MAX
            print(f"  [{elapsed_min}m] rate={rate:.0f}/s router_p99={p99*1000:.1f}ms "
                  f"nats_qd={qd} {'✅' if qd_ok else '❌'}")
            samples.append({
                "elapsed_min": elapsed_min,
                "elapsed_sec": elapsed_sec,
                "rate": rate,
                "router_p99_ms": p99 * 1000,
                "nats_qd": qd,
            })
            if not qd_ok:
                fail_(f"NATS queue depth {qd} exceeds max {NATS_QUEUE_DEPTH_MAX}")
        except AssertionError as e:
            fail_(f"bench soak breach at {elapsed_min}m: {e}")

    return samples


def step3_broker_ceiling_check(samples: list[dict]) -> None:
    """Final broker ceiling assertions over the full soak window."""
    print("\nStep 3 — broker ceiling check")
    # All samples already passed the router p99 check; do a final aggregate
    # assertion over the full window to confirm the ceiling held.
    p99_final = assert_router_p99_under(P99_ROUTER_THRESHOLD,
                                          window_seconds=BENCH_DURATION_MIN * 60)
    pass_(f"router p99 over full {BENCH_DURATION_MIN}min window: "
          f"{p99_final*1000:.1f}ms (threshold: {P99_ROUTER_THRESHOLD*1000:.1f}ms)")

    qd = nats_queue_depth()
    if qd < NATS_QUEUE_DEPTH_MAX:
        pass_(f"NATS queue depth: {qd} (max allowed: {NATS_QUEUE_DEPTH_MAX})")
    else:
        fail_(f"NATS queue depth {qd} exceeds max {NATS_QUEUE_DEPTH_MAX}")


def step4_teardown(proc) -> None:
    """Bring down the bench cluster."""
    print("\nStep 4 — teardown")
    r = subprocess.run(
        ["docker", "compose", "--profile", "bench", "down", "-v"],
        capture_output=True,
    )
    if r.returncode == 0:
        pass_("bench cluster torn down")
    else:
        warn_(f"teardown returned {r.returncode}: {r.stderr.decode().strip()}")


def print_summary(samples: list[dict]) -> None:
    """Print a summary table of the bench run."""
    print("\n=== M10-Bench Soak Summary ===")
    print(f"Duration:       {BENCH_DURATION_MIN} min")
    print(f"Target rate:   {CLUSTER_TARGET}/s (±{RATE_TOLERANCE*100:.0f}%)")
    print(f"Router p99 threshold: {P99_ROUTER_THRESHOLD*1000:.0f}ms")
    print(f"NATS qd max:  {NATS_QUEUE_DEPTH_MAX}")
    print()
    if samples:
        print(f"{'Time':>6}  {'Rate/s':>8}  {'rt_p99(ms)':>11}  {'nats_qd':>8}")
        print("-" * 42)
        for s in samples:
            print(f"{s['elapsed_min']:>5}m  {s['rate']:>8.0f}  "
                  f"{s['router_p99_ms']:>11.1f}  {s['nats_qd']:>8}")
    print()
    print("🎉 M10-bench smoke: all checks complete.")


def main():
    print(f"M10-Bench — target {CLUSTER_TARGET}/s for {BENCH_DURATION_MIN} min")
    print(f"Router p99 threshold: {P99_ROUTER_THRESHOLD*1000:.0f}ms | "
          f"NATS qd max: {NATS_QUEUE_DEPTH_MAX}")
    print()

    try:
        step1_preflight()
        proc = step2_start_loadgen()
        samples = step2_monitor_soak()
        step3_broker_ceiling_check(samples)
        step4_teardown(proc)
        print_summary(samples)
    except AssertionError as e:
        print(f"\n💥 M10-bench smoke FAILED: {e}")
        sys.exit(1)
    except KeyboardInterrupt:
        print("\nInterrupted.")
        sys.exit(1)


if __name__ == "__main__":
    main()