# Familiar - Self-Hosted AI Companion Platform
# Copyright (c) 2026 George Scott Foley
# ORCID: 0009-0006-4957-0540
# Email: Georgescottfoley@proton.me
# Licensed under the MIT License - see LICENSE file for details

"""
Structured Outputs (v1.4.0 Phase 4 - Enhancement #17)

Guaranteed schema-compliant LLM outputs with validation and retry:
- Pydantic schema validation
- JSON extraction from responses
- Automatic retry on validation failure
- Schema-aware prompting
- Multiple output formats (JSON, XML, YAML)
- Partial/streaming structured outputs

Usage:
    from familiar.core import StructuredGenerator, OutputSchema
    from pydantic import BaseModel

    class Analysis(BaseModel):
        summary: str
        confidence: float
        topics: list[str]

    generator = StructuredGenerator(llm_provider)
    result = generator.generate(
        prompt="Analyze this text...",
        schema=Analysis,
        max_retries=3
    )
    # result is guaranteed to be a valid Analysis instance
"""

import asyncio
import json
import logging
import re
import threading
from dataclasses import dataclass, field
from datetime import datetime, timezone
from enum import Enum
from typing import (
    Any,
    Callable,
    Dict,
    Generic,
    List,
    Optional,
    Tuple,
    Type,
    TypeVar,
    Union,
    get_args,
    get_origin,
    get_type_hints,
)

logger = logging.getLogger(__name__)

# Type variable for schema types
T = TypeVar("T")

# Try to import pydantic
try:
    from pydantic import BaseModel, Field, ValidationError

    PYDANTIC_AVAILABLE = True
except ImportError:
    PYDANTIC_AVAILABLE = False
    BaseModel = None
    ValidationError = Exception
    Field = None


# ============================================================
# CONSTANTS
# ============================================================


def _utcnow() -> datetime:
    """Get current UTC time as timezone-aware datetime."""
    return datetime.now(timezone.utc)


# JSON extraction patterns
JSON_PATTERNS = [
    # Code blocks
    re.compile(r"```json\s*([\s\S]*?)\s*```", re.IGNORECASE),
    re.compile(r"```\s*([\s\S]*?)\s*```"),
    # Bare JSON objects/arrays
    re.compile(r"(\{[\s\S]*\})", re.MULTILINE),
    re.compile(r"(\[[\s\S]*\])", re.MULTILINE),
]


# ============================================================
# ENUMS
# ============================================================


class OutputFormat(str, Enum):
    """Supported output formats."""

    JSON = "json"
    XML = "xml"
    YAML = "yaml"
    MARKDOWN = "markdown"
    PLAIN = "plain"


class ValidationStrategy(str, Enum):
    """How to handle validation failures."""

    STRICT = "strict"  # Raise on first failure
    RETRY = "retry"  # Retry with feedback
    COERCE = "coerce"  # Try to fix/coerce values
    PARTIAL = "partial"  # Accept partial matches


class ExtractionMethod(str, Enum):
    """How to extract structured data from response."""

    JSON_BLOCK = "json_block"  # Look for ```json blocks
    FIRST_JSON = "first_json"  # Find first JSON object
    FULL_RESPONSE = "full"  # Parse entire response as JSON
    TAGGED = "tagged"  # Look for <output> tags


# ============================================================
# EXCEPTIONS
# ============================================================


class StructuredOutputError(Exception):
    """Base exception for structured output errors."""

    pass


class SchemaValidationError(StructuredOutputError):
    """Schema validation failed."""

    def __init__(
        self,
        message: str,
        errors: Optional[List[Dict[str, Any]]] = None,
        raw_output: Optional[str] = None,
    ):
        super().__init__(message)
        self.errors = errors or []
        self.raw_output = raw_output


class ExtractionError(StructuredOutputError):
    """Failed to extract structured data from response."""

    def __init__(self, message: str, raw_output: Optional[str] = None):
        super().__init__(message)
        self.raw_output = raw_output


class MaxRetriesExceeded(StructuredOutputError):
    """Maximum retries exceeded."""

    def __init__(
        self,
        message: str,
        attempts: int,
        last_error: Optional[Exception] = None,
    ):
        super().__init__(message)
        self.attempts = attempts
        self.last_error = last_error


