#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
OpenAI 兼容接口(vLLM / TGI / SGLang / Ascend-vLLM ...) 推理性能压测脚本
=====================================================================

支持一次跑多个服务端（不同端口 / 不同模型 / 不同上下文上限），结果合并到一张
CSV，字段与性能看护表格完全对齐：

  硬件 | 测试版本 | 模型 | dtype | TP | 输入 | 输出 | 数据集条数 | 请求频率 | BS类型 |
  Concurrency | TTFT(ms,Prefill时延) | TPOT(ms,Decode时延) | E2E时间(ms) | QPS |
  输出吞吐 | 单芯输出吞吐 | 总吞吐 | 单卡E2E吞吐 | 自动化及看护责任田

指标口径（与表格一致）
----------------------
  E2E时间(ms)   = 该组用例的整体墙钟耗时（不是单请求时延）
  QPS           = 成功请求数 / E2E时间
  TTFT          = 各请求首 token 时延的均值
  TPOT          = 各请求 (单请求E2E - TTFT) / (输出token数 - 1) 的均值
  输出吞吐      = 总输出token / E2E时间          (tokens/s)
  总吞吐        = (总输入token + 总输出token) / E2E时间
  单芯输出吞吐  = 输出吞吐 / chips
  单卡E2E吞吐   = 总吞吐   / tp
  Concurrency   = 实测平均并发 = Σ(单请求E2E) / E2E时间

多模型 / 多端口
---------------
  1) 生成配置模板（已按 8000=Qwen3.6-27B/128k、8001=Qwen3-14B/40k 预填）：
         python bench_openai.py --gen-config bench.json
  2) 按需改 bench.json 里的 cases，然后：
         python bench_openai.py --config bench.json -o result.csv
  每个 target 有自己的 base_url / 模型 / max_model_len / tp / chips / 用例集；脚本会
  按各自的上下文上限自动跳过（或裁剪）放不下的用例，结果写进同一张表，用「模型」
  列区分。单端口临时测试仍可以完全走命令行参数，不需要配置文件。

长序列注意事项（10K 输入 / 50K 输出这类）
-----------------------------------------
  * 强烈建议装 aiohttp：50K 输出=单请求 5 万个 SSE 分片，线程模式 GIL 会污染测量。
        pip install aiohttp
  * 长度单位 K = 1024。10k→10240，50k→51200。要 10000/50000 就直接写数字。
  * 输入+输出 必须 ≤ 服务端 max_model_len，脚本会预检并按 --on-overflow 处理。
  * 用 rounds=1 让「数据集条数 = 并发」，别用默认 4，否则一组用例要跑几小时。
  * --case-timeout 给单组用例设墙钟上限，超时会真正中止在途请求（已生成 token 仍
    计入吞吐，用「中止数」列标记）。
  * 长输出下 decode 会随 KV 增长变慢：看 TPOT_前1000 / TPOT_后1000 / TPOT衰减倍数。

依赖
----
  纯标准库即可运行（线程模式）。装了 aiohttp 会自动走 asyncio 模式。
