#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""作业12 · 单次延迟（Windows 友好）"""
from __future__ import annotations

import argparse
import csv
import json
import os
import time
import urllib.request
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 main() -> int:
    ap = argparse.ArgumentParser()
    ap.add_argument("--rounds", type=int, default=3)
    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))
    args = ap.parse_args()
    url = args.url or (DEFAULT_HIGRESS if args.via_higress else DEFAULT_LITELLM)

    csv_path = out_dir() / f"latency_single_{datetime.now().strftime('%Y%m%d_%H%M%S')}.csv"
    print(f"[INFO] URL={url} rounds={args.rounds}")
    rows = []
    payload = {
        "model": args.model,
        "messages": [{"role": "user", "content": "只回复一个词：pong"}],
        "max_tokens": 16,
        "stream": False,
        "temperature": 0,
    }
    data = json.dumps(payload, ensure_ascii=False).encode("utf-8")

    for i in range(1, args.rounds + 1):
        req = urllib.request.Request(
            url,
            data=data,
            headers={
                "Authorization": f"Bearer {args.key}",
                "Content-Type": "application/json",
            },
            method="POST",
        )
        t0 = time.perf_counter()
        try:
            with urllib.request.urlopen(req, timeout=180) as resp:
                body = resp.read().decode("utf-8", errors="replace")
                elapsed = time.perf_counter() - t0
                code = int(getattr(resp, "status", 200))
                try:
                    snip = (
                        (json.loads(body).get("choices") or [{}])[0]
                        .get("message", {})
                        .get("content", "")
                    )[:60]
                except Exception:
                    snip = body[:60]
        except Exception as e:
            elapsed = time.perf_counter() - t0
            code = 0
            snip = str(e)[:60]
        print(f"[R{i}] http={code} time={elapsed:.3f}s snippet={snip!r}")
        rows.append(
            {
                "round": i,
                "http_code": code,
                "time_total_s": round(elapsed, 4),
                "url": url,
                "model": args.model,
            }
        )

    with csv_path.open("w", encoding="utf-8", newline="") as f:
        w = csv.DictWriter(
            f, fieldnames=["round", "http_code", "time_total_s", "url", "model"]
        )
        w.writeheader()
        w.writerows(rows)
    ok = [r for r in rows if r["http_code"] == 200]
    ts = [float(r["time_total_s"]) for r in ok]
    print(f"成功 {len(ok)}/{len(rows)}")
    if ts:
        print(f"延迟 min={min(ts):.3f}s avg={sum(ts)/len(ts):.3f}s max={max(ts):.3f}s")
    print(f"[OK] {csv_path}")
    return 0


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