# ============================================================
# SCHEMA UTILITIES
# ============================================================


def get_json_schema(model: Type) -> Dict[str, Any]:
    """
    Get JSON schema from a Pydantic model or type hints.

    Args:
        model: Pydantic model class or dataclass

    Returns:
        JSON Schema dict
    """
    if PYDANTIC_AVAILABLE and hasattr(model, "model_json_schema"):
        # Pydantic v2
        return model.model_json_schema()
    elif PYDANTIC_AVAILABLE and hasattr(model, "schema"):
        # Pydantic v1
        return model.schema()
    else:
        # Build schema from type hints
        return _build_schema_from_hints(model)


def _build_schema_from_hints(cls: Type) -> Dict[str, Any]:
    """Build JSON schema from type hints."""
    hints = get_type_hints(cls) if hasattr(cls, "__annotations__") else {}

    properties = {}
    required = []

    for name, hint in hints.items():
        prop_schema = _type_to_schema(hint)
        properties[name] = prop_schema

        # Check if required (no default)
        if hasattr(cls, name):
            default = getattr(cls, name, ...)
            if default is ...:
                required.append(name)
        else:
            required.append(name)

    return {
        "type": "object",
        "properties": properties,
        "required": required,
    }


def _type_to_schema(hint: Type) -> Dict[str, Any]:
    """Convert Python type hint to JSON schema."""
    origin = get_origin(hint)
    args = get_args(hint)

    # Handle Optional
    if origin is Union:
        non_none = [a for a in args if a is not type(None)]
        if len(non_none) == 1:
            return _type_to_schema(non_none[0])
        return {"anyOf": [_type_to_schema(a) for a in non_none]}

    # Handle List
    if origin is list:
        item_type = args[0] if args else Any
        return {"type": "array", "items": _type_to_schema(item_type)}

    # Handle Dict
    if origin is dict:
        return {"type": "object"}

    # Handle basic types
    type_map = {
        str: {"type": "string"},
        int: {"type": "integer"},
        float: {"type": "number"},
        bool: {"type": "boolean"},
        type(None): {"type": "null"},
    }

    if hint in type_map:
        return type_map[hint]

    # Handle Enum
    if isinstance(hint, type) and issubclass(hint, Enum):
        return {"type": "string", "enum": [e.value for e in hint]}

    # Default
    return {"type": "object"}


def schema_to_prompt_description(schema: Dict[str, Any], indent: int = 0) -> str:
    """
    Convert JSON schema to human-readable description for prompts.

    Args:
        schema: JSON Schema dict
        indent: Indentation level

    Returns:
        Human-readable schema description
    """
    prefix = "  " * indent
    lines = []

    schema_type = schema.get("type", "object")

    if schema_type == "object":
        properties = schema.get("properties", {})
        required = set(schema.get("required", []))

        for name, prop in properties.items():
            req_marker = "*" if name in required else ""
            prop_type = prop.get("type", "any")
            desc = prop.get("description", "")

            if prop_type == "array":
                items = prop.get("items", {})
                item_type = items.get("type", "any")
                lines.append(f"{prefix}- {name}{req_marker}: array of {item_type}")
            elif prop_type == "object":
                lines.append(f"{prefix}- {name}{req_marker}: object")
                if "properties" in prop:
                    lines.append(schema_to_prompt_description(prop, indent + 1))
            else:
                enum_values = prop.get("enum")
                if enum_values:
                    lines.append(f"{prefix}- {name}{req_marker}: one of {enum_values}")
                else:
                    lines.append(f"{prefix}- {name}{req_marker}: {prop_type}")

            if desc:
                lines.append(f"{prefix}  ({desc})")

    elif schema_type == "array":
        items = schema.get("items", {})
        lines.append(f"{prefix}Array of:")
        lines.append(schema_to_prompt_description(items, indent + 1))

    return "\n".join(lines)


# ============================================================
# JSON EXTRACTION
# ============================================================


