"""WebSocket 长连接并发压测（Issue #94）

支持 100 - 5000 并发连接梯度，统计建连成功率、消息往返延迟（P50/P95/P99）。

运行：
    python tests/performance/websocket_stress.py --url ws://127.0.0.1:8000/ws \
        --concurrency 500 --duration 60
"""

from __future__ import annotations

import argparse
import asyncio
import statistics
import time

try:
    import websockets
except ImportError:  # pragma: no cover
    websockets = None


class Stats:
    def __init__(self) -> None:
        self.connected = 0
        self.failed = 0
        self.latencies: list[float] = []

    def summary(self) -> str:
        if not self.latencies:
            lat = "无消息样本"
        else:
            ls = sorted(self.latencies)
            p = lambda q: ls[min(len(ls) - 1, int(q * len(ls)))]
            lat = (f"P50={p(0.5)*1000:.1f}ms P95={p(0.95)*1000:.1f}ms "
                   f"P99={p(0.99)*1000:.1f}ms avg={statistics.mean(ls)*1000:.1f}ms")
        return (f"建连成功={self.connected} 失败={self.failed} | 往返延迟 {lat}")


async def worker(url: str, duration: float, stats: Stats) -> None:
    try:
        async with websockets.connect(url, open_timeout=10) as ws:
            stats.connected += 1
            end = time.time() + duration
            while time.time() < end:
                t0 = time.perf_counter()
                await ws.send('{"type":"ping"}')
                try:
                    await asyncio.wait_for(ws.recv(), timeout=5)
                    stats.latencies.append(time.perf_counter() - t0)
                except asyncio.TimeoutError:
                    pass
                await asyncio.sleep(1)
    except Exception:
        stats.failed += 1


async def main_async(url: str, concurrency: int, duration: float, ramp: float) -> None:
    stats = Stats()
    tasks = []
    for i in range(concurrency):
        tasks.append(asyncio.create_task(worker(url, duration, stats)))
        await asyncio.sleep(ramp)  # 平滑加压，避免瞬时 SYN 风暴
    await asyncio.gather(*tasks)
    print(stats.summary())


def main() -> None:
    ap = argparse.ArgumentParser()
    ap.add_argument("--url", default="ws://127.0.0.1:8000/ws")
    ap.add_argument("--concurrency", type=int, default=100, help="并发连接数(100-5000)")
    ap.add_argument("--duration", type=float, default=60, help="每连接压测时长(秒)")
    ap.add_argument("--ramp", type=float, default=0.01, help="建连间隔(秒)")
    args = ap.parse_args()
    if websockets is None:
        raise SystemExit("请先 pip install websockets")
    if not 1 <= args.concurrency <= 5000:
        raise SystemExit("concurrency 需在 1-5000 之间")
    asyncio.run(main_async(args.url, args.concurrency, args.duration, args.ramp))


if __name__ == "__main__":
    main()