#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""作业13 · 逐步加压直到出现 429，输出「安全并发建议」

用法：
  # 直打模拟 primary（不经过 LiteLLM 也能学）
  python hw13_03_压测到429.py --url http://127.0.0.1:18101/v1/chat/completions --model mock-primary

  # 经测试 LiteLLM（虚拟模型 primary）
  python hw13_03_压测到429.py --url http://127.0.0.1:4010/v1/chat/completions \
    --model primary --key sk-hw13-test-4010
"""
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


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


def one_call(url: str, key: str, model: str, timeout: float) -> dict:
    payload = {
        "model": model,
        "messages": [{"role": "user", "content": "ping"}],
        "max_tokens": 8,
        "stream": False,
        "temperature": 0,
    }
    data = json.dumps(payload).encode("utf-8")
    req = urllib.request.Request(
        url,
        data=data,
        method="POST",
        headers={
            "Authorization": f"Bearer {key}",
            "Content-Type": "application/json",
        },
    )
    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))
            return {"http_code": code, "latency_s": round(elapsed, 4), "error": "" if code == 200 else body[:160]}
    except urllib.error.HTTPError as e:
        elapsed = time.perf_counter() - t0
        try:
            detail = e.read().decode("utf-8", errors="replace")[:160]
        except Exception:
            detail = str(e)
        return {"http_code": int(e.code), "latency_s": round(elapsed, 4), "error": detail}
    except Exception as e:
        elapsed = time.perf_counter() - t0
        return {"http_code": 0, "latency_s": round(elapsed, 4), "error": f"{type(e).__name__}: {e}"[:160]}


def run_level(url: str, key: str, model: str, concurrency: int, n: int, timeout: float) -> dict:
    rows: list[dict] = []
    with ThreadPoolExecutor(max_workers=concurrency) as ex:
        futs = [ex.submit(one_call, url, key, model, timeout) for _ in range(n)]
        for fut in as_completed(futs):
            rows.append(fut.result())
    codes: dict[int, int] = {}
    for r in rows:
        codes[r["http_code"]] = codes.get(r["http_code"], 0) + 1
    ok = codes.get(200, 0)
    n429 = codes.get(429, 0)
    n5xx = sum(v for k, v in codes.items() if k >= 500)
    n0 = codes.get(0, 0)
    lats = [r["latency_s"] for r in rows if r["http_code"] == 200]
    return {
        "concurrency": concurrency,
        "n": n,
        "ok": ok,
        "n429": n429,
        "n5xx": n5xx,
        "n_other": n - ok - n429 - n5xx - n0,
        "n0": n0,
        "ok_rate": round(ok / n, 4) if n else 0,
        "rate_429": round(n429 / n, 4) if n else 0,
        "latency_avg": round(sum(lats) / len(lats), 4) if lats else None,
        "codes": codes,
    }


def suggest_safe(levels: list[dict], margin: float = 0.7) -> dict:
    """取「无 429 且成功率>=95%」的最大并发，再乘安全余量。"""
    good = [lv for lv in levels if lv["n429"] == 0 and lv["ok_rate"] >= 0.95 and lv["n0"] == 0]
    if not good:
        return {
            "safe_concurrency": 1,
            "basis": "未找到稳定档位，建议并发=1 并先查上游/鉴权",
            "max_stable": 0,
        }
    max_stable = max(lv["concurrency"] for lv in good)
    safe = max(1, int(max_stable * margin))
    first_429 = next((lv["concurrency"] for lv in levels if lv["n429"] > 0), None)
    return {
        "safe_concurrency": safe,
        "max_stable": max_stable,
        "first_429_at": first_429,
        "basis": (
            f"最大稳定并发={max_stable}（429=0 且成功率>=95%）；"
            f"安全建议= floor({max_stable}*{margin}) = {safe}"
            + (f"；首次出现429的并发={first_429}" if first_429 else "")
        ),
    }


def main() -> int:
    ap = argparse.ArgumentParser(description="作业13 加压至429")
    ap.add_argument("--url", default="http://127.0.0.1:18101/v1/chat/completions")
    ap.add_argument("--model", default="mock-primary")
    ap.add_argument("--key", default="sk-mock")
    ap.add_argument("--levels", default="1,2,4,6,8,12,16", help="并发阶梯，逗号分隔")
    ap.add_argument("--requests-per-level", type=int, default=20, help="每档请求数")
    ap.add_argument("--timeout", type=float, default=30.0)
    ap.add_argument("--margin", type=float, default=0.7, help="安全余量系数")
    ap.add_argument("--stop-on-429", action="store_true", default=True)
    ap.add_argument("--no-stop-on-429", action="store_true", help="即使出429也跑完所有档")
    args = ap.parse_args()
    stop_on_429 = not args.no_stop_on_429

    levels = [int(x.strip()) for x in args.levels.split(",") if x.strip()]
    stamp = datetime.now().strftime("%Y%m%d_%H%M%S")
    csv_path = out_dir() / f"ramp_{stamp}.csv"
    json_path = out_dir() / f"ramp_{stamp}.json"

    print(f"[INFO] url={args.url}")
    print(f"[INFO] model={args.model} levels={levels} n/level={args.requests_per_level}")
    results: list[dict] = []

    for c in levels:
        print(f"\n=== 并发={c} 请求数={args.requests_per_level} ===")
        lv = run_level(args.url, args.key, args.model, c, args.requests_per_level, args.timeout)
        results.append(lv)
        print(
            f"ok={lv['ok']}/{lv['n']} ok_rate={lv['ok_rate']:.0%} "
            f"429={lv['n429']} 5xx={lv['n5xx']} other_fail={lv['n0']} "
            f"latency_avg={lv['latency_avg']}"
        )
        print(f"codes={lv['codes']}")
        if stop_on_429 and lv["n429"] > 0:
            print("[INFO] 已出现 429，停止继续加压（可用 --no-stop-on-429 跑完）")
            break

    advice = suggest_safe(results, margin=args.margin)
    print("\n========== 安全并发建议 ==========")
    print(f"safe_concurrency = {advice['safe_concurrency']}")
    print(f"说明：{advice['basis']}")
    print("==================================\n")

    with csv_path.open("w", encoding="utf-8", newline="") as f:
        w = csv.DictWriter(
            f,
            fieldnames=[
                "concurrency",
                "n",
                "ok",
                "n429",
                "n5xx",
                "n0",
                "ok_rate",
                "rate_429",
                "latency_avg",
            ],
        )
        w.writeheader()
        for lv in results:
            w.writerow({k: lv[k] for k in w.fieldnames})

    payload = {
        "url": args.url,
        "model": args.model,
        "levels": results,
        "advice": advice,
        "finished_at": datetime.now().isoformat(timespec="seconds"),
    }
    json_path.write_text(json.dumps(payload, ensure_ascii=False, indent=2), encoding="utf-8")
    print(f"[OK] CSV={csv_path}")
    print(f"[OK] JSON={json_path}")
    return 0


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