# -*- coding: utf-8 -*-
"""MCP Gateway Schema Definitions.

Copyright 2025
SPDX-License-Identifier: Apache-2.0
Authors: Mihai Criveti

This module provides Pydantic models for request/response validation in the MCP Gateway.
It implements schemas for:
- Tool registration and invocation
- Resource management and subscriptions
- Prompt templates and arguments
- Gateway federation
- RPC message formats
- Event messages
- Admin interface

The schemas ensure proper validation according to the MCP specification while adding
gateway-specific extensions for federation support.
"""

# Standard
import base64
from datetime import datetime, timezone
import json
import logging
from typing import Any, Dict, List, Literal, Optional, Union

# Third-Party
from pydantic import AnyHttpUrl, BaseModel, ConfigDict, Field, field_serializer, field_validator, model_validator, ValidationInfo

# First-Party
from mcpgateway.models import ImageContent
from mcpgateway.models import Prompt as MCPPrompt
from mcpgateway.models import Resource as MCPResource
from mcpgateway.models import ResourceContent, TextContent
from mcpgateway.models import Tool as MCPTool
from mcpgateway.utils.services_auth import decode_auth, encode_auth

logger = logging.getLogger(__name__)


def to_camel_case(s: str) -> str:
    """
    Convert a string from snake_case to camelCase.

    Args:
        s (str): The string to be converted, which is assumed to be in snake_case.

    Returns:
        str: The string converted to camelCase.

    Example:
        >>> to_camel_case("hello_world_example")
        'helloWorldExample'
    """
    return "".join(word.capitalize() if i else word for i, word in enumerate(s.split("_")))


def encode_datetime(v: datetime) -> str:
    """
    Convert a datetime object to an ISO 8601 formatted string.

    Args:
        v (datetime): The datetime object to be encoded.

    Returns:
        str: The ISO 8601 formatted string representation of the datetime object.

    Example:
        >>> encode_datetime(datetime(2023, 5, 22, 14, 30, 0))
        '2023-05-22T14:30:00'
    """
    return v.isoformat()


# --- Base Model ---
class BaseModelWithConfigDict(BaseModel):
    """Base model with common configuration.

    Provides:
    - ORM mode for SQLAlchemy integration
    - JSON encoders for datetime handling
    - Automatic conversion from snake_case to camelCase for output
    """

    model_config = ConfigDict(
        from_attributes=True,
        alias_generator=to_camel_case,
        populate_by_name=True,
        use_enum_values=True,
        extra="ignore",
        json_schema_extra={"nullable": True},
    )

    def to_dict(self, use_alias: bool = False) -> Dict[str, Any]:
        """
        Converts the model instance into a dictionary representation.

        Args:
            use_alias (bool): Whether to use aliases for field names (default is False). If True,
                               field names will be converted using the alias generator function.

        Returns:
            Dict[str, Any]: A dictionary where keys are field names and values are corresponding field values,
                             with any nested models recursively converted to dictionaries.
        """
        output = {}
        for key, value in self.dict(by_alias=use_alias).items():
            output[key] = value if not isinstance(value, BaseModel) else value.to_dict(use_alias)
        return output


# --- Metrics Schemas ---


class ToolMetrics(BaseModelWithConfigDict):
    """
    Represents the performance and execution statistics for a tool.

    Attributes:
        total_executions (int): Total number of tool invocations.
        successful_executions (int): Number of successful tool invocations.
        failed_executions (int): Number of failed tool invocations.
        failure_rate (float): Failure rate (failed invocations / total invocations).
        min_response_time (Optional[float]): Minimum response time in seconds.
        max_response_time (Optional[float]): Maximum response time in seconds.
        avg_response_time (Optional[float]): Average response time in seconds.
        last_execution_time (Optional[datetime]): Timestamp of the most recent invocation.
    """

    total_executions: int = Field(..., description="Total number of tool invocations")
    successful_executions: int = Field(..., description="Number of successful tool invocations")
    failed_executions: int = Field(..., description="Number of failed tool invocations")
    failure_rate: float = Field(..., description="Failure rate (failed invocations / total invocations)")
    min_response_time: Optional[float] = Field(None, description="Minimum response time in seconds")
    max_response_time: Optional[float] = Field(None, description="Maximum response time in seconds")
    avg_response_time: Optional[float] = Field(None, description="Average response time in seconds")
    last_execution_time: Optional[datetime] = Field(None, description="Timestamp of the most recent invocation")


class ResourceMetrics(BaseModelWithConfigDict):
    """
    Represents the performance and execution statistics for a resource.

    Attributes:
        total_executions (int): Total number of resource invocations.
        successful_executions (int): Number of successful resource invocations.
        failed_executions (int): Number of failed resource invocations.
        failure_rate (float): Failure rate (failed invocations / total invocations).
        min_response_time (Optional[float]): Minimum response time in seconds.
        max_response_time (Optional[float]): Maximum response time in seconds.
        avg_response_time (Optional[float]): Average response time in seconds.
        last_execution_time (Optional[datetime]): Timestamp of the most recent invocation.
    """

    total_executions: int = Field(..., description="Total number of resource invocations")
    successful_executions: int = Field(..., description="Number of successful resource invocations")
    failed_executions: int = Field(..., description="Number of failed resource invocations")
    failure_rate: float = Field(..., description="Failure rate (failed invocations / total invocations)")
    min_response_time: Optional[float] = Field(None, description="Minimum response time in seconds")
    max_response_time: Optional[float] = Field(None, description="Maximum response time in seconds")
    avg_response_time: Optional[float] = Field(None, description="Average response time in seconds")
    last_execution_time: Optional[datetime] = Field(None, description="Timestamp of the most recent invocation")


