"""
Structured Output Support - JSON Schema Generation from Signatures

Generates JSON schemas for reliable LLM outputs with OpenAI Structured Outputs API support.

Features:
- Automatic schema generation from Kaizen signatures
- OpenAI Structured Outputs API integration (strict mode)
- Legacy json_object format support
- 100% format compliance with strict mode (gpt-4o-2024-08-06+)
- Output validation against schema
- No regex parsing needed

Usage with BaseAgent:
    >>> from kaizen.core.base_agent import BaseAgent, BaseAgentConfig
    >>> from kaizen.signatures import Signature, InputField, OutputField
    >>> from kaizen.core.structured_output import create_structured_output_config
    >>>
    >>> class QASignature(Signature):
    ...     question: str = InputField(desc="User question")
    ...     answer: str = OutputField(desc="Answer")
    ...     confidence: float = OutputField(desc="Confidence 0-1")
    >>>
    >>> # Create structured output config
    >>> response_format = create_structured_output_config(QASignature(), strict=True)
    >>>
    >>> # Pass to BaseAgent via generation_config
    >>> config = BaseAgentConfig(
    ...     llm_provider="openai",
    ...     model="gpt-4o-2024-08-06",  # Required for strict mode
    ...     generation_config={"response_format": response_format}
    ... )
    >>> agent = BaseAgent(config=config, signature=QASignature())
    >>> result = agent.run(question="What is AI?")  # 100% schema compliance!

Supported Models (strict=True):
- gpt-4o-2024-08-06 and later
- gpt-4o-mini-2024-07-18 and later

Legacy Models (strict=False):
- All OpenAI models (best-effort compliance ~70-85%)
"""

import json
from typing import Any, Dict, List, Literal, Type, get_args, get_origin


