#!/usr/bin/env python3
"""Which server-side default makes Nex behave well for clients that pass no chat_template_kwargs (the gateway case)?
Follow-up to nex_tools_tpl.py: with the `high` shim, a no-kwargs one-word request came back entirely in
reasoning_content (the model never closed its forced-open think block). Two candidate defaults, greedy probes:
  C1  high shim + `--chat-template-kwargs {"enable_thinking": false}`      (thinking off unless a client asks)
  C2  v2 shim (explicit enable_thinking=false wins; otherwise enable_thinking maps to high unless reasoning_effort is
      set) + `--chat-template-kwargs {"reasoning_effort": "medium"}`      (upstream's adaptive default)
Diagnostic only -> results/nex_seat_default_probe.json."""
import json, os, subprocess, sys, time
from types import SimpleNamespace

os.environ["AGNES_BIN"] = "/opt/llama-rocm/rocmfpx-724/build-hipvk/bin"
sys.path.insert(0, "/mnt/models/nex-n2.5-mini")
import nex_harness as H  # noqa: E402

W = H.W
SRC = open(f"{W}/hf/chat_template.jinja", "rb").read()
SHIM_V2 = ("{%- if enable_thinking is defined and not enable_thinking %}{%- set reasoning_effort = 'none' %}"
           "{%- elif reasoning_effort is not defined and enable_thinking is defined %}"
           "{%- set reasoning_effort = 'high' %}{%- endif %}\n")
TPL_V2 = f"{W}/tpl/chat_template_enable_thinking_v2.jinja"
open(TPL_V2, "wb").write(SHIM_V2.encode() + SRC)
CONFIGS = [
    ("C1", [f"--chat-template-file", f"{W}/tpl/chat_template_enable_thinking.jinja",
            "--chat-template-kwargs", json.dumps({"enable_thinking": False})]),
    ("C2", ["--chat-template-file", TPL_V2, "--chat-template-kwargs", json.dumps({"reasoning_effort": "medium"})]),
]
PROMPTS = [("correct-decline", "What is 17 times 23? Answer directly."),
           ("single-word", "Reply with the single word: ready"),
           ("multi-arg", "What's the weather in Paris in celsius?")]
KW = [("no-kwargs", None), ("enable_thinking=true", {"enable_thinking": True}),
      ("enable_thinking=false", {"enable_thinking": False})]
_orig = subprocess.Popen
report = {"shim_v2": SHIM_V2, "configs": {}}
for name, extra in CONFIGS:
    def _popen(cmd, *a, _extra=extra, **k):
        if cmd and str(cmd[0]).endswith("llama-server"):
            cmd = list(cmd) + list(_extra)
        return _orig(cmd, *a, **k)
    subprocess.Popen = _popen
    a = SimpleNamespace(model=f"{W}/out/Nex-N2.5-mini-Q4_0_ROCMFP4_STRIX_LEAN.gguf", dev="ROCm0", ctx=16384,
                        draft=None, mtp_infile=False, nmax=4, pmin=0.0, strict=False,
                        serverlog=f"{W}/logs/probe_seat_default_{name}.log")
    try:
        s = H.Server(a, 18653)
    finally:
        subprocess.Popen = _orig
    rec = {"extra": extra, "results": {}}
    try:
        for kn, kw in KW:
            for pn, prompt in PROMPTS:
                body = {"messages": [{"role": "user", "content": prompt}], "tools": H.TOOLS, "tool_choice": "auto",
                        "temperature": 0, "top_k": 1, "max_tokens": 2048}
                if kw is not None:
                    body["chat_template_kwargs"] = kw
                try:
                    m = H.post(18653, "/v1/chat/completions", body)["choices"][0]["message"]
                    c = m.get("content") or ""
                    r = {"content": c[:120], "reasoning_len": len(m.get("reasoning_content") or ""),
                         "tool_calls": [(t.get("function") or {}).get("name") for t in m.get("tool_calls") or []],
                         "leaks": [x for x in H.LEAK if x in c]}
                except Exception as e:  # noqa: BLE001
                    r = {"error": repr(e)[:300]}
                rec["results"][f"{kn}|{pn}"] = r
                print(name, kn, pn, json.dumps(r)[:220], flush=True)
    finally:
        s.stop()
    report["configs"][name] = rec
    time.sleep(3)
json.dump(report, open(f"{W}/results/nex_seat_default_probe.json", "w"), indent=1)
print("NEX_SEAT_DEFAULT_PROBE_DONE")