class ServerMetrics(BaseModelWithConfigDict):
    """
    Represents the performance and execution statistics for a server.

    Attributes:
        total_executions (int): Total number of server invocations.
        successful_executions (int): Number of successful server invocations.
        failed_executions (int): Number of failed server invocations.
        failure_rate (float): Failure rate (failed invocations / total invocations).
        min_response_time (Optional[float]): Minimum response time in seconds.
        max_response_time (Optional[float]): Maximum response time in seconds.
        avg_response_time (Optional[float]): Average response time in seconds.
        last_execution_time (Optional[datetime]): Timestamp of the most recent invocation.
    """

    total_executions: int = Field(..., description="Total number of server invocations")
    successful_executions: int = Field(..., description="Number of successful server invocations")
    failed_executions: int = Field(..., description="Number of failed server invocations")
    failure_rate: float = Field(..., description="Failure rate (failed invocations / total invocations)")
    min_response_time: Optional[float] = Field(None, description="Minimum response time in seconds")
    max_response_time: Optional[float] = Field(None, description="Maximum response time in seconds")
    avg_response_time: Optional[float] = Field(None, description="Average response time in seconds")
    last_execution_time: Optional[datetime] = Field(None, description="Timestamp of the most recent invocation")


class PromptMetrics(BaseModelWithConfigDict):
    """
    Represents the performance and execution statistics for a prompt.

    Attributes:
        total_executions (int): Total number of prompt invocations.
        successful_executions (int): Number of successful prompt invocations.
        failed_executions (int): Number of failed prompt invocations.
        failure_rate (float): Failure rate (failed invocations / total invocations).
        min_response_time (Optional[float]): Minimum response time in seconds.
        max_response_time (Optional[float]): Maximum response time in seconds.
        avg_response_time (Optional[float]): Average response time in seconds.
        last_execution_time (Optional[datetime]): Timestamp of the most recent invocation.
    """

    total_executions: int = Field(..., description="Total number of prompt invocations")
    successful_executions: int = Field(..., description="Number of successful prompt invocations")
    failed_executions: int = Field(..., description="Number of failed prompt invocations")
    failure_rate: float = Field(..., description="Failure rate (failed invocations / total invocations)")
    min_response_time: Optional[float] = Field(None, description="Minimum response time in seconds")
    max_response_time: Optional[float] = Field(None, description="Maximum response time in seconds")
    avg_response_time: Optional[float] = Field(None, description="Average response time in seconds")
    last_execution_time: Optional[datetime] = Field(None, description="Timestamp of the most recent invocation")


# --- JSON Path API modifier Schema


class JsonPathModifier(BaseModelWithConfigDict):
    """Schema for JSONPath queries.

    Provides the structure for parsing JSONPath queries and optional mapping.
    """

    jsonpath: Optional[str] = Field(None, description="JSONPath expression for querying JSON data.")
    mapping: Optional[Dict[str, str]] = Field(None, description="Mapping of fields from original data to output.")


# --- Tool Schemas ---
# Authentication model
class AuthenticationValues(BaseModelWithConfigDict):
    """Schema for all Authentications.
    Provides the authentication values for different types of authentication.
    """

    auth_type: Optional[str] = Field(None, description="Type of authentication: basic, bearer, headers or None")
    auth_value: Optional[str] = Field(None, description="Encoded Authentication values")

    # Only For tool read and view tool
    username: str = Field("", description="Username for basic authentication")
    password: str = Field("", description="Password for basic authentication")
    token: str = Field("", description="Bearer token for authentication")
    auth_header_key: str = Field("", description="Key for custom headers authentication")
    auth_header_value: str = Field("", description="Value for custom headers authentication")


