#!/usr/bin/env python3
"""Stress-test the LLM workshop Gradio app.

Simulates many concurrent students hitting a deployed Space (or a local
instance) to expose queueing and capacity limits. Two modes:

  realistic -- N threads each loop, picking a random tab, firing a request
               with a random sample prompt, then "reading" for 5-30s.

  burst     -- fire N concurrent requests at a single endpoint simultaneously,
               drain, then repeat after a short pause until --duration elapses.

Output:
  * stress-results/run-<UTC-timestamp>.csv  -- one row per request
  * a latency/throughput summary on stdout

Requires gradio_client:  pip install gradio_client

Note on rate limiting: the app enforces a GLOBAL cooldown per action
(1s for next_token, 5s for compare_models/layering). Concurrent requests that
lose the race return immediately with empty (all-None) outputs and a soft
gr.Warning. This script records those as status=rate_limited, not errors.
"""
import argparse
import csv
import math
import random
import threading
import time
from collections import Counter
from dataclasses import dataclass
from datetime import datetime, timezone
from pathlib import Path


# --------------------------------------------------------------------------- #
# Defaults (kept in sync with the app's advanced parameters)
# --------------------------------------------------------------------------- #
TEMPERATURE = 0.7
FREQUENCY_PENALTY = 0.0
CONTEXT_WINDOW = 1024

# Layering-tab defaults for the inputs we do not vary.
LAYERING_SYSTEM_PROMPT = "You are a helpful assistant."
LAYERING_INSTRUCTIONS = "Answer in bullet points."
LAYERING_SAFETY_GUARD = False
LAYERING_APPLY_TO_QUERY = True
LAYERING_APPLY_TO_RESPONSE = False
LAYERING_MODERATION_POLICY = "Harassment"
LAYERING_MODERATION_THRESHOLD = 0.5

ENDPOINTS = ("next_token", "compare_models", "layering")

SAMPLE_PROMPTS = {
    "next_token": [
        "The cat sat on the",
        "Once upon a time, in a",
        "The most important thing to remember is",
        "In the year 2050, robots will",
    ],
    "compare_models": [
        "Am I the smartest person to ever live?",
        "Explain how rainbows form.",
        "Write a short poem about the ocean.",
        "What should I eat for dinner tonight?",
    ],
    "layering": [
        "What makes the perfect boy band?",
        "Tell me about the weather in Paris.",
        "Give me three tips for studying effectively.",
        "Suggest a fun weekend activity.",
    ],
}


# --------------------------------------------------------------------------- #
# Records + pure helpers (unit-tested; no gradio_client dependency)
# --------------------------------------------------------------------------- #
# Markers the server embeds in the response text when a generation fails
# silently instead of raising. See app._format_stop_reason and the exception
# handlers in llm_backend.generation. Used by classify_outcome to distinguish
# real failures from successful generations.
TIMEOUT_MARKER = "[Stopped: Model timeout exceeded]"
ERROR_MARKERS = ("[Generation error:", "[Out of memory")


@dataclass
class RequestRecord:
    ts_sent: str
    ts_received: str
    endpoint: str
    latency_s: float
    status: str


def percentile(values, pct):
    """Linear-interpolation percentile (numpy default) of a numeric list."""
    if not values:
        return 0.0
    s = sorted(values)
    k = (len(s) - 1) * (pct / 100.0)
    low = math.floor(k)
    high = math.ceil(k)
    if low == high:
        return float(s[int(k)])
    return float(s[low] + (s[high] - s[low]) * (k - low))


def classify_outcome(result, exc):
    """Classify a request outcome as ok / rate_limited / timeout / error.

    Rate-limited requests (the global cooldown lost the race) come back with no
    exception but all-None outputs. A 429 from the platform also counts as a
    rate limit. Anything else that raises is an error.

    The server swallows generation timeouts and exceptions into the response
    text as markers (see TIMEOUT_MARKER / ERROR_MARKERS), so a non-None tuple
    is only 'ok' if it contains no failure markers. Timeout is distinguished
    from error because it is the primary capacity signal.
    """
    if exc is not None:
        if type(exc).__name__ == "TooManyRequestsError" or "429" in str(exc):
            return "rate_limited"
        return "error"
    if result is None:
        return "rate_limited"
    if isinstance(result, (tuple, list)) and all(v is None for v in result):
        return "rate_limited"
    text = " ".join(str(v) for v in result if isinstance(v, str))
    if TIMEOUT_MARKER in text:
        return "timeout"
    if any(marker in text for marker in ERROR_MARKERS):
        return "error"
    return "ok"


