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


OPENAI_COMPAT_TESTS = [
    {
        "name": "official_disable",
        "model": "qwen/qwen3.6-35b-a3b",
        "enable_thinking": False,
    },
    {
        "name": "official_enable",
        "model": "qwen/qwen3.6-35b-a3b",
        "enable_thinking": True,
    },
    {
        "name": "hauhau_disable",
        "model": "qwen3.6-35b-a3b-uncensored-hauhaucs-aggressive",
        "enable_thinking": False,
    },
    {
        "name": "hauhau_enable",
        "model": "qwen3.6-35b-a3b-uncensored-hauhaucs-aggressive",
        "enable_thinking": True,
    },
]

NATIVE_TESTS = [
    {
        "name": "official_reasoning_off",
        "model": "qwen/qwen3.6-35b-a3b",
        "reasoning": "off",
    },
    {
        "name": "official_reasoning_on",
        "model": "qwen/qwen3.6-35b-a3b",
        "reasoning": "on",
    },
    {
        "name": "hauhau_reasoning_off",
        "model": "qwen3.6-35b-a3b-uncensored-hauhaucs-aggressive",
        "reasoning": "off",
    },
    {
        "name": "hauhau_reasoning_on",
        "model": "qwen3.6-35b-a3b-uncensored-hauhaucs-aggressive",
        "reasoning": "on",
    },
    {
        "name": "wrapped_reasoning_off",
        "model": "local/qwen3.6-35b-a3b-hauhaucs-aggressive",
        "reasoning": "off",
    },
    {
        "name": "wrapped_reasoning_on",
        "model": "local/qwen3.6-35b-a3b-hauhaucs-aggressive",
        "reasoning": "on",
    },
]


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


def get_json(url: str, timeout: int) -> tuple[int, str]:
    request = urllib.request.Request(url, method="GET")
    try:
        with urllib.request.urlopen(request, timeout=timeout) as response:
            return response.status, response.read().decode("utf-8", errors="replace")
    except urllib.error.HTTPError as exc:
        return exc.code, exc.read().decode("utf-8", errors="replace")


def parse_json_maybe(raw: str):
    try:
        return json.loads(raw)
    except json.JSONDecodeError:
        return None


def summarize_response(test: dict, status_code: int, raw_text: str) -> dict:
    parsed = parse_json_maybe(raw_text)
    summary = {
        "test": test["name"],
        "model": test["model"],
        "enable_thinking": test["enable_thinking"],
        "http_status": status_code,
        "json_valid": parsed is not None,
        "error": None,
        "finish_reason": None,
        "content": None,
        "reasoning_content": None,
        "contains_think_tag_in_content": None,
        "contains_think_tag_in_reasoning": None,
        "usage": None,
    }

    if parsed is None:
        summary["error"] = "response_not_json"
        return summary

    if "error" in parsed:
        summary["error"] = parsed["error"]

    choices = parsed.get("choices") or []
    if not choices:
        return summary

    choice = choices[0]
    message = choice.get("message") or {}
    content = message.get("content")
    reasoning = message.get("reasoning_content")

    summary["finish_reason"] = choice.get("finish_reason")
    summary["content"] = content
    summary["reasoning_content"] = reasoning
    summary["contains_think_tag_in_content"] = (
        isinstance(content, str) and ("<think>" in content or "</think>" in content)
    )
    summary["contains_think_tag_in_reasoning"] = (
        isinstance(reasoning, str) and ("<think>" in reasoning or "</think>" in reasoning)
    )
    summary["usage"] = parsed.get("usage")
    return summary


def summarize_native_response(test: dict, status_code: int, raw_text: str) -> dict:
    parsed = parse_json_maybe(raw_text)
    summary = {
        "test": test["name"],
        "model": test["model"],
        "reasoning": test["reasoning"],
        "http_status": status_code,
        "json_valid": parsed is not None,
        "error": None,
        "message_content": [],
        "reasoning_content": [],
        "output_types": [],
        "stats": None,
    }
    if parsed is None:
        summary["error"] = "response_not_json"
        return summary

    if "error" in parsed:
        summary["error"] = parsed["error"]

    output = parsed.get("output") or []
    summary["output_types"] = [item.get("type") for item in output]
    for item in output:
        if item.get("type") == "message":
            summary["message_content"].append(item.get("content"))
        elif item.get("type") == "reasoning":
            summary["reasoning_content"].append(item.get("content"))
    summary["stats"] = parsed.get("stats")
    return summary


def ensure_dir(path: Path) -> None:
    path.mkdir(parents=True, exist_ok=True)


def write_text(path: Path, text: str) -> None:
    path.write_text(text, encoding="utf-8")


def write_json(path: Path, obj) -> None:
    path.write_text(json.dumps(obj, ensure_ascii=False, indent=2), encoding="utf-8")