class ToolCreate(BaseModelWithConfigDict):
    """Schema for creating a new tool.

    Supports both MCP-compliant tools and REST integrations. Validates:
      - Unique tool name
      - Valid endpoint URL
      - Valid JSON Schema for input validation
      - Integration type: 'MCP' for MCP-compliant tools or 'REST' for REST integrations
      - Request type (For REST-GET,POST,PUT,DELETE and for MCP-SSE,STDIO,STREAMABLEHTTP)
      - Optional authentication credentials: BasicAuth or BearerTokenAuth or HeadersAuth.
    """

    name: str = Field(..., description="Unique name for the tool")
    url: Union[str, AnyHttpUrl] = Field(None, description="Tool endpoint URL")
    description: Optional[str] = Field(None, description="Tool description")
    request_type: Literal["GET", "POST", "PUT", "DELETE", "SSE", "STDIO", "STREAMABLEHTTP"] = Field("SSE", description="HTTP method to be used for invoking the tool")
    integration_type: Literal["MCP", "REST"] = Field("MCP", description="Tool integration type: 'MCP' for MCP-compliant tools, 'REST' for REST integrations")
    headers: Optional[Dict[str, str]] = Field(None, description="Additional headers to send when invoking the tool")
    input_schema: Optional[Dict[str, Any]] = Field(
        default_factory=lambda: {"type": "object", "properties": {}},
        description="JSON Schema for validating tool parameters",
    )
    annotations: Optional[Dict[str, Any]] = Field(
        default_factory=dict,
        description="Tool annotations for behavior hints (title, readOnlyHint, destructiveHint, idempotentHint, openWorldHint)",
    )
    jsonpath_filter: Optional[str] = Field(default="", description="JSON modification filter")
    auth: Optional[AuthenticationValues] = Field(None, description="Authentication credentials (Basic or Bearer Token or custom headers) if required")
    gateway_id: Optional[str] = Field(None, description="id of gateway for the tool")

    @model_validator(mode="before")
    def assemble_auth(cls, values: Dict[str, Any]) -> Dict[str, Any]:
        """
        Assemble authentication information from separate keys if provided.

        Looks for keys "auth_type", "auth_username", "auth_password", "auth_token", "auth_header_key" and "auth_header_value".
        Constructs the "auth" field as a dictionary suitable for BasicAuth or BearerTokenAuth or HeadersAuth.

        Args:
            values: Dict with authentication information

        Returns:
            Dict: Reformatedd values dict
        """
        logger.debug(
            "Assembling auth in ToolCreate with raw values",
            extra={
                "auth_type": values.get("auth_type"),
                "auth_username": values.get("auth_username"),
                "auth_password": values.get("auth_password"),
                "auth_token": values.get("auth_token"),
                "auth_header_key": values.get("auth_header_key"),
                "auth_header_value": values.get("auth_header_value"),
            },
        )

        auth_type = values.get("auth_type")
        if auth_type:
            if auth_type.lower() == "basic":
                creds = base64.b64encode(f"{values.get('auth_username', '')}:{values.get('auth_password', '')}".encode("utf-8")).decode()
                encoded_auth = encode_auth({"Authorization": f"Basic {creds}"})
                values["auth"] = {"auth_type": "basic", "auth_value": encoded_auth}
            elif auth_type.lower() == "bearer":
                encoded_auth = encode_auth({"Authorization": f"Bearer {values.get('auth_token', '')}"})
                values["auth"] = {"auth_type": "bearer", "auth_value": encoded_auth}
            elif auth_type.lower() == "authheaders":
                encoded_auth = encode_auth({values.get("auth_header_key", ""): values.get("auth_header_value", "")})
                values["auth"] = {"auth_type": "authheaders", "auth_value": encoded_auth}
        return values


class ToolUpdate(BaseModelWithConfigDict):
    """Schema for updating an existing tool.

    Similar to ToolCreate but all fields are optional to allow partial updates.
    """

    name: Optional[str] = Field(None, description="Unique name for the tool")
    url: Optional[Union[str, AnyHttpUrl]] = Field(None, description="Tool endpoint URL")
    description: Optional[str] = Field(None, description="Tool description")
    request_type: Optional[Literal["GET", "POST", "PUT", "DELETE", "SSE", "STDIO", "STREAMABLEHTTP"]] = Field(None, description="HTTP method to be used for invoking the tool")
    integration_type: Optional[Literal["MCP", "REST"]] = Field(None, description="Tool integration type")
    headers: Optional[Dict[str, str]] = Field(None, description="Additional headers to send when invoking the tool")
    input_schema: Optional[Dict[str, Any]] = Field(None, description="JSON Schema for validating tool parameters")
    annotations: Optional[Dict[str, Any]] = Field(None, description="Tool annotations for behavior hints")
    jsonpath_filter: Optional[str] = Field(None, description="JSON path filter for rpc tool calls")
    auth: Optional[AuthenticationValues] = Field(None, description="Authentication credentials (Basic or Bearer Token or custom headers) if required")
    gateway_id: Optional[str] = Field(None, description="id of gateway for the tool")

    @model_validator(mode="before")
    def assemble_auth(cls, values: Dict[str, Any]) -> Dict[str, Any]:
        """
        Assemble authentication information from separate keys if provided.

        Looks for keys "auth_type", "auth_username", "auth_password", "auth_token", "auth_header_key" and "auth_header_value".
        Constructs the "auth" field as a dictionary suitable for BasicAuth or BearerTokenAuth or HeadersAuth.

        Args:
            values: Dict with authentication information

        Returns:
            Dict: Reformatedd values dict
        """

        logger.debug(
            "Assembling auth in ToolCreate with raw values",
            extra={
                "auth_type": values.get("auth_type"),
                "auth_username": values.get("auth_username"),
                "auth_password": values.get("auth_password"),
                "auth_token": values.get("auth_token"),
                "auth_header_key": values.get("auth_header_key"),
                "auth_header_value": values.get("auth_header_value"),
            },
        )

        auth_type = values.get("auth_type")
        if auth_type:
            if auth_type.lower() == "basic":
                creds = base64.b64encode(f"{values.get('auth_username', '')}:{values.get('auth_password', '')}".encode("utf-8")).decode()
                encoded_auth = encode_auth({"Authorization": f"Basic {creds}"})
                values["auth"] = {"auth_type": "basic", "auth_value": encoded_auth}
            elif auth_type.lower() == "bearer":
                encoded_auth = encode_auth({"Authorization": f"Bearer {values.get('auth_token', '')}"})
                values["auth"] = {"auth_type": "bearer", "auth_value": encoded_auth}
            elif auth_type.lower() == "authheaders":
                encoded_auth = encode_auth({values.get("auth_header_key", ""): values.get("auth_header_value", "")})
                values["auth"] = {"auth_type": "authheaders", "auth_value": encoded_auth}
        return values


