#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""作业12 · 并发压测（Windows / 207 通用，不依赖 xargs -P）

用法：
  python hw12_02_并发压测.py
  python hw12_02_并发压测.py --concurrency 12 --n 36
  python hw12_02_并发压测.py --via-higress
"""
from __future__ import annotations

import argparse
import csv
import json
import os
import time
import urllib.error
import urllib.request
from concurrent.futures import ThreadPoolExecutor, as_completed
from datetime import datetime
from pathlib import Path

DEFAULT_LITELLM = "http://192.168.110.207:4000/v1/chat/completions"
DEFAULT_HIGRESS = "http://192.168.110.207:18080/llm/v1/chat/completions"
DEFAULT_KEY = "sk-litellm-bocheng-207"
DEFAULT_MODEL = "qwen3.6-27b"


def out_dir() -> Path:
    env = os.environ.get("HW12_OUT")
    if env:
        p = Path(env)
        p.mkdir(parents=True, exist_ok=True)
        return p
    repo = Path(__file__).resolve().parents[2] / "data" / "homework12"
    candidates: list[Path] = [repo]
    opt = Path("/opt/workbuddy/data/homework12")
    if os.name != "nt" and Path("/opt/workbuddy").is_dir():
        candidates.insert(0, opt)
    for cand in candidates:
        try:
            cand.mkdir(parents=True, exist_ok=True)
            return cand
        except OSError:
            continue
    p = Path.cwd() / "data" / "homework12"
    p.mkdir(parents=True, exist_ok=True)
    return p


def one_call(i: int, url: str, key: str, model: str, timeout: float) -> dict:
    payload = {
        "model": model,
        "messages": [{"role": "user", "content": "用一句话介绍什么是网关"}],
        "max_tokens": 32,
        "stream": False,
        "temperature": 0,
    }
    data = json.dumps(payload, ensure_ascii=False).encode("utf-8")
    req = urllib.request.Request(
        url,
        data=data,
        headers={
            "Authorization": f"Bearer {key}",
            "Content-Type": "application/json",
        },
        method="POST",
    )
    t0 = time.perf_counter()
    try:
        with urllib.request.urlopen(req, timeout=timeout) as resp:
            body = resp.read().decode("utf-8", errors="replace")
            elapsed = time.perf_counter() - t0
            code = int(getattr(resp, "status", 200))
            err = "" if code == 200 else body[:120]
            return {"idx": i, "http_code": code, "time_total_s": round(elapsed, 4), "error": err}
    except urllib.error.HTTPError as e:
        elapsed = time.perf_counter() - t0
        try:
            detail = e.read().decode("utf-8", errors="replace")[:120]
        except Exception:
            detail = str(e)
        return {"idx": i, "http_code": int(e.code), "time_total_s": round(elapsed, 4), "error": detail}
    except Exception as e:
        elapsed = time.perf_counter() - t0
        return {
            "idx": i,
            "http_code": 0,
            "time_total_s": round(elapsed, 4),
            "error": f"{type(e).__name__}: {e}"[:120],
        }


def pct(sorted_vals: list[float], q: float) -> float:
    if not sorted_vals:
        return 0.0
    i = min(len(sorted_vals) - 1, max(0, int(round(q * (len(sorted_vals) - 1)))))
    return sorted_vals[i]


def main() -> int:
    ap = argparse.ArgumentParser()
    ap.add_argument("--concurrency", type=int, default=8)
    ap.add_argument("--n", type=int, default=24)
    ap.add_argument("--via-higress", action="store_true")
    ap.add_argument("--url", default="")
    ap.add_argument("--key", default=os.environ.get("LITELLM_MASTER_KEY", DEFAULT_KEY))
    ap.add_argument("--model", default=os.environ.get("LITELLM_CHAT_MODEL", DEFAULT_MODEL))
    ap.add_argument("--timeout", type=float, default=300.0)
    args = ap.parse_args()

    if args.url:
        url = args.url
    elif args.via_higress:
        url = DEFAULT_HIGRESS
    else:
        url = DEFAULT_LITELLM

    stamp = datetime.now().strftime("%Y%m%d_%H%M%S")
    csv_path = out_dir() / f"latency_conc_{stamp}.csv"
    print(f"[INFO] URL={url} concurrency={args.concurrency} n={args.n}")
    print(f"[INFO] CSV={csv_path}")

    rows: list[dict] = []
    with ThreadPoolExecutor(max_workers=args.concurrency) as ex:
        futs = [
            ex.submit(one_call, i, url, args.key, args.model, args.timeout)
            for i in range(1, args.n + 1)
        ]
        for fut in as_completed(futs):
            r = fut.result()
            rows.append(r)
            print(f"[#{r['idx']}] http={r['http_code']} time={r['time_total_s']}s {r['error'] and 'err='+r['error']}")

    rows.sort(key=lambda x: x["idx"])
    with csv_path.open("w", encoding="utf-8", newline="") as f:
        w = csv.DictWriter(f, fieldnames=["idx", "http_code", "time_total_s", "error"])
        w.writeheader()
        w.writerows(rows)

    ok = [r for r in rows if r["http_code"] == 200]
    ts = sorted(float(r["time_total_s"]) for r in ok)
    from collections import Counter

    codes = Counter(str(r["http_code"]) for r in rows)
    print("=== 汇总 ===")
    print("HTTP 码分布:", dict(codes))
    print(f"成功 {len(ok)}/{len(rows)}")
    if ts:
        print(
            f"延迟 s: min={min(ts):.3f} p50={pct(ts, 0.5):.3f} "
            f"p95={pct(ts, 0.95):.3f} max={max(ts):.3f}"
        )
    print("请对照作业文档填写「现象→原因→处理」五联表。")
    print(f"[OK] {csv_path}")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())