def extract_json(
    text: str,
    method: ExtractionMethod = ExtractionMethod.FIRST_JSON,
) -> Tuple[Optional[Any], Optional[str]]:
    """
    Extract JSON from text response.

    Args:
        text: Raw text potentially containing JSON
        method: Extraction method to use

    Returns:
        Tuple of (parsed_json, error_message)
    """
    if not text or not text.strip():
        return None, "Empty response"

    text = text.strip()

    # Try based on method
    if method == ExtractionMethod.FULL_RESPONSE:
        try:
            return json.loads(text), None
        except json.JSONDecodeError as e:
            return None, f"JSON parse error: {e}"

    if method == ExtractionMethod.TAGGED:
        # Look for <output>...</output> or similar tags
        tag_patterns = [
            re.compile(r"<output>\s*([\s\S]*?)\s*</output>", re.IGNORECASE),
            re.compile(r"<json>\s*([\s\S]*?)\s*</json>", re.IGNORECASE),
            re.compile(r"<result>\s*([\s\S]*?)\s*</result>", re.IGNORECASE),
        ]

        for pattern in tag_patterns:
            match = pattern.search(text)
            if match:
                try:
                    return json.loads(match.group(1)), None
                except json.JSONDecodeError:
                    continue

    # JSON_BLOCK or FIRST_JSON - try patterns
    for pattern in JSON_PATTERNS:
        matches = pattern.findall(text)
        for match in matches:
            try:
                # Clean up the match
                cleaned = match.strip()
                if cleaned:
                    parsed = json.loads(cleaned)
                    return parsed, None
            except json.JSONDecodeError:
                continue

    # Last resort: try to find JSON-like content
    # Look for content between first { and last }
    first_brace = text.find("{")
    last_brace = text.rfind("}")

    if first_brace != -1 and last_brace != -1 and last_brace > first_brace:
        try:
            candidate = text[first_brace : last_brace + 1]
            return json.loads(candidate), None
        except json.JSONDecodeError:
            pass

    # Try array
    first_bracket = text.find("[")
    last_bracket = text.rfind("]")

    if first_bracket != -1 and last_bracket != -1 and last_bracket > first_bracket:
        try:
            candidate = text[first_bracket : last_bracket + 1]
            return json.loads(candidate), None
        except json.JSONDecodeError:
            pass

    return None, "No valid JSON found in response"


def repair_json(text: str) -> str:
    """
    Attempt to repair common JSON issues.

    Args:
        text: Potentially malformed JSON

    Returns:
        Repaired JSON string
    """
    # Remove trailing commas before } or ]
    text = re.sub(r",\s*([}\]])", r"\1", text)

    # Fix single quotes to double quotes (careful with nested)
    # Only if no double quotes present
    if '"' not in text and "'" in text:
        text = text.replace("'", '"')

    # Fix unquoted keys (simple cases)
    text = re.sub(r"(\{|,)\s*([a-zA-Z_][a-zA-Z0-9_]*)\s*:", r'\1"\2":', text)

    # Fix True/False/None to JSON equivalents
    text = re.sub(r"\bTrue\b", "true", text)
    text = re.sub(r"\bFalse\b", "false", text)
    text = re.sub(r"\bNone\b", "null", text)

    return text


# ============================================================
# VALIDATION
# ============================================================


@dataclass
class ValidationResult:
    """Result of schema validation."""

    valid: bool
    data: Optional[Any] = None
    errors: List[Dict[str, Any]] = field(default_factory=list)
    raw_output: Optional[str] = None

    @property
    def error_summary(self) -> str:
        """Get human-readable error summary."""
        if not self.errors:
            return ""

        summaries = []
        for error in self.errors[:5]:  # Limit to first 5
            loc = error.get("loc", [])
            msg = error.get("msg", "Unknown error")
            path = ".".join(str(part) for part in loc) if loc else "root"
            summaries.append(f"- {path}: {msg}")

        return "\n".join(summaries)


