#!/usr/bin/env python3
"""Test what an OpenAI-compatible server actually supports, rather than what it claims.

Purpose: send three probes to a server - a model listing, a tool-calling request and a
    JSON-schema structured-output request - and record for each whether the server accepted
    it, whether the answer had the shape the OpenAI Chat Completions API specifies, and the
    exact reason when it did not. One JSON line per server is appended to the lab notebook.
Platform: all (Python standard library only)
Minimum memory: 12 GB
Assumes: Python 3.9 or later; a server already running and reachable at --base-url with the
    model named by --model loaded. Nothing is installed and nothing is written except the
    notebook line.

Usage: python3 feature-probe.py --engine llama.cpp --base-url http://127.0.0.1:8080/v1 \
           --model Qwen3-8B-Q4_K_M --engine-version "0.4.0" --labbook labbook.md
       python3 feature-probe.py --engine exllamav3 --base-url http://127.0.0.1:5000/v1 \
           --model Qwen3-8B-exl3-4.0bpw --api-key "$TABBY_KEY" --print-only

What counts as supported:
  * tool calling: the reply carries message.tool_calls, the first call names
    get_current_weather, and its arguments are a JSON object with a "city" string. A reply
    whose text contains "<tool_call>" is reported separately: the model produced a call in
    its own format and the server did not convert it, which is a server configuration
    finding, not a model failure.
  * structured output: the request uses response_format {"type": "json_schema", ...}, the
    prompt does not mention JSON or the field names, and the reply parses as a JSON object
    with exactly the three required keys and the right types. A server that ignores
    response_format gets a sentence of prose back and fails, which is the point: only a
    server that constrains generation to the schema can pass reliably.
Both probe prompts end with Qwen3's documented "/no_think" switch (change it with
--prompt-suffix) so that thinking does not consume the token budget. A reply that still
starts with an empty <think></think> block has that block removed before parsing, and the
record says so.
"""

from __future__ import annotations

import argparse
import json
import re
import time
import urllib.error
import urllib.request
from pathlib import Path

LAB = "part-08/lab-same-model-every-engine"

WEATHER_TOOL = {
    "type": "function",
    "function": {
        "name": "get_current_weather",
        "description": "Get the current weather in a named city.",
        "parameters": {
            "type": "object",
            "properties": {
                "city": {"type": "string", "description": "City name, for example Lisbon"},
                "unit": {"type": "string", "enum": ["celsius", "fahrenheit"]},
            },
            "required": ["city"],
        },
    },
}

PERSON_SCHEMA = {
    "type": "object",
    "properties": {
        "name": {"type": "string"},
        "born": {"type": "integer"},
        "field": {"type": "string"},
    },
    "required": ["name", "born", "field"],
    "additionalProperties": False,
}

THINK_BLOCK = re.compile(r"^\s*<think>.*?</think>\s*", re.DOTALL)


def call(base_url: str, path: str, api_key: str, body, timeout: float):
    """One request. Returns (status, parsed body or text) and turns errors into the same shape."""
    url = base_url.rstrip("/") + path
    headers = {"Content-Type": "application/json", "Accept": "application/json"}
    if api_key:
        headers["Authorization"] = f"Bearer {api_key}"
    data = json.dumps(body).encode("utf-8") if body is not None else None
    method = "POST" if body is not None else "GET"
    request = urllib.request.Request(url, data=data, headers=headers, method=method)
    try:
        with urllib.request.urlopen(request, timeout=timeout) as response:  # noqa: S310 - local server
            raw = response.read().decode("utf-8", "replace")
            try:
                return response.status, json.loads(raw)
            except json.JSONDecodeError:
                return response.status, raw[:400]
    except urllib.error.HTTPError as exc:
        return exc.code, exc.read().decode("utf-8", "replace")[:400]
    except (urllib.error.URLError, TimeoutError, OSError) as exc:
        return 0, str(exc)[:400]


def first_message(response: dict) -> tuple:
    choice = (response.get("choices") or [{}])[0] or {}
    return choice.get("message") or {}, choice.get("finish_reason")


def reasoning_field(message: dict):
    for key in ("reasoning_content", "reasoning"):
        if message.get(key):
            return key
    return None


def probe_models(args) -> dict:
    status, body = call(args.base_url, "/models", args.api_key, None, args.timeout)
    ids = []
    if status == 200 and isinstance(body, dict):
        ids = [m.get("id") for m in body.get("data") or [] if isinstance(m, dict)]
    result = {"status": status, "supported": status == 200, "model_ids": ids[:10]}
    if status != 200:
        result["detail"] = f"status {status}: {str(body)[:160]}"
    return result