class StructuredOutputGenerator:
    """
    Generates JSON schemas from Kaizen signatures.

    Supports OpenAI's structured output format for reliable JSON responses.
    """

    @staticmethod
    def signature_to_json_schema(signature: Any) -> Dict[str, Any]:
        """
        Convert signature to JSON schema.

        Args:
            signature: Kaizen Signature instance

        Returns:
            Dict: JSON schema for structured output

        Example:
            >>> schema = StructuredOutputGenerator.signature_to_json_schema(qa_signature)
            >>> # Use with OpenAI: model='gpt-4-turbo-preview', response_format={"type": "json_object", "schema": schema}
        """
        schema = {
            "type": "object",
            "properties": {},
            "required": [],
            "additionalProperties": False,
        }

        # Extract output fields from signature
        if hasattr(signature, "output_fields"):
            for field_name, field_info in signature.output_fields.items():
                field_type = field_info.get("type", str)
                # Support both 'desc' parameter and 'description' in metadata
                field_desc = field_info.get("desc", "") or field_info.get(
                    "metadata", {}
                ).get("description", "")

                # FIX BUG #1: Detect Literal types and extract enum values
                origin = get_origin(field_type)
                if origin is Literal:
                    # Extract allowed values from Literal["A", "B", "C"]
                    enum_values = list(get_args(field_type))
                    field_schema = {
                        "type": "string",
                        "description": field_desc,
                        "enum": enum_values,
                    }
                else:
                    # Map Python types to JSON schema types
                    json_type = StructuredOutputGenerator._python_type_to_json_type(
                        field_type
                    )

                    # Build field schema
                    field_schema = {
                        "type": json_type,
                        "description": field_desc,
                    }

                    # Fix 3: Add items property for arrays
                    if json_type == "array":
                        field_schema["items"] = {"type": "string"}

                    # Fix 4: Extract validation constraints from metadata
                    # Note: metadata is double-nested as metadata.metadata.validation
                    metadata = field_info.get("metadata", {}).get("metadata", {})
                    validation = metadata.get("validation", {})

                    # Add enum constraint if present (from metadata)
                    if "enum" in validation:
                        field_schema["enum"] = validation["enum"]

                    # Add numeric constraints if present
                    if "min" in validation:
                        field_schema["minimum"] = validation["min"]
                    if "max" in validation:
                        field_schema["maximum"] = validation["max"]

                    # Add pattern constraint if present
                    if "pattern" in validation:
                        field_schema["pattern"] = validation["pattern"]

                schema["properties"][field_name] = field_schema

                # All output fields are required
                schema["required"].append(field_name)

        return schema

    @staticmethod
    def _python_type_to_json_type(python_type: Type) -> str:
        """Map Python type to JSON schema type."""
        type_mapping = {
            str: "string",
            int: "integer",
            float: "number",
            bool: "boolean",
            list: "array",
            dict: "object",
        }

        return type_mapping.get(python_type, "string")

    @staticmethod
    def generate_system_prompt_with_schema(signature: Any) -> str:
        """
        Generate system prompt with embedded JSON schema.

        Args:
            signature: Kaizen Signature instance

        Returns:
            str: System prompt with schema instructions

        Example:
            >>> prompt = StructuredOutputGenerator.generate_system_prompt_with_schema(signature)
        """
        prompt_parts = []

        # Add signature description
        if hasattr(signature, "description") and signature.description:
            prompt_parts.append(signature.description)
        elif hasattr(signature, "name") and signature.name:
            prompt_parts.append(f"Task: {signature.name}")

        # Add input field descriptions
        if hasattr(signature, "input_fields") and signature.input_fields:
            prompt_parts.append("\nExpected Inputs:")
            for field_name, field_info in signature.input_fields.items():
                if isinstance(field_info, dict):
                    # Support both 'desc' parameter and 'description' in metadata
                    field_desc = field_info.get("desc", "") or field_info.get(
                        "metadata", {}
                    ).get("description", "")
                    if field_desc:
                        prompt_parts.append(f"  - {field_name}: {field_desc}")

        # Add output field descriptions with types
        if hasattr(signature, "output_fields") and signature.output_fields:
            prompt_parts.append("\nRequired Outputs:")
            for field_name, field_info in signature.output_fields.items():
                if isinstance(field_info, dict):
                    field_type = field_info.get("type", str).__name__
                    # Support both 'desc' parameter and 'description' in metadata
                    field_desc = field_info.get("desc", "") or field_info.get(
                        "metadata", {}
                    ).get("description", "")
                    prompt_parts.append(
                        f"  - {field_name} ({field_type}): {field_desc}"
                    )

        # Add JSON schema
        schema = StructuredOutputGenerator.signature_to_json_schema(signature)
        prompt_parts.append("\n---")
        prompt_parts.append(
            "\nYou MUST respond with a valid JSON object matching this exact schema:"
        )
        prompt_parts.append(f"```json\n{json.dumps(schema, indent=2)}\n```")
        prompt_parts.append("\nDo NOT include any text outside the JSON object.")
        prompt_parts.append(
            "Ensure all required fields are present with correct types."
        )

        return "\n".join(prompt_parts)

    @staticmethod
    def validate_output(
        output: Dict[str, Any], signature: Any
    ) -> tuple[bool, List[str]]:
        """
        Validate output against signature schema.

        Args:
            output: LLM output to validate
            signature: Kaizen Signature instance

        Returns:
            tuple: (is_valid, list of errors)

        Example:
            >>> is_valid, errors = StructuredOutputGenerator.validate_output(result, signature)
            >>> if not is_valid:
            ...     print(f"Validation errors: {errors}")
        """
        errors = []

        if not hasattr(signature, "output_fields"):
            return True, []

        # Check all required fields present
        for field_name, field_info in signature.output_fields.items():
            if field_name not in output:
                errors.append(f"Missing required field: {field_name}")
                continue

            # Check type
            expected_type = field_info.get("type", str)
            actual_value = output[field_name]

            if not isinstance(actual_value, expected_type):
                # Special case: int/float are interchangeable for numeric types
                if expected_type == float and isinstance(actual_value, int):
                    pass  # OK
                elif expected_type == int and isinstance(actual_value, float):
                    pass  # OK
                else:
                    errors.append(
                        f"Type mismatch for {field_name}: "
                        f"expected {expected_type.__name__}, got {type(actual_value).__name__}"
                    )

        return len(errors) == 0, errors


# Convenience function
def create_structured_output_config(
    signature: Any, strict: bool = True, name: str = "response"
) -> Dict[str, Any]:
    """
    Create OpenAI-compatible structured output configuration.

    Supports both OpenAI Structured Outputs formats:
    - New format (strict=True): 100% reliability with gpt-4o-2024-08-06+
    - Legacy format (strict=False): Best-effort with older models

    Args:
        signature: Kaizen Signature instance
        strict: Use strict mode (new format) for 100% reliability.
                Requires gpt-4o-2024-08-06 or gpt-4o-mini-2024-07-18+.
                Default: True
        name: Schema name for strict mode. Default: "response"

    Returns:
        Dict: Config for OpenAI API response_format parameter

    Example (Strict Mode - Recommended):
        >>> config = create_structured_output_config(qa_signature, strict=True)
        >>> # Use with: model="gpt-4o-2024-08-06", response_format=config
        >>> # Guaranteed 100% schema compliance

    Example (Legacy Mode):
        >>> config = create_structured_output_config(qa_signature, strict=False)
        >>> # Use with any model, best-effort schema compliance (~70-85%)
    """
    schema = StructuredOutputGenerator.signature_to_json_schema(signature)

    if strict:
        # New OpenAI Structured Outputs format (gpt-4o-2024-08-06+)
        # Provides 100% reliability via constrained sampling
        return {
            "type": "json_schema",
            "json_schema": {"name": name, "strict": True, "schema": schema},
        }
    else:
        # Legacy format (prompt-based, best-effort)
        return {"type": "json_object", "schema": schema}
