#!/usr/bin/env python3
"""Measure a streamed reply from the local Qwen server, without third-party packages."""
import argparse
import json
import time
import urllib.error
import urllib.request


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("prompt", help="The exact prompt to send")
    parser.add_argument("--max-tokens", type=int, default=512)
    args = parser.parse_args()
    if args.max_tokens < 1:
        parser.error("--max-tokens must be positive")

    payload = {
        "model": "Qwen3.5-2B",
        "messages": [
            {"role": "system", "content": "Answer in the user's language. Keep answers concise and be honest about uncertainty."},
            {"role": "user", "content": args.prompt},
        ],
        "temperature": 0.7,
        "max_tokens": args.max_tokens,
        "tools": [],
        "tool_choice": "none",
        "stream": True,
        "stream_options": {"include_usage": True},
        "timings_per_token": True,
    }
    request = urllib.request.Request(
        "http://127.0.0.1:8092/v1/chat/completions",
        data=json.dumps(payload, ensure_ascii=False).encode("utf-8"),
        headers={"Content-Type": "application/json", "Accept": "text/event-stream"},
    )
    http = urllib.request.build_opener(urllib.request.ProxyHandler({}))
    first_text = None
    last_text = None
    usage = None
    timings = None
    finish_reason = None
    saw_done = False
    answer = []
    started = time.perf_counter()
    try:
        with http.open(request, timeout=300) as response:
            for raw_line in response:
                line = raw_line.decode("utf-8").strip()
                if not line.startswith("data:"):
                    continue
                data = line[5:].strip()
                if data == "[DONE]":
                    saw_done = True
                    break
                event = json.loads(data)
                if event.get("error"):
                    raise RuntimeError(event["error"])
                if event.get("usage"):
                    usage = event["usage"]
                if event.get("timings"):
                    timings = event["timings"]
                for choice in event.get("choices", []):
                    content = choice.get("delta", {}).get("content")
                    if content:
                        now = time.perf_counter() - started
                        if first_text is None:
                            first_text = now
                        last_text = now
                        answer.append(content)
                        print(content, end="", flush=True)
                    if choice.get("finish_reason"):
                        finish_reason = choice["finish_reason"]
    except urllib.error.HTTPError as exc:
        raise SystemExit(f"Local server returned HTTP {exc.code}: {exc.read().decode('utf-8', 'replace')}") from exc
    except urllib.error.URLError as exc:
        raise SystemExit(f"Cannot reach the local server: {exc.reason}") from exc
    total = time.perf_counter() - started
    if not saw_done or finish_reason is None:
        raise SystemExit("Incomplete stream: missing completion marker or finish reason; timing sample is incomplete")
    if not answer:
        raise SystemExit("The server returned no answer text")
    summary = {
        "first_answer_seconds": round(first_text, 3),
        "last_answer_seconds": round(last_text, 3),
        "total_stream_seconds": round(total, 3),
        "finish_reason": finish_reason,
        "usage": usage,
        "server_timings": timings,
    }
    print("\n\n" + json.dumps(summary, ensure_ascii=False, indent=2))


if __name__ == "__main__":
    main()