# Familiar - Self-Hosted AI Companion Platform
# Copyright (c) 2026 George Scott Foley
# ORCID: 0009-0006-4957-0540
# Email: Georgescottfoley@proton.me
# Licensed under the MIT License - see LICENSE file for details

"""
MCP Configuration

Defines configuration structures for MCP server connections.
Supports loading from YAML, environment variables, and programmatic config.
"""

import logging
import os
import re
from dataclasses import dataclass, field
from enum import Enum
from pathlib import Path
from typing import Any, Dict, List, Optional, Union

import yaml

logger = logging.getLogger(__name__)


class MCPTransport(str, Enum):
    """Supported MCP transport types."""

    STDIO = "stdio"  # Standard input/output (most common)
    HTTP = "http"  # Streamable HTTP
    SSE = "sse"  # Server-Sent Events (deprecated)
    WEBSOCKET = "websocket"  # WebSocket (future)


@dataclass
class MCPServerConfig:
    """
    Configuration for a single MCP server.

    Attributes:
        name: Unique identifier for this server (used as tool prefix)
        transport: Transport type (stdio, http, sse)

        # For stdio transport:
        command: Command to execute (e.g., "npx", "uvx", "python")
        args: Arguments to pass to command
        cwd: Working directory for the command

        # For http/sse transport:
        url: Server URL
        headers: HTTP headers to include

        # Common:
        env: Environment variables to set
        env_inherit: Whether to inherit current environment (default: True)
        timeout: Connection timeout in seconds
        enabled: Whether this server is enabled
        auto_connect: Whether to connect on startup

        # Security:
        trust_level: Trust level for tools from this server
        allowed_tools: If set, only allow these tools (whitelist)
        blocked_tools: If set, block these tools (blacklist)
    """

    name: str
    transport: MCPTransport = MCPTransport.STDIO

    # Stdio transport
    command: Optional[str] = None
    args: List[str] = field(default_factory=list)
    cwd: Optional[str] = None

    # HTTP transport
    url: Optional[str] = None
    headers: Dict[str, str] = field(default_factory=dict)

    # Common
    env: Dict[str, str] = field(default_factory=dict)
    env_inherit: bool = True
    timeout: float = 30.0
    enabled: bool = True
    auto_connect: bool = True

    # Security
    trust_level: str = "standard"  # "untrusted", "standard", "trusted"
    allowed_tools: Optional[List[str]] = None
    blocked_tools: Optional[List[str]] = None
    requires_confirmation: bool = False  # Require user confirmation for all tools

    # Metadata
    description: Optional[str] = None
    version: Optional[str] = None

    def __post_init__(self):
        """Validate configuration."""
        # Convert string transport to enum
        if isinstance(self.transport, str):
            self.transport = MCPTransport(self.transport.lower())

        # Validate transport-specific requirements
        if self.transport == MCPTransport.STDIO:
            if not self.command:
                raise ValueError(f"MCP server '{self.name}': stdio transport requires 'command'")
        elif self.transport in (MCPTransport.HTTP, MCPTransport.SSE):
            if not self.url:
                raise ValueError(
                    f"MCP server '{self.name}': {self.transport.value} transport requires 'url'"
                )

        # Expand environment variables in config values
        self._expand_env_vars()

    def _expand_env_vars(self):
        """Expand ${VAR} patterns in string values."""
        pattern = re.compile(r"\$\{([^}]+)\}")

        def expand(value: str) -> str:
            def replacer(match):
                var_name = match.group(1)
                return os.environ.get(var_name, match.group(0))

            return pattern.sub(replacer, value)

        # Expand in various fields
        if self.command:
            self.command = expand(self.command)

        self.args = [expand(arg) for arg in self.args]

        if self.url:
            self.url = expand(self.url)

        if self.cwd:
            self.cwd = expand(self.cwd)

        # Expand in env dict
        self.env = {k: expand(v) for k, v in self.env.items()}

        # Expand in headers
        self.headers = {k: expand(v) for k, v in self.headers.items()}

    def get_env(self) -> Dict[str, str]:
        """Get full environment for server process."""
        if self.env_inherit:
            env = os.environ.copy()
            env.update(self.env)
            return env
        return self.env.copy()

    def is_tool_allowed(self, tool_name: str) -> bool:
        """Check if a tool is allowed based on whitelist/blacklist."""
        # Check blacklist first
        if self.blocked_tools and tool_name in self.blocked_tools:
            return False

        # Check whitelist if set
        if self.allowed_tools is not None:
            return tool_name in self.allowed_tools

        return True

    @classmethod
    def from_dict(cls, data: Dict[str, Any]) -> "MCPServerConfig":
        """Create config from dictionary."""
        return cls(**data)

    def to_dict(self) -> Dict[str, Any]:
        """Convert to dictionary."""
        return {
            "name": self.name,
            "transport": self.transport.value,
            "command": self.command,
            "args": self.args,
            "cwd": self.cwd,
            "url": self.url,
            "headers": self.headers,
            "env": self.env,
            "env_inherit": self.env_inherit,
            "timeout": self.timeout,
            "enabled": self.enabled,
            "auto_connect": self.auto_connect,
            "trust_level": self.trust_level,
            "allowed_tools": self.allowed_tools,
            "blocked_tools": self.blocked_tools,
            "requires_confirmation": self.requires_confirmation,
            "description": self.description,
            "version": self.version,
        }


