#!/usr/bin/env python3
"""Get a validated JSON object out of a local model, with a schema and a repair loop.

Purpose: demonstrate the three layers that make structured output reliable on a local
    server: a JSON Schema sent with the request so the sampler is constrained, Pydantic
    validation of what comes back, and a bounded repair loop that shows the model its own
    error. Two tasks are included, extraction and classification, plus a --free-form mode
    that sends no schema so the difference is visible rather than asserted.
Platform: all (pure Python over HTTP; the server may be on any track or on another machine)
Minimum memory: 8 GB on the machine running the model; this script needs almost none
Assumes: Python 3.9 or later, pydantic 2.x installed in the active environment, and an
    OpenAI-compatible endpoint (llama-server from Part 6, or the gateway from Part 9)
    reachable at --base-url.

Usage: python3 structured-extract.py --base-url http://127.0.0.1:8080/v1 --model qwen3-8b --demo
       python3 structured-extract.py --base-url http://127.0.0.1:8080/v1 --model qwen3-8b \
           --task classify --input message.txt --labbook labbook.md
       python3 structured-extract.py --base-url http://127.0.0.1:8080/v1 --model qwen3-8b \
           --demo --free-form        # no schema: see how often it still parses
"""

from __future__ import annotations

import argparse
import json
import sys
import time
import urllib.error
import urllib.request
from pathlib import Path
from typing import List, Optional

try:
    from pydantic import BaseModel, Field, ValidationError
except ImportError:  # pragma: no cover - environment check, not logic
    sys.exit("pydantic is not installed. Run: uv pip install 'pydantic>=2'")


# --------------------------------------------------------------------------------------
# The two schemas. These are the contract: the same class produces the JSON Schema that
# constrains the sampler and the validator that checks the reply, so the two cannot drift.
# --------------------------------------------------------------------------------------

class ModelRecord(BaseModel):
    """One open-weight model as described in a paragraph of release notes."""

    name: str = Field(description="The model name exactly as written in the text")
    publisher: Optional[str] = Field(default=None, description="Who released it, or null")
    total_params_b: Optional[float] = Field(
        default=None, description="Total parameters in billions, or null if not stated"
    )
    licence: Optional[str] = Field(default=None, description="Licence name, or null if not stated")
    quantisations: List[str] = Field(
        default_factory=list, description="Quantisation formats named in the text"
    )


class Extraction(BaseModel):
    """The top-level object for the extraction task."""

    models: List[ModelRecord]


class Classification(BaseModel):
    """A single-label classification with a short justification."""

    label: str = Field(description="One of: bug, question, feature-request, documentation, other")
    confidence: float = Field(ge=0.0, le=1.0, description="0 to 1")
    rationale: str = Field(description="One sentence, at most 25 words")


LABELS = ["bug", "question", "feature-request", "documentation", "other"]

TASKS = {
    "extract": {
        "schema_model": Extraction,
        "system": (
            "You extract structured records from release notes. "
            "Report only what the text states. Use null for anything it does not state."
        ),
        "demo_input": (
            "This week's roundup. Alibaba published Qwen3-8B, a dense 8.2B model under "
            "Apache-2.0, with community GGUF builds at Q4_K_M and Q8_0. OpenAI released "
            "gpt-oss-20b, a mixture-of-experts model with 21B total parameters shipped with "
            "native MXFP4 weights, also Apache-2.0. A third release was mentioned on a forum "
            "but the licence was not given."
        ),
    },
    "classify": {
        "schema_model": Classification,
        "system": (
            "You classify an incoming message into exactly one of these labels: "
            + ", ".join(LABELS)
            + ". Choose 'other' when none of the rest fits."
        ),
        "demo_input": (
            "After I upgrade the server my long conversations start returning empty replies "
            "once they pass about eight thousand tokens. Nothing in the log looks wrong."
        ),
    },
}


def post_chat(base_url: str, api_key: Optional[str], payload: dict, timeout: int) -> dict:
    """One POST to /chat/completions. Raises RuntimeError with a readable message."""
    url = base_url.rstrip("/") + "/chat/completions"
    body = json.dumps(payload).encode("utf-8")
    headers = {"Content-Type": "application/json"}
    if api_key:
        headers["Authorization"] = f"Bearer {api_key}"
    request = urllib.request.Request(url, data=body, headers=headers, method="POST")
    try:
        with urllib.request.urlopen(request, timeout=timeout) as response:
            return json.loads(response.read().decode("utf-8"))
    except urllib.error.HTTPError as exc:
        detail = exc.read().decode("utf-8", "replace")[:500]
        raise RuntimeError(f"{url} returned HTTP {exc.code}: {detail}") from exc
    except urllib.error.URLError as exc:
        raise RuntimeError(f"cannot reach {url}: {exc.reason}") from exc