class ToolRead(BaseModelWithConfigDict):
    """Schema for reading tool information.

    Includes all tool fields plus:
    - Database ID
    - Creation/update timestamps
    - enabled: If Tool is enabled or disabled.
    - reachable: If Tool is reachable or not.
    - Gateway ID for federation
    - Execution count indicating the number of times the tool has been executed.
    - Metrics: Aggregated metrics for the tool invocations.
    - Request type and authentication settings.
    """

    id: str
    original_name: str
    url: Optional[str]
    description: Optional[str]
    request_type: str
    integration_type: str
    headers: Optional[Dict[str, str]]
    input_schema: Dict[str, Any]
    annotations: Optional[Dict[str, Any]]
    jsonpath_filter: Optional[str]
    auth: Optional[AuthenticationValues]
    created_at: datetime
    updated_at: datetime
    enabled: bool
    reachable: bool
    gateway_id: Optional[str]
    execution_count: int
    metrics: ToolMetrics
    name: str
    gateway_slug: str
    original_name_slug: str


class ToolInvocation(BaseModelWithConfigDict):
    """Schema for tool invocation requests.

    Captures:
    - Tool name to invoke
    - Arguments matching tool's input schema
    """

    name: str = Field(..., description="Name of tool to invoke")
    arguments: Dict[str, Any] = Field(default_factory=dict, description="Arguments matching tool's input schema")


class ToolResult(BaseModelWithConfigDict):
    """Schema for tool invocation results.

    Supports:
    - Multiple content types (text/image)
    - Error reporting
    - Optional error messages
    """

    content: List[Union[TextContent, ImageContent]]
    is_error: bool = False
    error_message: Optional[str] = None


class ResourceCreate(BaseModelWithConfigDict):
    """Schema for creating a new resource.

    Supports:
    - Static resources with URI
    - Template resources with parameters
    - Both text and binary content
    """

    uri: str = Field(..., description="Unique URI for the resource")
    name: str = Field(..., description="Human-readable resource name")
    description: Optional[str] = Field(None, description="Resource description")
    mime_type: Optional[str] = Field(None, description="Resource MIME type")
    template: Optional[str] = Field(None, description="URI template for parameterized resources")
    content: Union[str, bytes] = Field(..., description="Resource content (text or binary)")


class ResourceUpdate(BaseModelWithConfigDict):
    """Schema for updating an existing resource.

    Similar to ResourceCreate but URI is not required and all fields are optional.
    """

    name: Optional[str] = Field(None, description="Human-readable resource name")
    description: Optional[str] = Field(None, description="Resource description")
    mime_type: Optional[str] = Field(None, description="Resource MIME type")
    template: Optional[str] = Field(None, description="URI template for parameterized resources")
    content: Optional[Union[str, bytes]] = Field(None, description="Resource content (text or binary)")


class ResourceRead(BaseModelWithConfigDict):
    """Schema for reading resource information.

    Includes all resource fields plus:
    - Database ID
    - Content size
    - Creation/update timestamps
    - Active status
    - Metrics: Aggregated metrics for the resource invocations.
    """

    id: int
    uri: str
    name: str
    description: Optional[str]
    mime_type: Optional[str]
    size: Optional[int]
    created_at: datetime
    updated_at: datetime
    is_active: bool
    metrics: ResourceMetrics


class ResourceSubscription(BaseModelWithConfigDict):
    """Schema for resource subscriptions.

    Tracks:
    - Resource URI being subscribed to
    - Unique subscriber identifier
    """

    uri: str = Field(..., description="URI of resource to subscribe to")
    subscriber_id: str = Field(..., description="Unique subscriber identifier")


class ResourceNotification(BaseModelWithConfigDict):
    """Schema for resource update notifications.

    Contains:
    - Resource URI
    - Updated content
    - Update timestamp
    """

    uri: str
    content: ResourceContent
    timestamp: datetime = Field(default_factory=lambda: datetime.now(timezone.utc))

    @field_serializer("timestamp")
    def serialize_timestamp(self, dt: datetime) -> str:
        """Serialize the `timestamp` field as an ISO 8601 string with UTC timezone.

        Converts the given datetime to UTC and returns it in ISO 8601 format,
        replacing the "+00:00" suffix with "Z" to indicate UTC explicitly.

        Args:
            dt (datetime): The datetime object to serialize.

        Returns:
            str: ISO 8601 formatted string in UTC, ending with 'Z'.
        """
        return dt.astimezone(timezone.utc).isoformat().replace("+00:00", "Z")


# --- Prompt Schemas ---


class PromptArgument(BaseModelWithConfigDict):
    """Schema for prompt template arguments.

    Defines:
    - Argument name
    - Optional description
    - Required flag
    """

    name: str = Field(..., description="Argument name")
    description: Optional[str] = Field(None, description="Argument description")
    required: bool = Field(default=False, description="Whether argument is required")

    model_config: ConfigDict = ConfigDict(
        **{
            # start with every key from the base
            **BaseModelWithConfigDict.model_config,
            # override only json_schema_extra by merging the two dicts:
            "json_schema_extra": {
                **BaseModelWithConfigDict.model_config.get("json_schema_extra", {}),
                "example": {
                    "name": "language",
                    "description": "Programming language",
                    "required": True,
                },
            },
        }
    )


