#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
8020 接口压测（在 249 本机跑，避免网络变量）。

阶梯并发（step）：每个并发档跑 per-step 秒，请求 multipart 上传同一张图并带 refresh=1
（击穿缓存，逼真实推理）；统计每档 RPS / p50 / p90 / p99 / 状态码直方图。
结尾打一次 /health 拿队列与缓存计数。产出"单机最大可用并发"结论：p99 不超时
（<50s）且 429 少量出现的最高档。

用法（任意有 requests 的 env）:
  python tools/loadtest_8020.py --image /path/to/多卡图.jpg --steps 1,2,4,8,16,32 --per-step 60
"""
import os
import sys
import time
import argparse
import threading
import statistics
from concurrent.futures import ThreadPoolExecutor, as_completed

sys.stdout.reconfigure(encoding="utf-8")

import requests


def _pctl(sorted_vals, q):
    if not sorted_vals:
        return None
    k = min(len(sorted_vals) - 1, max(0, int(round(q / 100.0 * (len(sorted_vals) - 1)))))
    return sorted_vals[k]


def _worker(url, img_path, stop_ts, results, lock):
    with open(img_path, "rb") as f:
        blob = f.read()
    sess = requests.Session()
    i = 0
    while time.time() < stop_ts:
        i += 1
        t0 = time.perf_counter()
        status = None
        try:
            r = sess.post(url, files={"image": ("loadtest.jpg", blob)},
                          data={"refresh": "1"}, timeout=90)
            status = r.status_code
        except Exception:
            status = -1
        dt = (time.perf_counter() - t0) * 1000
        with lock:
            results.append((status, dt))


def _step(name, concurrency, duration, url, img_path):
    results = []
    lock = threading.Lock()
    stop_ts = time.time() + duration
    with ThreadPoolExecutor(max_workers=concurrency) as ex:
        futs = [ex.submit(_worker, url, img_path, stop_ts, results, lock)
                for _ in range(concurrency)]
        for f in as_completed(futs):
            f.result()
    lat = sorted(d for _s, d in results)
    codes = {}
    for s, _d in results:
        codes[s] = codes.get(s, 0) + 1
    rps = len(results) / duration
    print(f"[{name}] 请求 {len(results)}  RPS={rps:.2f}  "
          f"p50={_pctl(lat, 50):.0f}ms  p90={_pctl(lat, 90):.0f}ms  p99={_pctl(lat, 99):.0f}ms  "
          f"max={lat[-1] if lat else 0:.0f}ms  状态码={dict(sorted(codes.items()))}", flush=True)
    return {"n": len(results), "rps": rps, "p99": _pctl(lat, 99), "codes": codes}


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--url", default="http://127.0.0.1:8020/match_fields")
    ap.add_argument("--image", required=True, help="压测用图（多卡图更能测出真实负载）")
    ap.add_argument("--steps", default="1,2,4,8,16,32")
    ap.add_argument("--per-step", type=int, default=60, help="每档时长(秒)")
    args = ap.parse_args()

    print(f"[LoadTest] {args.url}  image={args.image}  steps={args.steps}  "
          f"per-step={args.per_step}s", flush=True)

    # 预热一发（不带 refresh），确认接口通
    with open(args.image, "rb") as f:
        r = requests.post(args.url, files={"image": ("warmup.jpg", f.read())}, timeout=120)
    print(f"[LoadTest] 预热请求 status={r.status_code} body前80字={r.text[:80]!r}", flush=True)
    if r.status_code != 200:
        sys.exit("预热失败，终止压测")

    summary = []
    for c in [int(x) for x in args.steps.split(",") if x.strip()]:
        summary.append((c, _step(f"c={c}", c, args.per_step, args.url, args.image)))

    print("\n[Summary] 并发 → RPS / p99 / 状态码", flush=True)
    for c, s in summary:
        print(f"  c={c:<3} RPS={s['rps']:.2f}  p99={s['p99']:.0f}ms  {s['codes']}", flush=True)

    try:
        h = requests.get(args.url.rsplit("/", 1)[0] + "/health", timeout=5).json()
        print(f"\n[Health] queue={h.get('queue')}  stats={h.get('stats')}  "
              f"cache={ {k: h.get('cache', {}).get(k) for k in ('size', 'hits', 'misses')} }", flush=True)
    except Exception:
        pass

    print("\n[结论口径] 取 p99 < 50s 且 429 占比低的最高并发档为可用上限；"
          "RPS 进入平台期即 GPU 吞吐上限。", flush=True)


if __name__ == "__main__":
    main()