#!/usr/bin/env python3
import argparse
import json
import urllib.error
import urllib.request
from datetime import datetime
from pathlib import Path


PROBE_FIELDS = [
    ("baseline", {}),
    ("enable_thinking_false", {"enable_thinking": False}),
    ("reasoning_off_string", {"reasoning": "off"}),
    ("reasoning_false_bool", {"reasoning": False}),
    ("extra_body_enable_thinking_false", {"extra_body": {"enable_thinking": False}}),
    ("extra_body_reasoning_off", {"extra_body": {"reasoning": "off"}}),
    ("chat_template_kwargs_enable_thinking_false", {"chat_template_kwargs": {"enable_thinking": False}}),
    ("chat_template_kwargs_preserve_false", {"chat_template_kwargs": {"preserve_thinking": False, "enable_thinking": False}}),
    ("include_reasoning_false", {"include_reasoning": False}),
    ("separate_reasoning_false", {"separate_reasoning": False}),
]


def post_json(url: str, payload: dict, timeout: int) -> tuple[int, str]:
    body = json.dumps(payload, ensure_ascii=False).encode("utf-8")
    req = urllib.request.Request(
        url,
        data=body,
        headers={"Content-Type": "application/json"},
        method="POST",
    )
    try:
        with urllib.request.urlopen(req, timeout=timeout) as resp:
            return resp.status, resp.read().decode("utf-8", errors="replace")
    except urllib.error.HTTPError as exc:
        return exc.code, exc.read().decode("utf-8", errors="replace")


def main() -> int:
    parser = argparse.ArgumentParser()
    parser.add_argument("--port", type=int, default=7860)
    parser.add_argument(
        "--model",
        default="local/qwen3.6-35b-a3b-hauhaucs-aggressive",
    )
    parser.add_argument(
        "--prompt",
        default="Compute 317 * 29. Give the final answer only.",
    )
    parser.add_argument("--output-root", default="artifacts/lmstudio-openai-reasoning-probe")
    parser.add_argument("--max-tokens", type=int, default=64)
    parser.add_argument("--timeout", type=int, default=300)
    args = parser.parse_args()

    run_dir = Path(args.output_root) / datetime.now().strftime("%Y%m%d-%H%M%S")
    run_dir.mkdir(parents=True, exist_ok=True)

    results = []
    url = f"http://127.0.0.1:{args.port}/v1/chat/completions"
    for index, (name, fields) in enumerate(PROBE_FIELDS, start=1):
        payload = {
            "model": args.model,
            "messages": [{"role": "user", "content": args.prompt}],
            "temperature": 0.2,
            "max_tokens": args.max_tokens,
            "stream": False,
        }
        payload.update(fields)
        prefix = f"{index:02d}_{name}"
        (run_dir / f"{prefix}.request.json").write_text(
            json.dumps(payload, ensure_ascii=False, indent=2) + "\n",
            encoding="utf-8",
        )

        status_code, raw = post_json(url, payload, timeout=args.timeout)
        (run_dir / f"{prefix}.response.raw.json").write_text(raw, encoding="utf-8")

        parsed = None
        try:
            parsed = json.loads(raw)
        except json.JSONDecodeError:
            pass

        msg = ((parsed or {}).get("choices") or [{}])[0].get("message") or {}
        summary = {
            "name": name,
            "status_code": status_code,
            "fields": fields,
            "has_reasoning_content": "reasoning_content" in msg and msg.get("reasoning_content") is not None,
            "content": msg.get("content"),
            "reasoning_content": msg.get("reasoning_content"),
            "error": (parsed or {}).get("error"),
            "usage": (parsed or {}).get("usage"),
        }
        (run_dir / f"{prefix}.summary.json").write_text(
            json.dumps(summary, ensure_ascii=False, indent=2) + "\n",
            encoding="utf-8",
        )
        results.append(summary)

    (run_dir / "all_summaries.json").write_text(
        json.dumps({"results": results}, ensure_ascii=False, indent=2) + "\n",
        encoding="utf-8",
    )
    print(run_dir)
    return 0


if __name__ == "__main__":
    raise SystemExit(main())