class PromptCreate(BaseModelWithConfigDict):
    """Schema for creating a new prompt template.

    Includes:
    - Template name and description
    - Template text
    - Expected arguments
    """

    name: str = Field(..., description="Unique name for the prompt")
    description: Optional[str] = Field(None, description="Prompt description")
    template: str = Field(..., description="Prompt template text")
    arguments: List[PromptArgument] = Field(default_factory=list, description="List of arguments for the template")


class PromptUpdate(BaseModelWithConfigDict):
    """Schema for updating an existing prompt.

    Similar to PromptCreate but all fields are optional to allow partial updates.
    """

    name: Optional[str] = Field(None, description="Unique name for the prompt")
    description: Optional[str] = Field(None, description="Prompt description")
    template: Optional[str] = Field(None, description="Prompt template text")
    arguments: Optional[List[PromptArgument]] = Field(None, description="List of arguments for the template")


class PromptRead(BaseModelWithConfigDict):
    """Schema for reading prompt information.

    Includes all prompt fields plus:
    - Database ID
    - Creation/update timestamps
    - Active status
    - Metrics: Aggregated metrics for the prompt invocations.
    """

    id: int
    name: str
    description: Optional[str]
    template: str
    arguments: List[PromptArgument]
    created_at: datetime
    updated_at: datetime
    is_active: bool
    metrics: PromptMetrics


class PromptInvocation(BaseModelWithConfigDict):
    """Schema for prompt invocation requests.

    Contains:
    - Prompt name to use
    - Arguments for template rendering
    """

    name: str = Field(..., description="Name of prompt to use")
    arguments: Dict[str, str] = Field(default_factory=dict, description="Arguments for template rendering")


# --- Gateway Schemas ---


class GatewayCreate(BaseModelWithConfigDict):
    """Schema for registering a new federation gateway.

    Captures:
    - Gateway name
    - Endpoint URL
    - Optional description
    - Authentication type: basic, bearer, headers
    - Authentication credentials: username/password or token or custom headers
    """

    name: str = Field(..., description="Unique name for the gateway")
    url: Union[str, AnyHttpUrl] = Field(..., description="Gateway endpoint URL")
    description: Optional[str] = Field(None, description="Gateway description")
    transport: str = Field(default="SSE", description="Transport used by MCP server: SSE or STREAMABLEHTTP")

    # Authorizations
    auth_type: Optional[str] = Field(None, description="Type of authentication: basic, bearer, headers, or none")
    # Fields for various types of authentication
    auth_username: Optional[str] = Field(None, description="Username for basic authentication")
    auth_password: Optional[str] = Field(None, description="Password for basic authentication")
    auth_token: Optional[str] = Field(None, description="Token for bearer authentication")
    auth_header_key: Optional[str] = Field(None, description="Key for custom headers authentication")
    auth_header_value: Optional[str] = Field(None, description="Value for custom headers authentication")

    # Adding `auth_value` as an alias for better access post-validation
    auth_value: Optional[str] = Field(None, validate_default=True)

    @field_validator("url", mode="before")
    def ensure_url_scheme(cls, v: str) -> str:
        """
        Ensure URL has an http/https scheme.

        Args:
            v: Input url

        Returns:
            str: URL with correct schema

        """
        if isinstance(v, str) and not (v.startswith("http://") or v.startswith("https://")):
            return f"http://{v}"
        return v

    @field_validator("auth_value", mode="before")
    def create_auth_value(cls, v, info):
        """
        This validator will run before the model is fully instantiated (mode="before")
        It will process the auth fields based on auth_type and generate auth_value.

        Args:
            v: Input url
            info: ValidationInfo containing auth_type

        Returns:
            str: Auth value
        """
        data = info.data
        auth_type = data.get("auth_type")

        if (auth_type is None) or (auth_type == ""):
            return v  # If no auth_type is provided, no need to create auth_value

        # Process the auth fields and generate auth_value based on auth_type
        auth_value = cls._process_auth_fields(info)

        return auth_value

    @staticmethod
    def _process_auth_fields(info: ValidationInfo) -> Optional[Dict[str, Any]]:
        """
        Processes the input authentication fields and returns the correct auth_value.
        This method is called based on the selected auth_type.

        Args:
            info: ValidationInfo containing auth fields

        Returns:
            Dict with encoded auth

        Raises:
            ValueError: If auth_type is invalid
        """
        data = info.data
        auth_type = data.get("auth_type")

        if auth_type == "basic":
            # For basic authentication, both username and password must be present
            username = data.get("auth_username")
            password = data.get("auth_password")

            if not username or not password:
                raise ValueError("For 'basic' auth, both 'auth_username' and 'auth_password' must be provided.")

            creds = base64.b64encode(f"{username}:{password}".encode("utf-8")).decode()
            return encode_auth({"Authorization": f"Basic {creds}"})

        if auth_type == "bearer":
            # For bearer authentication, only token is required
            token = data.get("auth_token")

            if not token:
                raise ValueError("For 'bearer' auth, 'auth_token' must be provided.")

            return encode_auth({"Authorization": f"Bearer {token}"})

        if auth_type == "authheaders":
            # For headers authentication, both key and value must be present
            header_key = data.get("auth_header_key")
            header_value = data.get("auth_header_value")

            if not header_key or not header_value:
                raise ValueError("For 'headers' auth, both 'auth_header_key' and 'auth_header_value' must be provided.")

            return encode_auth({header_key: header_value})

        raise ValueError("Invalid 'auth_type'. Must be one of: basic, bearer, or headers.")


