import json, time, sys, requests, jsonschema
from openai import OpenAI
M = "qwen2.5-1.5b-local"; B = "http://127.0.0.1:11434"
out = {"model": M, "started": time.strftime("%H:%M:%S")}
# C2: 10 known-answer prompts via native API
Q = [("What is 17*23? Reply with only the number.", ["391"]),
 ("What is the capital of Australia? One word.", ["canberra"]),
 ("What is the chemical symbol for gold? Reply with only the symbol.", ["au"]),
 ("How many days are in a leap year? Only the number.", ["366"]),
 ("What is 2 to the power of 10? Only the number.", ["1024"]),
 ("Who wrote the novel 'Pride and Prejudice'? Only the name.", ["jane austen"]),
 ("What is the largest planet in our solar system? One word.", ["jupiter"]),
 ("What is the common name of H2O? One word.", ["water"]),
 ("What is the square root of 144? Only the number.", ["12"]),
 ("What is the boiling point of water at sea level in degrees Celsius? Only the number.", ["100"])]
c2 = []
for q, gt in Q:
    t = time.time()
    r = requests.post(B + "/api/generate", json={"model": M, "prompt": q, "stream": False, "options": {"temperature": 0, "seed": 1}}, timeout=300).json()
    ans = r.get("response", "")
    c2.append({"q": q, "gt": gt, "answer": ans.strip(), "correct": any(g in ans.lower() for g in gt),
               "wall_s": round(time.time() - t, 2), "prompt_tokens": r.get("prompt_eval_count"), "eval_tokens": r.get("eval_count"),
               "tok_per_s": round(r["eval_count"] / (r["eval_duration"] / 1e9), 1) if r.get("eval_duration") else None,
               "load_s": round(r.get("load_duration", 0) / 1e9, 2)})
    print("C2", c2[-1]["correct"], c2[-1]["answer"][:60], c2[-1]["tok_per_s"], flush=True)
out["c2"] = c2
# speed: one longer generation
r = requests.post(B + "/api/generate", json={"model": M, "prompt": "Write a 200-word explanation of how a refrigerator works.", "stream": False, "options": {"temperature": 0, "seed": 1, "num_predict": 300}}, timeout=600).json()
out["speed_long"] = {"eval_tokens": r["eval_count"], "tok_per_s": round(r["eval_count"] / (r["eval_duration"] / 1e9), 1),
                     "prompt_tok_per_s": round(r["prompt_eval_count"] / (r["prompt_eval_duration"] / 1e9), 1), "total_s": round(r["total_duration"] / 1e9, 2)}
print("speed", out["speed_long"], flush=True)
# C3: OpenAI client
c3 = {}
cl = OpenAI(base_url=B + "/v1/", api_key="ollama")
try:
    x = cl.chat.completions.create(model=M, messages=[{"role": "user", "content": "What is 17*23? Reply with only the number."}], temperature=0)
    c3["a_chat"] = {"ok": "391" in x.choices[0].message.content, "content": x.choices[0].message.content, "usage": x.usage.model_dump() if x.usage else None, "finish": x.choices[0].finish_reason}
except Exception as e: c3["a_chat"] = {"ok": False, "error": repr(e)[:300]}
try:
    chunks = []; 
    for ch in cl.chat.completions.create(model=M, messages=[{"role": "user", "content": "Name three primary colors."}], temperature=0, stream=True):
        if ch.choices and ch.choices[0].delta.content: chunks.append(ch.choices[0].delta.content)
    c3["b_stream"] = {"ok": len(chunks) > 1, "chunks": len(chunks), "text": "".join(chunks)}
except Exception as e: c3["b_stream"] = {"ok": False, "error": repr(e)[:300]}
try:
    ids = [m.id for m in cl.models.list().data]
    c3["c_models"] = {"ok": any(i.startswith(M) for i in ids), "ids": ids}
except Exception as e: c3["c_models"] = {"ok": False, "error": repr(e)[:300]}
try:
    x = OpenAI(base_url=B + "/v1/", api_key="x").chat.completions.create(model=M, messages=[{"role": "user", "content": "Say OK."}], temperature=0)
    c3["d_dummy_key"] = {"ok": bool(x.choices[0].message.content), "content": x.choices[0].message.content}
except Exception as e: c3["d_dummy_key"] = {"ok": False, "error": repr(e)[:300]}
out["c3"] = c3; print("C3", {k: v["ok"] for k, v in c3.items()}, flush=True)
# C5: structured outputs
schema = {"type": "object", "properties": {"name": {"type": "string"}, "year": {"type": "integer"}, "country": {"type": "string"}}, "required": ["name", "year", "country"]}
c5 = []
for yr in [1879, 1452, 1564, 1809, 1912, 1643, 1867, 1756, 1918, 1769]:
    r = requests.post(B + "/api/chat", json={"model": M, "messages": [{"role": "user", "content": f"Give one famous person born in {yr}, with their birth year and country of birth."}], "format": schema, "stream": False, "options": {"temperature": 0}}, timeout=300).json()
    txt = r.get("message", {}).get("content", ""); ok = False; err = None
    try: jsonschema.validate(json.loads(txt), schema); ok = True
    except Exception as e: err = repr(e)[:200]
    c5.append({"year_asked": yr, "raw": txt, "valid": ok, "err": err}); print("C5", ok, txt[:80], flush=True)
out["c5"] = c5; out["ended"] = time.strftime("%H:%M:%S")
json.dump(out, open(sys.argv[1], "w"), indent=1, ensure_ascii=False)