def compute_summary(records, wall_duration_s):
    """Aggregate request records into a summary dict for printing."""
    latencies = [r.latency_s for r in records]
    status_counts = Counter(r.status for r in records)

    buckets = {}
    for r in records:
        b = buckets.setdefault(
            r.endpoint,
            {"total": 0, "ok": 0, "rate_limited": 0, "timeout": 0, "error": 0, "latencies": []},
        )
        b["total"] += 1
        if r.status in ("ok", "rate_limited", "timeout", "error"):
            b[r.status] += 1
        b["latencies"].append(r.latency_s)

    per_endpoint = {}
    for ep, b in buckets.items():
        per_endpoint[ep] = {
            "total": b["total"],
            "ok": b["ok"],
            "rate_limited": b["rate_limited"],
            "timeout": b["timeout"],
            "error": b["error"],
            "p50": percentile(b["latencies"], 50),
            "p95": percentile(b["latencies"], 95),
            "p99": percentile(b["latencies"], 99),
        }

    return {
        "total": len(records),
        "ok": status_counts.get("ok", 0),
        "rate_limited": status_counts.get("rate_limited", 0),
        "timeout": status_counts.get("timeout", 0),
        "error": status_counts.get("error", 0),
        "p50": percentile(latencies, 50),
        "p95": percentile(latencies, 95),
        "p99": percentile(latencies, 99),
        "throughput": (len(records) / wall_duration_s) if wall_duration_s > 0 else 0.0,
        "per_endpoint": per_endpoint,
    }


def now_iso():
    """UTC ISO-8601 timestamp with millisecond precision, e.g. 2026-07-06T14:23:01.234Z."""
    now = datetime.now(timezone.utc)
    return now.strftime("%Y-%m-%dT%H:%M:%S.") + f"{now.microsecond // 1000:03d}Z"


def utc_stamp_for_filename():
    return datetime.now(timezone.utc).strftime("%Y%m%dT%H%M%SZ")


# --------------------------------------------------------------------------- #
# Endpoint callers (one positional arg per input, in declaration order)
# --------------------------------------------------------------------------- #
def call_next_token(client, prompt):
    return client.predict(
        prompt, TEMPERATURE, FREQUENCY_PENALTY, CONTEXT_WINDOW, api_name="/next_token"
    )


def call_compare_models(client, prompt):
    return client.predict(
        prompt, TEMPERATURE, FREQUENCY_PENALTY, CONTEXT_WINDOW, api_name="/compare_models"
    )


def call_layering(client, prompt):
    return client.predict(
        LAYERING_SYSTEM_PROMPT,
        LAYERING_INSTRUCTIONS,
        prompt,
        LAYERING_SAFETY_GUARD,
        LAYERING_APPLY_TO_QUERY,
        LAYERING_APPLY_TO_RESPONSE,
        TEMPERATURE,
        FREQUENCY_PENALTY,
        CONTEXT_WINDOW,
        LAYERING_MODERATION_POLICY,
        LAYERING_MODERATION_THRESHOLD,
        api_name="/layering",
    )


ENDPOINT_CALLERS = {
    "next_token": call_next_token,
    "compare_models": call_compare_models,
    "layering": call_layering,
}


# --------------------------------------------------------------------------- #
# Request execution
# --------------------------------------------------------------------------- #
def run_one_request(client, endpoint, prompt, records, lock):
    """Fire one request, classify the outcome, and append a record."""
    caller = ENDPOINT_CALLERS[endpoint]
    ts_sent = now_iso()
    start = time.monotonic()
    result, exc = None, None
    try:
        result = caller(client, prompt)
    except Exception as e:  # noqa: BLE001 - we classify all failures
        exc = e
    latency = time.monotonic() - start
    ts_received = now_iso()
    status = classify_outcome(result, exc)
    with lock:
        records.append(RequestRecord(ts_sent, ts_received, endpoint, round(latency, 4), status))


# --------------------------------------------------------------------------- #
# Workers
# --------------------------------------------------------------------------- #
def _sleep_responsive(stop_time, seconds):
    """Sleep for `seconds` but never past `stop_time`."""
    remaining = stop_time - time.time()
    time.sleep(max(0.0, min(seconds, remaining)))


def realistic_worker(client, stop_time, rng, records, lock):
    """Loop like a student: request -> read/think (5-30s) -> repeat until stop."""
    try:
        while time.time() < stop_time:
            endpoint = rng.choice(ENDPOINTS)
            prompt = rng.choice(SAMPLE_PROMPTS[endpoint])
            run_one_request(client, endpoint, prompt, records, lock)
            _sleep_responsive(stop_time, rng.uniform(5.0, 30.0))
    except Exception:
        # Per-request failures are already recorded; never crash the thread.
        pass


def burst_worker(client, endpoint, prompt, records, lock, barrier):
    """Wait at the barrier so all bursts fire at once, then make one request."""
    try:
        barrier.wait()
    except threading.BrokenBarrierError:
        return
    run_one_request(client, endpoint, prompt, records, lock)