class GatewayUpdate(BaseModelWithConfigDict):
    """Schema for updating an existing federation gateway.

    Similar to GatewayCreate but all fields are optional to allow partial updates.
    """

    name: Optional[str] = Field(None, description="Unique name for the gateway")
    url: Optional[Union[str, AnyHttpUrl]] = Field(None, description="Gateway endpoint URL")
    description: Optional[str] = Field(None, description="Gateway description")
    transport: str = Field(default="SSE", description="Transport used by MCP server: SSE or STREAMABLEHTTP")

    name: Optional[str] = Field(None, description="Unique name for the prompt")
    # Authorizations
    auth_type: Optional[str] = Field(None, description="auth_type: basic, bearer, headers or None")
    auth_username: Optional[str] = Field(None, description="username for basic authentication")
    auth_password: Optional[str] = Field(None, description="password for basic authentication")
    auth_token: Optional[str] = Field(None, description="token for bearer authentication")
    auth_header_key: Optional[str] = Field(None, description="key for custom headers authentication")
    auth_header_value: Optional[str] = Field(None, description="vallue for custom headers authentication")

    # Adding `auth_value` as an alias for better access post-validation
    auth_value: Optional[str] = Field(None, validate_default=True)

    @field_validator("url", mode="before")
    def ensure_url_scheme(cls, v: Optional[str]) -> Optional[str]:
        """
        Ensure URL has an http/https scheme.

        Args:
            v: Input URL

        Returns:
            str: Validated URL
        """
        if isinstance(v, str) and not (v.startswith("http://") or v.startswith("https://")):
            return f"http://{v}"
        return v

    @field_validator("auth_value", mode="before")
    def create_auth_value(cls, v, info):
        """
        This validator will run before the model is fully instantiated (mode="before")
        It will process the auth fields based on auth_type and generate auth_value.

        Args:
            v: Input URL
            info: ValidationInfo containing auth_type

        Returns:
            str: Auth value or URL
        """
        data = info.data
        auth_type = data.get("auth_type")

        if (auth_type is None) or (auth_type == ""):
            return v  # If no auth_type is provided, no need to create auth_value

        # Process the auth fields and generate auth_value based on auth_type
        auth_value = cls._process_auth_fields(info)

        return auth_value

    @staticmethod
    def _process_auth_fields(values: Dict[str, Any]) -> Optional[Dict[str, Any]]:
        """
        Processes the input authentication fields and returns the correct auth_value.
        This method is called based on the selected auth_type.

        Args:
            values: Dict container auth information auth_type, auth_username, auth_password, auth_token, auth_header_key and auth_header_value

        Returns:
            dict: Encoded auth information

        Raises:
            ValueError: If auth type is invalid
        """
        auth_type = values.get("auth_type")

        if auth_type == "basic":
            # For basic authentication, both username and password must be present
            username = values.get("auth_username")
            password = values.get("auth_password")

            if not username or not password:
                raise ValueError("For 'basic' auth, both 'auth_username' and 'auth_password' must be provided.")

            creds = base64.b64encode(f"{username}:{password}".encode("utf-8")).decode()
            return encode_auth({"Authorization": f"Basic {creds}"})

        if auth_type == "bearer":
            # For bearer authentication, only token is required
            token = values.get("auth_token")

            if not token:
                raise ValueError("For 'bearer' auth, 'auth_token' must be provided.")

            return encode_auth({"Authorization": f"Bearer {token}"})

        if auth_type == "authheaders":
            # For headers authentication, both key and value must be present
            header_key = values.get("auth_header_key")
            header_value = values.get("auth_header_value")

            if not header_key or not header_value:
                raise ValueError("For 'headers' auth, both 'auth_header_key' and 'auth_header_value' must be provided.")

            return encode_auth({header_key: header_value})

        raise ValueError("Invalid 'auth_type'. Must be one of: basic, bearer, or headers.")


