"""
Self-contained async load test — no extra tools needed.

Starts a uvicorn server in a background thread, fires N concurrent slow
requests, and polls /metrics every 200ms to print a live bar chart of
runway_requests_in_flight as the queue builds and drains.

Modes
-----
burst  (default)
    Fire N requests all at once and watch the in-flight bar chart.
    Good for confirming the metric is working.

    uv run examples/load_test.py
    uv run examples/load_test.py --requests 100 --delay 2.0

sweep
    Step through increasing concurrency levels and measure p50/p95/p99
    latency at each level. Finds where your server saturates and prints
    the recommended autoscaling threshold.

    uv run examples/load_test.py --mode sweep
    uv run examples/load_test.py --mode sweep --workload cpu
    uv run examples/load_test.py --mode sweep --max-concurrency 200

Workloads
---------
async (default)
    asyncio.sleep — simulates pure I/O-bound work (DB queries, API calls).
    Will NOT saturate at any concurrency — that is expected and correct.
    Use this to confirm the metric tracks in-flight accurately.

cpu
    Actual CPU computation — simulates sync/CPU-bound work.
    Will saturate when the thread pool fills up.
    Use this to find the real saturation point of CPU-heavy endpoints.
"""

import argparse
import asyncio
import math
import sys
import threading
import time
from typing import List

import httpx
import uvicorn
from fastapi import FastAPI
from fastapi.concurrency import run_in_threadpool

from asgi_runway import setup

# ---------------------------------------------------------------------------
# Inline app (mirrors examples/app.py)
# ---------------------------------------------------------------------------

_app = FastAPI()
setup(
    _app,
    route_groups=[
        (r"^/api/infer", "inference"),
        (r"^/api/embed", "embedding"),
    ],
)


@_app.get("/api/infer")
async def infer(delay: float = 1.0):
    await asyncio.sleep(delay)
    return {"result": "done"}


@_app.get("/api/embed")
async def embed(delay: float = 0.1):
    await asyncio.sleep(delay)
    return {"result": "done"}


@_app.get("/api/cpu")
async def cpu_work(duration_ms: int = 100):
    """
    Simulate CPU-bound work by running a blocking computation in the thread pool.
    This WILL saturate — unlike asyncio.sleep — because it consumes a real OS thread
    for the full duration_ms, exhausting the thread pool under high concurrency.
    """
    def _compute(ms: int) -> int:
        # Keep the CPU busy for ~ms milliseconds.
        deadline = time.perf_counter() + ms / 1000
        n = 0
        while time.perf_counter() < deadline:
            n += sum(i * i for i in range(500))
        return n

    result = await run_in_threadpool(_compute, duration_ms)
    return {"result": result}


@_app.get("/healthz")
async def health():
    return {"status": "ok"}


# ---------------------------------------------------------------------------
# Load test
# ---------------------------------------------------------------------------

PORT = 8765
BASE_URL = f"http://127.0.0.1:{PORT}"
BAR_WIDTH = 48
PALETTE = [" ", "▏", "▎", "▍", "▌", "▋", "▊", "▉", "█"]


def _parse_gauge(text: str, name: str) -> float:
    """Extract a gauge value from Prometheus text format."""
    for line in text.splitlines():
        if line.startswith(f"{name} ") and not line.startswith("#"):
            try:
                return float(line.split()[-1])
            except ValueError:
                pass
    return 0.0


def _bar(value: int, peak: int) -> str:
    ratio = value / max(peak, 1)
    filled_cells = ratio * BAR_WIDTH
    full = int(filled_cells)
    partial_idx = int((filled_cells - full) * 8)
    bar = "█" * full
    if full < BAR_WIDTH:
        bar += PALETTE[partial_idx]
        bar += " " * (BAR_WIDTH - full - 1)
    color = "\033[91m" if ratio > 0.75 else "\033[93m" if ratio > 0.4 else "\033[92m"
    reset = "\033[0m"
    return f"{color}│{bar}│{reset} {value:>4}"


async def _wait_for_server(client: httpx.AsyncClient, retries: int = 30) -> None:
    for _ in range(retries):
        try:
            await client.get(f"{BASE_URL}/healthz", timeout=1.0)
            return
        except Exception:
            await asyncio.sleep(0.15)
    raise RuntimeError("Server did not start in time.")


async def _poll_loop(
    client: httpx.AsyncClient,
    stop: asyncio.Event,
    peak: list,
    interval: float = 0.2,
) -> None:
    while not stop.is_set():
        try:
            resp = await client.get(f"{BASE_URL}/metrics", timeout=2.0)
            val = int(_parse_gauge(resp.text, "runway_requests_in_flight"))
            if val > peak[0]:
                peak[0] = val
            bar = _bar(val, peak[0])
            print(f"\r  in-flight {bar}", end="", flush=True)
        except Exception:
            pass
        await asyncio.sleep(interval)


