from typing import Any, Dict, List, Optional, Type

from litellm._logging import verbose_proxy_logger


class CustomOpenAPISpec:
    """
    Handler for customizing OpenAPI specifications with Pydantic models
    for documentation purposes without runtime validation.
    """
    
    CHAT_COMPLETION_PATHS = [
        "/v1/chat/completions",
        "/chat/completions", 
        "/engines/{model}/chat/completions",
        "/openai/deployments/{model}/chat/completions"
    ]
    
    EMBEDDING_PATHS = [
        "/v1/embeddings",
        "/embeddings",
        "/engines/{model}/embeddings", 
        "/openai/deployments/{model}/embeddings"
    ]
    
    RESPONSES_API_PATHS = [
        "/v1/responses",
        "/responses"
    ]
    
    @staticmethod
    def get_pydantic_schema(model_class) -> Optional[Dict[str, Any]]:
        """
        Get JSON schema from a Pydantic model, handling both v1 and v2 APIs.
        
        Args:
            model_class: Pydantic model class
            
        Returns:
            JSON schema dict or None if failed
        """
        try:
            # Try Pydantic v2 method first
            return model_class.model_json_schema()  # type: ignore
        except AttributeError:
            try:
                # Fallback to Pydantic v1 method
                return model_class.schema()  # type: ignore
            except AttributeError:
                # If both methods fail, return None
                return None
        except Exception as e:
            # FastAPI 0.120+ may fail schema generation for certain types (e.g., openai.Timeout)
            # Log the error and return None to skip schema generation for this model
            verbose_proxy_logger.debug(f"Failed to generate schema for {model_class}: {e}")
            return None
    
    @staticmethod
    def add_schema_to_components(openapi_schema: Dict[str, Any], schema_name: str, schema_def: Dict[str, Any]) -> None:
        """
        Add a schema definition to the OpenAPI components/schemas section.
        
        Args:
            openapi_schema: The OpenAPI schema dict to modify
            schema_name: Name for the schema component
            schema_def: The schema definition
        """
        # Ensure components/schemas structure exists
        if "components" not in openapi_schema:
            openapi_schema["components"] = {}
        if "schemas" not in openapi_schema["components"]:
            openapi_schema["components"]["schemas"] = {}
        
        # Add the schema
        CustomOpenAPISpec._move_defs_to_components(openapi_schema, {schema_name: schema_def})
    
    @staticmethod
    def add_request_body_to_paths(openapi_schema: Dict[str, Any], paths: List[str], schema_ref: str) -> None:
        """
        Add request body with expanded form fields for better Swagger UI display.
        This keeps the request body but expands it to show individual fields in the UI.
        
        Args:
            openapi_schema: The OpenAPI schema dict to modify
            paths: List of paths to update
            schema_ref: Reference to the schema component (e.g., "#/components/schemas/ModelName")
        """
        for path in paths:
            if path in openapi_schema.get("paths", {}) and "post" in openapi_schema["paths"][path]:
                # Get the actual schema to extract ALL field definitions
                schema_name = schema_ref.split("/")[-1]  # Extract "ProxyChatCompletionRequest" from the ref
                actual_schema = openapi_schema.get("components", {}).get("schemas", {}).get(schema_name, {})
                schema_properties = actual_schema.get("properties", {})
                required_fields = actual_schema.get("required", [])
                
                # Extract $defs and add them to components/schemas
                # This fixes Pydantic v2 $defs not being resolvable in Swagger/OpenAPI
                if "$defs" in actual_schema:
                    CustomOpenAPISpec._move_defs_to_components(openapi_schema, actual_schema["$defs"])
                
                # Create an expanded inline schema instead of just a $ref
                # This makes Swagger UI show all individual fields in the request body editor
                expanded_schema = {
                    "type": "object",
                    "required": required_fields,
                    "properties": {}
                }
                
                # Add all properties with their full definitions
                for field_name, field_def in schema_properties.items():
                    expanded_field = CustomOpenAPISpec._expand_field_definition(field_def)
                    
                    # Rewrite $defs references to use components/schemas instead
                    expanded_field = CustomOpenAPISpec._rewrite_defs_refs(expanded_field)
                    
                    # Add a simple example for the messages field
                    if field_name == "messages":
                        expanded_field["example"] = [
                            {"role": "user", "content": "Hello, how are you?"}
                        ]
                    
                    expanded_schema["properties"][field_name] = expanded_field
                
                # Set the request body with the expanded schema
                openapi_schema["paths"][path]["post"]["requestBody"] = {
                    "required": True,
                    "content": {
                        "application/json": {
                            "schema": expanded_schema
                        }
                    }
                }
                
                # Keep any existing parameters (like path parameters) but remove conflicting query params
                if "parameters" in openapi_schema["paths"][path]["post"]:
                    existing_params = openapi_schema["paths"][path]["post"]["parameters"]
                    # Only keep path parameters, remove query params that conflict with request body
                    filtered_params = [
                        param for param in existing_params 
                        if param.get("in") == "path"
                    ]
                    openapi_schema["paths"][path]["post"]["parameters"] = filtered_params
    
    @staticmethod
    def _move_defs_to_components(openapi_schema: Dict[str, Any], defs: Dict[str, Any]) -> None:
        """
        Move $defs from Pydantic v2 schema to OpenAPI components/schemas.
        This makes the definitions resolvable in Swagger/OpenAPI viewers.
        
        Args:
            openapi_schema: The OpenAPI schema dict to modify
            defs: The $defs dictionary from Pydantic schema
        """
        if not defs:
            return
        
        # Ensure components/schemas exists
        if "components" not in openapi_schema:
            openapi_schema["components"] = {}
        if "schemas" not in openapi_schema["components"]:
            openapi_schema["components"]["schemas"] = {}
        
        # Add each definition to components/schemas
        for def_name, def_schema in defs.items():
            # Recursively rewrite any nested $defs references within this definition
            rewritten_def = CustomOpenAPISpec._rewrite_defs_refs(def_schema)
            openapi_schema["components"]["schemas"][def_name] = rewritten_def
            
            # If this definition also has $defs, process them recursively
            if "$defs" in def_schema:
                CustomOpenAPISpec._move_defs_to_components(openapi_schema, def_schema["$defs"])
    
    @staticmethod
    def _rewrite_defs_refs(schema: Any) -> Any:
        """
        Recursively rewrite $ref values from #/$defs/... to #/components/schemas/...
        This converts Pydantic v2 references to OpenAPI-compatible references.
        
        Args:
            schema: Schema object to process (can be dict, list, or primitive)
            
        Returns:
            Schema with rewritten references
        """
        if isinstance(schema, dict):
            result = {}
            for key, value in schema.items():
                if key == "$ref" and isinstance(value, str) and value.startswith("#/$defs/"):
                    # Rewrite the reference to use components/schemas
                    def_name = value.replace("#/$defs/", "")
                    result[key] = f"#/components/schemas/{def_name}"
                elif key == "$defs":
                    # Remove $defs from the schema since they're moved to components
                    continue
                else:
                    # Recursively process nested structures
                    result[key] = CustomOpenAPISpec._rewrite_defs_refs(value)
            return result
        elif isinstance(schema, list):
            return [CustomOpenAPISpec._rewrite_defs_refs(item) for item in schema]
        else:
            return schema
    
    @staticmethod
    def _extract_field_schema(field_def: Dict[str, Any]) -> Dict[str, Any]:
        """
        Extract a simple schema from a Pydantic field definition for parameter display.
        
        Args:
            field_def: Pydantic field definition
            
        Returns:
            Simplified schema for OpenAPI parameter
        """
        # Handle simple types
        if "type" in field_def:
            return {"type": field_def["type"]}
        
        # Handle anyOf (Optional fields in Pydantic v2)
        if "anyOf" in field_def:
            any_of = field_def["anyOf"]
            # Find the non-null type
            for option in any_of:
                if option.get("type") != "null":
                    return option
            # Fallback to string if all else fails
            return {"type": "string"}
        
        # Default fallback
        return {"type": "string"}
    
    @staticmethod
    def _expand_field_definition(field_def: Dict[str, Any]) -> Dict[str, Any]:
        """
        Expand a Pydantic field definition for inline use in OpenAPI schema.
        This creates a full field definition that Swagger UI can render as individual form fields.
        
        Args:
            field_def: Pydantic field definition
            
        Returns:
            Expanded field definition for OpenAPI schema
        """
        # Return the field definition as-is since Pydantic already provides proper schemas
        return field_def.copy()
    
    @staticmethod
    def add_request_schema(
        openapi_schema: Dict[str, Any], 
        model_class: Type, 
        schema_name: str, 
        paths: List[str],
        operation_name: str
    ) -> Dict[str, Any]:
        """
        Generic method to add a request schema to OpenAPI specification.
        
        Args:
            openapi_schema: The OpenAPI schema dict to modify
            model_class: The Pydantic model class to get schema from
            schema_name: Name for the schema component
            paths: List of paths to add the request body to
            operation_name: Name of the operation for logging (e.g., "chat completion", "embedding")
            
        Returns:
            Modified OpenAPI schema
        """
        try:
            # Get the schema for the model class
            request_schema = CustomOpenAPISpec.get_pydantic_schema(model_class)
            
            # Only proceed if we successfully got the schema
            if request_schema is not None:
                # Add schema to components
                CustomOpenAPISpec.add_schema_to_components(openapi_schema, schema_name, request_schema)
                
                # Add request body to specified endpoints
                CustomOpenAPISpec.add_request_body_to_paths(
                    openapi_schema, 
                    paths, 
                    f"#/components/schemas/{schema_name}"
                )
                
                verbose_proxy_logger.debug(f"Successfully added {schema_name} schema to OpenAPI spec")
            else:
                verbose_proxy_logger.debug(f"Could not get schema for {schema_name}")
                
        except Exception as e:
            # If schema addition fails, continue without it
            verbose_proxy_logger.debug(f"Failed to add {operation_name} request schema: {str(e)}")
        
        return openapi_schema
    
    @staticmethod
    def add_chat_completion_request_schema(openapi_schema: Dict[str, Any]) -> Dict[str, Any]:
        """
        Add ProxyChatCompletionRequest schema to chat completion endpoints for documentation.
        This shows the request body in Swagger without runtime validation.
        
        Args:
            openapi_schema: The OpenAPI schema dict to modify
            
        Returns:
            Modified OpenAPI schema
        """
        try:
            from litellm.proxy._types import ProxyChatCompletionRequest
            
            return CustomOpenAPISpec.add_request_schema(
                openapi_schema=openapi_schema,
                model_class=ProxyChatCompletionRequest,
                schema_name="ProxyChatCompletionRequest",
                paths=CustomOpenAPISpec.CHAT_COMPLETION_PATHS,
                operation_name="chat completion"
            )
        except ImportError as e:
            verbose_proxy_logger.debug(f"Failed to import ProxyChatCompletionRequest: {str(e)}")
            return openapi_schema
    
    @staticmethod
    def add_embedding_request_schema(openapi_schema: Dict[str, Any]) -> Dict[str, Any]:
        """
        Add EmbeddingRequest schema to embedding endpoints for documentation.
        This shows the request body in Swagger without runtime validation.
        
        Args:
            openapi_schema: The OpenAPI schema dict to modify
            
        Returns:
            Modified OpenAPI schema
        """
        try:
            from litellm.types.embedding import EmbeddingRequest
            
            return CustomOpenAPISpec.add_request_schema(
                openapi_schema=openapi_schema,
                model_class=EmbeddingRequest,
                schema_name="EmbeddingRequest",
                paths=CustomOpenAPISpec.EMBEDDING_PATHS,
                operation_name="embedding"
            )
        except ImportError as e:
            verbose_proxy_logger.debug(f"Failed to import EmbeddingRequest: {str(e)}")
            return openapi_schema
    
    @staticmethod
    def add_responses_api_request_schema(openapi_schema: Dict[str, Any]) -> Dict[str, Any]:
        """
        Add ResponsesAPIRequestParams schema to responses API endpoints for documentation.
        This shows the request body in Swagger without runtime validation.
        
        Args:
            openapi_schema: The OpenAPI schema dict to modify
            
        Returns:
            Modified OpenAPI schema
        """
        try:
            from litellm.types.llms.openai import ResponsesAPIRequestParams
            
            return CustomOpenAPISpec.add_request_schema(
                openapi_schema=openapi_schema,
                model_class=ResponsesAPIRequestParams,
                schema_name="ResponsesAPIRequestParams",
                paths=CustomOpenAPISpec.RESPONSES_API_PATHS,
                operation_name="responses API"
            )
        except ImportError as e:
            verbose_proxy_logger.debug(f"Failed to import ResponsesAPIRequestParams: {str(e)}")
            return openapi_schema
    
    @staticmethod
    def add_llm_api_request_schema_body(openapi_schema: Dict[str, Any]) -> Dict[str, Any]:
        """
        Add LLM API request schema bodies to OpenAPI specification for documentation.
        
        Args:
            openapi_schema: The base OpenAPI schema
            
        Returns:
            OpenAPI schema with added request body schemas
        """
        # Add chat completion request schema
        openapi_schema = CustomOpenAPISpec.add_chat_completion_request_schema(openapi_schema)
        
        # Add embedding request schema
        openapi_schema = CustomOpenAPISpec.add_embedding_request_schema(openapi_schema)
        
        # Add responses API request schema
        openapi_schema = CustomOpenAPISpec.add_responses_api_request_schema(openapi_schema)
        
        return openapi_schema