"""
MCP (Model Context Protocol) server management for Oh My Codex.
"""

import asyncio
import os
from typing import Optional, Dict, Any, List
from dataclasses import dataclass


@dataclass
class MCPServerConfig:
    """Configuration for an MCP server."""
    name: str
    command: str
    args: List[str]
    env: Optional[Dict[str, str]] = None
    description: str = ""


# Pre-configured MCP servers
MCP_SERVERS = {
    "codex": MCPServerConfig(
        name="Codex CLI",
        command="npx",
        args=["-y", "codex", "mcp"],
        description="OpenAI Codex CLI as MCP server - file editing, shell execution",
    ),
    
    "filesystem": MCPServerConfig(
        name="Filesystem",
        command="npx",
        args=["-y", "@anthropic/mcp-server-filesystem"],
        description="File system access with sandboxing",
    ),
    
    "github": MCPServerConfig(
        name="GitHub",
        command="npx", 
        args=["-y", "@anthropic/mcp-server-github"],
        env={"GITHUB_TOKEN": os.environ.get("GITHUB_TOKEN", "")},
        description="GitHub API access - repos, issues, PRs",
    ),
    
    "postgres": MCPServerConfig(
        name="PostgreSQL",
        command="npx",
        args=["-y", "@anthropic/mcp-server-postgres"],
        description="PostgreSQL database access",
    ),
    
    "brave": MCPServerConfig(
        name="Brave Search",
        command="npx",
        args=["-y", "@anthropic/mcp-server-brave-search"],
        env={"BRAVE_API_KEY": os.environ.get("BRAVE_API_KEY", "")},
        description="Web search via Brave Search API",
    ),
    
    "memory": MCPServerConfig(
        name="Memory",
        command="npx",
        args=["-y", "@anthropic/mcp-server-memory"],
        description="Persistent memory/knowledge graph",
    ),
}


class MCPManager:
    """Manages MCP server connections."""
    
    def __init__(self):
        self.active_servers: Dict[str, Any] = {}
    
    def get_server_config(self, name: str) -> Optional[MCPServerConfig]:
        """Get configuration for a named MCP server."""
        return MCP_SERVERS.get(name)
    
    def list_available(self) -> List[str]:
        """List available MCP server names."""
        return list(MCP_SERVERS.keys())
    
    def get_codex_mcp_params(self) -> dict:
        """Get params for Codex MCP server (most common use case)."""
        config = MCP_SERVERS["codex"]
        return {
            "command": config.command,
            "args": config.args,
        }
    
    async def start_server(self, name: str) -> Any:
        """Start an MCP server (requires agents SDK)."""
        try:
            from agents.mcp import MCPServerStdio
        except ImportError:
            raise ImportError("openai-agents-sdk required. pip install openai-agents-sdk")
        
        config = self.get_server_config(name)
        if not config:
            raise ValueError(f"Unknown MCP server: {name}")
        
        params = {
            "command": config.command,
            "args": config.args,
        }
        if config.env:
            params["env"] = config.env
        
        server = MCPServerStdio(name=config.name, params=params)
        self.active_servers[name] = server
        return server
    
    async def stop_all(self):
        """Stop all active MCP servers."""
        for name, server in self.active_servers.items():
            try:
                await server.__aexit__(None, None, None)
            except Exception:
                pass
        self.active_servers.clear()


def generate_config_yaml(servers: List[str] = None) -> str:
    """Generate config.yaml content for specified MCP servers."""
    servers = servers or ["codex"]

    lines = [
        "# Oh My Codex - MCP Configuration",
        "# Generated by orchestrator/mcp.py",
        "",
        "mcp:",
    ]

    for name in servers:
        config = MCP_SERVERS.get(name)
        if not config:
            continue

        lines.append(f"")
        lines.append(f"  {name}:")
        lines.append(f"    command: \"{config.command}\"")
        lines.append(f"    args: {config.args}")

        if config.env:
            lines.append(f"    env:")
            for k, v in config.env.items():
                lines.append(f"      {k}: \"{v}\"")

    return "\n".join(lines)