async def run(n_requests: int, delay: float, workers: int) -> None:
    # --- Start server ---
    config = uvicorn.Config(
        _app,
        host="127.0.0.1",
        port=PORT,
        workers=workers if workers > 1 else None,
        log_level="error",
    )
    server = uvicorn.Server(config)
    thread = threading.Thread(target=server.run, daemon=True)
    thread.start()

    async with httpx.AsyncClient() as client:
        print(f"\n  Starting server on {BASE_URL}... ", end="", flush=True)
        await _wait_for_server(client)
        print("ready.\n")

        print(f"  Firing {n_requests} concurrent requests  (each takes {delay}s)\n")
        print(f"  {'─' * (BAR_WIDTH + 16)}")

        stop = asyncio.Event()
        peak = [0]
        poll_task = asyncio.create_task(_poll_loop(client, stop, peak))

        # Fire all requests concurrently
        t0 = time.perf_counter()
        results = await asyncio.gather(
            *[
                client.get(f"{BASE_URL}/api/infer?delay={delay}", timeout=delay + 10)
                for _ in range(n_requests)
            ],
            return_exceptions=True,
        )
        elapsed = time.perf_counter() - t0

        stop.set()
        await poll_task

        ok = sum(
            1
            for r in results
            if isinstance(r, httpx.Response) and r.status_code == 200
        )
        errors = n_requests - ok

        print(f"\n  {'─' * (BAR_WIDTH + 16)}")
        print(f"\n  Results:")
        print(f"    Requests : {n_requests}  ({ok} OK, {errors} errors)")
        print(f"    Duration : {elapsed:.2f}s")
        print(f"    Peak in-flight : {peak[0]}")
        print(
            f"    Throughput : {ok / elapsed:.1f} req/s"
            f"  (expected: {n_requests / delay:.1f} req/s with {delay}s delay)\n"
        )

    server.should_exit = True
    thread.join(timeout=3)


# ---------------------------------------------------------------------------
# Sweep mode — find the saturation point
# ---------------------------------------------------------------------------

def _percentile(data: List[float], p: float) -> float:
    if not data:
        return 0.0
    s = sorted(data)
    idx = max(0, int(math.ceil(p / 100 * len(s))) - 1)
    return s[idx]


async def _sustained_load(
    client: httpx.AsyncClient,
    url: str,
    concurrency: int,
    duration: float,
) -> List[float]:
    """
    Maintain exactly `concurrency` in-flight requests for `duration` seconds.

    Each worker fires requests back-to-back (as soon as one finishes, the next
    starts) so the concurrency level stays constant throughout the window —
    the same pattern as real production load.

    Returns all latency samples collected across all workers.
    """
    latencies: List[float] = []
    deadline = time.perf_counter() + duration

    async def _worker() -> None:
        while time.perf_counter() < deadline:
            t0 = time.perf_counter()
            try:
                await client.get(url, timeout=30.0)
            except Exception:
                pass
            latencies.append(time.perf_counter() - t0)

    await asyncio.gather(*[_worker() for _ in range(concurrency)])
    return latencies