def validate_against_schema(
    data: Any,
    schema: Union[Type[T], Dict[str, Any]],
    coerce: bool = False,
) -> ValidationResult:
    """
    Validate data against a schema.

    Args:
        data: Data to validate
        schema: Pydantic model or JSON schema dict
        coerce: Whether to attempt type coercion

    Returns:
        ValidationResult with valid flag and any errors
    """
    # If schema is a Pydantic model
    if PYDANTIC_AVAILABLE and isinstance(schema, type) and issubclass(schema, BaseModel):
        try:
            if hasattr(schema, "model_validate"):
                # Pydantic v2
                instance = schema.model_validate(data)
            else:
                # Pydantic v1
                instance = schema(**data)

            return ValidationResult(valid=True, data=instance)

        except ValidationError as e:
            errors = []
            if hasattr(e, "errors"):
                errors = e.errors() if callable(e.errors) else e.errors

            return ValidationResult(
                valid=False,
                errors=errors,
                raw_output=str(data),
            )

    # JSON Schema validation (basic)
    if isinstance(schema, dict):
        errors = _validate_json_schema(data, schema)
        if errors:
            return ValidationResult(valid=False, errors=errors)
        return ValidationResult(valid=True, data=data)

    # Fallback - just check it's dict-like
    if isinstance(data, dict):
        return ValidationResult(valid=True, data=data)

    return ValidationResult(
        valid=False,
        errors=[{"loc": [], "msg": f"Expected object, got {type(data).__name__}"}],
    )


def _validate_json_schema(
    data: Any, schema: Dict[str, Any], path: List[str] = None
) -> List[Dict[str, Any]]:
    """Basic JSON schema validation."""
    path = path or []
    errors = []

    schema_type = schema.get("type")

    # Type check
    type_map = {
        "string": str,
        "integer": int,
        "number": (int, float),
        "boolean": bool,
        "array": list,
        "object": dict,
        "null": type(None),
    }

    if schema_type and schema_type in type_map:
        expected = type_map[schema_type]
        if not isinstance(data, expected):
            errors.append(
                {
                    "loc": path,
                    "msg": f"Expected {schema_type}, got {type(data).__name__}",
                    "type": "type_error",
                }
            )
            return errors

    # Object validation
    if schema_type == "object" and isinstance(data, dict):
        properties = schema.get("properties", {})
        required = set(schema.get("required", []))

        # Check required fields
        for req in required:
            if req not in data:
                errors.append(
                    {
                        "loc": path + [req],
                        "msg": "Field required",
                        "type": "missing",
                    }
                )

        # Validate each property
        for key, value in data.items():
            if key in properties:
                errors.extend(_validate_json_schema(value, properties[key], path + [key]))

    # Array validation
    if schema_type == "array" and isinstance(data, list):
        items_schema = schema.get("items", {})
        for i, item in enumerate(data):
            errors.extend(_validate_json_schema(item, items_schema, path + [str(i)]))

    # Enum validation
    if "enum" in schema:
        if data not in schema["enum"]:
            errors.append(
                {
                    "loc": path,
                    "msg": f"Value must be one of {schema['enum']}",
                    "type": "enum",
                }
            )

    return errors


# ============================================================
# PROMPT BUILDING
# ============================================================


def build_structured_prompt(
    prompt: str,
    schema: Union[Type, Dict[str, Any]],
    format: OutputFormat = OutputFormat.JSON,
    examples: Optional[List[Dict[str, Any]]] = None,
) -> str:
    """
    Build a prompt that instructs the LLM to output structured data.

    Args:
        prompt: Base prompt/instruction
        schema: Schema to follow
        format: Output format
        examples: Optional examples

    Returns:
        Enhanced prompt with schema instructions
    """
    parts = [prompt]

    # Get schema description
    if isinstance(schema, type):
        json_schema = get_json_schema(schema)
    else:
        json_schema = schema

    schema_desc = schema_to_prompt_description(json_schema)

    # Add format instructions
    parts.append(f"\n\nRespond with valid {format.value.upper()} matching this schema:")
    parts.append(f"\n{schema_desc}")

    # Add examples if provided
    if examples:
        parts.append("\n\nExamples:")
        for i, example in enumerate(examples[:3], 1):
            parts.append(f"\nExample {i}:")
            parts.append(f"```json\n{json.dumps(example, indent=2)}\n```")

    # Add format reminder
    if format == OutputFormat.JSON:
        parts.append("\n\nRespond ONLY with the JSON object, no additional text.")
        parts.append("Use ```json code blocks.")

    return "".join(parts)


