"""
HTTP 端到端 RPC 压测脚本（chongming 网关全链路）

链路: 压测客户端(asyncio+httpx) → FastAPI 网关 :8000 → Zenoh Query/Reply → worker 进程

前置启动（均后台运行，需在仓库根目录）:
  1. Zenoh router:  uv run python utils/worker-rs/scripts/zenoh_router.py
  2. 网关(无热重载): uv run uvicorn chongming.app:app --host 0.0.0.0 --port 8000 --no-access-log
  3. Worker:         uv run chongming-worker --type demo --count 2

运行:
  uv run python test/bench_http.py --concurrency 1,10,50,100,200 --duration 30

指标: RPS(总请求/秒)、p50/p95/p99/max 延迟(成功请求, ms)、错误率(非 2xx/超时/连接异常)、
      结束后拉取网关侧 GET /api/v1/stats 做服务端对照。
"""
import argparse
import asyncio
import json
import statistics
import sys
import time
from dataclasses import dataclass, field

import httpx

DEFAULT_URL = "http://localhost:8000/api/v1/demo/hello?name=bench"


@dataclass
class Bench:
    """单档压测共享状态（asyncio 单线程内访问，无需加锁）"""
    start: float = 0.0
    total: int = 0          # 全部完成的请求（含错误/超时）
    ok: int = 0             # 2xx 成功请求
    timeouts: int = 0
    errors: int = 0         # 非 2xx / 连接异常
    status_dist: dict = field(default_factory=dict)
    lats_ms: list = field(default_factory=list)  # 成功请求延迟(ms)


def percentile(sorted_lats: list, p: float) -> float:
    """有序列表上的线性插值分位数（p 为 0~1）"""
    if not sorted_lats:
        return 0.0
    idx = (len(sorted_lats) - 1) * p
    lo = int(idx)
    hi = min(lo + 1, len(sorted_lats) - 1)
    frac = idx - lo
    return sorted_lats[lo] * (1 - frac) + sorted_lats[hi] * frac


async def worker(client: httpx.AsyncClient, url: str, bench: Bench, end_time: float) -> None:
    loop = asyncio.get_running_loop()
    while True:
        t0 = loop.time()
        if t0 >= end_time:
            return
        try:
            resp = await client.get(url, timeout=10.0)
        except httpx.TimeoutException:
            bench.timeouts += 1
            continue
        except httpx.HTTPError:
            bench.errors += 1
            continue
        lat_ms = (loop.time() - t0) * 1000
        if resp.status_code == 200:
            bench.ok += 1
            bench.lats_ms.append(lat_ms)
        else:
            bench.errors += 1
            bench.status_dist[resp.status_code] = bench.status_dist.get(resp.status_code, 0) + 1
        bench.total += 1


async def progress_reporter(bench: Bench, duration: float) -> None:
    """每 5s 打印进度（RPS/成功/错误），供长压测观测"""
    while True:
        await asyncio.sleep(5)
        elapsed = time.perf_counter() - bench.start
        if elapsed <= 0:
            continue
        print(f"    [{elapsed:5.0f}s] RPS={bench.total / elapsed:7.1f} "
              f"ok={bench.ok} err={bench.errors} timeout={bench.timeouts}", flush=True)


async def run_level(client: httpx.AsyncClient, url: str, concurrency: int,
                    duration: float, warmup: float) -> dict:
    """单档并发：warmup 预热（结果丢弃）→ duration 计时统计"""
    bench = Bench()

    async def phase(end_time: float) -> None:
        tasks = [asyncio.create_task(worker(client, url, bench, end_time))
                 for _ in range(concurrency)]
        await asyncio.gather(*tasks)

    loop = asyncio.get_running_loop()
    print(f"  预热 {warmup:.0f}s ...", flush=True)
    bench.start = time.perf_counter()
    await phase(loop.time() + warmup)

    bench = Bench()  # 丢弃预热结果，正式计时
    bench.start = time.perf_counter()
    print(f"  压测 {duration:.0f}s (并发 {concurrency}) ...", flush=True)
    reporter = asyncio.create_task(progress_reporter(bench, duration))
    await phase(loop.time() + duration)
    reporter.cancel()
    await asyncio.gather(reporter, return_exceptions=True)

    elapsed = time.perf_counter() - bench.start
    sorted_lats = sorted(bench.lats_ms)
    total = bench.total
    return {
        "concurrency": concurrency,
        "duration_s": round(elapsed, 2),
        "total": total,
        "ok": bench.ok,
        "errors": bench.errors,
        "timeouts": bench.timeouts,
        "error_rate_pct": round((bench.errors + bench.timeouts) / total * 100, 2) if total else 0.0,
        "rps": round(total / elapsed, 1) if elapsed else 0.0,
        "p50_ms": round(percentile(sorted_lats, 0.50), 2),
        "p95_ms": round(percentile(sorted_lats, 0.95), 2),
        "p99_ms": round(percentile(sorted_lats, 0.99), 2),
        "max_ms": round(sorted_lats[-1], 2) if sorted_lats else 0.0,
        "avg_ms": round(statistics.fmean(bench.lats_ms), 2) if bench.lats_ms else 0.0,
        "status_dist": dict(sorted(bench.status_dist.items())),
    }


