#!/usr/bin/env python3
"""Regression / reliability check for structured (schema-constrained) output.

Exercises a llama.cpp-compatible /v1/chat/completions endpoint (either the
llama-server directly, or through the FastAPI gateway in this repo) with:

  A. an ordinary, unconstrained completion
  B. response_format={"type": "json_object"}
  C. response_format={"type": "json_schema", ...} with a simple schema
  D. the same, with a realistic tool-call repair schema, fed a malformed
     tool-call and checked against the expected repaired structure
  E. test D repeated 10x, checking for zero sampler-initialization failures

Usage:
    python3 scripts/validate_structured_output.py --base-url http://127.0.0.1:8080/v1
    python3 scripts/validate_structured_output.py \\
        --base-url https://leon4gr45-minicpm.hf.space/v1 --api-key-env API_KEY

The API key (if any) is read from the environment variable named by
--api-key-env and is never printed or logged.
"""

from __future__ import annotations

import argparse
import json
import os
import sys
import time
import urllib.error
import urllib.request
from dataclasses import dataclass, field

SCHEMA_C = {
    "type": "object",
    "properties": {
        "repairable": {"type": "boolean"},
        "name": {"type": "string"},
        "arguments": {"type": "object"},
    },
    "required": ["repairable"],
    "additionalProperties": False,
}

SCHEMA_D = {
    "type": "object",
    "properties": {
        "repairable": {"type": "boolean"},
        "name": {"type": "string", "enum": ["bash"]},
        "arguments": {
            "type": "object",
            "properties": {"command": {"type": "string"}},
            "required": ["command"],
            "additionalProperties": False,
        },
    },
    "required": ["repairable", "name", "arguments"],
    "additionalProperties": False,
}

MALFORMED_TOOL_CALL = '<|tool_call>call:bash{command:<|"|>ls -la<|"|>}<tool_call|>'

REPAIR_SYSTEM_PROMPT = (
    "You repair malformed tool calls. A tool call is repairable if its "
    "intended tool name and arguments can be confidently recovered from the "
    "malformed text, even if the syntax (quotes, brackets, tags) is wrong. "
    "Given a malformed tool call, output the repaired call as JSON matching "
    "the schema."
)


@dataclass
class CallResult:
    ok: bool
    status: int
    body: str
    latency_s: float
    content: str | None = None
    parsed: object | None = None
    error: str | None = None
    usage: dict = field(default_factory=dict)
    timings: dict = field(default_factory=dict)


def call(base_url: str, api_key: str | None, payload: dict, timeout: float = 60.0) -> CallResult:
    url = base_url.rstrip("/") + "/chat/completions"
    headers = {"Content-Type": "application/json"}
    if api_key:
        headers["Authorization"] = f"Bearer {api_key}"
    data = json.dumps(payload).encode()
    req = urllib.request.Request(url, data=data, headers=headers, method="POST")

    start = time.monotonic()
    try:
        with urllib.request.urlopen(req, timeout=timeout) as resp:
            body = resp.read().decode()
            status = resp.status
    except urllib.error.HTTPError as exc:
        body = exc.read().decode()
        status = exc.code
    except urllib.error.URLError as exc:
        return CallResult(ok=False, status=0, body="", latency_s=time.monotonic() - start, error=str(exc))
    latency_s = time.monotonic() - start

    if status != 200:
        return CallResult(ok=False, status=status, body=body, latency_s=latency_s, error=f"HTTP {status}: {body}")

    try:
        response_json = json.loads(body)
        content = response_json["choices"][0]["message"]["content"]
        usage = response_json.get("usage", {})
        timings = response_json.get("timings", {})
    except (KeyError, IndexError, json.JSONDecodeError) as exc:
        return CallResult(ok=False, status=status, body=body, latency_s=latency_s, error=f"malformed response envelope: {exc}")

    return CallResult(ok=True, status=status, body=body, latency_s=latency_s, content=content, usage=usage, timings=timings)


def try_parse_json(text: str):
    try:
        return json.loads(text)
    except json.JSONDecodeError:
        # Some models wrap JSON in ```json fences despite instructions not to.
        stripped = text.strip()
        if stripped.startswith("```"):
            stripped = stripped.strip("`")
            if stripped.lower().startswith("json"):
                stripped = stripped[4:]
            try:
                return json.loads(stripped.strip())
            except json.JSONDecodeError:
                return None
        return None


def validates_schema_c(parsed) -> bool:
    return isinstance(parsed, dict) and isinstance(parsed.get("repairable"), bool)


def validates_schema_d_repair(parsed) -> bool:
    if not isinstance(parsed, dict):
        return False
    if parsed.get("repairable") is not True:
        return False
    if parsed.get("name") != "bash":
        return False
    args = parsed.get("arguments")
    return isinstance(args, dict) and args.get("command") == "ls -la"


def test_ordinary_completion(base_url: str, api_key: str | None, model: str) -> tuple[bool, str]:
    payload = {"model": model, "stream": False, "messages": [{"role": "user", "content": "Say hi in one word."}]}
    result = call(base_url, api_key, payload)
    if not result.ok:
        return False, result.error or "unknown failure"
    if not result.content or not result.content.strip():
        return False, "empty assistant content"
    return True, f"latency={result.latency_s:.2f}s content={result.content!r}"


