#!/usr/bin/env python3
"""Structured JSON output: schema-constrained LLM generation + validation loop.

Demonstrates: schema-in-prompt, raw JSON extraction (removing markdown fences),
validation with Pydantic, and a bounded retry loop for repair.
"""
from __future__ import annotations

import json
import re
import sys
from typing import Any

try:
    from pydantic import BaseModel, Field, ValidationError
except ImportError:
    print("Install pydantic: pip install pydantic")
    sys.exit(1)


# --- 1. Define the output schema (Pydantic = single source of truth) ---

class ArticleSummary(BaseModel):
    title: str = Field(description="Article title, max 100 chars")
    key_points: list[str] = Field(
        min_length=1, max_length=5,
        description="1-5 key takeaways, each 10-200 chars"
    )
    sentiment: str = Field(
        pattern="^(positive|negative|neutral)$",
        description="Overall sentiment"
    )
    word_count: int = Field(ge=0, le=100000, description="Estimated word count")


# --- 2. Schema-in-prompt helper ---

def schema_prompt(schema_json: str) -> str:
    return f"""
Output MUST be valid JSON matching this schema exactly. No markdown, no commentary.
Return ONLY the JSON object.

Schema:
{schema_json}

Example:
{{"title": "Example Article", "key_points": ["Point 1", "Point 2"],
  "sentiment": "neutral", "word_count": 5000}}
"""


# --- 3. Extract raw JSON from model output (handles markdown fences) ---

def extract_json(text: str) -> str:
    """Strip markdown ```json fences and surrounding whitespace."""
    # Try to find JSON inside ```json ... ``` blocks
    m = re.search(r'```(?:json)?\s*\n?(.*?)\n?```', text, re.DOTALL)
    if m:
        return m.group(1).strip()
    # Try to find a JSON object
    m = re.search(r'\{.*\}', text, re.DOTALL)
    if m:
        return m.group(0).strip()
    return text.strip()


# --- 4. Validate + repair loop ---

def validate_or_repair(
    raw_output: str,
    schema: type[BaseModel],
    retries: int = 2,
) -> tuple[BaseModel | None, list[str]]:
    """
    Try to parse and validate raw LLM output.
    Returns (model, error_log). On failure, return error_log for the LLM to repair.
    """
    errors: list[str] = []
    for attempt in range(1 + retries):
        try:
            cleaned = extract_json(raw_output)
            data = json.loads(cleaned)
            model = schema.model_validate(data)
            return model, errors
        except json.JSONDecodeError as e:
            errors.append(f"Attempt {attempt}: JSON parse error at pos {e.pos}: {e.msg}")
        except ValidationError as e:
            errors.append(f"Attempt {attempt}: Validation error: {e}")
        except Exception as e:
            errors.append(f"Attempt {attempt}: {e}")
    return None, errors


# --- 5. Demo ---

def main() -> None:
    # Simulate LLM outputs — good, fixable, and broken
    test_outputs = [
        # Good: clean JSON
        '{"title": "LLMs in Production", "key_points": ["Scale matters", "Latency critical"], '
        '"sentiment": "positive", "word_count": 3000}',
        # Fixable: wrapped in markdown fences
        '```json\n{"title": "RAG Best Practices", '
        '"key_points": ["Chunking", "Retrieval"], '
        '"sentiment": "neutral", "word_count": 1500}\n```',
        # Broken: wrong sentiment enum
        '{"title": "Bad Example", "key_points": ["One"], "sentiment": "angry", "word_count": -1}',
    ]

    for i, raw in enumerate(test_outputs, 1):
        print(f"\n--- Test {i} ---")
        print(f"Raw: {raw[:80]}...")
        model, errors = validate_or_repair(raw, ArticleSummary)
        if model:
            print(f"  OK: title={model.title}, sentiment={model.sentiment}, "
                  f"points={len(model.key_points)}")
        else:
            print(f"  FAILED after retries:")
            for err in errors:
                print(f"    {err}")


if __name__ == "__main__":
    main()