"""Schema normalization utilities for OpenAI strict mode compliance.

This module provides functions to prepare Pydantic schemas for OpenAI's strict mode
by fixing Pydantic v2 patterns and ensuring OpenAI requirements are met.
"""

from typing import Dict, Any


def normalize_pydantic_optional_fields(schema: Dict[str, Any]) -> Dict[str, Any]:
    """
    Fix Pydantic v2's Optional field schemas that use anyOf with [type, null].
    OpenAI requires every property to have a 'type' key directly.
    
    Pydantic v2 generates: {"anyOf": [{"type": "string"}, {"type": "null"}], "description": "..."}
    OpenAI expects: {"type": "string", "description": "..."}
    
    This function modifies the schema in-place and also returns it for convenience.
    
    Args:
        schema: JSON schema dict generated by Pydantic's model_json_schema()
        
    Returns:
        The same schema dict (modified in-place) for chaining
    """
    if isinstance(schema, dict):
        # Check if this is an Optional field with anyOf pattern
        if "anyOf" in schema and "type" not in schema:
            any_of = schema["anyOf"]
            # Look for pattern: [{"type": "<some_type>"}, {"type": "null"}]
            if isinstance(any_of, list) and len(any_of) == 2:
                type_schemas = [item for item in any_of if isinstance(item, dict) and item.get("type") != "null"]
                if len(type_schemas) == 1:
                    # Extract the actual type schema (not null)
                    actual_type_schema = type_schemas[0]
                    # Remove anyOf and merge the actual type
                    schema.pop("anyOf")
                    schema.update(actual_type_schema)
                    # Keep other fields like description, default, title
        
        # Recursively process nested schemas
        for key, value in list(schema.items()):
            if key == "properties" and isinstance(value, dict):
                for prop_value in value.values():
                    normalize_pydantic_optional_fields(prop_value)
            elif isinstance(value, dict):
                normalize_pydantic_optional_fields(value)
            elif isinstance(value, list):
                for item in value:
                    if isinstance(item, dict):
                        normalize_pydantic_optional_fields(item)
    
    return schema


def ensure_openai_strict_compliance(schema: Dict[str, Any]) -> Dict[str, Any]:
    """
    Apply OpenAI strict mode requirements to a JSON schema.
    
    OpenAI strict mode requires:
    1. All objects must have "additionalProperties": false
    2. All objects must have "required" array listing ALL properties (not just required fields)
    
    This function modifies the schema in-place and also returns it for convenience.
    
    Args:
        schema: JSON schema dict to make compliant
        
    Returns:
        The same schema dict (modified in-place) for chaining
    """
    if isinstance(schema, dict):
        if schema.get("type") == "object":
            # Requirement 1: No additional properties
            schema["additionalProperties"] = False
            
            # Requirement 2: All properties must be in required array
            if "properties" in schema:
                schema["required"] = list(schema["properties"].keys())
        
        # Recursively process nested schemas
        for key, value in schema.items():
            if key == "properties" and isinstance(value, dict):
                for prop_value in value.values():
                    ensure_openai_strict_compliance(prop_value)
            elif isinstance(value, dict):
                ensure_openai_strict_compliance(value)
            elif isinstance(value, list):
                for item in value:
                    if isinstance(item, dict):
                        ensure_openai_strict_compliance(item)
    
    return schema


def prepare_schema_for_openai(schema_dict: Dict[str, Any]) -> Dict[str, Any]:
    """
    Prepare a Pydantic-generated JSON schema for OpenAI strict mode.
    
    This is the main entry point that applies all necessary transformations:
    1. Fixes Pydantic v2 Optional field anyOf patterns
    2. Adds OpenAI strict mode requirements (additionalProperties, complete required arrays)
    
    Args:
        schema_dict: JSON schema dict from Pydantic's model_json_schema()
        
    Returns:
        The prepared schema dict ready for OpenAI strict mode
        
    Example:
        ```python
        from pydantic import BaseModel
        from ultragpt.schemas import prepare_schema_for_openai
        
        class MySchema(BaseModel):
            name: str
            age: Optional[int] = None
            
        raw_schema = MySchema.model_json_schema()
        prepared_schema = prepare_schema_for_openai(raw_schema)
        ```
    """
    # Make a copy to avoid modifying the input
    import copy
    schema = copy.deepcopy(schema_dict)
    
    # Apply transformations
    normalize_pydantic_optional_fields(schema)
    ensure_openai_strict_compliance(schema)
    
    return schema


