"""No-think latency instrument — measures the latency win of suppressing reasoning.

Runs each probe in THINK vs NO_THINK mode against the endpoint and compares
wall time, completion tokens, and thinking length. The NO_THINK mode uses
the `chat_template_kwargs.enable_thinking` field (for backends that support
it) and/or the empty-think assistant prefill trick.

Usage (via the bench.py CLI):
  python bench.py no-think run

Env (via .env): ENDPOINT, MODEL, API_KEY.
"""
import json
import os
import re
import sys
import time

from bench.core import client

PROBES = [
    ("triage-classify",
     "You are a security triage assistant. Classify this finding as one of "
     "{FALSE_POSITIVE, LOW, MEDIUM, HIGH, CRITICAL} and give a one-sentence reason. "
     "Finding: CVE in a dev-only test dependency (pytest plugin) not shipped in the "
     "production image; CVSS 9.8 RCE. Answer with 'SEVERITY: <x>' then the reason."),
    ("mail-triage",
     "Extract sender, intent, urgency (low/med/high) from this email as strict JSON with "
     "keys sender,intent,urgency and nothing else: 'From: billing@acme.com — Subject: "
     "Invoice #4471 overdue, service suspension in 48h.'"),
    ("classify-short",
     "Classify the sentiment of this ticket as POSITIVE/NEUTRAL/NEGATIVE and nothing else: "
     "'Your update broke our export and support hasn't replied in three days.'"),
    ("arithmetic",
     "A backup runs every 15 minutes. How many backups happen in one day? Give just the number."),
]


def _split_think(msg, content):
    rc = msg.get("reasoning_content")
    if rc:
        return rc, content
    m = re.search(chr(60) + "think" + chr(62) + "(.*?)" + chr(60) + "/think" + chr(62) + "(.*)$", content, re.DOTALL)
    if m:
        return m.group(1).strip(), m.group(2).strip()
    return "", content


def run():
    print("### no-think latency bench vs %s ###" % client.MODEL, flush=True)
    print("%-16s %-9s %8s %8s %8s   %s" % ("probe", "mode", "wall_s", "comp_tok", "think_w", "answer"), flush=True)
    for name, base in PROBES:
        for mode in ("think", "no_think"):
            no_think = (mode == "no_think")
            msgs = [{"role": "user", "content": base}]
            if no_think and not client.NOTHINK:
                msgs = msgs + [{"role": "assistant", "content": chr(60) + "think" + chr(62) + chr(60) + "/think" + chr(62)}]
            try:
                resp, dt = client.call_model(
                    msgs, max_tokens=4000,
                    temperature=0.3,
                )
                ch = resp["choices"][0]
                msg = ch.get("message", {})
                content = (msg.get("content") or "").strip()
                think, ans = _split_think(msg, content)
                ct = (resp.get("usage") or {}).get("completion_tokens")
                ans1 = ans.replace("\n", " ")[:70]
                print("%-16s %-9s %8.1f %8s %8d   %s" % (name, mode, dt, ct, len(think.split()), ans1), flush=True)
            except Exception as e:
                print("%-16s %-9s  ERROR %r" % (name, mode, e), flush=True)
    print("DONE.", flush=True)




def main():
    mode = sys.argv[1] if len(sys.argv) > 1 else "run"
    if mode == "run":
        run()
    else:
        print(f"usage: python bench.py no-think [run]")
        sys.exit(1)


if __name__ == "__main__":
    main()