"""Measure real vLLM streamed token delivery, with exact returned token IDs."""
import argparse
import json
import time
import urllib.request
from pathlib import Path

PROMPTS = [
    'Explain how a hash table works in Python, including collisions, resizing, and time complexity. Write a detailed tutorial of at least 800 words.',
    'Write a complete Python implementation of an LRU cache with type annotations, a detailed explanation, and examples. Include at least 800 words of explanation after the code.',
]
for repetitions in (90, 210, 410):
    PROMPTS.append(
        'This is background information for a technical documentation exercise. The system stores records and processes requests.\n' * repetitions
        + 'Now write a detailed tutorial of at least 800 words explaining database indexes and their tradeoffs.')


def request(base_url, model, prompt, maximum):
    body = {'model':model, 'messages':[{'role':'user','content':prompt}],
        'temperature':0, 'max_tokens':maximum, 'chat_template_kwargs':{'enable_thinking':False},
        'stream':True, 'stream_options':{'include_usage':True}, 'return_token_ids':True}
    req = urllib.request.Request(base_url.rstrip('/')+'/chat/completions',
        data=json.dumps(body).encode(), headers={'Content-Type':'application/json'})
    start = time.perf_counter()
    events, content, reasoning = [], [], []
    usage = None
    with urllib.request.urlopen(req, timeout=1800) as response:
        for line in response:
            if not line.startswith(b'data: '):
                continue
            payload = line[6:].strip()
            if payload == b'[DONE]':
                break
            chunk = json.loads(payload)
            if 'error' in chunk:
                raise RuntimeError(chunk['error'])
            if chunk.get('usage') is not None:
                usage = chunk['usage']
            for choice in chunk.get('choices',[]):
                tokens = choice.get('token_ids') or []
                if tokens:
                    events.append({'elapsed_seconds':time.perf_counter()-start,'token_ids':tokens})
                delta = choice.get('delta',{})
                content.append(delta.get('content') or '')
                reasoning.append(delta.get('reasoning') or delta.get('reasoning_content') or '')
    wall = time.perf_counter()-start
    count = sum(len(event['token_ids']) for event in events)
    assert usage and count == usage['completion_tokens'], (count,usage)
    assert len(events) >= 2, 'At least two token-bearing SSE events are required.'
    first_count = len(events[0]['token_ids'])
    seconds = events[-1]['elapsed_seconds']-events[0]['elapsed_seconds']
    return {'request':body,'usage':usage,'content':''.join(content),'reasoning':''.join(reasoning),
        'wall_seconds':wall,'time_to_first_token_seconds':events[0]['elapsed_seconds'],
        'decode_tokens':count-first_count,'decode_seconds':seconds,
        'decode_tokens_per_second':(count-first_count)/seconds,
        'end_to_end_tokens_per_second':count/wall,'first_event_token_count':first_count,
        'token_events':events}


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument('--base-url',default='http://127.0.0.1:8000/v1')
    parser.add_argument('--model',default='Naive-N0.5-Flash-NVFP4')
    parser.add_argument('--tokens',type=int,default=256)
    parser.add_argument('--cases',type=int,default=5)
    parser.add_argument('--output',type=Path,required=True)
    args = parser.parse_args()
    request(args.base_url,args.model,'Explain why the sky is blue in several sentences.',32)
    results=[]
    for prompt in PROMPTS[:args.cases]:
        result=request(args.base_url,args.model,prompt,args.tokens)
        assert result['usage']['completion_tokens'] >= min(128,args.tokens)
        results.append(result)
        print(json.dumps({key:result[key] for key in ('usage','time_to_first_token_seconds',
            'decode_tokens_per_second','end_to_end_tokens_per_second')}),flush=True)
        report={'engine':'vLLM','measurement':'Client wall time between token-bearing SSE events, using exact token IDs; first event excluded from decode rate',
            'single_active_request':True,'cases':results,'complete':len(results)==args.cases,
            'target_met':len(results)==args.cases and all(r['decode_tokens_per_second']>=50 for r in results)}
        args.output.parent.mkdir(parents=True,exist_ok=True)
        args.output.write_text(json.dumps(report,indent=2)+'\n')


if __name__ == '__main__':
    main()