class GatewayRead(BaseModelWithConfigDict):
    """Schema for reading gateway information.

    Includes all gateway fields plus:
    - Database ID
    - Capabilities dictionary
    - Creation/update timestamps
    - enabled status
    - reachable status
    - Last seen timestamp
    - Authentication type: basic, bearer, headers
    - Authentication value: username/password or token or custom headers

    Auto Populated fields:
    - Authentication username: for basic auth
    - Authentication password: for basic auth
    - Authentication token: for bearer auth
    - Authentication header key: for headers auth
    - Authentication header value: for headers auth
    """

    id: str = Field(None, description="Unique ID of the gateway")
    name: str = Field(..., description="Unique name for the gateway")
    url: str = Field(..., description="Gateway endpoint URL")
    description: Optional[str] = Field(None, description="Gateway description")
    transport: str = Field(default="SSE", description="Transport used by MCP server: SSE or STREAMABLEHTTP")
    capabilities: Dict[str, Any] = Field(default_factory=dict, description="Gateway capabilities")
    created_at: datetime = Field(default_factory=lambda: datetime.now(timezone.utc), description="Creation timestamp")
    updated_at: datetime = Field(default_factory=lambda: datetime.now(timezone.utc), description="Last update timestamp")
    enabled: bool = Field(default=True, description="Is the gateway enabled?")
    reachable: bool = Field(default=True, description="Is the gateway reachable/online?")

    last_seen: Optional[datetime] = Field(default_factory=lambda: datetime.now(timezone.utc), description="Last seen timestamp")

    # Authorizations
    auth_type: Optional[str] = Field(None, description="auth_type: basic, bearer, headers or None")
    auth_value: Optional[str] = Field(None, description="auth value: username/password or token or custom headers")

    # auth_value will populate the following fields
    auth_username: Optional[str] = Field(None, description="username for basic authentication")
    auth_password: Optional[str] = Field(None, description="password for basic authentication")
    auth_token: Optional[str] = Field(None, description="token for bearer authentication")
    auth_header_key: Optional[str] = Field(None, description="key for custom headers authentication")
    auth_header_value: Optional[str] = Field(None, description="vallue for custom headers authentication")

    slug: str = Field(None, description="Slug for gateway endpoint URL")

    # This will be the main method to automatically populate fields
    @model_validator(mode="after")
    def _populate_auth(cls, values: Dict[str, Any]) -> Dict[str, Any]:
        auth_type = values.auth_type
        auth_value_encoded = values.auth_value
        auth_value = decode_auth(auth_value_encoded)
        if auth_type == "basic":
            u = auth_value.get("username")
            p = auth_value.get("password")
            if not u or not p:
                raise ValueError("basic auth requires both username and password")
            values.auth_username, values.auth_password = u, p

        elif auth_type == "bearer":
            auth = auth_value.get("Authorization")
            if not (isinstance(auth, str) and auth.startswith("Bearer ")):
                raise ValueError("bearer auth requires an Authorization header of the form 'Bearer <token>'")
            values.auth_token = auth.removeprefix("Bearer ")

        elif auth_type == "authheaders":
            # must be exactly one header
            if len(auth_value) != 1:
                raise ValueError("authheaders requires exactly one key/value pair")
            k, v = next(iter(auth_value.items()))
            values.auth_header_key, values.auth_header_value = k, v

        return values


class FederatedTool(BaseModelWithConfigDict):
    """Schema for tools provided by federated gateways.

    Contains:
    - Tool definition
    - Source gateway information
    """

    tool: MCPTool
    gateway_id: str
    gateway_name: str
    gateway_url: str


class FederatedResource(BaseModelWithConfigDict):
    """Schema for resources from federated gateways.

    Contains:
    - Resource definition
    - Source gateway information
    """

    resource: MCPResource
    gateway_id: str
    gateway_name: str
    gateway_url: str


class FederatedPrompt(BaseModelWithConfigDict):
    """Schema for prompts from federated gateways.

    Contains:
    - Prompt definition
    - Source gateway information
    """

    prompt: MCPPrompt
    gateway_id: str
    gateway_name: str
    gateway_url: str


# --- RPC Schemas ---


class RPCRequest(BaseModelWithConfigDict):
    """Schema for JSON-RPC 2.0 requests.

    Validates:
    - Protocol version
    - Method name
    - Optional parameters
    - Optional request ID
    """

    jsonrpc: Literal["2.0"]
    method: str
    params: Optional[Dict[str, Any]] = None
    id: Optional[Union[int, str]] = None


class RPCResponse(BaseModelWithConfigDict):
    """Schema for JSON-RPC 2.0 responses.

    Contains:
    - Protocol version
    - Result or error
    - Request ID
    """

    jsonrpc: Literal["2.0"]
    result: Optional[Any] = None
    error: Optional[Dict[str, Any]] = None
    id: Optional[Union[int, str]] = None


# --- Event and Admin Schemas ---


class EventMessage(BaseModelWithConfigDict):
    """Schema for SSE event messages.

    Includes:
    - Event type
    - Event data payload
    - Event timestamp
    """

    type: str = Field(..., description="Event type (tool_added, resource_updated, etc)")
    data: Dict[str, Any] = Field(..., description="Event payload")
    timestamp: datetime = Field(default_factory=lambda: datetime.now(timezone.utc))

    @field_serializer("timestamp")
    def serialize_timestamp(self, dt: datetime) -> str:
        """
        Serialize the `timestamp` field as an ISO 8601 string with UTC timezone.

        Converts the given datetime to UTC and returns it in ISO 8601 format,
        replacing the "+00:00" suffix with "Z" to indicate UTC explicitly.

        Args:
            dt (datetime): The datetime object to serialize.

        Returns:
            str: ISO 8601 formatted string in UTC, ending with 'Z'.
        """
        return dt.astimezone(timezone.utc).isoformat().replace("+00:00", "Z")


