#!/usr/bin/env python3
"""Measure client-observed first content, not internal model TTFT. Stdlib only."""
import argparse
import json
import os
import time
import urllib.request


def events(response):
    lines = []
    for raw in response:
        line = raw.decode('utf-8').rstrip('\r\n')
        if not line:
            if lines:
                yield '\n'.join(lines)
                lines = []
        elif line.startswith('data:'):
            value = line[5:]
            lines.append(value[1:] if value.startswith(' ') else value)
    if lines:
        raise ValueError('Truncated SSE event (no terminating blank line)')


def measure(url, model, prompt, max_tokens=64, timeout=120, api_key=''):
    payload = json.dumps({'model': model, 'messages': [{'role': 'user', 'content': prompt}],
                          'stream': True, 'max_tokens': max_tokens, 'temperature': 0}).encode()
    headers = {'Content-Type': 'application/json', 'Accept': 'text/event-stream'}
    if api_key:
        headers['Authorization'] = 'Bearer ' + api_key
    request = urllib.request.Request(url, data=payload, headers=headers)
    start = time.perf_counter()
    first = None
    content_chunks = 0
    chars = 0
    finished = False
    # No implicit environment proxy; this diagnostic targets an explicitly chosen server.
    opener = urllib.request.build_opener(urllib.request.ProxyHandler({}))
    with opener.open(request, timeout=timeout) as response:
        if 'text/event-stream' not in response.headers.get('Content-Type', '').lower():
            raise ValueError('Expected text/event-stream, not a buffered JSON/HTML response')
        for event in events(response):
            if event == '[DONE]':
                finished = True
                break
            data = json.loads(event)
            if 'error' in data:
                raise ValueError('Server returned an error event; inspect sanitized server logs')
            for choice in data.get('choices', []):
                if choice.get('index', 0) != 0:
                    continue
                content = choice.get('delta', {}).get('content')
                if isinstance(content, str) and content:
                    if first is None:
                        first = time.perf_counter() - start
                    content_chunks += 1
                    chars += len(content)
    elapsed = time.perf_counter() - start
    if not finished:
        raise ValueError('Stream ended before [DONE]; result not accepted as complete')
    return {'first_content_ms': round(first * 1000, 2) if first is not None else None,
            'total_ms': round(elapsed * 1000, 2), 'content_chunks': content_chunks,
            'output_characters': chars, 'complete': True}


def main():
    p = argparse.ArgumentParser(description=__doc__)
    p.add_argument('--url', default='http://127.0.0.1:8080/v1/chat/completions')
    p.add_argument('--model', required=True)
    p.add_argument('--prompt', default='Explain a cache hit in one short sentence.')
    p.add_argument('--max-tokens', type=int, default=64)
    p.add_argument('--timeout', type=float, default=120)
    p.add_argument('--runs', type=int, default=3)
    args = p.parse_args()
    if not 1 <= args.runs <= 20 or args.max_tokens < 1 or args.timeout <= 0:
        p.error('Use 1–20 runs, positive token limit and timeout')
    for i in range(args.runs):
        result = measure(args.url, args.model, args.prompt, args.max_tokens, args.timeout,
                         os.environ.get('LLAMA_API_KEY', ''))
        print(json.dumps({'run': i + 1, **result}), flush=True)


if __name__ == '__main__':
    main()