@dataclass
class MCPConfig:
    """
    Configuration for MCP integration.

    Attributes:
        enabled: Whether MCP integration is enabled
        servers: List of MCP server configurations
        tool_prefix_separator: Separator between server name and tool name
        connect_timeout: Timeout for initial connections
        reconnect_on_failure: Whether to auto-reconnect on failure
        max_reconnect_attempts: Maximum reconnection attempts
        log_level: Logging level for MCP operations
    """

    enabled: bool = True
    servers: List[MCPServerConfig] = field(default_factory=list)
    tool_prefix_separator: str = ":"  # e.g., "github:create_issue"
    connect_timeout: float = 30.0
    reconnect_on_failure: bool = True
    max_reconnect_attempts: int = 3
    reconnect_delay: float = 5.0
    log_level: str = "INFO"

    # Global security settings
    require_confirmation_untrusted: bool = True

    @classmethod
    def from_dict(cls, data: Dict[str, Any]) -> "MCPConfig":
        """Create config from dictionary."""
        servers = [
            MCPServerConfig.from_dict(s) if isinstance(s, dict) else s
            for s in data.get("servers", [])
        ]
        return cls(
            enabled=data.get("enabled", True),
            servers=servers,
            tool_prefix_separator=data.get("tool_prefix_separator", ":"),
            connect_timeout=data.get("connect_timeout", 30.0),
            reconnect_on_failure=data.get("reconnect_on_failure", True),
            max_reconnect_attempts=data.get("max_reconnect_attempts", 3),
            reconnect_delay=data.get("reconnect_delay", 5.0),
            log_level=data.get("log_level", "INFO"),
            require_confirmation_untrusted=data.get("require_confirmation_untrusted", True),
        )

    @classmethod
    def from_yaml(cls, path: Union[str, Path]) -> "MCPConfig":
        """Load configuration from YAML file."""
        path = Path(path)
        if not path.exists():
            logger.warning(f"MCP config file not found: {path}")
            return cls()

        with open(path) as f:
            data = yaml.safe_load(f) or {}

        # Support both top-level and nested under 'mcp' key
        if "mcp" in data:
            data = data["mcp"]

        return cls.from_dict(data)

    def to_dict(self) -> Dict[str, Any]:
        """Convert to dictionary."""
        return {
            "enabled": self.enabled,
            "servers": [s.to_dict() for s in self.servers],
            "tool_prefix_separator": self.tool_prefix_separator,
            "connect_timeout": self.connect_timeout,
            "reconnect_on_failure": self.reconnect_on_failure,
            "max_reconnect_attempts": self.max_reconnect_attempts,
            "reconnect_delay": self.reconnect_delay,
            "log_level": self.log_level,
            "require_confirmation_untrusted": self.require_confirmation_untrusted,
        }

    def save_yaml(self, path: Union[str, Path]):
        """Save configuration to YAML file."""
        path = Path(path)
        path.parent.mkdir(parents=True, exist_ok=True)

        with open(path, "w") as f:
            yaml.dump({"mcp": self.to_dict()}, f, default_flow_style=False)

    def get_server(self, name: str) -> Optional[MCPServerConfig]:
        """Get server config by name."""
        for server in self.servers:
            if server.name == name:
                return server
        return None

    def add_server(self, server: MCPServerConfig):
        """Add a server configuration."""
        # Check for duplicate names
        existing = self.get_server(server.name)
        if existing:
            raise ValueError(f"Server '{server.name}' already exists")
        self.servers.append(server)

    def remove_server(self, name: str) -> bool:
        """Remove a server configuration by name."""
        for i, server in enumerate(self.servers):
            if server.name == name:
                del self.servers[i]
                return True
        return False


