import json, time, urllib.request
PORT = "8081"
MODEL = "Qwen3.8-27B-Uncensored-HauhauCS-Aggressive-Q4_K_P.gguf"

def call(messages, max_tokens=600, tools=None, label=""):
    body = {"model": MODEL, "messages": messages, "max_tokens": max_tokens, "stream": True}
    if tools: body["tools"] = tools
    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; err=None
    try:
        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
    except Exception as e:
        err=str(e)[:100]
    tot=time.time()-t0; n=n_gen+n_reason
    tc = f" tool={toolcall.get('name')}{toolcall.get('arguments')}" if toolcall else ""
    print(f"{label}: ttft={ttf:.2f}s total={tot:.2f}s out={n} (reason={n_reason}) {n/max(tot-ttf,0.01):.1f} t/s{tc}{(' ERR:'+err) if err else ''}")

# realistic Hermes-ish system prompt ~6k tokens
sysprompt = "You are Hermes Agent, an AI assistant. " + ("You have tools for web search, terminal, file editing and memory. " * 180) + "Reply concisely."
# turn 1: short chat
call([{"role":"system","content":sysprompt},{"role":"user","content":"简短介绍一下你自己。"}], max_tokens=300, label="[turn1 short chat]")
# turn 2: same system + growing history (simulate 12k total)
hist = [{"role":"system","content":sysprompt},
        {"role":"user","content":"帮我查一下今天有什么新闻。"},
        {"role":"assistant","content":"(tool call)"},
        {"role":"user","content":"(result: AI发布新产品)"},
        {"role":"assistant","content":"今天有AI新产品的新闻。"}*3 if False else {"role":"assistant","content":"今天有AI新产品的新闻。"},
        {"role":"user","content":"展开讲讲第一条。"}]
call(hist, max_tokens=400, label="[turn2 12k ctx]  ")
# turn 3: creative / long answer
call([{"role":"system","content":sysprompt},{"role":"user","content":"写一首关于秋天的五言绝句，再解释每句意思。"}], max_tokens=600, label="[turn3 creative]   ")
# tool call
TOOLS=[{"type":"function","function":{"name":"terminal","parameters":{"type":"object","properties":{"command":{"type":"string"}},"required":["command"]}}}]
call([{"role":"system","content":"You are an agent with a terminal tool. Use it for shell commands."},
      {"role":"user","content":"Run: list files in /tmp and count them"}], tools=TOOLS, max_tokens=200, label="[toolcall]         ")
# accuracy probes
for q,exp in [("9.11 和 9.9 哪个大？只回答数字","9.9"),("小明比小红高，小红比小刚高，谁最矮？只回答名字","小刚"),("一个房间10人，一半离开，3人进来，1人出去。现在几人？只回答数字","7")]:
    r=None
    body={"model":MODEL,"messages":[{"role":"user","content":q}],"max_tokens":300}
    req=urllib.request.Request(f"http://localhost:{PORT}/v1/chat/completions",data=json.dumps(body).encode(),headers={"Content-Type":"application/json"})
    d=json.loads(urllib.request.urlopen(req,timeout=180).read())
    m=d["choices"][0]["message"]; out=(m.get("content") or "")+(m.get("reasoning_content") or "")
    ok = exp in out
    print(f"[probe] {ok} expect={exp} content={(m.get('content') or '')[-60:]!r}")