"""轻量级压测 — health + stats

用法:
  NEXUS_API_KEY=your-key python tests/quick_load.py 50 30
  NEXUS_TEST_BASE=http://127.0.0.1:8600  (可选)
"""
import asyncio
import httpx
import os
import sys
import time

BASE = os.environ.get("NEXUS_TEST_BASE", "http://127.0.0.1:8600")
API_KEY = os.environ.get("NEXUS_API_KEY", "")
oks = [0]
fails = [0]
times = []


async def health(c, d):
    end = time.time() + d
    while time.time() < end:
        t0 = time.perf_counter()
        try:
            r = await c.get(f"{BASE}/health", timeout=5)
            if r.status_code == 200:
                oks[0] += 1
            else:
                fails[0] += 1
        except Exception:
            fails[0] += 1
        times.append(time.perf_counter() - t0)
        await asyncio.sleep(0.001)


async def stats_w(c, d):
    if not API_KEY:
        return
    end = time.time() + d
    while time.time() < end:
        t0 = time.perf_counter()
        try:
            r = await c.get(
                f"{BASE}/api/servers/stats",
                headers={"X-API-Key": API_KEY},
                timeout=10,
            )
            if r.status_code == 200:
                oks[0] += 1
            else:
                fails[0] += 1
        except Exception:
            fails[0] += 1
        times.append(time.perf_counter() - t0)
        await asyncio.sleep(0.05)


async def main(n, d):
    if not API_KEY:
        print("WARN: NEXUS_API_KEY unset — stats workers skipped", file=sys.stderr)
    async with httpx.AsyncClient(limits=httpx.Limits(max_connections=n + 100)) as c:
        tasks = [health(c, d) for _ in range(int(n * 0.7))]
        if API_KEY:
            tasks += [stats_w(c, d) for _ in range(max(1, int(n * 0.3)))]
        t0 = time.time()
        await asyncio.gather(*tasks, return_exceptions=True)
        elapsed = time.time() - t0
        total = oks[0] + fails[0]
        if not total:
            print("No requests completed")
            return
        t = sorted(times)
        nt = len(t)
        print(
            f"并发={n} | RPS={total/elapsed:.0f} | P50={t[int(nt*0.5)]*1000:.1f}ms | "
            f"P95={t[min(int(nt*0.95),nt-1)]*1000:.1f}ms | "
            f"P99={t[min(int(nt*0.99),nt-1)]*1000:.1f}ms | "
            f"ok={oks[0]} fail={fails[0]} err={(fails[0]/total*100):.1f}%"
        )


if __name__ == "__main__":
    if len(sys.argv) < 3:
        print("Usage: NEXUS_API_KEY=... python tests/quick_load.py <concurrency> <seconds>")
        sys.exit(1)
    asyncio.run(main(int(sys.argv[1]), int(sys.argv[2])))