# Pre-configured server templates for common services
COMMON_SERVERS = {
    "filesystem": MCPServerConfig(
        name="filesystem",
        transport=MCPTransport.STDIO,
        command="uvx",
        args=["mcp-server-filesystem"],
        description="File system operations",
    ),
    "github": MCPServerConfig(
        name="github",
        transport=MCPTransport.STDIO,
        command="npx",
        args=["-y", "@modelcontextprotocol/server-github"],
        env={"GITHUB_TOKEN": "${GITHUB_TOKEN}"},
        description="GitHub repository operations",
    ),
    "postgres": MCPServerConfig(
        name="postgres",
        transport=MCPTransport.STDIO,
        command="uvx",
        args=["mcp-server-postgres"],
        env={"DATABASE_URL": "${DATABASE_URL}"},
        description="PostgreSQL database operations",
    ),
    "sqlite": MCPServerConfig(
        name="sqlite",
        transport=MCPTransport.STDIO,
        command="uvx",
        args=["mcp-server-sqlite"],
        description="SQLite database operations",
    ),
    "slack": MCPServerConfig(
        name="slack",
        transport=MCPTransport.STDIO,
        command="npx",
        args=["-y", "@anthropic/mcp-server-slack"],
        env={"SLACK_BOT_TOKEN": "${SLACK_BOT_TOKEN}"},
        description="Slack messaging",
    ),
    "puppeteer": MCPServerConfig(
        name="puppeteer",
        transport=MCPTransport.STDIO,
        command="npx",
        args=["-y", "@anthropic/mcp-server-puppeteer"],
        description="Browser automation",
    ),
    "memory": MCPServerConfig(
        name="memory",
        transport=MCPTransport.STDIO,
        command="npx",
        args=["-y", "@anthropic/mcp-server-memory"],
        description="Key-value memory storage",
    ),
}


def get_common_server(name: str, **overrides) -> MCPServerConfig:
    """
    Get a pre-configured server template with optional overrides.

    Args:
        name: Server template name (e.g., "github", "filesystem")
        **overrides: Config values to override

    Returns:
        Configured MCPServerConfig

    Example:
        server = get_common_server("github", env={"GITHUB_TOKEN": "my_token"})
    """
    if name not in COMMON_SERVERS:
        available = ", ".join(COMMON_SERVERS.keys())
        raise ValueError(f"Unknown server template '{name}'. Available: {available}")

    # Create a copy with overrides
    base = COMMON_SERVERS[name].to_dict()
    base.update(overrides)
    return MCPServerConfig.from_dict(base)