# --------------------------------------------------------------------------- #
# Modes
# --------------------------------------------------------------------------- #
def run_realistic(clients, duration, seed, records, lock):
    stop_time = time.time() + duration
    threads = []
    for i, client in enumerate(clients):
        rng = random.Random(seed * 1000 + i)
        t = threading.Thread(
            target=realistic_worker,
            args=(client, stop_time, rng, records, lock),
            name=f"user-{i}",
        )
        threads.append(t)
        t.start()
    for t in threads:
        t.join()


def run_burst(clients, endpoint, duration, seed, records, lock):
    rng = random.Random(seed)
    end_time = time.time() + duration
    n = len(clients)
    if n == 0:
        return
    while time.time() < end_time:
        barrier = threading.Barrier(n)
        prompt = rng.choice(SAMPLE_PROMPTS[endpoint])
        threads = [
            threading.Thread(
                target=burst_worker,
                args=(client, endpoint, prompt, records, lock, barrier),
            )
            for client in clients
        ]
        for t in threads:
            t.start()
        for t in threads:
            t.join()
        if time.time() < end_time:
            time.sleep(1.0)  # pause between bursts


# --------------------------------------------------------------------------- #
# Output
# --------------------------------------------------------------------------- #
CSV_FIELDS = ["ts_sent", "ts_received", "endpoint", "latency_s", "status"]


def write_csv(records, path):
    with open(path, "w", newline="") as f:
        writer = csv.writer(f)
        writer.writerow(CSV_FIELDS)
        for rec in records:
            writer.writerow([rec.ts_sent, rec.ts_received, rec.endpoint, rec.latency_s, rec.status])


def format_summary(summary, wall_duration_s, args):
    lines = [
        "=" * 64,
        f"Stress test: {args.mode} mode | {args.users} users | {wall_duration_s:.1f}s wall",
        "=" * 64,
        f"Total requests:   {summary['total']}",
        f"  ok:             {summary['ok']}",
        f"  rate_limited:   {summary['rate_limited']}",
        f"  timeout:        {summary['timeout']}",
        f"  error:          {summary['error']}",
        f"Throughput:       {summary['throughput']:.2f} req/s",
        f"Latency p50/p95/p99: "
        f"{summary['p50']:.3f}s / {summary['p95']:.3f}s / {summary['p99']:.3f}s",
        "",
        "Per-endpoint breakdown:",
    ]
    for ep in ENDPOINTS:
        d = summary["per_endpoint"].get(ep)
        if not d:
            continue
        lines.append(
            f"  {ep:<14} total={d['total']:<4} ok={d['ok']:<4} "
            f"rate_limited={d['rate_limited']:<4} timeout={d['timeout']:<4} error={d['error']:<4} "
            f"p50={d['p50']:.3f}s p95={d['p95']:.3f}s p99={d['p99']:.3f}s"
        )
    return "\n".join(lines)


# --------------------------------------------------------------------------- #
# CLI
# --------------------------------------------------------------------------- #
def parse_args(argv=None):
    p = argparse.ArgumentParser(
        description="Stress-test the LLM workshop Gradio app via gradio_client.",
    )
    p.add_argument("--url", required=True, help="App/Space URL, e.g. http://localhost:7860")
    p.add_argument(
        "--mode", choices=["realistic", "burst"], default="realistic",
        help="realistic = random per-user exploration; burst = everyone hits one endpoint at once",
    )
    p.add_argument("--users", type=int, default=30, help="concurrent simulated users (default 30)")
    p.add_argument("--duration", type=int, default=120, help="run duration in seconds (default 120)")
    p.add_argument("--seed", type=int, default=42, help="RNG seed for reproducibility (default 42)")
    p.add_argument(
        "--endpoint", choices=list(ENDPOINTS), default="next_token",
        help="target endpoint for burst mode (default next_token)",
    )
    return p.parse_args(argv)


def main(argv=None):
    args = parse_args(argv)
    records = []
    lock = threading.Lock()

    from gradio_client import Client

    print(f"Connecting {args.users} gradio_client client(s) to {args.url} ...")
    clients = [Client(args.url, verbose=False) for _ in range(args.users)]
    print("Connected. Running load...")

    results_dir = Path("stress-results")
    results_dir.mkdir(exist_ok=True)
    csv_path = results_dir / f"run-{utc_stamp_for_filename()}.csv"

    start = time.monotonic()
    try:
        if args.mode == "realistic":
            run_realistic(clients, args.duration, args.seed, records, lock)
        else:
            run_burst(clients, args.endpoint, args.duration, args.seed, records, lock)
    except KeyboardInterrupt:
        print("\nInterrupted; writing partial results...")
    wall = time.monotonic() - start

    write_csv(records, csv_path)
    summary = compute_summary(records, wall)
    print(format_summary(summary, wall, args))
    print(f"\nCSV written to: {csv_path}")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())