def probe_tool_calling(args) -> dict:
    body = {
        "model": args.model,
        "messages": [{"role": "user", "content":
                      "What is the weather in Lisbon right now? Use the tool." + args.prompt_suffix}],
        "tools": [WEATHER_TOOL],
        "tool_choice": "auto",
        "max_tokens": args.max_tokens,
        "temperature": 0,
    }
    status, response = call(args.base_url, "/chat/completions", args.api_key, body, args.timeout)
    result = {"status": status, "supported": False, "called": None, "arguments": None}
    if status != 200 or not isinstance(response, dict):
        result["detail"] = f"status {status}: {str(response)[:160]}"
        return result
    message, finish = first_message(response)
    result["finish_reason"] = finish
    result["reasoning_field"] = reasoning_field(message)
    calls = message.get("tool_calls") or []
    if not calls:
        text = message.get("content") or ""
        if "<tool_call>" in text:
            result["detail"] = "the model wrote a <tool_call> block as text; the server did not parse it"
        elif finish == "length":
            result["detail"] = "ran out of tokens before any tool call (finish_reason length)"
        else:
            result["detail"] = "answered without a tool_calls field: " + text.strip()[:80]
        return result
    function = (calls[0] or {}).get("function") or {}
    result["called"] = function.get("name")
    arguments = function.get("arguments")
    try:
        parsed = json.loads(arguments) if isinstance(arguments, str) else arguments
    except json.JSONDecodeError:
        result["detail"] = "tool_calls present but the arguments were not valid JSON"
        return result
    result["arguments"] = parsed
    ok = (result["called"] == "get_current_weather" and isinstance(parsed, dict)
          and isinstance(parsed.get("city"), str))
    result["supported"] = ok
    if not ok:
        result["detail"] = "tool_calls present but not the expected function and city argument"
    return result


def probe_structured_output(args) -> dict:
    body = {
        "model": args.model,
        "messages": [{"role": "user", "content":
                      "Name one physicist who won a Nobel Prize." + args.prompt_suffix}],
        "response_format": {
            "type": "json_schema",
            "json_schema": {"name": "person", "strict": True, "schema": PERSON_SCHEMA},
        },
        "max_tokens": args.max_tokens,
        "temperature": 0,
    }
    status, response = call(args.base_url, "/chat/completions", args.api_key, body, args.timeout)
    result = {"status": status, "supported": False, "value": None, "think_block_removed": False}
    if status != 200 or not isinstance(response, dict):
        result["detail"] = f"status {status}: {str(response)[:160]}"
        return result
    message, finish = first_message(response)
    result["finish_reason"] = finish
    result["reasoning_field"] = reasoning_field(message)
    content = message.get("content") or ""
    stripped = THINK_BLOCK.sub("", content, count=1)
    result["think_block_removed"] = stripped != content
    try:
        value = json.loads(stripped)
    except json.JSONDecodeError:
        result["detail"] = "accepted the request but the content was not JSON: " + stripped.strip()[:80]
        return result
    result["value"] = value
    if not isinstance(value, dict):
        result["detail"] = "valid JSON, but not an object"
        return result
    problems = []
    missing = [k for k in PERSON_SCHEMA["required"] if k not in value]
    extra = [k for k in value if k not in PERSON_SCHEMA["properties"]]
    if missing:
        problems.append(f"missing keys {missing}")
    if extra:
        problems.append(f"extra keys {extra}")
    if "born" in value and not (isinstance(value["born"], int) and not isinstance(value["born"], bool)):
        problems.append("born is not an integer")
    for key in ("name", "field"):
        if key in value and not isinstance(value[key], str):
            problems.append(f"{key} is not a string")
    result["supported"] = not problems
    if problems:
        result["detail"] = "valid JSON that does not match the schema: " + "; ".join(problems)
    return result


def main() -> int:
    parser = argparse.ArgumentParser(description=__doc__.splitlines()[0])
    parser.add_argument("--engine", required=True, help="engine name as it should appear in the notebook")
    parser.add_argument("--base-url", default="http://127.0.0.1:8080/v1", help="OpenAI-compatible base URL")
    parser.add_argument("--model", required=True, help="model name to send in the request")
    parser.add_argument("--api-key", default="", help="bearer token, when the server wants one")
    parser.add_argument("--engine-version", default="", help="the version string you read from the engine")
    parser.add_argument("--host", default="", help="short description of the machine")
    parser.add_argument("--prompt-suffix", default=" /no_think", help="appended to both probe prompts")
    parser.add_argument("--max-tokens", type=int, default=1024, help="tokens to allow per probe")
    parser.add_argument("--timeout", type=float, default=300.0, help="seconds to wait for a response")
    parser.add_argument("--labbook", default="labbook.md", help="notebook to append to")
    parser.add_argument("--print-only", action="store_true", help="print the result, record nothing")
    args = parser.parse_args()

    record = {
        "lab": LAB,
        "probe": "features",
        "engine": args.engine,
        "engine_version": args.engine_version,
        "host": args.host,
        "model": args.model,
        "base_url": args.base_url,
        "models_endpoint": probe_models(args),
        "tool_calling": probe_tool_calling(args),
        "structured_output": probe_structured_output(args),
        "measured_on": time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()),
    }

    def mark(section: str, shown) -> str:
        entry = record[section]
        if entry.get("supported"):
            return "yes " + json.dumps(shown)
        return "no (" + str(entry.get("detail")) + ")"

    models = record["models_endpoint"]
    tools = record["tool_calling"]
    schema = record["structured_output"]
    print(f"    {args.engine}: /v1/models {mark('models_endpoint', models['model_ids'])}")
    print(f"    {args.engine}: tool calling {mark('tool_calling', {tools['called']: tools['arguments']})}")
    print(f"    {args.engine}: JSON schema {mark('structured_output', schema['value'])}")

    if args.print_only:
        return 0
    notebook = Path(args.labbook)
    if not notebook.exists():
        print(f"    {notebook} does not exist; creating it")
    with notebook.open("a", encoding="utf-8") as handle:
        handle.write(json.dumps(record) + "\n")
    print(f"    recorded 1 line in {notebook}")
    return 0


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