def build_payload(args, messages: list, schema_model) -> dict:
    payload = {
        "model": args.model,
        "messages": messages,
        "temperature": args.temperature,
        "max_tokens": args.max_tokens,
    }
    if args.seed is not None:
        payload["seed"] = args.seed
    if not args.free_form:
        payload["response_format"] = {
            "type": "json_schema",
            "json_schema": {
                "name": schema_model.__name__.lower(),
                "schema": schema_model.model_json_schema(),
                "strict": True,
            },
        }
    return payload


def run(args) -> dict:
    task = TASKS[args.task]
    schema_model = task["schema_model"]
    text = task["demo_input"] if args.demo else Path(args.input).read_text(encoding="utf-8")

    messages = [
        {"role": "system", "content": task["system"]},
        {"role": "user", "content": text},
    ]

    started = time.time()
    last_error = None
    for attempt in range(1, args.retries + 2):
        payload = build_payload(args, messages, schema_model)
        reply = post_chat(args.base_url, args.api_key, payload, args.timeout)
        content = reply["choices"][0]["message"]["content"] or ""
        try:
            parsed = schema_model.model_validate_json(content)
        except ValidationError as exc:
            last_error = f"validation failed: {exc.error_count()} error(s)"
            complaint = json.dumps(exc.errors(include_url=False)[:5], default=str)
        except json.JSONDecodeError as exc:
            last_error = f"not JSON at all: {exc}"
            complaint = str(exc)
        else:
            elapsed = time.time() - started
            print(json.dumps(parsed.model_dump(), indent=2))
            print(f"\nvalid on attempt {attempt} of {args.retries + 1}  ({elapsed:.1f} s)")
            return {
                "ok": True,
                "attempts": attempt,
                "seconds": round(elapsed, 2),
                "result": parsed.model_dump(),
            }

        # The repair turn: the model is shown its own output and the validator's complaint.
        print(f"attempt {attempt}: {last_error}", file=sys.stderr)
        messages = messages + [
            {"role": "assistant", "content": content},
            {
                "role": "user",
                "content": (
                    "That reply did not validate against the schema. The validator reported: "
                    f"{complaint}. Return the corrected JSON object and nothing else."
                ),
            },
        ]

    elapsed = time.time() - started
    print(f"gave up after {args.retries + 1} attempt(s): {last_error}", file=sys.stderr)
    return {"ok": False, "attempts": args.retries + 1, "seconds": round(elapsed, 2), "error": last_error}


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__.splitlines()[0])
    parser.add_argument("--base-url", default="http://127.0.0.1:8080/v1",
                        help="OpenAI-compatible base URL: llama-server, or the Part 9 gateway")
    parser.add_argument("--api-key", default=None, help="sent as a bearer token when set")
    parser.add_argument("--model", required=True, help="model name or alias the server answers to")
    parser.add_argument("--task", choices=sorted(TASKS), default="extract")
    parser.add_argument("--input", default=None, help="file to read; omit with --demo")
    parser.add_argument("--demo", action="store_true", help="use the built-in sample text")
    parser.add_argument("--free-form", action="store_true",
                        help="send no schema, so the reply is unconstrained (teaching mode)")
    parser.add_argument("--retries", type=int, default=2, help="repair attempts after the first")
    parser.add_argument("--temperature", type=float, default=0.0)
    parser.add_argument("--max-tokens", type=int, default=1024)
    parser.add_argument("--seed", type=int, default=None)
    parser.add_argument("--timeout", type=int, default=180)
    parser.add_argument("--labbook", default=None, help="append one JSON line per run to this file")
    args = parser.parse_args()

    if not args.demo and not args.input:
        parser.error("give --input FILE or --demo")
    if args.input and not Path(args.input).is_file():
        parser.error(f"--input {args.input} does not exist")

    outcome = run(args)

    if args.labbook:
        record = {
            "lab": "part-10/structured-extract",
            "task": args.task,
            "model": args.model,
            "base_url": args.base_url,
            "constrained": not args.free_form,
            "temperature": args.temperature,
            "seed": args.seed,
            "ok": outcome["ok"],
            "attempts": outcome["attempts"],
            "seconds": outcome["seconds"],
            "date": time.strftime("%Y-%m-%d"),
        }
        with Path(args.labbook).open("a", encoding="utf-8") as handle:
            handle.write(json.dumps(record) + "\n")
        print(f"recorded in {args.labbook}")

    sys.exit(0 if outcome["ok"] else 1)


if __name__ == "__main__":
    main()