def sanitize_tool_parameters_schema(raw_schema: Dict[str, Any]) -> Dict[str, Any]:
    """
    Sanitize a tool parameters schema for OpenAI's tool calling API.
    
    OpenAI's tool schema validator has strict requirements (separate from strict mode):
    1. Root must be {"type": "object", "properties": {...}, "required": [...], "additionalProperties": false}
    2. Every nested object MUST have additionalProperties: false
    3. Every nested object MUST have required array with ALL property keys
    4. Arrays MUST have "items" schema
    5. $ref CANNOT have sibling keywords (default, description, title, etc)
    6. No "default" keyword anywhere (causes 400 BadRequestError)
    7. anyOf/oneOf should be normalized to simple types where possible
    
    This function:
    - Ensures root is proper object schema
    - Fixes Pydantic Optional anyOf patterns (using normalize_pydantic_optional_fields)
    - Strips "default" keyword from all properties
    - If a dict has "$ref", keeps ONLY "$ref" (removes description, title, etc.)
    - Ensures all objects have additionalProperties: false
    - Ensures all objects have complete required arrays
    - Ensures all arrays have items schema
    - Does NOT add additionalProperties or force required at root (for backward compatibility)
    
    Use this for tool schemas before bind_tools(), not for output schemas.
    
    Args:
        raw_schema: JSON schema dict generated by Pydantic's model_json_schema()
        
    Returns:
        Sanitized schema dict ready for OpenAI tool calling
        
    Example:
        ```python
        from pydantic import BaseModel, Field
        from typing import Optional
        from ultragpt.schemas import sanitize_tool_parameters_schema
        
        class ToolParams(BaseModel):
            name: str
            config: Optional[ConfigModel] = Field(default=None, description="Config")
            
        raw_schema = ToolParams.model_json_schema()
        clean_schema = sanitize_tool_parameters_schema(raw_schema)
        # Now safe to use with OpenAI bind_tools()
        ```
    """
    import copy
    schema = copy.deepcopy(raw_schema)
    
    # Step 1: Ensure root is proper object schema
    if "type" not in schema:
        schema["type"] = "object"
    if schema.get("type") == "object":
        if "properties" not in schema:
            schema["properties"] = {}
        # Note: We don't force required/additionalProperties at root for backward compat
        # Individual providers or callers can add these if needed
    
    # Step 2: Fix Pydantic Optional anyOf patterns first
    normalize_pydantic_optional_fields(schema)
    
    # Step 3: Walk the schema and apply OpenAI tool-specific rules recursively
    def walk(node: Any, parent_key: str = "") -> None:
        if isinstance(node, dict):
            # Rule: If this node has $ref, keep ONLY $ref (strip all sibling keywords)
            if "$ref" in node:
                ref_val = node["$ref"]
                node.clear()
                node["$ref"] = ref_val
                return  # Don't recurse into $ref nodes
            
            # Rule: Remove "default" keyword (OpenAI rejects it in tool schemas)
            if "default" in node:
                node.pop("default", None)
            
            # Rule: If this is an object type, ensure it has required + additionalProperties
            if node.get("type") == "object":
                # Ensure properties exists
                if "properties" not in node:
                    node["properties"] = {}
                
                # Ensure required array exists and contains ALL properties
                # (OpenAI requires this even for "optional" fields)
                if node["properties"]:
                    node["required"] = list(node["properties"].keys())
                elif "required" not in node:
                    node["required"] = []
                
                # Ensure additionalProperties is false
                node["additionalProperties"] = False
            
            # Rule: If this is an array type, ensure it has items
            if node.get("type") == "array":
                if "items" not in node:
                    # Add a permissive items schema as fallback
                    node["items"] = {"type": "string"}
            
            # Recurse into nested structures
            for k, v in list(node.items()):
                walk(v, parent_key=k)
        
        elif isinstance(node, list):
            for i, item in enumerate(node):
                walk(item, parent_key=f"{parent_key}[{i}]")
    
    walk(schema)
    return schema


__all__ = [
    "normalize_pydantic_optional_fields",
    "ensure_openai_strict_compliance",
    "prepare_schema_for_openai",
    "sanitize_tool_parameters_schema",
]
