import json
from typing import Any, Dict, List, Union

from litellm.constants import DEFAULT_MAX_RECURSE_DEPTH


def normalize_json_schema_types(schema: Union[Dict[str, Any], List[Any], Any], depth: int = 0, max_depth: int = DEFAULT_MAX_RECURSE_DEPTH) -> Union[Dict[str, Any], List[Any], Any]:
    """
    Normalize JSON schema types from uppercase to lowercase format.
    
    Some providers (like certain Google services) use uppercase types like 'BOOLEAN', 'STRING', 'ARRAY', 'OBJECT'
    but standard JSON Schema requires lowercase: 'boolean', 'string', 'array', 'object'
    
    This function recursively normalizes all type fields in a schema to lowercase.
    
    Args:
        schema: The schema to normalize (dict, list, or other)
        depth: Current recursion depth
        max_depth: Maximum recursion depth to prevent infinite loops
        
    Returns:
        The normalized schema with lowercase types
    """
    # Prevent infinite recursion
    if depth >= max_depth:
        return schema
        
    if not isinstance(schema, (dict, list)):
        return schema
        
    # Type mapping from uppercase to lowercase
    type_mapping = {
        'BOOLEAN': 'boolean',
        'STRING': 'string', 
        'ARRAY': 'array',
        'OBJECT': 'object',
        'NUMBER': 'number',
        'INTEGER': 'integer',
        'NULL': 'null'
    }
    
    if isinstance(schema, list):
        return [normalize_json_schema_types(item, depth + 1, max_depth) for item in schema]
    
    if isinstance(schema, dict):
        normalized_schema: Dict[str, Any] = {}
        
        for key, value in schema.items():
            if key == 'type' and isinstance(value, str) and value in type_mapping:
                normalized_schema[key] = type_mapping[value]
            elif key == 'properties' and isinstance(value, dict):
                # Recursively normalize properties
                normalized_schema[key] = {
                    prop_key: normalize_json_schema_types(prop_value, depth + 1, max_depth)
                    for prop_key, prop_value in value.items()
                }
            elif key == 'items' and isinstance(value, (dict, list)):
                # Recursively normalize array items
                normalized_schema[key] = normalize_json_schema_types(value, depth + 1, max_depth)
            elif isinstance(value, (dict, list)):
                # Recursively normalize any nested dict or list
                normalized_schema[key] = normalize_json_schema_types(value, depth + 1, max_depth)
            else:
                normalized_schema[key] = value
                
        return normalized_schema
    
    return schema


def normalize_tool_schema(tool: Dict[str, Any]) -> Dict[str, Any]:
    """
    Normalize a tool's parameter schema to use standard JSON Schema lowercase types.
    
    Args:
        tool: The tool definition containing function parameters
        
    Returns:
        The tool with normalized schema types
    """
    if not isinstance(tool, dict):
        return tool
        
    normalized_tool = tool.copy()
    
    # Normalize function parameters if present
    if 'function' in tool and isinstance(tool['function'], dict):
        normalized_tool['function'] = tool['function'].copy()
        if 'parameters' in tool['function']:
            normalized_tool['function']['parameters'] = normalize_json_schema_types(
                tool['function']['parameters']
            )
    
    return normalized_tool


def validate_schema(schema: dict, response: str):
    """
    Validate if the returned json response follows the schema.

    Params:
    - schema - dict: JSON schema
    - response - str: Received json response as string.
    """
    from jsonschema import ValidationError, validate

    from litellm import JSONSchemaValidationError

    try:
        response_dict = json.loads(response)
    except json.JSONDecodeError:
        raise JSONSchemaValidationError(
            model="", llm_provider="", raw_response=response, schema=json.dumps(schema)
        )

    try:
        validate(response_dict, schema=schema)
    except ValidationError:
        raise JSONSchemaValidationError(
            model="", llm_provider="", raw_response=response, schema=json.dumps(schema)
        )