#!/usr/bin/env python3
"""Probe what --reasoning-budget actually enforces.

Sends a prompt hard enough to blow past the budget, then measures the thinking
block exactly via the server's /tokenize endpoint (word counts are too rough to
tell 4096 from 4500). Also tries a per-request override to distinguish a hard
cap from a default.
"""
import json, argparse, urllib.request

# Deliberately open-ended search problem: forces long systematic enumeration
# rather than a one-line recall answer.
HARD = (
    "Bestimme alle Tripel positiver ganzer Zahlen (a,b,c) mit a<=b<=c, "
    "a+b+c = 30 und a*b*c durch 30 teilbar, fuer die zusaetzlich a^2+b^2 = c^2+7 gilt. "
    "Pruefe systematisch alle Faelle durch und begruende jeden Ausschluss einzeln."
)


def post(port, path, payload, timeout=3600):
    req = urllib.request.Request(f"http://localhost:{port}{path}",
                                 data=json.dumps(payload).encode(),
                                 headers={"Content-Type": "application/json"})
    with urllib.request.urlopen(req, timeout=timeout) as r:
        return json.load(r)


def ntok(port, text):
    if not text:
        return 0
    return len(post(port, "/tokenize", {"content": text}).get("tokens", []))


def probe(port, label, extra):
    body = {"messages": [{"role": "user", "content": HARD}], "temperature": 0.3}
    body.update(extra)
    d = post(port, "/v1/chat/completions", body)
    ch = d["choices"][0]
    msg = ch["message"]
    think = msg.get("reasoning_content") or ""
    answer = msg.get("content") or ""
    t_think = ntok(port, think)
    t_ans = ntok(port, answer)
    print(f"[{label}]", flush=True)
    print(f"  reasoning tokens : {t_think}", flush=True)
    print(f"  answer tokens    : {t_ans}", flush=True)
    print(f"  completion_tokens: {d.get('usage', {}).get('completion_tokens')}", flush=True)
    print(f"  finish_reason    : {ch.get('finish_reason')}", flush=True)
    print(f"  answer non-empty : {bool(answer.strip())}", flush=True)
    print(f"  thinking tail    : ...{think[-160:]!r}", flush=True)
    return t_think


if __name__ == "__main__":
    ap = argparse.ArgumentParser()
    ap.add_argument("--port", type=int, default=18008)
    ap.add_argument("--case", choices=["default", "override", "off"], default="default")
    ap.add_argument("--budget", type=int, default=16384)
    a = ap.parse_args()

    if a.case == "default":
        probe(a.port, "server default", {})
    elif a.case == "override":
        probe(a.port, f"per-request reasoning_budget={a.budget}",
              {"reasoning_budget": a.budget})
    else:
        probe(a.port, "thinking disabled",
              {"chat_template_kwargs": {"enable_thinking": False}})