def build_retry_prompt(
    original_prompt: str,
    previous_output: str,
    validation_errors: str,
) -> str:
    """
    Build a prompt for retry after validation failure.

    Args:
        original_prompt: The original prompt
        previous_output: What the LLM returned
        validation_errors: What was wrong with it

    Returns:
        Retry prompt with error feedback
    """
    return f"""{original_prompt}

Your previous response had validation errors:
{validation_errors}

Previous output:
{previous_output[:500]}{"..." if len(previous_output) > 500 else ""}

Please fix these errors and respond with valid JSON matching the schema exactly."""


# ============================================================
# STRUCTURED GENERATOR
# ============================================================


@dataclass
class GenerationResult(Generic[T]):
    """Result of structured generation."""

    success: bool
    data: Optional[T] = None
    raw_output: Optional[str] = None
    attempts: int = 1
    errors: List[str] = field(default_factory=list)
    elapsed_ms: float = 0.0

    def __bool__(self) -> bool:
        return self.success

    def unwrap(self) -> T:
        """Get data or raise if failed."""
        if not self.success or self.data is None:
            raise SchemaValidationError(
                f"Generation failed after {self.attempts} attempts: {self.errors[-1] if self.errors else 'Unknown error'}"
            )
        return self.data


class StructuredGenerator:
    """
    Generates structured outputs with schema validation.

    Features:
    - Automatic JSON extraction from LLM responses
    - Schema validation with Pydantic or JSON Schema
    - Retry with error feedback
    - Multiple extraction methods

    Usage:
        from pydantic import BaseModel

        class Analysis(BaseModel):
            summary: str
            confidence: float
            topics: list[str]

        generator = StructuredGenerator(llm_provider)
        result = generator.generate(
            prompt="Analyze: The stock market rose 2%...",
            schema=Analysis,
        )

        if result.success:
            print(result.data.summary)
            print(result.data.confidence)
    """

    def __init__(
        self,
        llm_provider: Callable,
        default_model: Optional[str] = None,
        max_retries: int = 3,
        extraction_method: ExtractionMethod = ExtractionMethod.FIRST_JSON,
        validation_strategy: ValidationStrategy = ValidationStrategy.RETRY,
        repair_json: bool = True,
    ):
        """
        Initialize generator.

        Args:
            llm_provider: Function that calls LLM (messages, **kwargs) -> response
            default_model: Default model to use
            max_retries: Maximum retry attempts
            extraction_method: How to extract JSON from responses
            validation_strategy: How to handle validation failures
            repair_json: Whether to attempt JSON repair
        """
        self.llm_provider = llm_provider
        self.default_model = default_model
        self.max_retries = max_retries
        self.extraction_method = extraction_method
        self.validation_strategy = validation_strategy
        self.repair_json_enabled = repair_json

        self._is_async = asyncio.iscoroutinefunction(llm_provider)

    def generate(
        self,
        prompt: str,
        schema: Union[Type[T], Dict[str, Any]],
        examples: Optional[List[Dict[str, Any]]] = None,
        max_retries: Optional[int] = None,
        model: Optional[str] = None,
        **llm_kwargs,
    ) -> GenerationResult[T]:
        """
        Generate structured output (sync).

        Args:
            prompt: The prompt/instruction
            schema: Pydantic model or JSON schema
            examples: Optional examples for few-shot
            max_retries: Override default max retries
            model: Override default model
            **llm_kwargs: Additional LLM arguments

        Returns:
            GenerationResult with parsed data
        """
        if self._is_async:
            return asyncio.run(
                self.generate_async(
                    prompt=prompt,
                    schema=schema,
                    examples=examples,
                    max_retries=max_retries,
                    model=model,
                    **llm_kwargs,
                )
            )

        return self._generate_sync(
            prompt=prompt,
            schema=schema,
            examples=examples,
            max_retries=max_retries,
            model=model,
            **llm_kwargs,
        )

    async def generate_async(
        self,
        prompt: str,
        schema: Union[Type[T], Dict[str, Any]],
        examples: Optional[List[Dict[str, Any]]] = None,
        max_retries: Optional[int] = None,
        model: Optional[str] = None,
        **llm_kwargs,
    ) -> GenerationResult[T]:
        """
        Generate structured output (async).

        Args:
            prompt: The prompt/instruction
            schema: Pydantic model or JSON schema
            examples: Optional examples for few-shot
            max_retries: Override default max retries
            model: Override default model
            **llm_kwargs: Additional LLM arguments

        Returns:
            GenerationResult with parsed data
        """
        import time

        start_time = time.time()

        max_retries = max_retries if max_retries is not None else self.max_retries
        model = model or self.default_model

        # Build structured prompt
        structured_prompt = build_structured_prompt(
            prompt=prompt,
            schema=schema,
            examples=examples,
        )

        messages = [{"role": "user", "content": structured_prompt}]
        errors = []
        last_output = ""

        for attempt in range(max_retries + 1):
            try:
                # Call LLM
                if self._is_async:
                    response = await self.llm_provider(messages=messages, model=model, **llm_kwargs)
                else:
                    response = self.llm_provider(messages=messages, model=model, **llm_kwargs)

                # Extract text from response
                raw_output = self._extract_text(response)
                last_output = raw_output

                # Extract JSON
                json_data, extract_error = extract_json(raw_output, self.extraction_method)

                if json_data is None:
                    # Try repair if enabled
                    if self.repair_json_enabled:
                        repaired = repair_json(raw_output)
                        json_data, extract_error = extract_json(
                            repaired, ExtractionMethod.FIRST_JSON
                        )

                    if json_data is None:
                        errors.append(f"Extraction failed: {extract_error}")

                        if self.validation_strategy == ValidationStrategy.STRICT:
                            break

                        if attempt >= max_retries:
                            break

                        # Build retry prompt
                        messages = [
                            {
                                "role": "user",
                                "content": build_retry_prompt(
                                    structured_prompt,
                                    raw_output,
                                    f"Could not extract valid JSON: {extract_error}",
                                ),
                            }
                        ]
                        continue

                # Validate against schema
                validation = validate_against_schema(json_data, schema)

                if validation.valid:
                    elapsed_ms = (time.time() - start_time) * 1000
                    return GenerationResult(
                        success=True,
                        data=validation.data,
                        raw_output=raw_output,
                        attempts=attempt + 1,
                        elapsed_ms=elapsed_ms,
                    )

                # Validation failed
                errors.append(f"Validation: {validation.error_summary}")

                if self.validation_strategy == ValidationStrategy.STRICT:
                    break

                if self.validation_strategy == ValidationStrategy.COERCE:
                    # Try to coerce the data
                    coerced = self._coerce_data(json_data, schema)
                    if coerced is not None:
                        elapsed_ms = (time.time() - start_time) * 1000
                        return GenerationResult(
                            success=True,
                            data=coerced,
                            raw_output=raw_output,
                            attempts=attempt + 1,
                            elapsed_ms=elapsed_ms,
                        )

                # Build retry prompt
                messages = [
                    {
                        "role": "user",
                        "content": build_retry_prompt(
                            structured_prompt, raw_output, validation.error_summary
                        ),
                    }
                ]

            except Exception as e:
                errors.append(f"Error: {str(e)}")
                if self.validation_strategy == ValidationStrategy.STRICT:
                    break

        # All retries exhausted
        elapsed_ms = (time.time() - start_time) * 1000
        return GenerationResult(
            success=False,
            raw_output=last_output,
            attempts=attempt + 1,
            errors=errors,
            elapsed_ms=elapsed_ms,
        )

    def _generate_sync(
        self,
        prompt: str,
        schema: Union[Type[T], Dict[str, Any]],
        examples: Optional[List[Dict[str, Any]]] = None,
        max_retries: Optional[int] = None,
        model: Optional[str] = None,
        **llm_kwargs,
    ) -> GenerationResult[T]:
        """Sync generation implementation."""
        import time

        start_time = time.time()

        max_retries = max_retries if max_retries is not None else self.max_retries
        model = model or self.default_model

        # Build structured prompt
        structured_prompt = build_structured_prompt(
            prompt=prompt,
            schema=schema,
            examples=examples,
        )

        messages = [{"role": "user", "content": structured_prompt}]
        errors = []
        last_output = ""

        for attempt in range(max_retries + 1):
            try:
                # Call LLM
                response = self.llm_provider(messages=messages, model=model, **llm_kwargs)

                # Extract text from response
                raw_output = self._extract_text(response)
                last_output = raw_output

                # Extract JSON
                json_data, extract_error = extract_json(raw_output, self.extraction_method)

                if json_data is None and self.repair_json_enabled:
                    repaired = repair_json(raw_output)
                    json_data, extract_error = extract_json(repaired, ExtractionMethod.FIRST_JSON)

                if json_data is None:
                    errors.append(f"Extraction: {extract_error}")

                    if self.validation_strategy == ValidationStrategy.STRICT:
                        break

                    if attempt >= max_retries:
                        break

                    messages = [
                        {
                            "role": "user",
                            "content": build_retry_prompt(
                                structured_prompt,
                                raw_output,
                                f"Could not extract valid JSON: {extract_error}",
                            ),
                        }
                    ]
                    continue

                # Validate
                validation = validate_against_schema(json_data, schema)

                if validation.valid:
                    elapsed_ms = (time.time() - start_time) * 1000
                    return GenerationResult(
                        success=True,
                        data=validation.data,
                        raw_output=raw_output,
                        attempts=attempt + 1,
                        elapsed_ms=elapsed_ms,
                    )

                errors.append(f"Validation: {validation.error_summary}")

                if self.validation_strategy == ValidationStrategy.STRICT:
                    break

                messages = [
                    {
                        "role": "user",
                        "content": build_retry_prompt(
                            structured_prompt, raw_output, validation.error_summary
                        ),
                    }
                ]

            except Exception as e:
                errors.append(f"Error: {str(e)}")
                if self.validation_strategy == ValidationStrategy.STRICT:
                    break

        elapsed_ms = (time.time() - start_time) * 1000
        return GenerationResult(
            success=False,
            raw_output=last_output,
            attempts=attempt + 1,
            errors=errors,
            elapsed_ms=elapsed_ms,
        )

    def _extract_text(self, response: Any) -> str:
        """Extract text content from LLM response."""
        if isinstance(response, str):
            return response

        if isinstance(response, dict):
            # Common response formats
            if "content" in response:
                content = response["content"]
                if isinstance(content, str):
                    return content
                if isinstance(content, list):
                    # Anthropic format
                    texts = [
                        c.get("text", "")
                        for c in content
                        if isinstance(c, dict) and c.get("type") == "text"
                    ]
                    return "".join(texts)
            if "text" in response:
                return response["text"]
            if "message" in response:
                return self._extract_text(response["message"])

        # Has .text attribute
        if hasattr(response, "text"):
            return response.text

        # Has .content attribute
        if hasattr(response, "content"):
            content = response.content
            if isinstance(content, str):
                return content
            if isinstance(content, list):
                texts = [c.text if hasattr(c, "text") else str(c) for c in content]
                return "".join(texts)

        return str(response)

    def _coerce_data(
        self, data: Dict[str, Any], schema: Union[Type[T], Dict[str, Any]]
    ) -> Optional[T]:
        """Attempt to coerce data to match schema."""
        if not isinstance(data, dict):
            return None

        if not PYDANTIC_AVAILABLE:
            return None

        if isinstance(schema, type) and issubclass(schema, BaseModel):
            try:
                # Get schema info
                if hasattr(schema, "model_fields"):
                    fields = schema.model_fields
                else:
                    fields = schema.__fields__

                coerced = {}
                for field_name, field_info in fields.items():
                    if field_name in data:
                        value = data[field_name]
                        # Try basic type coercion
                        coerced[field_name] = self._coerce_value(value, field_info)
                    elif hasattr(field_info, "default") and field_info.default is not None:
                        coerced[field_name] = field_info.default

                # Try to create instance
                if hasattr(schema, "model_validate"):
                    return schema.model_validate(coerced)
                else:
                    return schema(**coerced)

            except Exception:
                pass

        return None

    def _coerce_value(self, value: Any, field_info: Any) -> Any:
        """Coerce a single value based on field info."""
        # Get target type
        if hasattr(field_info, "annotation"):
            target_type = field_info.annotation
        elif hasattr(field_info, "outer_type_"):
            target_type = field_info.outer_type_
        else:
            return value

        # Basic coercions
        if target_type is str and not isinstance(value, str):
            return str(value)
        if target_type is int and isinstance(value, (float, str)):
            try:
                return int(float(value))
            except (ValueError, TypeError):
                pass
        if target_type is float and isinstance(value, (int, str)):
            try:
                return float(value)
            except (ValueError, TypeError):
                pass
        if target_type is bool:
            if isinstance(value, str):
                return value.lower() in ("true", "yes", "1")
            return bool(value)

        return value


