"""MCP (Model Context Protocol) client for integrating external tool servers.

Uses the official `mcp` Python SDK for proper protocol handling.

Configuration (.mcp.json):
    {
        "mcpServers": {
            "server-name": {
                "command": "npx",           # Required: executable to run
                "args": ["-y", "@pkg"],     # Optional: command arguments
                "env": {"KEY": "value"},    # Optional: environment variables
                "description": "...",       # Optional: human-readable description
                "enabled": true             # Optional: set to false to disable
            }
        }
    }

Supported server types:
    - NPM packages: {"command": "npx", "args": ["-y", "@scope/package"]}
    - Python packages: {"command": "uvx", "args": ["package-name"]}
    - Local scripts: {"command": "node", "args": ["./path/to/server.js"]}
    - Docker: {"command": "docker", "args": ["run", "-i", "image:tag"]}
"""

from __future__ import annotations

import asyncio
import json
import os
import shutil
import threading
from dataclasses import dataclass
from pathlib import Path
from typing import Any

from logging_config import get_logger
from mcp import ClientSession, StdioServerParameters
from mcp.client.stdio import stdio_client

from skills import Skill, SkillParameter

log = get_logger(__name__)


@dataclass
class MCPServerConfig:
    """Configuration for an MCP server."""

    name: str
    command: str
    args: list[str]
    env: dict[str, str] | None = None
    description: str = ""
    enabled: bool = True

    @classmethod
    def from_dict(cls, name: str, data: dict[str, Any]) -> MCPServerConfig:
        """Create from config dictionary."""
        return cls(
            name=name,
            command=data["command"],
            args=data.get("args", []),
            env=data.get("env"),
            description=data.get("description", ""),
            enabled=data.get("enabled", True),
        )

    def validate(self) -> list[str]:
        """Validate the configuration. Returns list of errors."""
        errors = []

        if not self.command:
            errors.append(f"[{self.name}] 'command' is required")
        elif not shutil.which(self.command):
            errors.append(f"[{self.name}] command '{self.command}' not found in PATH")

        return errors


class MCPClient:
    """Async MCP client using the official SDK."""

    def __init__(self, config: MCPServerConfig) -> None:
        self.config = config
        self.name = config.name
        self._session: ClientSession | None = None
        self._stdio_context: Any = None
        self._session_context: Any = None
        self._log = log.bind(component="mcp_client", server=config.name)

    async def connect(self) -> None:
        """Connect to the MCP server."""
        # Merge environment variables
        env = {**os.environ}
        if self.config.env:
            env.update(self.config.env)

        server_params = StdioServerParameters(
            command=self.config.command,
            args=self.config.args,
            env=env,
        )

        self._log.debug(
            "connecting_to_mcp_server",
            command=self.config.command,
            args=self.config.args,
        )

        # Create the stdio client context
        self._stdio_context = stdio_client(server_params)
        read_stream, write_stream = await self._stdio_context.__aenter__()

        # Create and initialize session
        self._session = ClientSession(read_stream, write_stream)
        self._session_context = self._session
        await self._session_context.__aenter__()

        # Initialize the connection
        result = await asyncio.wait_for(
            self._session.initialize(),
            timeout=30,
        )

        self._log.debug(
            "mcp_server_connected",
            server_info=result.serverInfo.model_dump() if result.serverInfo else None,
        )

    async def list_tools(self) -> list[dict]:
        """List available tools from the server."""
        if not self._session:
            raise RuntimeError("Not connected")

        result = await self._session.list_tools()
        tools = []
        for tool in result.tools:
            tool_dict = {
                "name": tool.name,
                "description": tool.description or "",
            }
            if hasattr(tool, "inputSchema") and tool.inputSchema:
                tool_dict["inputSchema"] = tool.inputSchema
            tools.append(tool_dict)
        return tools

    async def call_tool(self, name: str, arguments: dict[str, Any]) -> str:
        """Call a tool and return the result."""
        if not self._session:
            raise RuntimeError("Not connected")

        result = await self._session.call_tool(name, arguments)

        # Extract text content from result
        parts = []
        for content in result.content:
            if content.type == "text":
                parts.append(content.text)
            elif content.type == "image":
                parts.append(f"[Image: {content.mimeType}]")
            else:
                parts.append(str(content))

        return "\n".join(parts) if parts else json.dumps({"ok": True})

    async def close(self) -> None:
        """Close the connection."""
        # Note: We don't call __aexit__ here because the contexts must be
        # exited in the same task they were entered. The subprocess will
        # be terminated when Python exits or when we explicitly kill it.
        self._session = None
        self._session_context = None
        self._stdio_context = None


