"""5.7 本地模型测速：单流速度，以及多路并发时每一路和总共的速度。

对任何兼容 OpenAI 接口的本地服务都能用：llama.cpp 的 llama-server、MLX 的 mlx_lm.server、vLLM。
每个请求都让模型从 1 往上数数，用 max_tokens 卡住输出长度，这样每路的输出 token 数基本一样。

用法（先按 README 起好服务）：
    python bench.py --url http://127.0.0.1:8080 --label "llama.cpp Q4_K_M"
    python bench.py --url http://127.0.0.1:8000 --label "vLLM bf16" --concurrency 1 4 8 16 32
"""
import argparse
import asyncio
import json
import random
import statistics
import time
from datetime import datetime
from pathlib import Path

import httpx

FILLER = ('The quick brown fox jumps over the lazy dog while the model server keeps answering requests. ' * 40)


async def one_request(client, url, model, max_tokens):
    """发一个流式请求。返回 (首 token 秒数, 总秒数, 输出 token 数)。"""
    body = {
        'model': model,
        'messages': [{'role': 'system', 'content': f'[{random.random()}] ' + FILLER},   # 约 800 token 的输入；随机前缀避免命中缓存
                     {'role': 'user', 'content': 'Count from 1 upward, separated by commas. Do not stop.'}],
        'max_tokens': max_tokens, 'temperature': 0, 'stream': True,
        'stream_options': {'include_usage': True},
        'chat_template_kwargs': {'enable_thinking': False},
    }
    t0 = time.perf_counter()
    first, chunks, usage = None, 0, None
    async with client.stream('POST', url + '/v1/chat/completions', json=body) as r:
        r.raise_for_status()
        async for line in r.aiter_lines():
            if not line.startswith('data:') or line.endswith('[DONE]'):
                continue
            d = json.loads(line[5:])
            if d.get('usage'):
                usage = d['usage']
            delta = (d.get('choices') or [{}])[0].get('delta') or {}
            if delta.get('content') or delta.get('reasoning_content'):
                chunks += 1
                if first is None:
                    first = time.perf_counter() - t0
    total = time.perf_counter() - t0
    n = (usage or {}).get('completion_tokens') or chunks     # 服务没报 usage 时，用收到的块数近似 token 数
    return first, total, n


async def level(url, model, n, max_tokens):
    async with httpx.AsyncClient(timeout=600) as client:
        t0 = time.perf_counter()
        rs = await asyncio.gather(*[one_request(client, url, model, max_tokens) for _ in range(n)])
        wall = time.perf_counter() - t0
    per_stream = [(k - 1) / (tot - f) for f, tot, k in rs if f is not None and tot > f and k > 1]
    return dict(concurrency=n, ttft_median=statistics.median(f for f, _, _ in rs),
                per_stream_tok_s=statistics.median(per_stream), total_tok_s=sum(k for _, _, k in rs) / wall,
                tokens=[k for _, _, k in rs])


async def main():
    ap = argparse.ArgumentParser()
    ap.add_argument('--url', default='http://127.0.0.1:8080')
    ap.add_argument('--model', help='服务里的模型名；不给就取 /v1/models 里的第一个')
    ap.add_argument('--label', default='', help='写进结果里，方便对比，例如 "MLX 4bit"')
    ap.add_argument('--max-tokens', type=int, default=256)
    ap.add_argument('--concurrency', type=int, nargs='+', default=[1, 4, 8, 16])
    ap.add_argument('--out', default='results')
    args = ap.parse_args()
    url = args.url.rstrip('/')
    model = args.model or httpx.get(url + '/v1/models').json()['data'][0]['id']
    shown = model.rstrip('/').split('/')[-1] if model.startswith('/') else model   # llama-server 报的是文件路径，只留文件名

    print(f'{args.label or shown}（{url}），每个请求约 800 token 输入、{args.max_tokens} token 输出，{datetime.now():%Y-%m-%d %H:%M}')
    await level(url, model, 1, 32)                                   # 预热一次，不计入结果
    rows = []
    print('\n| 并发路数 | 首 token 延迟（中位数） | 每一路的生成速度 | 所有路加起来 |')
    print('|---|---|---|---|')
    for n in args.concurrency:
        r = await level(url, model, n, args.max_tokens)
        rows.append(r)
        print(f'| {n} | {r["ttft_median"]:.2f} 秒 | {r["per_stream_tok_s"]:.0f} token/秒 | {r["total_tok_s"]:.0f} token/秒 |')
    if any(min(r['tokens']) < args.max_tokens * 0.9 for r in rows):
        print('\n注意：有的请求没生成满 max_tokens 就停了，速度会偏低或不可比。')
    Path(args.out).mkdir(exist_ok=True)
    stem = (args.label or shown).replace('/', '_').replace(' ', '_')
    path = Path(args.out) / f'{stem}_{datetime.now():%Y%m%d_%H%M}.json'
    path.write_text(json.dumps(dict(label=args.label, model=shown, max_tokens=args.max_tokens, rows=rows), indent=1), encoding='utf-8')
    print(f'\n已存到 {path}')


if __name__ == '__main__':
    asyncio.run(main())