# Copyright (c) 2025 CoReason, Inc.
#
# This software is proprietary and dual-licensed.
# Licensed under the Prosperity Public License 3.0 (the "License").
# A copy of the license is available at https://prosperitylicense.com/versions/3.0.0
# For details, see the LICENSE file.
# Commercial use beyond a 30-day trial requires a separate license.
#
# Source Code: https://github.com/CoReason-AI/coreason_mcp

"""Configuration models for the CoReason MCP Gateway.

This module defines the Pydantic models used to configure the gateway, including
server connections and global settings.
"""

from enum import Enum
from typing import Dict, List, Optional

from pydantic import BaseModel, Field, model_validator


class TransportType(str, Enum):
    """Enumeration of supported MCP transport types."""

    STDIO = "stdio"
    SSE = "sse"


class McpServerConfig(BaseModel):
    """Configuration for a single MCP server connection.

    Defines how the gateway connects to a specific MCP server, supporting
    both SSE (Server-Sent Events) and Stdio transports.

    Attributes:
        name: Unique name for the server (e.g. 'brave_search').
        transport: The transport type to use (stdio or sse).
        url: URL for SSE transport (e.g. 'http://localhost:8080/sse').
        command: Command to execute for Stdio transport (e.g. 'npx').
        args: Arguments for the command (Stdio).
        env: Environment variables for the command (Stdio), used for credential injection.
        auth_headers: Authentication headers for SSE transport.
    """

    name: str = Field(..., description="Unique name for the server (e.g. 'brave_search').")
    transport: Optional[TransportType] = Field(
        None, description="Transport type (auto-detected if missing, but explicit is preferred)."
    )

    # SSE Specific
    url: Optional[str] = Field(None, description="URL for SSE transport (e.g. 'http://localhost:8080/sse').")
    auth_headers: Dict[str, str] = Field(default_factory=dict, description="Authentication headers for SSE transport.")
    # Keep auth_token for backward compatibility if needed, but PRD says auth_headers
    auth_token: Optional[str] = Field(
        None, description="Legacy: Authentication token for SSE transport (deprecated, use auth_headers)."
    )

    # Stdio Specific
    command: Optional[str] = Field(None, description="Command to execute for Stdio transport (e.g. 'npx').")
    args: List[str] = Field(default_factory=list, description="Arguments for the command (Stdio).")
    env: Dict[str, str] = Field(
        default_factory=dict,
        description="Environment variables for the command (Stdio), used for credential injection.",
    )

    @model_validator(mode="after")
    def check_transport(self) -> "McpServerConfig":
        """Validates that exactly one transport method (SSE or Stdio) is configured correctly.

        Raises:
            ValueError: If neither or both transport configurations are present, or if
                the explicit transport type doesn't match the provided fields.
        """
        if self.transport == TransportType.SSE or (self.url and not self.command):
            self.transport = TransportType.SSE
            if not self.url:
                raise ValueError("SSE Transport requires 'url'.")
        elif self.transport == TransportType.STDIO or (self.command and not self.url):
            self.transport = TransportType.STDIO
            if not self.command:
                raise ValueError("Stdio Transport requires 'command'.")
        else:
            # Ambiguous or missing
            if not self.url and not self.command:
                raise ValueError("Either 'url' (SSE) or 'command' (Stdio) must be provided.")
            if self.url and self.command:
                raise ValueError("Cannot provide both 'url' and 'command'. Choose one transport type.")

        return self


class McpConfig(BaseModel):
    """Global configuration for the MCP Host.

    Attributes:
        servers: List of server configurations.
        timeout: Default timeout for requests in seconds.
    """

    servers: List[McpServerConfig] = Field(..., description="List of server configurations.")
    timeout: float = Field(30.0, description="Default timeout for requests in seconds.")