def test_json_object(base_url: str, api_key: str | None, model: str) -> tuple[bool, str]:
    payload = {
        "model": model,
        "stream": False,
        "messages": [{"role": "user", "content": 'Return {"ok":true} as JSON.'}],
        "response_format": {"type": "json_object"},
    }
    result = call(base_url, api_key, payload)
    if not result.ok:
        return False, result.error or "unknown failure"
    parsed = try_parse_json(result.content)
    if parsed is None or not isinstance(parsed, dict):
        return False, f"not valid JSON object: {result.content!r}"
    return True, f"latency={result.latency_s:.2f}s parsed={parsed}"


def test_json_schema_simple(base_url: str, api_key: str | None, model: str) -> tuple[bool, str]:
    payload = {
        "model": model,
        "stream": False,
        "messages": [{"role": "user", "content": "Return a JSON object describing whether this is repairable."}],
        "response_format": {
            "type": "json_schema",
            "json_schema": {"name": "tool_repair_simple", "schema": SCHEMA_C},
        },
    }
    result = call(base_url, api_key, payload)
    if not result.ok:
        return False, result.error or "unknown failure"
    parsed = try_parse_json(result.content)
    if parsed is None:
        return False, f"not valid JSON: {result.content!r}"
    if not validates_schema_c(parsed):
        return False, f"does not validate against schema: {parsed!r}"
    return True, f"latency={result.latency_s:.2f}s parsed={parsed}"


def build_repair_payload(model: str) -> dict:
    return {
        "model": model,
        "stream": False,
        # Greedy decoding: this is a verifier/repair role where deterministic,
        # repeatable output matters more than sampling diversity, and it
        # keeps this reliability check meaningful (a schema grammar can force
        # *valid* JSON at any temperature, but whether the model's content
        # judgement is repeatable is a separate, real property to check).
        "temperature": 0,
        "messages": [
            {"role": "system", "content": REPAIR_SYSTEM_PROMPT},
            {"role": "user", "content": f"Repair this malformed tool call:\n{MALFORMED_TOOL_CALL}"},
        ],
        "response_format": {
            "type": "json_schema",
            "json_schema": {"name": "tool_repair", "strict": True, "schema": SCHEMA_D},
        },
    }


def test_json_schema_repair(base_url: str, api_key: str | None, model: str) -> tuple[bool, str, CallResult | None]:
    payload = build_repair_payload(model)
    result = call(base_url, api_key, payload)
    if not result.ok:
        return False, result.error or "unknown failure", result
    parsed = try_parse_json(result.content)
    if parsed is None:
        return False, f"not valid JSON: {result.content!r}", result
    if not validates_schema_d_repair(parsed):
        return False, f"does not validate against repair schema: {parsed!r}", result
    return True, f"latency={result.latency_s:.2f}s parsed={parsed}", result


def test_repeated_repair(base_url: str, api_key: str | None, model: str, n: int = 10) -> tuple[bool, str]:
    payload = build_repair_payload(model)
    passed = 0
    sampler_failures = 0
    details = []
    for i in range(n):
        result = call(base_url, api_key, payload)
        if not result.ok:
            details.append(f"#{i}: HTTP {result.status} {result.error}")
            if result.error and "sampler" in result.error.lower():
                sampler_failures += 1
            continue
        parsed = try_parse_json(result.content)
        if parsed is None or not validates_schema_d_repair(parsed):
            details.append(f"#{i}: schema-invalid content={result.content!r}")
            continue
        passed += 1
    ok = passed == n
    summary = f"{passed}/{n} passed, {sampler_failures} sampler-init failures"
    if details:
        summary += "; " + "; ".join(details[:5])
    return ok, summary


def main() -> int:
    parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    parser.add_argument("--base-url", required=True, help="Base URL including /v1, e.g. http://127.0.0.1:8080/v1")
    parser.add_argument("--model", default="minicpm4-0.5b")
    parser.add_argument("--api-key-env", default=None, help="Env var name holding a bearer API key (never printed)")
    parser.add_argument("--repeat", type=int, default=10, help="Number of repair calls for the reliability test")
    args = parser.parse_args()

    api_key = os.environ.get(args.api_key_env) if args.api_key_env else None

    results: list[tuple[str, bool, str]] = []

    ok, detail = test_ordinary_completion(args.base_url, api_key, args.model)
    results.append(("ordinary_completion", ok, detail))

    ok, detail = test_json_object(args.base_url, api_key, args.model)
    results.append(("json_object", ok, detail))

    ok, detail = test_json_schema_simple(args.base_url, api_key, args.model)
    results.append(("json_schema_simple", ok, detail))

    ok, detail, last_result = test_json_schema_repair(args.base_url, api_key, args.model)
    results.append(("json_schema_repair", ok, detail))

    ok, detail = test_repeated_repair(args.base_url, api_key, args.model, n=args.repeat)
    results.append((f"repeated_repair_{args.repeat}x", ok, detail))

    print()
    width = max(len(name) for name, _, _ in results) + 4
    failures = 0
    for name, ok, detail in results:
        status = "PASS" if ok else "FAIL"
        if not ok:
            failures += 1
        print(f"{name.ljust(width)}{status}")
        if not ok or "--verbose" in sys.argv:
            print(f"{'':{width}}{detail}")

    if last_result is not None and last_result.ok:
        print()
        print("Performance (last repair call):")
        print(f"  prompt_tokens     = {last_result.usage.get('prompt_tokens')}")
        print(f"  completion_tokens = {last_result.usage.get('completion_tokens')}")
        print(f"  prompt_ms         = {last_result.timings.get('prompt_ms')}")
        print(f"  predicted_ms      = {last_result.timings.get('predicted_ms')}")
        print(f"  wall_time_s       = {last_result.latency_s:.3f}")

    print()
    print(f"failures: {failures}")
    return 1 if failures else 0


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