from __future__ import annotations

import os
from typing import Annotated, Literal

from pydantic import BaseModel, ConfigDict, Field


class MCPServerBase(BaseModel):
    """Base model for MCP server configuration."""

    type: str
    """Type discriminator for MCP server configurations."""

    enabled: bool = True
    """Whether this server is currently enabled."""

    environment: dict[str, str] | None = None
    """Environment variables to pass to the server process."""

    model_config = ConfigDict(use_attribute_docstrings=True, extra="forbid")

    def get_env_vars(self) -> dict[str, str]:
        """Get environment variables for the server process."""
        env = os.environ.copy()
        if self.environment:
            env.update(self.environment)
        env["PYTHONIOENCODING"] = "utf-8"
        return env


class StdioMCPServer(MCPServerBase):
    """MCP server started via stdio.

    Uses subprocess communication through standard input/output streams.
    """

    type: Literal["stdio"] = Field("stdio", init=False)
    """Stdio server coniguration."""

    command: str
    """Command to execute (e.g. "pipx", "python", "node")."""

    args: list[str] = Field(default_factory=list)
    """Command arguments (e.g. ["run", "some-server", "--debug"])."""

    @classmethod
    def from_string(cls, command: str) -> StdioMCPServer:
        """Create a MCP server from a command string."""
        cmd, args = command.split(maxsplit=1)
        return cls(command=cmd, args=args.split())


class SSEMCPServer(MCPServerBase):
    """MCP server using Server-Sent Events transport.

    Connects to a server over HTTP with SSE for real-time communication.
    """

    type: Literal["sse"] = Field("sse", init=False)
    """SSE server configuration."""

    url: str
    """URL of the SSE server endpoint."""


MCPServerConfig = Annotated[StdioMCPServer | SSEMCPServer, Field(discriminator="type")]