class AdminToolCreate(BaseModelWithConfigDict):
    """Schema for creating tools via admin UI.

    Handles:
    - Basic tool information
    - JSON string inputs for headers/schema
    """

    name: str
    url: str
    description: Optional[str] = None
    integration_type: str = "MCP"
    headers: Optional[str] = None  # JSON string
    input_schema: Optional[str] = None  # JSON string

    @field_validator("headers", "input_schema")
    def validate_json(cls, v: Optional[str]) -> Optional[Dict[str, Any]]:
        """
        Validate and parse JSON string inputs.

        Args:
            v: Input string

        Returns:
            dict: Output JSON version of v

        Raises:
            ValueError: When unable to convert to JSON
        """
        if not v:
            return None
        try:
            return json.loads(v)
        except json.JSONDecodeError:
            raise ValueError("Invalid JSON")


class AdminGatewayCreate(BaseModelWithConfigDict):
    """Schema for creating gateways via admin UI.

    Captures:
    - Gateway name
    - Endpoint URL
    - Optional description
    """

    name: str
    url: str
    description: Optional[str] = None


# --- New Schemas for Status Toggle Operations ---


class StatusToggleRequest(BaseModelWithConfigDict):
    """Request schema for toggling active status."""

    activate: bool = Field(..., description="Whether to activate (true) or deactivate (false) the item")


class StatusToggleResponse(BaseModelWithConfigDict):
    """Response schema for status toggle operations."""

    id: int
    name: str
    is_active: bool
    message: str = Field(..., description="Success message")


# --- Optional Filter Parameters for Listing Operations ---


class ListFilters(BaseModelWithConfigDict):
    """Filtering options for list operations."""

    include_inactive: bool = Field(False, description="Whether to include inactive items in the results")


# --- Server Schemas ---


class ServerCreate(BaseModelWithConfigDict):
    """Schema for creating a new server.

    Attributes:
        name: The server's name (required).
        description: Optional text description.
        icon: Optional URL for the server's icon.
        associated_tools: Optional list of tool IDs (as strings).
        associated_resources: Optional list of resource IDs (as strings).
        associated_prompts: Optional list of prompt IDs (as strings).
    """

    name: str = Field(..., description="The server's name")
    description: Optional[str] = Field(None, description="Server description")
    icon: Optional[str] = Field(None, description="URL for the server's icon")
    associated_tools: Optional[List[str]] = Field(None, description="Comma-separated tool IDs")
    associated_resources: Optional[List[str]] = Field(None, description="Comma-separated resource IDs")
    associated_prompts: Optional[List[str]] = Field(None, description="Comma-separated prompt IDs")

    @field_validator("associated_tools", "associated_resources", "associated_prompts", mode="before")
    def split_comma_separated(cls, v):
        """
        Splits a comma-separated string into a list of strings if needed.

        Args:
            v: Input string

        Returns:
            list: Comma separated array of input string
        """
        if isinstance(v, str):
            return [item.strip() for item in v.split(",") if item.strip()]
        return v


class ServerUpdate(BaseModelWithConfigDict):
    """Schema for updating an existing server.

    All fields are optional to allow partial updates.
    """

    name: Optional[str] = Field(None, description="The server's name")
    description: Optional[str] = Field(None, description="Server description")
    icon: Optional[str] = Field(None, description="URL for the server's icon")
    associated_tools: Optional[List[str]] = Field(None, description="Comma-separated tool IDs")
    associated_resources: Optional[List[str]] = Field(None, description="Comma-separated resource IDs")
    associated_prompts: Optional[List[str]] = Field(None, description="Comma-separated prompt IDs")

    @field_validator("associated_tools", "associated_resources", "associated_prompts", mode="before")
    def split_comma_separated(cls, v):
        """
        Splits a comma-separated string into a list of strings if needed.

        Args:
            v: Input string

        Returns:
            list: Comma separated array of input string
        """
        if isinstance(v, str):
            return [item.strip() for item in v.split(",") if item.strip()]
        return v


class ServerRead(BaseModelWithConfigDict):
    """Schema for reading server information.

    Includes all server fields plus:
    - Database ID
    - Associated tool, resource, and prompt IDs
    - Creation/update timestamps
    - Active status
    - Metrics: Aggregated metrics for the server invocations.
    """

    id: str
    name: str
    description: Optional[str]
    icon: Optional[str]
    created_at: datetime
    updated_at: datetime
    is_active: bool
    associated_tools: List[str] = []
    associated_resources: List[int] = []
    associated_prompts: List[int] = []
    metrics: ServerMetrics

    @model_validator(mode="before")
    def populate_associated_ids(cls, values):
        """
        Pre-validation method that converts associated objects to their 'id'.

        This method checks 'associated_tools', 'associated_resources', and
        'associated_prompts' in the input and replaces each object with its `id`
        if present.

        Args:
            values (dict): The input values.

        Returns:
            dict: Updated values with object ids, or the original values if no
            changes are made.
        """
        # If values is not a dict (e.g. it's a Server instance), convert it
        if not isinstance(values, dict):
            try:
                values = vars(values)
            except Exception:
                return values
        if "associated_tools" in values and values["associated_tools"]:
            values["associated_tools"] = [tool.id if hasattr(tool, "id") else tool for tool in values["associated_tools"]]
        if "associated_resources" in values and values["associated_resources"]:
            values["associated_resources"] = [res.id if hasattr(res, "id") else res for res in values["associated_resources"]]
        if "associated_prompts" in values and values["associated_prompts"]:
            values["associated_prompts"] = [prompt.id if hasattr(prompt, "id") else prompt for prompt in values["associated_prompts"]]
        return values
