#!/usr/bin/env python3
"""Three replay demos for a running OpenAI-compatible local model (stdlib only).

These are publication examples, not new Thor benchmark measurements.
Run on the Thor or through an SSH tunnel. Use --dry-run to inspect requests.
"""
import argparse
import json
import os
import sys
import time
import urllib.error
import urllib.request


def payload(args):
    data = {
        "model": args.model,
        "temperature": 0,
        "max_tokens": args.max_tokens,
        "chat_template_kwargs": {"enable_thinking": args.thinking},
    }
    if args.demo == "chat":
        data.update(messages=[{"role": "user", "content": "用三句话解释 MoE。"}],
                    stream=True, stream_options={"include_usage": True})
    elif args.demo == "tool":
        data.update(
            messages=[{"role": "user", "content": "What's the weather in Paris? Use the tool."}],
            tools=[{"type": "function", "function": {
                "name": "get_weather", "description": "Get weather for a city",
                "parameters": {"type": "object", "properties": {"city": {"type": "string"}},
                               "required": ["city"], "additionalProperties": False},
            }}], tool_choice="auto")
    else:
        data.update(
            messages=[{"role": "user", "content": "Give basic facts about Tokyo as JSON."}],
            response_format={"type": "json_schema", "json_schema": {
                "name": "city", "schema": {
                    "type": "object",
                    "properties": {"city": {"type": "string"}, "country": {"type": "string"},
                                   "population_million": {"type": "number"}},
                    "required": ["city", "country", "population_million"],
                    "additionalProperties": False,
                },
            }})
    return data


def post(args, data):
    headers = {"Content-Type": "application/json"}
    if os.environ.get("MODEL_API_KEY"):
        headers["Authorization"] = "Bearer " + os.environ["MODEL_API_KEY"]
    return urllib.request.urlopen(urllib.request.Request(
        args.base_url.rstrip("/") + "/chat/completions",
        data=json.dumps(data).encode(), headers=headers), timeout=600)


def run(args):
    data = payload(args)
    if args.dry_run:
        print(json.dumps(data, ensure_ascii=False, indent=2))
        return 0
    started = time.perf_counter()
    with post(args, data) as response:
        if args.demo == "chat":
            first = None
            first_answer = None
            usage = None
            finish = None
            for raw in response:
                line = raw.decode().strip()
                if not line.startswith("data:"):
                    continue
                event = line[5:].strip()
                if event == "[DONE]":
                    break
                event = json.loads(event)
                usage = event.get("usage") or usage
                for choice in event.get("choices", []):
                    delta = choice.get("delta") or {}
                    reasoning = delta.get("reasoning") or delta.get("reasoning_content") or ""
                    answer = delta.get("content") or ""
                    if reasoning or answer:
                        first = first if first is not None else time.perf_counter() - started
                    if answer and first_answer is None:
                        first_answer = time.perf_counter() - started
                    if reasoning:
                        print(reasoning, end="", file=sys.stderr, flush=True)
                    if answer:
                        print(answer, end="", flush=True)
                    finish = choice.get("finish_reason") or finish
            print()
            print(json.dumps({"first_output_s": first, "first_answer_s": first_answer,
                              "elapsed_s": time.perf_counter() - started,
                              "finish_reason": finish, "usage": usage}, ensure_ascii=False), file=sys.stderr)
            if finish != "stop":
                print("Output did not finish normally; inspect finish_reason/token budget.", file=sys.stderr)
                return 1
            return 0
        result = json.load(response)
    choice = result["choices"][0]
    msg = choice["message"]
    print(json.dumps(result, ensure_ascii=False, indent=2))
    if args.demo == "tool":
        calls = msg.get("tool_calls") or []
        if not calls:
            raise ValueError("No tool_calls returned")
        for call in calls:
            fn = call["function"]
            params = json.loads(fn["arguments"])
            if fn["name"] != "get_weather" or str(params.get("city", "")).lower() != "paris":
                raise ValueError("Unexpected function or city")
        print("PASS: get_weather(city=Paris) parsed. No weather service was called.")
    else:
        obj = json.loads(msg.get("content") or "")
        expected = {"city", "country", "population_million"}
        valid = (isinstance(obj, dict) and set(obj) == expected
                 and isinstance(obj["city"], str) and isinstance(obj["country"], str)
                 and type(obj["population_million"]) in (int, float))
        if not valid or choice.get("finish_reason") != "stop":
            raise ValueError("Incomplete output or wrong keys/types")
        print("PASS: complete JSON with expected keys and types; factual accuracy not checked.")
    return 0


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("demo", choices=["chat", "tool", "schema"])
    parser.add_argument("--base-url", default="http://127.0.0.1:8000/v1")
    parser.add_argument("--model", default="qwen38-flash-next")
    parser.add_argument("--max-tokens", type=int, default=256)
    parser.add_argument("--thinking", action="store_true")
    parser.add_argument("--dry-run", action="store_true")
    args = parser.parse_args()
    try:
        return run(args)
    except (urllib.error.URLError, ValueError, KeyError, TypeError, IndexError) as exc:
        print(f"FAIL: {exc}", file=sys.stderr)
        return 1


if __name__ == "__main__":
    raise SystemExit(main())