# ============================================================
# CONVENIENCE FUNCTIONS
# ============================================================

# Global generator instance
_generator: Optional[StructuredGenerator] = None
_generator_lock = threading.Lock()


def get_structured_generator() -> Optional[StructuredGenerator]:
    """Get the global StructuredGenerator instance."""
    with _generator_lock:
        return _generator


def set_structured_generator(generator: StructuredGenerator):
    """Set the global StructuredGenerator instance."""
    global _generator
    with _generator_lock:
        _generator = generator


def reset_structured_generator():
    """Reset the global StructuredGenerator instance."""
    global _generator
    with _generator_lock:
        _generator = None


def create_structured_generator(
    llm_provider: Callable, max_retries: int = 3, **kwargs
) -> StructuredGenerator:
    """Create a configured StructuredGenerator."""
    return StructuredGenerator(llm_provider=llm_provider, max_retries=max_retries, **kwargs)


def generate_structured(
    prompt: str,
    schema: Union[Type[T], Dict[str, Any]],
    llm_provider: Optional[Callable] = None,
    **kwargs,
) -> GenerationResult[T]:
    """
    Convenience function for one-off structured generation.

    Args:
        prompt: The prompt
        schema: Schema to validate against
        llm_provider: LLM provider (uses global if not provided)
        **kwargs: Additional arguments

    Returns:
        GenerationResult
    """
    if llm_provider is None:
        generator = get_structured_generator()
        if generator is None:
            raise ValueError("No LLM provider specified and no global generator set")
    else:
        generator = StructuredGenerator(llm_provider)

    return generator.generate(prompt, schema, **kwargs)


