import json, sys, time, urllib.request

PORT = sys.argv[1]
MODEL = "Qwen3.8-27B-Uncensored-HauhauCS-Aggressive-Q4_K_P.gguf"

def call(messages, max_tokens=400, tools=None, temperature=None):
    body = {"model": MODEL, "messages": messages, "max_tokens": max_tokens, "stream": True}
    if tools: body["tools"] = tools
    if temperature is not None: body["temperature"] = temperature
    req = urllib.request.Request(f"http://localhost:{PORT}/v1/chat/completions",
        data=json.dumps(body).encode(), headers={"Content-Type": "application/json"})
    t0=time.time(); ttf=None; n_gen=0; n_reason=0; toolcall=None
    with urllib.request.urlopen(req, timeout=300) as r:
        for line in r:
            line=line.decode().strip()
            if not line.startswith("data: ") or line[6:]=="[DONE]": continue
            d=json.loads(line[6:]); delta=d["choices"][0].get("delta",{})
            c=delta.get("content"); rc=delta.get("reasoning_content")
            if rc: n_reason+=1
            if c: n_gen+=1
            if delta.get("tool_calls"): toolcall=delta["tool_calls"][0].get("function",{})
            if (c or rc) and ttf is None: ttf=time.time()-t0
    tot=time.time()-t0; n=n_gen+n_reason
    return dict(ttft=ttf, total=tot, out=n, reason=n_reason, speed=n/max(tot-ttf,0.01), toolcall=toolcall)

# --- latency tests ---
r = call([{"role":"user","content":"你好，1+1等于几？简短回答。"}], max_tokens=150)
print(f"[short]    ttft={r['ttft']:.2f}s total={r['total']:.2f}s out={r['out']} (reason={r['reason']}) {r['speed']:.1f} t/s")

r = call([{"role":"user","content":"写一首关于秋天的五言绝句。"}], max_tokens=400)
print(f"[creative] ttft={r['ttft']:.2f}s total={r['total']:.2f}s out={r['out']} (reason={r['reason']}) {r['speed']:.1f} t/s")

# --- accuracy probes (expected answers) ---
PROBES = [
    ("一个房间10人，一半离开，3人进来，1人出去。现在几人？(final number)", "7"),
    ("9.11 和 9.9 哪个大？(which is larger, answer 9.11 or 9.9)", "9.9"),
    ("小明比小红高，小红比小刚高，谁最矮？(answer 小明/小红/小刚)", "小刚"),
]
correct=0
for q,exp in PROBES:
    r = call([{"role":"user","content":q}], max_tokens=300)
    last=(r.get('reasoning') or '')
    ok = exp.lower() in (json.dumps(r,ensure_ascii=False)).lower()
    correct += ok
    print(f"[probe] {ok} (expect {exp}) out={r['out']} reason={r['reason']} total={r['total']:.1f}s")
print(f"[accuracy] {correct}/{len(PROBES)}")

# --- tool call accuracy ---
TOOLS=[{"type":"function","function":{"name":"get_weather","parameters":{"type":"object","properties":{"city":{"type":"string"}},"required":["city"]}}}]
r = call([{"role":"system","content":"You are a helpful assistant with tools. Call get_weather for weather questions."},
          {"role":"user","content":"What is the weather in Beijing right now?"}], tools=TOOLS, max_tokens=200)
tc=r["toolcall"] or {}
tool_ok = tc.get("name")=="get_weather" and "Beijing" in (tc.get("arguments") or "")
print(f"[toolcall] {tool_ok} name={tc.get('name')} args={tc.get('arguments')}")