"""5.7 Local model speed benchmark: single-stream speed, plus per-stream and total throughput under concurrency.

Works against any OpenAI-compatible local server: llama.cpp's llama-server, MLX's mlx_lm.server, or vLLM.
Every request asks the model to count upward from 1 and caps output length with max_tokens so each stream does equal work.

Usage (start the server first as shown in README.md):
    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):
    """Send one streaming request. Return (time to first token in seconds, total seconds, output tokens)."""
    body = {
        'model': model,
        'messages': [{'role': 'system', 'content': f'[{random.random()}] ' + FILLER},   # ~800-token input; random prefix avoids cache hits
                     {'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     # Fall back to chunk count if the server omits usage
    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='Model ID on the server; defaults to the first entry in /v1/models')
    ap.add_argument('--label', default='', help='Label saved in the output for comparison, e.g. "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 reports a file path; keep only the filename

    print(f'{args.label or shown} ({url}), ~800 input tokens and {args.max_tokens} output tokens per request, {datetime.now():%Y-%m-%d %H:%M}')
    await level(url, model, 1, 32)                                   # Warm-up run, excluded from results
    rows = []
    print('\n| Concurrent streams | Time to first token (median) | Per-stream decode speed | Total throughput |')
    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} s | {r["per_stream_tok_s"]:.0f} tok/s | {r["total_tok_s"]:.0f} tok/s |')
    if any(min(r['tokens']) < args.max_tokens * 0.9 for r in rows):
        print('\nNote: some requests stopped before reaching max_tokens; speeds may skew low or not be comparable.')
    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'\nSaved to {path}')


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