# ============================================================
# OUTPUT SCHEMA HELPER
# ============================================================


class OutputSchema:
    """
    Helper for defining output schemas without Pydantic.

    Usage:
        schema = OutputSchema({
            "summary": {"type": "string", "description": "Brief summary"},
            "score": {"type": "number", "minimum": 0, "maximum": 100},
            "tags": {"type": "array", "items": {"type": "string"}},
        }, required=["summary", "score"])
    """

    def __init__(
        self,
        properties: Dict[str, Dict[str, Any]],
        required: Optional[List[str]] = None,
        description: Optional[str] = None,
    ):
        self.properties = properties
        self.required = required or list(properties.keys())
        self.description = description

    def to_json_schema(self) -> Dict[str, Any]:
        """Convert to JSON Schema."""
        schema = {
            "type": "object",
            "properties": self.properties,
            "required": self.required,
        }
        if self.description:
            schema["description"] = self.description
        return schema

    def validate(self, data: Any) -> ValidationResult:
        """Validate data against this schema."""
        return validate_against_schema(data, self.to_json_schema())


# ============================================================
# EXPORTS
# ============================================================

__all__ = [
    # Core classes
    "StructuredGenerator",
    "GenerationResult",
    "ValidationResult",
    "OutputSchema",
    # Enums
    "OutputFormat",
    "ValidationStrategy",
    "ExtractionMethod",
    # Exceptions
    "StructuredOutputError",
    "SchemaValidationError",
    "ExtractionError",
    "MaxRetriesExceeded",
    # Schema utilities
    "get_json_schema",
    "schema_to_prompt_description",
    # JSON utilities
    "extract_json",
    "repair_json",
    # Validation
    "validate_against_schema",
    # Prompt building
    "build_structured_prompt",
    "build_retry_prompt",
    # Convenience functions
    "get_structured_generator",
    "set_structured_generator",
    "reset_structured_generator",
    "create_structured_generator",
    "generate_structured",
    # Constants
    "PYDANTIC_AVAILABLE",
]

# Compatibility alias
StructuredOutputParser = OutputSchema
