"""5.7 Local model speed benchmark (in-process MLX edition; Apple Silicon Mac only).

Calls mlx-lm directly in Python without an HTTP server: runs single-stream and N-way parallel batches (batch_generate),
printing the same table as bench.py for side-by-side comparison with llama.cpp and vLLM.

Usage:
    python bench_mlx.py --model mlx-community/Qwen3.5-2B-4bit --label "MLX 4bit"
    python bench_mlx.py --model mlx-community/Qwen3.5-2B-8bit --label "MLX 8bit"
"""
import argparse
import json
import time
from datetime import datetime
from pathlib import Path

from mlx_lm import batch_generate, load

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


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument('--model', default='mlx-community/Qwen3.5-2B-4bit')
    ap.add_argument('--label', default='')
    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()

    model, tok = load(args.model)
    msgs = [{'role': 'system', 'content': FILLER},
            {'role': 'user', 'content': 'Count from 1 upward, separated by commas. Do not stop.'}]
    ids = tok.apply_chat_template(msgs, add_generation_prompt=True, enable_thinking=False)
    print(f'{args.label or args.model} (in-process MLX), {len(ids)} input tokens and {args.max_tokens} output tokens per request, {datetime.now():%Y-%m-%d %H:%M}')
    batch_generate(model, tok, [ids], max_tokens=16)                     # Warm-up

    rows = []
    print('\n| Concurrent streams | Time to first token (batch prefill) | Per-stream decode speed | Total throughput |')
    print('|---|---|---|---|')
    for n in args.concurrency:
        t0 = time.perf_counter()
        r = batch_generate(model, tok, [ids] * n, max_tokens=args.max_tokens, return_token_ids=True)
        wall = time.perf_counter() - t0
        tokens = [len(t) for t in r.token_ids]
        prefill = r.stats.prompt_time
        per_stream = (sum(tokens) / n) / max(wall - prefill, 1e-9)
        rows.append(dict(concurrency=n, ttft=prefill, per_stream_tok_s=per_stream, total_tok_s=sum(tokens) / wall,
                         tokens=tokens, peak_memory_gb=r.stats.peak_memory))
        print(f'| {n} | {prefill:.2f} s | {per_stream:.0f} tok/s | {sum(tokens) / wall:.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 args.model).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=args.model, max_tokens=args.max_tokens, rows=rows), indent=1), encoding='utf-8')
    print(f'\nSaved to {path}')


if __name__ == '__main__':
    main()