def main() -> int:
    parser = argparse.ArgumentParser()
    parser.add_argument("--host", default="127.0.0.1")
    parser.add_argument("--port", type=int, default=7860)
    parser.add_argument("--timeout", type=int, default=1800)
    parser.add_argument("--max-tokens", type=int, default=256)
    parser.add_argument(
        "--prompt",
        default="Compute 317 * 29. Give the final answer only.",
    )
    parser.add_argument(
        "--output-root",
        default="artifacts/lmstudio-thinking-probe",
    )
    args = parser.parse_args()

    timestamp = datetime.now().strftime("%Y%m%d-%H%M%S")
    run_dir = Path(args.output_root) / timestamp
    ensure_dir(run_dir)

    base_url = f"http://{args.host}:{args.port}/v1"
    metadata = {
        "base_url": base_url,
        "started_at": datetime.now().isoformat(),
        "cwd": os.getcwd(),
        "python": sys.version,
        "openai_compat_tests": OPENAI_COMPAT_TESTS,
        "native_tests": NATIVE_TESTS,
        "prompt": args.prompt,
        "timeout": args.timeout,
        "max_tokens": args.max_tokens,
    }
    write_json(run_dir / "run_metadata.json", metadata)

    status_code, models_raw = get_json(f"{base_url}/models", timeout=30)
    write_text(run_dir / "models.raw.json", models_raw)
    write_json(
        run_dir / "models.summary.json",
        {
            "http_status": status_code,
            "json_valid": parse_json_maybe(models_raw) is not None,
        },
    )

    all_summaries = []
    for index, test in enumerate(OPENAI_COMPAT_TESTS, start=1):
        payload = {
            "model": test["model"],
            "messages": [{"role": "user", "content": args.prompt}],
            "temperature": 0.2,
            "max_tokens": args.max_tokens,
            "stream": False,
            "enable_thinking": test["enable_thinking"],
        }

        prefix = f"{index:02d}_{test['name']}"
        write_json(run_dir / f"{prefix}.request.json", payload)

        started = time.time()
        status_code, raw_text = post_json(
            f"{base_url}/chat/completions",
            payload,
            timeout=args.timeout,
        )
        elapsed = time.time() - started

        write_text(run_dir / f"{prefix}.response.raw.json", raw_text)
        summary = summarize_response(test, status_code, raw_text)
        summary["elapsed_seconds"] = round(elapsed, 3)
        write_json(run_dir / f"{prefix}.summary.json", summary)
        all_summaries.append(summary)

    native_summaries = []
    for index, test in enumerate(NATIVE_TESTS, start=1):
        payload = {
            "model": test["model"],
            "input": args.prompt,
            "temperature": 0.2,
            "max_output_tokens": args.max_tokens,
            "reasoning": test["reasoning"],
            "store": False,
        }
        prefix = f"native_{index:02d}_{test['name']}"
        write_json(run_dir / f"{prefix}.request.json", payload)

        started = time.time()
        status_code, raw_text = post_json(
            f"http://{args.host}:{args.port}/api/v1/chat",
            payload,
            timeout=args.timeout,
        )
        elapsed = time.time() - started

        write_text(run_dir / f"{prefix}.response.raw.json", raw_text)
        summary = summarize_native_response(test, status_code, raw_text)
        summary["elapsed_seconds"] = round(elapsed, 3)
        write_json(run_dir / f"{prefix}.summary.json", summary)
        native_summaries.append(summary)

    report_lines = [
        "# LM Studio Thinking Probe",
        "",
        f"- Base URL: `{base_url}`",
        f"- Prompt: `{args.prompt}`",
        f"- Run dir: `{run_dir}`",
        "",
        "| Test | HTTP | enable_thinking | finish | reasoning_content | think tags in content | content |",
        "| --- | --- | --- | --- | --- | --- | --- |",
    ]
    for summary in all_summaries:
        content = summary["content"]
        if isinstance(content, str):
            content = content.replace("\n", "\\n")
            content = content[:80]
        reasoning_present = summary["reasoning_content"] is not None
        report_lines.append(
            "| {test} | {http_status} | {enable_thinking} | {finish_reason} | {reasoning_present} | {contains_think_tag_in_content} | {content} |".format(
                test=summary["test"],
                http_status=summary["http_status"],
                enable_thinking=summary["enable_thinking"],
                finish_reason=summary["finish_reason"],
                reasoning_present=reasoning_present,
                contains_think_tag_in_content=summary["contains_think_tag_in_content"],
                content=content,
            )
        )

    report_lines.extend(
        [
            "",
            "## Native API",
            "",
            "| Test | HTTP | reasoning | reasoning items | message items | stats |",
            "| --- | --- | --- | --- | --- | --- |",
        ]
    )
    for summary in native_summaries:
        report_lines.append(
            "| {test} | {http_status} | {reasoning} | {reasoning_count} | {message_count} | {stats} |".format(
                test=summary["test"],
                http_status=summary["http_status"],
                reasoning=summary["reasoning"],
                reasoning_count=len(summary["reasoning_content"]),
                message_count=len(summary["message_content"]),
                stats=str(summary["stats"]).replace("\n", " "),
            )
        )

    write_text(run_dir / "report.md", "\n".join(report_lines) + "\n")
    write_json(
        run_dir / "all_summaries.json",
        {
            "completed_at": datetime.now().isoformat(),
            "openai_compat_results": all_summaries,
            "native_results": native_summaries,
        },
    )

    print(run_dir)
    return 0


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