async def fetch_stats_snapshot(client: httpx.AsyncClient) -> dict:
    """拉取网关侧 /api/v1/stats 快照（失败返回空 dict）"""
    try:
        resp = await client.get("http://localhost:8000/api/v1/stats", timeout=5.0)
        return resp.json() if resp.status_code == 200 else {}
    except Exception:
        return {}


async def main() -> int:
    parser = argparse.ArgumentParser(description="chongming HTTP 端到端 RPC 压测")
    parser.add_argument("--url", default=DEFAULT_URL, help="压测目标 URL")
    parser.add_argument("--concurrency", default="100", help="并发数，逗号分隔多档（如 1,10,50,100,200）")
    parser.add_argument("--duration", type=float, default=30.0, help="每档压测时长(秒)")
    parser.add_argument("--warmup", type=float, default=3.0, help="每档预热时长(秒)")
    parser.add_argument("--max-connections", type=int, default=100,
                        help="httpx 连接池上限（并发数超过此值时客户端会排队）")
    parser.add_argument("--output", default=None, help="结果保存为 JSON 文件")
    args = parser.parse_args()

    levels = [int(c.strip()) for c in args.concurrency.split(",") if c.strip()]
    print(f"压测目标: {args.url}")
    print(f"并发档位: {levels} | 每档 {args.duration:.0f}s + 预热 {args.warmup:.0f}s")

    results = []
    limits = httpx.Limits(max_connections=args.max_connections,
                          max_keepalive_connections=args.max_connections)
    async with httpx.AsyncClient(http2=False, limits=limits) as client:
        for level in levels:
            stats_before = await fetch_stats_snapshot(client)
            result = await run_level(client, args.url, level, args.duration, args.warmup)
            stats_after = await fetch_stats_snapshot(client)
            delta = max(0, stats_after.get("total_calls", 0) - stats_before.get("total_calls", 0))
            result["gateway_total_calls_delta"] = delta
            results.append(result)

    # ---- 汇总表格 ----
    print("\n" + "=" * 88)
    print(f"{'并发':>5} | {'RPS':>8} | {'avg(ms)':>8} | {'p50':>7} | {'p95':>7} | "
          f"{'p99':>7} | {'max':>7} | {'错误率%':>7} | 网关侧计数差")
    print("-" * 88)
    for r in results:
        print(f"{r['concurrency']:>5} | {r['rps']:>8.1f} | {r['avg_ms']:>8.2f} | {r['p50_ms']:>7.2f} | "
              f"{r['p95_ms']:>7.2f} | {r['p99_ms']:>7.2f} | {r['max_ms']:>7.2f} | "
              f"{r['error_rate_pct']:>7.2f} | {r.get('gateway_total_calls_delta', '-'):>10}")
    print("=" * 88)

    # ---- 网关侧统计对照 ----
    async with httpx.AsyncClient() as client:
        stats = await fetch_stats_snapshot(client)
    if stats.get("latency"):
        print("\n网关侧 /api/v1/stats 延迟统计（进程内聚合，全部档累计）:")
        for lat in stats["latency"]:
            print(f"  worker_type={lat['worker_type']:<8} count={lat['count']:<8} "
                  f"avg={lat['avg_ms']:<8.2f}ms max={lat['max_ms']:.2f}ms")
        print(f"  累计调用: {stats.get('total_calls', 0)}")
        for c in stats.get("calls", []):
            if c["count"]:
                print(f"  subject={c['subject']:<24} status={c['status']} count={c['count']}")
    else:
        print("\n（网关侧 /api/v1/stats 无数据或不可用）")

    if args.output:
        with open(args.output, "w") as f:
            json.dump(results, f, indent=2, ensure_ascii=False)
        print(f"\n结果已保存: {args.output}")

    return 0


if __name__ == "__main__":
    sys.exit(asyncio.run(main()))