async def sweep(
    workload: str,
    max_concurrency: int,
    duration_ms: int,
    step_duration: float,
    external_host: str | None = None,
    external_path: str | None = None,
) -> None:
    # --- Determine target URL and whether we need a local server ---
    if external_host:
        base = external_host.rstrip("/")
        path = (external_path or "/").lstrip("/")
        target_url = f"{base}/{path}"
        workload_desc = f"External  ({target_url})"
        local_server = False
    elif workload == "cpu":
        target_url = f"{BASE_URL}/api/cpu?duration_ms={duration_ms}"
        workload_desc = f"CPU-bound  ({duration_ms}ms of real computation per request)"
        local_server = True
    else:
        delay_s = duration_ms / 1000
        target_url = f"{BASE_URL}/api/infer?delay={delay_s}"
        workload_desc = f"Async I/O  ({delay_s}s asyncio.sleep per request)"
        local_server = True

    # Geometric steps: 1, 2, 4, 8 … up to max_concurrency
    steps: List[int] = []
    n = 1
    while n <= max_concurrency:
        steps.append(n)
        n = n * 2
    if steps[-1] != max_concurrency:
        steps.append(max_concurrency)

    server = thread = None
    if local_server:
        config = uvicorn.Config(_app, host="127.0.0.1", port=PORT, log_level="error")
        server = uvicorn.Server(config)
        thread = threading.Thread(target=server.run, daemon=True)
        thread.start()

    async with httpx.AsyncClient() as client:
        if local_server:
            print(f"\n  Starting server on {BASE_URL}... ", end="", flush=True)
            await _wait_for_server(client)
            print("ready.\n")
        else:
            print(f"\n  Targeting external server: {external_host}\n")

        print(f"  Workload      : {workload_desc}")
        print(f"  Steps         : {steps}")
        print(f"  Step duration : {step_duration}s  (~{len(steps) * step_duration:.0f}s total)\n")

        col = "{:>6}  {:>7}  {:>7}  {:>7}  {:>7}  {:>12}  {}"
        print(col.format("conc", "p50", "p95", "p99", "samples", "throughput", "status"))
        print("  " + "─" * 70)

        baseline_p95: float | None = None
        saturation_concurrency: int | None = None

        for concurrency in steps:
            latencies = await _sustained_load(client, target_url, concurrency, step_duration)
            p50 = _percentile(latencies, 50)
            p95 = _percentile(latencies, 95)
            p99 = _percentile(latencies, 99)
            n_samples = len(latencies)
            throughput = n_samples / step_duration

            if baseline_p95 is None:
                baseline_p95 = p95

            # Saturation = p95 has more than doubled from baseline
            is_saturated = p95 > baseline_p95 * 2.0
            is_degrading = p95 > baseline_p95 * 1.3

            if is_saturated:
                if saturation_concurrency is None:
                    saturation_concurrency = concurrency
                status = "\033[91m✗ saturated\033[0m"
            elif is_degrading:
                status = "\033[93m⚠ degrading\033[0m"
            else:
                status = "\033[92m✓ ok\033[0m"

            print(
                "  "
                + col.format(
                    concurrency,
                    f"{p50:.3f}s",
                    f"{p95:.3f}s",
                    f"{p99:.3f}s",
                    n_samples,
                    f"{throughput:.1f} req/s",
                    status,
                )
            )

        print("\n  " + "─" * 62)

        if saturation_concurrency:
            threshold = int(saturation_concurrency * 0.75)
            print(f"\n  Saturation point : ~{saturation_concurrency} concurrent requests")
            print(f"  Recommended autoscaling threshold : {threshold}  (75% of saturation)")
            print(f"\n  Little's Law check:")
            print(f"    If your SLA is p95 < {duration_ms}ms and you expect X req/s,")
            print(f"    each pod can handle ~{threshold} in-flight → deploy ceil(X / (1000/{duration_ms}) / {threshold}) pods")
        else:
            # No saturation found — typical for async I/O workloads
            print(f"\n  No saturation detected up to {max_concurrency} concurrent requests.")
            print(f"  This is expected for async I/O workloads (asyncio handles them cheaply).")
            print(f"\n  For async workloads, set your threshold using Little's Law instead:")
            print(f"    threshold = (target RPS per pod) × (expected p95 latency in seconds)")
            print(f"    e.g. 50 req/s × 0.2s p95 = 10  →  scale when in-flight > 10")

    if server:
        server.should_exit = True
    if thread:
        thread.join(timeout=3)


# ---------------------------------------------------------------------------
# Entry point
# ---------------------------------------------------------------------------

def main() -> None:
    parser = argparse.ArgumentParser(description="asgi-runway load test")
    parser.add_argument(
        "--mode", choices=["burst", "sweep"], default="burst",
        help="burst: fire all at once (default). sweep: find saturation point.",
    )
    # burst options
    parser.add_argument("--requests", type=int, default=50, help="concurrent requests (burst mode)")
    parser.add_argument("--delay", type=float, default=1.5, help="simulated work in seconds (burst mode)")
    parser.add_argument("--workers", type=int, default=1, help="uvicorn worker count")
    # sweep options
    parser.add_argument(
        "--workload", choices=["async", "cpu"], default="async",
        help="async: asyncio.sleep (I/O bound). cpu: real computation (CPU bound). Ignored when --host is set.",
    )
    parser.add_argument("--max-concurrency", type=int, default=64, help="max concurrency to test in sweep")
    parser.add_argument("--duration-ms", type=int, default=200, help="work duration per request in ms (built-in workloads only)")
    parser.add_argument(
        "--step-duration", type=float, default=10.0,
        help="seconds of sustained load per concurrency step (default: 10). "
             "Higher = more samples = more stable percentiles.",
    )
    parser.add_argument(
        "--host", default=None,
        help="Target an external server instead of the built-in one. "
             "e.g. http://staging.internal:8000  "
             "Use with --path to specify the endpoint.",
    )
    parser.add_argument(
        "--path", default="/",
        help="Endpoint path to hit when using --host. e.g. /api/infer  (default: /)",
    )
    args = parser.parse_args()

    try:
        if args.mode == "sweep":
            asyncio.run(sweep(args.workload, args.max_concurrency, args.duration_ms, args.step_duration, args.host, args.path))
        else:
            asyncio.run(run(args.requests, args.delay, args.workers))
    except KeyboardInterrupt:
        print("\n  Interrupted.")
        sys.exit(0)


if __name__ == "__main__":
    main()
