"""Synthetic streaming benchmark; run inside serving container, writes JSONL to stdout.

No external datasets or user prompts. Usage: python3 benchmark.py --label eager
TTFT includes reasoning tokens; first_content_s separately measures visible answer.
Decode tok/s is an approximation from total completion tokens and streaming time.
"""
import argparse
import concurrent.futures
import hashlib
import json
import random
import re
import threading
import time
import urllib.request
from pathlib import Path

parser = argparse.ArgumentParser()
parser.add_argument('--label', required=True)
parser.add_argument('--long-max-tokens', type=int, default=1024)
parser.add_argument('--phase', choices=['all', 'short', 'long'], default='all')
args = parser.parse_args()
BASE = 'http://127.0.0.1:8000'
KEY = Path('/run/secrets/qwen_api_key').read_text().strip()
HEADERS = {'Authorization': 'Bearer ' + KEY, 'Content-Type': 'application/json'}

def emit(record):
    print(json.dumps({'label': args.label, **record}, ensure_ascii=False), flush=True)

def metrics():
    request = urllib.request.Request(BASE + '/metrics', headers=HEADERS)
    with urllib.request.urlopen(request, timeout=10) as response:
        text = response.read().decode()
    values = {}
    for line in text.splitlines():
        if line.startswith('#') or 'prefix_cache_' not in line:
            continue
        name = line.split('{')[0].split()[0]
        if name.endswith('_total'):
            values[name] = values.get(name, 0) + float(line.split()[-1])
    return values

def run(name, prompt, expected=None, barrier=None, max_tokens=512):
    payload = {'model': 'qwen3.8-flash-next', 'messages': [{'role': 'user', 'content': prompt}],
               'temperature': 0, 'seed': 42, 'max_tokens': max_tokens,
               'reasoning_effort': 'low', 'stream': True,
               'stream_options': {'include_usage': True}}
    request = urllib.request.Request(BASE + '/v1/chat/completions',
                                     data=json.dumps(payload).encode(), headers=HEADERS)
    if barrier:
        barrier.wait()
    start = time.perf_counter()
    first = first_content = last = None
    content, reasoning = '', ''
    usage, finish = {}, None
    with urllib.request.urlopen(request, timeout=900) as response:
        for raw in response:
            line = raw.decode().strip()
            if not line.startswith('data: ') or line == 'data: [DONE]':
                continue
            item = json.loads(line[6:])
            if item.get('usage'):
                usage = item['usage']
            for choice in item.get('choices', []):
                delta = choice.get('delta', {})
                answer = delta.get('content') or ''
                thought = delta.get('reasoning') or delta.get('reasoning_content') or ''
                if answer or thought:
                    now = time.perf_counter()
                    first = first if first is not None else now
                    last = now
                if answer:
                    first_content = first_content if first_content is not None else now
                content += answer
                reasoning += thought
                finish = choice.get('finish_reason') or finish
    elapsed = time.perf_counter() - start
    count = usage.get('completion_tokens', 0)
    record = {'test': name, 'max_tokens': max_tokens, 'elapsed_s': round(elapsed, 4),
              'ttft_s': round(first-start, 4) if first else None,
              'first_content_s': round(first_content-start, 4) if first_content else None,
              'decode_tps_approx': round((count-1)/(last-first), 3) if last and first and last>first else None,
              'usage': usage, 'finish_reason': finish, 'content': content,
              'output_sha256': hashlib.sha256((reasoning+'\0'+content).encode()).hexdigest(),
              'prompt_sha256': hashlib.sha256(prompt.encode()).hexdigest(),
              'correct': bool(content.strip()) and (expected is None or expected in content)}
    emit(record)
    if not record['correct'] or finish != 'stop':
        emit({'event':'ANSWER_CHECK_FAILED','test':name})
    return record

emit({'event': 'start', 'phase': args.phase})
if args.phase in ('all', 'short'):
    run('warmup', '计算17乘19，只输出结果。', '323')
    prompt = '用中文写一段约150字的说明，解释数据库索引为什么能加快查询，以及它对写入有什么影响。'
    for i in range(3):
        run('short_'+str(i), prompt)
    run('code', '写一个Python函数is_even(n)，判断整数是否为偶数，只输出代码。', '%')
    for concurrency in [2, 4]:
        barrier = threading.Barrier(concurrency)
        start = time.perf_counter()
        with concurrent.futures.ThreadPoolExecutor(max_workers=concurrency) as pool:
            futures = [pool.submit(run, f'concurrent_{concurrency}_{i}',
                        f'场景编号{i}。'+prompt, None, barrier) for i in range(concurrency)]
            results = [f.result() for f in futures]
        seconds = time.perf_counter()-start
        emit({'test': 'concurrency_summary', 'concurrency': concurrency,
              'elapsed_s': round(seconds, 3),
              'aggregate_completion_tps': round(sum(r['usage']['completion_tokens'] for r in results)/seconds, 3)})
if args.phase in ('all', 'long'):
    from transformers import AutoTokenizer
    snapshot = '/root/.cache/huggingface/hub/models--nvidia--Qwen3.8-Flash-Next-NVFP4/snapshots/fc694b54fb0174e0913e6adf86691ef85a4ead47'
    tokenizer = AutoTokenizer.from_pretrained(snapshot, local_files_only=True)
    rng = random.Random(42)
    filler = '\n'.join(f'Record {i}: item={rng.randrange(100000,999999)}; status=archived; category={i%17}.' for i in range(6000))
    ids = tokenizer.encode(filler, add_special_tokens=False)
    for size in [8192, 32768]:
        text = tokenizer.decode(ids[:size-150])
        pos = len(text)//2
        text = text[:pos]+'\n唯一有效的校验口令是：青松739251。\n'+text[pos:]
        prompt = '从下面的档案里找出唯一有效的校验口令。只输出口令，不要解释。\n<档案>\n'+text+'\n</档案>\n请输出校验口令。'
        before = metrics()
        a = run(f'long_{size}_cold', prompt, '青松739251', max_tokens=args.long_max_tokens)
        b = run(f'long_{size}_repeat', prompt, '青松739251', max_tokens=args.long_max_tokens)
        after = metrics()
        emit({'test':f'prefix_{size}', 'same_output':a['output_sha256']==b['output_sha256'],
              'metrics_delta':{key:after[key]-before.get(key,0) for key in after}})
emit({'event':'BENCHMARK_COMPLETED'})