"""

import argparse
import copy
import csv
import json
import math
import os
import random
import statistics
import string
import sys
import threading
import time
import urllib.error
import urllib.request
from collections import deque
from concurrent.futures import ThreadPoolExecutor
from dataclasses import dataclass, field
from typing import List, Optional, Tuple

try:
    import aiohttp  # type: ignore
    import asyncio

    HAS_AIOHTTP = True
except Exception:  # pragma: no cover
    HAS_AIOHTTP = False


ITL_KEEP = 1000          # 首/尾各保留多少个 token 间隔用于分段 TPOT
CTX_MARGIN = 64          # 上下文预留：对话模板/特殊 token 的余量

# 进度心跳共享计数（近似值即可，不加锁）
PROG = {"tokens": 0, "done": 0, "total": 0, "t0": 0.0, "on": False}

# --case-timeout 触发后置位：所有在途请求立刻停止读流（已生成的 token 仍计入统计）
STOP = threading.Event()


# --------------------------------------------------------------------------- #
# 结果结构
# --------------------------------------------------------------------------- #
@dataclass
class ReqResult:
    ok: bool = False
    ttft: float = 0.0          # s
    latency: float = 0.0       # s，单请求 E2E
    prompt_tokens: int = 0
    output_tokens: int = 0
    itl_head: List[float] = field(default_factory=list)   # 前 ITL_KEEP 个 token 间隔
    itl_tail: List[float] = field(default_factory=list)   # 后 ITL_KEEP 个 token 间隔
    itl_n: int = 0
    itl_sum: float = 0.0
    itl_max: float = 0.0
    partial: bool = False      # 被 --case-timeout 中途截断
    error: str = ""


# --------------------------------------------------------------------------- #
# Prompt 构造（无 tokenizer，近似 + 在线校准）
# --------------------------------------------------------------------------- #
_WORD_POOL = [
    "system", "network", "compute", "memory", "kernel", "tensor", "matrix",
    "vector", "buffer", "stream", "device", "driver", "module", "layer",
    "weight", "signal", "sample", "record", "region", "server", "client",
    "packet", "thread", "socket", "engine", "runtime", "backend", "cluster",
    "pipeline", "storage", "monitor", "profile", "latency", "traffic",
    "request", "response", "channel", "gateway", "session", "control",
]


class PromptBuilder:
    """用随机词近似构造指定 token 长度的 prompt；用服务端 usage 在线校准。
    每个 target 单独校准一次（不同模型 tokenizer 不同）。"""

    def __init__(self, tokens_per_word: float = 1.0, overhead: int = 0, seed: int = 0):
        self.k = tokens_per_word     # 每个词的平均 token 数
        self.b = overhead            # 模板/特殊 token 固定开销
        self.rng = random.Random(seed)

    def words_for(self, target_tokens: int) -> int:
        return max(1, int(round((target_tokens - self.b) / max(self.k, 1e-6))))

    def make(self, target_tokens: int) -> str:
        return self.make_words(self.words_for(target_tokens))

    def make_words(self, n_words: int) -> str:
        # 每条 prompt 都不同，避免命中 prefix cache 造成虚高
        head = "".join(self.rng.choice(string.ascii_lowercase) for _ in range(8))
        body = " ".join(self.rng.choice(_WORD_POOL) for _ in range(max(0, n_words - 1)))
        return (head + " " + body).strip()


# --------------------------------------------------------------------------- #
# HTTP 基础
# --------------------------------------------------------------------------- #
def build_payload(cfg, prompt: str, out_len: int, stream: bool = True) -> Tuple[str, dict]:
    if cfg.api == "chat":
        url = cfg.base_url.rstrip("/") + "/chat/completions"
        payload = {"model": cfg.model, "messages": [{"role": "user", "content": prompt}]}
    else:
        url = cfg.base_url.rstrip("/") + "/completions"
        payload = {"model": cfg.model, "prompt": prompt}
    payload.update({"max_tokens": out_len, "temperature": cfg.temperature,
                    "top_p": cfg.top_p, "stream": stream})
    if stream:
        payload["stream_options"] = {"include_usage": True}
    if cfg.ignore_eos:
        payload["ignore_eos"] = True
        payload["min_tokens"] = out_len          # vLLM 支持，强制定长输出
    if cfg.extra_body:
        payload.update(cfg.extra_body)
    return url, payload


def headers_of(cfg) -> dict:
    h = {"Content-Type": "application/json", "Accept": "text/event-stream"}
    if cfg.api_key:
        h["Authorization"] = "Bearer " + cfg.api_key
    return h


def _extract_piece(obj: dict) -> str:
    """从一个 SSE chunk 中取出本次新增文本（兼容 chat / completions / reasoning）。"""
    choices = obj.get("choices") or []
    if not choices:
        return ""
    c = choices[0]
    if "delta" in c and isinstance(c["delta"], dict):
        d = c["delta"]
        return (d.get("content") or d.get("reasoning_content") or "") or ""
    return c.get("text") or ""


def new_state() -> dict:
    return {"ttft": None, "last": 0.0, "n_chunks": 0, "usage": None, "error": None,
            "head": [], "tail": deque(maxlen=ITL_KEEP), "isum": 0.0, "imax": 0.0,
            "aborted": False}


def _handle_line(line: str, state: dict, t0: float) -> bool:
    """处理一行 SSE，返回 False 表示流结束。"""
    if not line:
        return True
    if line.startswith("data:"):
        data = line[5:].strip()
    else:
        return True
    if data == "[DONE]":
        return False
    try:
        obj = json.loads(data)
    except json.JSONDecodeError:
        return True
    if obj.get("usage"):
        state["usage"] = obj["usage"]
    if obj.get("error"):
        state["error"] = str(obj["error"])
        return False
    piece = _extract_piece(obj)
    if piece:
        now = time.perf_counter()
        if state["ttft"] is None:
            state["ttft"] = now - t0
        else:
            gap = now - state["last"]
            state["isum"] += gap
            if gap > state["imax"]:
                state["imax"] = gap
            if len(state["head"]) < ITL_KEEP:
                state["head"].append(gap)
            state["tail"].append(gap)
        state["last"] = now
        state["n_chunks"] += 1
        if PROG["on"]:
            PROG["tokens"] += 1              # 近似计数，仅用于进度显示
    return True


def _finish(state: dict, t0: float) -> ReqResult:
    r = ReqResult()
    r.latency = time.perf_counter() - t0
    if state.get("error") or state["ttft"] is None:
        r.ok = False
        r.error = state.get("error") or (
            "aborted before first token" if state.get("aborted") else "no token received")
        return r
    r.partial = bool(state.get("aborted"))
    usage = state.get("usage") or {}
    r.ok = True
    r.ttft = state["ttft"]
    r.output_tokens = int(usage.get("completion_tokens") or state["n_chunks"])
    r.prompt_tokens = int(usage.get("prompt_tokens") or 0)
    r.itl_head = state["head"]
    r.itl_tail = list(state["tail"])
    r.itl_n = max(0, state["n_chunks"] - 1)
    r.itl_sum = state["isum"]
    r.itl_max = state["imax"]
    return r


# ------------------------------ 线程实现 ----------------------------------- #
def request_sync(cfg, prompt: str, out_len: int) -> ReqResult:
    url, payload = build_payload(cfg, prompt, out_len, stream=True)
    req = urllib.request.Request(url, data=json.dumps(payload).encode("utf-8"),
                                 headers=headers_of(cfg), method="POST")
    state = new_state()
    t0 = time.perf_counter()
    try:
        # urllib 的 timeout 是 socket 级读超时，正好等价于「多久没收到新 token」
        with urllib.request.urlopen(req, timeout=cfg.stall_timeout) as resp:
            for raw in resp:
                if not _handle_line(raw.decode("utf-8", "ignore").strip(), state, t0):
                    break
                if STOP.is_set():
                    state["aborted"] = True
                    break
                if cfg.timeout and time.perf_counter() - t0 > cfg.timeout:
                    state["error"] = f"request exceeded --timeout {cfg.timeout}s"
                    break
    except urllib.error.HTTPError as e:
        state["error"] = f"HTTP {e.code}: {e.read().decode('utf-8', 'ignore')[:200]}"
    except Exception as e:  # noqa: BLE001
        state["error"] = f"{type(e).__name__}: {e}"
    return _finish(state, t0)


def run_case_threads(cfg, prompts, out_len, concurrency, request_rate):
    results: List[Optional[ReqResult]] = [None] * len(prompts)
    sem = threading.Semaphore(concurrency)
    rng = random.Random(cfg.seed)

    def work(i: int):
        with sem:
            results[i] = request_sync(cfg, prompts[i], out_len)
        PROG["done"] += 1

    t0 = time.perf_counter()
    submitted = 0
    with ThreadPoolExecutor(max_workers=max(concurrency, 1) + 4) as pool:
        futs = []
        for i in range(len(prompts)):
            if STOP.is_set():
                print(f"    [warn] 达到 case_timeout，停止投放剩余 {len(prompts)-i} 条请求")
                break
            if request_rate > 0 and i > 0:
                time.sleep(rng.expovariate(request_rate))
            futs.append(pool.submit(work, i))
            submitted += 1
        for f in futs:
            f.result()
    return [r for r in results if r is not None], time.perf_counter() - t0, submitted


# ------------------------------ asyncio 实现 -------------------------------- #
async def _request_async(session, cfg, prompt, out_len) -> ReqResult:
    url, payload = build_payload(cfg, prompt, out_len, stream=True)
    state = new_state()
    t0 = time.perf_counter()
    try:
        async with session.post(url, json=payload, headers=headers_of(cfg)) as resp:
            if resp.status != 200:
                state["error"] = f"HTTP {resp.status}: {(await resp.text())[:200]}"
            else:
                async for raw in resp.content:
                    if not _handle_line(raw.decode("utf-8", "ignore").strip(), state, t0):
                        break
                    if STOP.is_set():
                        state["aborted"] = True
                        break
    except Exception as e:  # noqa: BLE001
        state["error"] = f"{type(e).__name__}: {e}"
    return _finish(state, t0)


async def _run_case_async(cfg, prompts, out_len, concurrency, request_rate):
    conn = aiohttp.TCPConnector(limit=0, ssl=False)
    # total=None：长输出请求可能跑几十分钟；用 sock_read 检测「卡住不出 token」
    timeout = aiohttp.ClientTimeout(total=cfg.timeout or None, sock_connect=30,
                                    sock_read=cfg.stall_timeout)
    sem = asyncio.Semaphore(concurrency)
    rng = random.Random(cfg.seed)

    async with aiohttp.ClientSession(connector=conn, timeout=timeout) as session:
        async def one(p):
            async with sem:
                r = await _request_async(session, cfg, p, out_len)
            PROG["done"] += 1
            return r

        t0 = time.perf_counter()
        tasks = []
        for i, p in enumerate(prompts):
            if STOP.is_set():
                print(f"    [warn] 达到 case_timeout，停止投放剩余 {len(prompts)-i} 条请求")
                break
            if request_rate > 0 and i > 0:
                await asyncio.sleep(rng.expovariate(request_rate))
            tasks.append(asyncio.create_task(one(p)))
        results = await asyncio.gather(*tasks)
        dur = time.perf_counter() - t0
    return list(results), dur, len(tasks)


def run_case(cfg, prompts, out_len, concurrency, request_rate):
    STOP.clear()
    watchdog = None
    if cfg.case_timeout:
        watchdog = threading.Timer(cfg.case_timeout, STOP.set)
        watchdog.daemon = True
        watchdog.start()
    try:
        if HAS_AIOHTTP and not cfg.force_threads:
            return asyncio.run(
                _run_case_async(cfg, prompts, out_len, concurrency, request_rate))
        return run_case_threads(cfg, prompts, out_len, concurrency, request_rate)
    finally:
        if watchdog:
            watchdog.cancel()
        STOP.clear()


# --------------------------------------------------------------------------- #
# 进度心跳
# --------------------------------------------------------------------------- #
class Heartbeat:
    def __init__(self, interval: float, total_req: int):
        self.interval = interval
        self.total = total_req
        self._stop = threading.Event()
        self._t = None

    def __enter__(self):
        if self.interval <= 0:
            return self
        PROG.update({"tokens": 0, "done": 0, "total": self.total,
                     "t0": time.perf_counter(), "on": True})
        self._t = threading.Thread(target=self._run, daemon=True)
        self._t.start()
        return self

    def _run(self):
        while not self._stop.wait(self.interval):
            el = time.perf_counter() - PROG["t0"]
            tok = PROG["tokens"]
            print(f"    ... {el:6.0f}s  完成 {PROG['done']}/{self.total}  "
                  f"已生成 {tok} tokens  瞬时输出吞吐 ≈ {tok/max(el,1e-6):.1f} tok/s",
                  flush=True)

    def __exit__(self, *exc):
        PROG["on"] = False
        self._stop.set()
        return False


# --------------------------------------------------------------------------- #
# 校准 / 预检
# --------------------------------------------------------------------------- #
def probe_prompt_tokens(cfg, prompt: str) -> Optional[int]:
    url, payload = build_payload(cfg, prompt, 1, stream=False)
    payload.pop("stream_options", None)
    payload.pop("min_tokens", None)
    hh = {"Content-Type": "application/json"}
    if cfg.api_key:
        hh["Authorization"] = "Bearer " + cfg.api_key
    req = urllib.request.Request(url, data=json.dumps(payload).encode("utf-8"),
                                 headers=hh, method="POST")
    try:
        with urllib.request.urlopen(req, timeout=cfg.stall_timeout) as resp:
            obj = json.loads(resp.read().decode("utf-8"))
        return int(obj["usage"]["prompt_tokens"])
    except Exception as e:  # noqa: BLE001
        print(f"[warn] prompt token 校准探测失败: {e}", file=sys.stderr)
        return None


def calibrate(cfg, pb: PromptBuilder) -> PromptBuilder:
    """两点线性拟合 tokens = k*words + b（每个模型单独校准）。"""
    if cfg.no_calibrate:
        return pb
    n1, n2 = 100, 500
    t1 = probe_prompt_tokens(cfg, pb.make_words(n1))
    t2 = probe_prompt_tokens(cfg, pb.make_words(n2))
    if not t1 or not t2 or t2 <= t1:
        print("[warn] 未能校准，使用默认近似 (1 词 ≈ 1 token)", file=sys.stderr)
        return pb
    k = (t2 - t1) / (n2 - n1)
    pb.k, pb.b = k, t1 - k * n1
    print(f"  [calib] tokens ≈ {pb.k:.4f} × words + {pb.b:.1f}")
    return pb


def make_prompts(cfg, pb: PromptBuilder, in_len: int, n: int) -> List[str]:
    """生成 n 条互不相同的 prompt（避免命中 prefix cache）。校准探测用一次性的临时
    prompt，不会污染正式压测样本。"""
    w = pb.words_for(in_len)
    if not cfg.no_calibrate:
        actual = probe_prompt_tokens(cfg, pb.make_words(w))     # 一次性探测样本
        if actual and abs(actual - in_len) > max(2, in_len * 0.01):
            w = max(1, int(round(w * (in_len - pb.b) / max(actual - pb.b, 1))))
            actual2 = probe_prompt_tokens(cfg, pb.make_words(w))
            print(f"  [calib] 输入目标 {in_len} → 首次 {actual} → 修正后 {actual2}")
    return [pb.make_words(w) for _ in range(n)]


def get_model_info(cfg) -> dict:
    url = cfg.base_url.rstrip("/") + "/models"
    req = urllib.request.Request(
        url, headers={"Authorization": "Bearer " + cfg.api_key} if cfg.api_key else {})
    with urllib.request.urlopen(req, timeout=30) as resp:
        obj = json.loads(resp.read().decode("utf-8"))
    return (obj.get("data") or [{}])[0]


def resolve_target(cfg) -> dict:
    """补全 model / max_model_len。config 里显式写了就以配置为准。"""
    info = {}
    if not cfg.model or not cfg.max_model_len:
        try:
            info = get_model_info(cfg)
        except Exception as e:  # noqa: BLE001
            print(f"  [warn] 读取 /models 失败: {e}", file=sys.stderr)
    if not cfg.model:
        cfg.model = info.get("id", "")
        print(f"  [info] 自动识别 model = {cfg.model}")
    if not cfg.max_model_len:
        cfg.max_model_len = info.get("max_model_len") or info.get("context_length") or 0
        if cfg.max_model_len:
            print(f"  [info] 服务端 max_model_len = {cfg.max_model_len}")
        else:
            print("  [warn] 未能获知 max_model_len，跳过上下文预检（配置里可手工写死）")
    return cfg


def filter_cases(cfg, cases):
    """按该 target 的上下文上限跳过 / 裁剪放不下的用例。"""
    if not cfg.max_model_len:
        return cases
    limit = int(cfg.max_model_len) - CTX_MARGIN
    kept = []
    for (i, o, c, n) in cases:
        if i + o <= limit:
            kept.append((i, o, c, n))
            continue
        if cfg.on_overflow == "run":
            print(f"  [warn] 用例 {i}/{o} 超出上下文 {cfg.max_model_len}，仍按配置执行（大概率报错）")
            kept.append((i, o, c, n))
        elif cfg.on_overflow == "clamp":
            o2 = limit - i
            if o2 <= 0:
                print(f"  [skip] 用例 {i}/{o}：输入本身就超过上下文 {cfg.max_model_len}")
            else:
                print(f"  [clamp] 用例 {i}/{o} → {i}/{o2}（受上下文 {cfg.max_model_len} 限制）")
                kept.append((i, o2, c, n))
        else:  # skip
            print(f"  [skip] 用例 {i}/{o}：输入+输出 {i+o} > 上下文 {cfg.max_model_len}")
    return kept


# --------------------------------------------------------------------------- #
# 统计
# --------------------------------------------------------------------------- #
def pct(xs: List[float], q: float) -> float:
    if not xs:
        return 0.0
    s = sorted(xs)
    return s[min(len(s) - 1, max(0, int(math.ceil(q / 100.0 * len(s))) - 1))]


def summarize(cfg, results, dur, in_len, out_len, concurrency, n_prompts, request_rate) -> dict:
    ok = [r for r in results if r.ok]
    fail = len(results) - len(ok)
    n = len(ok)
    if n == 0 or dur <= 0:
        return {"错误": "全部失败", "失败数": fail}

    ttfts = [r.ttft for r in ok]
    lats = [r.latency for r in ok]
    tpots = []
    for r in ok:
        if r.output_tokens > 1:
            tpots.append((r.latency - r.ttft) / (r.output_tokens - 1))
        elif r.itl_n:
            tpots.append(r.itl_sum / r.itl_n)
    head = [statistics.mean(r.itl_head) for r in ok if r.itl_head]
    tail = [statistics.mean(r.itl_tail) for r in ok if r.itl_tail]
    tot_out = sum(r.output_tokens for r in ok)
    tot_in = sum(r.prompt_tokens or in_len for r in ok)

    out_tps = tot_out / dur
    all_tps = (tot_in + tot_out) / dur

    return {
        "硬件": cfg.hardware, "测试版本": cfg.version, "模型": cfg.model_label or cfg.model,
        "dtype": cfg.dtype, "TP": cfg.tp, "输入": in_len, "输出": out_len,
        "数据集条数": n_prompts, "请求频率": request_rate, "BS类型": concurrency,
        "Concurrency": round(sum(lats) / dur, 4),
        "TTFT(ms,Prefill时延)": round(statistics.mean(ttfts) * 1000, 4),
        "TPOT(ms,Decode时延)": round(statistics.mean(tpots) * 1000, 4) if tpots else 0.0,
        "E2E时间(ms)": round(dur * 1000, 4),
        "QPS": round(n / dur, 4),
        "输出吞吐": round(out_tps, 4),
        "单芯输出吞吐": round(out_tps / cfg.chips, 4),
        "总吞吐": round(all_tps, 4),
        "单卡E2E吞吐": round(all_tps / cfg.tp, 4),
        "自动化及看护责任田": cfg.owner,
        # ---- 附加诊断字段 ----
        "服务地址": cfg.base_url,
        "成功数": n, "失败数": fail,
        "中止数(超时截断)": sum(1 for r in ok if r.partial),
        "实际输入token均值": round(tot_in / n, 1),
        "实际输出token均值": round(tot_out / n, 1),
        "单请求平均时延(ms)": round(statistics.mean(lats) * 1000, 2),
        "TTFT_P50(ms)": round(pct(ttfts, 50) * 1000, 2),
        "TTFT_P90(ms)": round(pct(ttfts, 90) * 1000, 2),
        "TTFT_P99(ms)": round(pct(ttfts, 99) * 1000, 2),
        "TPOT_P90(ms)": round(pct(tpots, 90) * 1000, 2) if tpots else 0.0,
        f"TPOT_前{ITL_KEEP}(ms)": round(statistics.mean(head) * 1000, 4) if head else 0.0,
        f"TPOT_后{ITL_KEEP}(ms)": round(statistics.mean(tail) * 1000, 4) if tail else 0.0,
        "TPOT衰减倍数": round(statistics.mean(tail) / statistics.mean(head), 3)
                        if head and tail and statistics.mean(head) > 0 else 0.0,
        "最大token间隔(ms)": round(max((r.itl_max for r in ok), default=0.0) * 1000, 2),
        "单请求E2E_P90(ms)": round(pct(lats, 90) * 1000, 2),
    }


MAIN_COLS = ["硬件", "测试版本", "模型", "dtype", "TP", "输入", "输出", "数据集条数", "请求频率",
             "BS类型", "Concurrency", "TTFT(ms,Prefill时延)", "TPOT(ms,Decode时延)",
             "E2E时间(ms)", "QPS", "输出吞吐", "单芯输出吞吐", "总吞吐", "单卡E2E吞吐",
             "自动化及看护责任田"]
EXTRA_COLS = ["服务地址", "成功数", "失败数", "中止数(超时截断)", "实际输入token均值",
              "实际输出token均值", "单请求平均时延(ms)", "TTFT_P50(ms)", "TTFT_P90(ms)",
              "TTFT_P99(ms)", "TPOT_P90(ms)", f"TPOT_前{ITL_KEEP}(ms)",
              f"TPOT_后{ITL_KEEP}(ms)", "TPOT衰减倍数", "最大token间隔(ms)", "单请求E2E_P90(ms)"]
PRINT_COLS = ["模型", "输入", "输出", "数据集条数", "BS类型", "Concurrency",
              "TTFT(ms,Prefill时延)", "TPOT(ms,Decode时延)", "E2E时间(ms)", "QPS",
              "输出吞吐", "单芯输出吞吐", "总吞吐", "单卡E2E吞吐"]


def _w(s) -> int:
    return sum(2 if ord(c) > 0x2E80 else 1 for c in str(s))


def print_table(rows, cols):
    if not rows:
        return
    widths = [max(_w(c), max(_w(r.get(c, "")) for r in rows)) for c in cols]
    def line(vals):
        return " | ".join(str(v) + " " * (w - _w(v)) for v, w in zip(vals, widths))
    print(line(cols))
    print("-+-".join("-" * w for w in widths))
    for r in rows:
        print(line([r.get(c, "") for c in cols]))


def write_csv(path, rows):
    cols = MAIN_COLS + EXTRA_COLS
    with open(path, "w", newline="", encoding="utf-8-sig") as f:
        w = csv.DictWriter(f, fieldnames=cols, extrasaction="ignore")
        w.writeheader()
        for r in rows:
            w.writerow(r)


# --------------------------------------------------------------------------- #
# 用例 / 配置
# --------------------------------------------------------------------------- #
def parse_len(s) -> int:
    """256 / 4k / 10K / 50k —— K = 1024。要 10000 就直接写 10000。"""
    s = str(s).strip().lower()
    if s.endswith("k"):
        return int(float(s[:-1]) * 1024)
    if s.endswith("m"):
        return int(float(s[:-1]) * 1024 * 1024)
    return int(s)


def fmt_secs(x: float) -> str:
    return f"{x/3600:.1f}h" if x >= 3600 else (f"{x/60:.1f}min" if x >= 60 else f"{x:.0f}s")


def _conc_list(cfg):
    cc = cfg.concurrency
    return [int(x) for x in (cc if isinstance(cc, list) else str(cc).split(",")) if str(x).strip()]


def cases_from_cli(cfg):
    """纯命令行方式（单端口）构造用例。"""
    cases = []
    if cfg.cases:
        for item in cfg.cases.split(","):
            parts = item.strip().split(":")
            i, o, c = parse_len(parts[0]), parse_len(parts[1]), int(parts[2])
            n = int(parts[3]) if len(parts) > 3 else (cfg.num_prompts or c * cfg.rounds)
            cases.append((i, o, c, n))
        return cases
    ins = [parse_len(x) for x in cfg.input_lens.split(",") if x.strip()]
    outs = [parse_len(x) for x in cfg.output_lens.split(",") if x.strip()]
    if len(outs) == 1:
        outs = outs * len(ins)
    if len(outs) != len(ins):
        sys.exit("[error] --input-lens 与 --output-lens 数量不一致")
    concs = _conc_list(cfg)
    for i, o in zip(ins, outs):
        for c in concs:
            cases.append((i, o, c, cfg.num_prompts or c * cfg.rounds))
    return cases


def cases_from_spec(cfg, spec):
    """配置文件里的 cases：支持 "10k:50k:4[:n]" 字符串，或
       {"in":"10k","out":"50k","concurrency":[1,4,8],"rounds":1,"n":4} 对象。"""
    if spec is None:
        return cases_from_cli(cfg)
    out = []
    for item in spec:
        if isinstance(item, str):
            parts = item.split(":")
            i, o = parse_len(parts[0]), parse_len(parts[1])
            concs = [int(parts[2])] if len(parts) > 2 else _conc_list(cfg)
            n_fixed = int(parts[3]) if len(parts) > 3 else None
            rounds = cfg.rounds
        else:
            i, o = parse_len(item["in"]), parse_len(item["out"])
            cc = item.get("concurrency", cfg.concurrency)
            concs = [int(x) for x in (cc if isinstance(cc, list) else str(cc).split(","))]
            n_fixed = item.get("n")
            rounds = int(item.get("rounds", cfg.rounds))
        for c in concs:
            out.append((i, o, c, int(n_fixed) if n_fixed else c * rounds))
    return out


_CFG_ONLY_KEYS = {"cases", "name", "targets", "defaults", "_comment"}


def load_config(path, base_cfg):
    with open(path, encoding="utf-8") as f:
        obj = json.load(f)
    defaults = obj.get("defaults", {})
    targets = []
    for t in obj.get("targets", []):
        tc = copy.copy(base_cfg)
        for k, v in list(defaults.items()) + list(t.items()):
            if k in _CFG_ONLY_KEYS:
                continue
            setattr(tc, k.replace("-", "_"), v)
        if isinstance(getattr(tc, "extra_body", None), str) and tc.extra_body:
            tc.extra_body = json.loads(tc.extra_body)
        tc.name = t.get("name") or t.get("model_label") or t.get("base_url", "")
        tc._case_spec = t.get("cases", defaults.get("cases"))
        targets.append(tc)
    if not targets:
        sys.exit(f"[error] {path} 里没有 targets")
    return targets


CONFIG_TEMPLATE = {
    "_comment": [
        "多端口/多模型压测配置。defaults 里的键就是命令行参数名（横线换成下划线都行）。",
        "target 里的同名键会覆盖 defaults。cases 支持两种写法：",
        "  字符串  \"输入:输出:并发[:条数]\"，例 \"10k:50k:4\"",
        "  对象    {\"in\":\"10k\",\"out\":\"50k\",\"concurrency\":[1,4,8],\"rounds\":1}",
        "长度单位 K=1024；输入+输出 超过该 target 的 max_model_len 会按 on_overflow 处理。",
        "条数默认 = 并发 × rounds；长序列用例建议 rounds=1。",
    ],
    "defaults": {
        "api": "chat",
        "rounds": 4,
        "concurrency": [1, 4, 8, 16],
        "warmup": 2,
        "sleep_between": 5,
        "progress_interval": 30,
        "stall_timeout": 600,
        "case_timeout": 0,
        "on_overflow": "skip",
        "hardware": "",
        "version": "",
        "dtype": "",
        "owner": "vLLM",
    },
    "targets": [
        {
            "name": "Qwen3.6-27B @8000",
            "base_url": "http://127.0.0.1:8000/v1",
            "model": "",
            "model_label": "Qwen3.6-27B",
            "max_model_len": 131072,
            "tp": 1,
            "chips": 1,
            "cases": [
                {"in": 256, "out": 256},
                {"in": 512, "out": 256},
                {"in": "1k", "out": "1k"},
                {"in": "2k", "out": "1k"},
                {"in": "4k", "out": "1k"},
                {"in": "10k", "out": "50k", "concurrency": [1, 4, 8], "rounds": 1},
                {"in": "32k", "out": "32k", "concurrency": [1, 4], "rounds": 1},
                {"in": "64k", "out": "16k", "concurrency": [1, 2], "rounds": 1},
            ],
        },
        {
            "name": "Qwen3-14B @8001",
            "base_url": "http://127.0.0.1:8001/v1",
            "model": "",
            "model_label": "Qwen3-14B",
            "max_model_len": 40960,
            "tp": 1,
            "chips": 1,
            "cases": [
                {"in": 256, "out": 256},
                {"in": 512, "out": 256},
                {"in": "1k", "out": "1k"},
                {"in": "2k", "out": "1k"},
                {"in": "4k", "out": "1k"},
                {"in": "10k", "out": "25k", "concurrency": [1, 4, 8], "rounds": 1},
                {"in": "20k", "out": "16k", "concurrency": [1, 2], "rounds": 1},
                {"in": "32k", "out": "7k", "concurrency": [1, 2], "rounds": 1},
            ],
        },
    ],
}


# --------------------------------------------------------------------------- #
# CLI
# --------------------------------------------------------------------------- #
def parse_args(argv=None):
    p = argparse.ArgumentParser(description="OpenAI 兼容接口推理性能压测（支持多端口多模型）",
                                formatter_class=argparse.ArgumentDefaultsHelpFormatter)
    p.add_argument("--config", default="", help="多端口/多模型配置文件(JSON)")
    p.add_argument("--gen-config", default="", help="生成配置模板到指定路径后退出")

    p.add_argument("--base-url", default="", help="单端口模式用；如 http://127.0.0.1:8000/v1")
    p.add_argument("--api-key", default=os.getenv("OPENAI_API_KEY", ""))
    p.add_argument("--model", default="", help="留空则自动从 /models 取第一个")
    p.add_argument("--model-label", default="", help="表格里显示的模型名")
    p.add_argument("--max-model-len", type=int, default=0,
                   help="该服务的上下文上限，0=自动从 /models 读")
    p.add_argument("--api", choices=["chat", "completions"], default="chat",
                   help="completions 无对话模板，测纯长度更准")
    p.add_argument("--on-overflow", choices=["skip", "clamp", "run"], default="skip",
                   help="输入+输出超过上下文时：跳过 / 裁剪输出长度 / 照跑")

    p.add_argument("--timeout", type=float, default=0.0,
                   help="单请求总超时(s)，0=不限（长输出建议 0，用 --stall-timeout 兜底）")
    p.add_argument("--stall-timeout", type=float, default=300.0,
                   help="多久收不到新 token 判定卡死(s)")
    p.add_argument("--case-timeout", type=float, default=0.0,
                   help="单组用例墙钟上限(s)，超时会中止在途请求，0=不限")

    p.add_argument("--input-lens", default="256,512,1024,2048,4k")
    p.add_argument("--output-lens", default="256,256,1024,1024,1k")
    p.add_argument("--concurrency", default="1,4,8,16")
    p.add_argument("--rounds", type=int, default=4, help="数据集条数 = 并发 × rounds（长序列建议 1）")
    p.add_argument("--num-prompts", type=int, default=0, help=">0 则固定条数，忽略 rounds")
    p.add_argument("--request-rate", type=float, default=0.0, help="req/s 泊松到达；0=不限速")
    p.add_argument("--cases", default="", help="精确指定用例 in:out:conc[:n]，例 10k:50k:1")
    p.add_argument("--warmup", type=int, default=2, help="每组用例前的预热请求数(不计入统计)")
    p.add_argument("--sleep-between", type=float, default=3.0, help="用例之间的冷却秒数")
    p.add_argument("--progress-interval", type=float, default=30.0, help="进度心跳间隔(s)，0=关")

    p.add_argument("--temperature", type=float, default=0.0)
    p.add_argument("--top-p", type=float, default=1.0)
    p.add_argument("--ignore-eos", action="store_true", default=True,
                   help="强制定长输出(vLLM 支持)，长输出必须开，否则模型早停")
    p.add_argument("--no-ignore-eos", dest="ignore_eos", action="store_false")
    p.add_argument("--extra-body", default="", help='额外 JSON，如 \'{"repetition_penalty":1.0}\'')
    p.add_argument("--no-calibrate", action="store_true", help="跳过输入长度在线校准")

    p.add_argument("--hardware", default="")
    p.add_argument("--version", default="")
    p.add_argument("--dtype", default="")
    p.add_argument("--tp", type=int, default=1, help="张量并行卡数，用于算单卡吞吐")
    p.add_argument("--chips", type=int, default=1, help="芯片数，用于算单芯吞吐")
    p.add_argument("--owner", default="")

    p.add_argument("-o", "--output", default="bench_result.csv")
    p.add_argument("--seed", type=int, default=1234)
    p.add_argument("--force-threads", action="store_true")
    cfg = p.parse_args(argv)
    if cfg.extra_body and isinstance(cfg.extra_body, str):
        cfg.extra_body = json.loads(cfg.extra_body)
    return cfg


def run_target(cfg, cases, rows, out_path) -> None:
    pb = calibrate(cfg, PromptBuilder(seed=cfg.seed))
    for idx, (in_len, out_len, conc, n) in enumerate(cases, 1):
        print(f"  [{idx}/{len(cases)}] 输入={in_len} 输出={out_len} 并发={conc} 条数={n} ...",
              flush=True)
        prompts = make_prompts(cfg, pb, in_len, n)
        if cfg.warmup > 0:
            warm = [pb.make(in_len) for _ in range(cfg.warmup)]   # 预热用独立 prompt
            run_case(cfg, warm, min(out_len, 16), 1, 0.0)

        with Heartbeat(cfg.progress_interval, n):
            results, dur, submitted = run_case(cfg, prompts, out_len, conc, cfg.request_rate)

        row = summarize(cfg, results, dur, in_len, out_len, conc, submitted, cfg.request_rate)
        if row.get("错误"):
            print(f"      !! 全部失败: {list({r.error for r in results if not r.ok})[:2]}")
            continue
        rows.append(row)
        print(f"      用时 {fmt_secs(dur)}  TTFT={row['TTFT(ms,Prefill时延)']}ms  "
              f"TPOT={row['TPOT(ms,Decode时延)']}ms  输出吞吐={row['输出吞吐']} tok/s  "
              f"QPS={row['QPS']}  失败={row['失败数']}")
        if row.get("中止数(超时截断)"):
            print(f"      [warn] {row['中止数(超时截断)']} 条请求被 case_timeout 截断，"
                  f"吞吐/TTFT/TPOT 仍有效，但目标输出长度未跑满")
        if row.get("TPOT衰减倍数", 0) and row["TPOT衰减倍数"] > 1.2:
            print(f"      [note] decode 衰减：后段 TPOT 是前段的 {row['TPOT衰减倍数']}× "
                  f"({row[f'TPOT_前{ITL_KEEP}(ms)']} → {row[f'TPOT_后{ITL_KEEP}(ms)']} ms)")

        write_csv(out_path, rows)            # 边跑边落盘，中途中断也不丢数据
        if idx < len(cases) and cfg.sleep_between > 0:
            time.sleep(cfg.sleep_between)


def main(argv=None):
    cfg = parse_args(argv)

    if cfg.gen_config:
        with open(cfg.gen_config, "w", encoding="utf-8") as f:
            json.dump(CONFIG_TEMPLATE, f, ensure_ascii=False, indent=2)
        print(f"[done] 配置模板已写入 {os.path.abspath(cfg.gen_config)}\n"
              f"       改完后运行: python {os.path.basename(sys.argv[0])} "
              f"--config {cfg.gen_config} -o result.csv")
        return 0

    if cfg.config:
        targets = load_config(cfg.config, cfg)
    else:
        if not cfg.base_url:
            sys.exit("[error] 需要 --base-url（单端口）或 --config（多端口）")
        cfg.name = cfg.model_label or cfg.base_url
        cfg._case_spec = None
        targets = [cfg]

    print(f"[info] 传输模式: {'asyncio/aiohttp' if HAS_AIOHTTP and not cfg.force_threads else '线程'}")
    print(f"[info] 共 {len(targets)} 个服务端\n")

    rows = []
    for ti, tcfg in enumerate(targets, 1):
        print(f"=== target {ti}/{len(targets)}: {tcfg.name}  ({tcfg.base_url}) ===")
        try:
            resolve_target(tcfg)
        except Exception as e:  # noqa: BLE001
            print(f"  [error] 连不上，跳过该 target: {e}", file=sys.stderr)
            continue
        cases = filter_cases(tcfg, cases_from_spec(tcfg, getattr(tcfg, "_case_spec", None)))
        if not cases:
            print("  [warn] 没有可执行的用例，跳过")
            continue
        if not HAS_AIOHTTP and max(o for _, o, _, _ in cases) > 8192:
            print("  [warn] 超长输出但未装 aiohttp：线程模式 GIL 会影响 TTFT/TPOT 精度，"
                  "建议 pip install aiohttp", file=sys.stderr)
        print(f"  [info] {len(cases)} 组用例")
        run_target(tcfg, cases, rows, cfg.output)
        print()

    print()
    print_table(rows, PRINT_COLS)
    print(f"\n[done] 结果已写入 {os.path.abspath(cfg.output)}"
          f"（含 {len(EXTRA_COLS)} 个附加诊断列）")
    return 0


if __name__ == "__main__":
    sys.exit(main())