"""Run every README demo against a local Ollama model and save the raw responses.

Usage: python scripts/run_demos.py [model]   (default: sagui-4b)
Writes examples/transcripts.json with the unedited output and Ollama's timing counters.
Standard library only.
"""
import base64
import json
import sys
import time
import urllib.request
from pathlib import Path

ROOT = Path(__file__).resolve().parent.parent
MODEL = sys.argv[1] if len(sys.argv) > 1 else "sagui-4b"
URL = "http://localhost:11434/api/chat"
WEATHER_TOOL = {
    "type": "function",
    "function": {
        "name": "get_weather",
        "description": "Get the current weather for a city",
        "parameters": {"type": "object", "properties": {"city": {"type": "string"}}, "required": ["city"]},
    },
}


def image(name):
    return base64.b64encode((ROOT / "examples" / name).read_bytes()).decode()


def chat(**body):
    req = urllib.request.Request(URL, json.dumps({"model": MODEL, "stream": False, **body}).encode(),
                                 {"Content-Type": "application/json"})
    with urllib.request.urlopen(req, timeout=900) as resp:
        return json.load(resp)


def record(results, name, **body):
    started = time.time()
    data = chat(**body)
    msg = data["message"]
    results.append({
        "demo": name,
        "think": body["think"],
        "prompt": body["messages"][-1].get("content"),
        "thinking": msg.get("thinking") or "",
        "content": msg.get("content") or "",
        "tool_calls": msg.get("tool_calls"),
        "tokens": data.get("eval_count"),
        "tokens_per_second": round(data.get("eval_count", 0) / (data.get("eval_duration", 1) / 1e9), 1),
        "seconds": round(time.time() - started, 1),
    })
    print(f"{name}: {results[-1]['tokens']} tokens, {results[-1]['tokens_per_second']} tok/s")
    return msg


def main():
    results = []
    record(results, "vision-fix-ts-error", think=False, messages=[{
        "role": "user",
        "content": "This is a screenshot from my editor. Explain the error in one sentence, then give the corrected line 4.",
        "images": [image("ts-error.png")],
    }])
    record(results, "vision-ui-to-code", think=True, messages=[{
        "role": "user",
        "content": "Recreate this screen as a single React + TypeScript component with inline styles. Code only.",
        "images": [image("signup-ui.png")],
    }])
    record(results, "code-typescript", think=False, messages=[{
        "role": "user",
        "content": "Write a TypeScript function groupBy<T, K extends PropertyKey>(items: T[], key: (item: T) => K): Record<K, T[]>. Code only.",
    }])
    record(results, "thinking-logic", think=True, messages=[{
        "role": "user",
        "content": "A train leaves at 14:35 and arrives at 17:10. How long is the trip in minutes?",
    }])

    ask = {"role": "user", "content": "What's the weather in Lisbon? Use the tool."}
    call = record(results, "tool-call", think=False, tools=[WEATHER_TOOL], messages=[ask])
    if call.get("tool_calls"):
        record(results, "tool-result", think=False, tools=[WEATHER_TOOL], messages=[
            ask, call, {"role": "tool", "tool_name": "get_weather", "content": '{"temp_c": 21, "sky": "sunny"}'},
        ])

    out = ROOT / "examples" / "transcripts.json"
    out.write_text(json.dumps({"model": MODEL, "results": results}, ensure_ascii=False, indent=1), encoding="utf-8")


if __name__ == "__main__":
    main()