class MCPManager:
    """Manages multiple MCP server connections and exposes their tools as Skills.

    This is a synchronous wrapper around async MCP operations for compatibility
    with the existing Chatixia architecture. Uses a persistent event loop to maintain
    MCP connections across tool calls.

    Example usage:
        manager = MCPManager()
        manager.connect_from_config(".mcp.json")
        skills = manager.get_skills()  # Returns dict of Skill objects
    """

    def __init__(self) -> None:
        self._clients: dict[str, MCPClient] = {}
        self._configs: dict[str, MCPServerConfig] = {}
        self._tools_cache: dict[str, list[dict]] = {}
        self._log = log.bind(component="mcp_manager")
        # Persistent event loop for MCP operations
        self._loop: asyncio.AbstractEventLoop | None = None
        self._loop_thread: threading.Thread | None = None

    def _ensure_loop(self) -> asyncio.AbstractEventLoop:
        """Ensure we have a running event loop in a background thread."""
        if self._loop is None or not self._loop.is_running():
            self._loop = asyncio.new_event_loop()

            def run_loop():
                asyncio.set_event_loop(self._loop)
                self._loop.run_forever()

            self._loop_thread = threading.Thread(target=run_loop, daemon=True)
            self._loop_thread.start()
            # Give the loop a moment to start
            import time

            time.sleep(0.05)

        return self._loop

    def _run_async(self, coro: Any) -> Any:
        """Run an async coroutine in the persistent event loop."""
        loop = self._ensure_loop()
        future = asyncio.run_coroutine_threadsafe(coro, loop)
        return future.result(timeout=120)

    def connect_server(self, config: MCPServerConfig) -> bool:
        """Connect to a single MCP server from config. Returns True on success."""
        if not config.enabled:
            self._log.debug("mcp_server_disabled", server=config.name)
            return False

        # Validate configuration
        errors = config.validate()
        if errors:
            for err in errors:
                self._log.warning("mcp_config_error", error=err)
            return False

        server_log = self._log.bind(server=config.name)

        async def _connect() -> None:
            client = MCPClient(config)
            try:
                await client.connect()
                self._clients[config.name] = client
                self._configs[config.name] = config

                # Cache tools
                tools = await client.list_tools()
                self._tools_cache[config.name] = tools
                server_log.info(
                    "mcp_server_ready",
                    tool_count=len(tools),
                    description=config.description,
                )
            except Exception as exc:
                server_log.warning("mcp_server_failed", error=str(exc))
                await client.close()
                raise

        try:
            self._run_async(_connect())
            return True
        except Exception as exc:
            server_log.warning("mcp_connect_error", error=str(exc))
            return False

    def connect(self, name: str, command: list[str], env: dict[str, str] | None = None) -> bool:
        """Start and initialize a single MCP server (legacy API)."""
        config = MCPServerConfig(
            name=name,
            command=command[0],
            args=command[1:] if len(command) > 1 else [],
            env=env,
        )
        return self.connect_server(config)

    def connect_from_config(self, config_path: str | Path) -> dict[str, bool]:
        """Load .mcp.json and connect to all configured servers.

        Returns a dict of {server_name: success} for each server.
        """
        path = Path(config_path)
        if not path.exists():
            self._log.debug("no_mcp_config", path=str(path))
            return {}

        try:
            with open(path) as f:
                config = json.load(f)
        except json.JSONDecodeError as exc:
            self._log.error("mcp_config_parse_error", path=str(path), error=str(exc))
            return {}

        results: dict[str, bool] = {}
        servers = config.get("mcpServers", {})

        if not servers:
            self._log.debug("no_mcp_servers_configured", path=str(path))
            return {}

        self._log.info("loading_mcp_servers", count=len(servers))

        for name, server_cfg in servers.items():
            try:
                server_config = MCPServerConfig.from_dict(name, server_cfg)
                results[name] = self.connect_server(server_config)
            except KeyError as exc:
                self._log.warning(
                    "mcp_config_missing_field",
                    server=name,
                    field=str(exc),
                )
                results[name] = False

        connected = sum(1 for v in results.values() if v)
        self._log.info("mcp_servers_loaded", connected=connected, total=len(results))

        return results

    def get_skills(self) -> dict[str, Skill]:
        """Return all MCP tools wrapped as Skill objects with handlers."""
        skills: dict[str, Skill] = {}

        for server_name, tools in self._tools_cache.items():
            client = self._clients.get(server_name)
            if not client:
                continue

            for tool in tools:
                tool_name = tool["name"]
                parameters = _schema_to_params(tool.get("inputSchema", {}))

                # Create a synchronous handler that wraps the async call
                def make_handler(_client: MCPClient, _tool_name: str, _mgr: MCPManager):
                    def handler(**kwargs: Any) -> str:
                        async def _call():
                            return await _client.call_tool(_tool_name, kwargs)

                        return _mgr._run_async(_call())

                    return handler

                skill = Skill(
                    name=tool_name,
                    description=tool.get("description", ""),
                    version="1.0.0",
                    prompt_template="",
                    parameters=parameters,
                    handler=make_handler(client, tool_name, self),
                )
                skills[tool_name] = skill
                self._log.debug(
                    "mcp_tool_registered",
                    server=server_name,
                    tool=tool_name,
                )

        return skills

    def list_servers(self) -> list[dict[str, Any]]:
        """List all connected MCP servers with their tool counts."""
        servers = []
        for name, config in self._configs.items():
            tools = self._tools_cache.get(name, [])
            servers.append(
                {
                    "name": name,
                    "description": config.description,
                    "command": config.command,
                    "tool_count": len(tools),
                    "tools": [t["name"] for t in tools],
                }
            )
        return servers

    def close(self) -> None:
        """Shut down all MCP server connections."""
        for name, client in self._clients.items():
            self._log.debug("closing_mcp_server", server=name)
            # Clear references - subprocesses terminate when Python exits
            client._session = None
            client._session_context = None
            client._stdio_context = None

        self._clients.clear()
        self._configs.clear()
        self._tools_cache.clear()


def _schema_to_params(schema: dict) -> list[SkillParameter]:
    """Convert a JSON Schema object into a list of SkillParameter."""
    params: list[SkillParameter] = []
    properties = schema.get("properties", {})
    required = set(schema.get("required", []))

    for name, prop in properties.items():
        params.append(
            SkillParameter(
                name=name,
                param_type=prop.get("type", "string"),
                description=prop.get("description", ""),
                required=name in required,
                default=prop.get("default"),
            )
        )
    return params
