"""
MCP Connection Manager

Handles spawning, connecting to, and communicating with external MCP servers.
This module is designed to be integrated into the AgenticWerx MCP client.

Ported from JavaScript implementation to Python with full persistence support.
"""

import asyncio
import json
import logging
import os
import subprocess
from datetime import datetime
from pathlib import Path
from typing import Any, Dict, List, Optional

logger = logging.getLogger(__name__)


class MCPServerConnection:
    """Manages a single external MCP server connection"""

    def __init__(self, server_config: Dict[str, Any]):
        self.name = server_config["name"]
        self.command = server_config["command"]
        self.args = server_config["args"]
        self.env = server_config.get("env", {})
        self.process: Optional[subprocess.Popen] = None
        self.is_initialized = False
        self.server_info: Optional[Dict] = None
        self.capabilities: Optional[Dict] = None
        self.tools: List[Dict] = []
        self.pending_requests: Dict[int, asyncio.Future] = {}
        self.request_id_counter = 1
        self.buffer = ""
        self._read_task: Optional[asyncio.Task] = None
        self._should_stop = False

    async def connect(self) -> Dict[str, Any]:
        """Start the external MCP server process and initialize connection"""
        try:
            # Prepare environment
            env = os.environ.copy()
            env.update(self.env)

            # Spawn the external MCP server
            self.process = subprocess.Popen(
                [self.command] + self.args,
                env=env,
                stdin=subprocess.PIPE,
                stdout=subprocess.PIPE,
                stderr=subprocess.PIPE,
                text=True,
                bufsize=1,  # Line buffered
            )

            logger.info(f"[{self.name}] Started process: {self.command} {' '.join(self.args)}")

            # Start background task to read stdout
            self._should_stop = False
            self._read_task = asyncio.create_task(self._read_stdout())

            # Send initialize request
            init_result = await self.send_initialize()
            self.is_initialized = True
            self.server_info = init_result.get("serverInfo")
            self.capabilities = init_result.get("capabilities")

            # Get list of available tools
            tools_result = await self.list_tools()
            self.tools = tools_result

            logger.info(
                f"[{self.name}] Connected successfully - {len(self.tools)} tools available"
            )

            return {
                "serverInfo": self.server_info,
                "capabilities": self.capabilities,
                "tools": self.tools,
            }

        except Exception as e:
            logger.error(f"[{self.name}] Failed to connect: {e}")
            self.disconnect()
            raise

    async def _read_stdout(self):
        """Read stdout in background task"""
        try:
            while not self._should_stop and self.process and self.process.poll() is None:
                # Read line from stdout (non-blocking via executor)
                line = await asyncio.get_event_loop().run_in_executor(
                    None, self.process.stdout.readline
                )

                if line:
                    self.handle_server_output(line)
                else:
                    # EOF reached
                    break

        except Exception as e:
            logger.error(f"[{self.name}] Error reading stdout: {e}")
        finally:
            logger.debug(f"[{self.name}] Stdout reader stopped")

    def handle_server_output(self, data: str):
        """Handle output from the external server"""
        self.buffer += data

        # Process complete JSON-RPC messages (newline-delimited)
        lines = self.buffer.split("\n")
        self.buffer = lines.pop() if lines else ""  # Keep incomplete line in buffer

        for line in lines:
            if not line.strip():
                continue

            try:
                message = json.loads(line)

                # Handle responses to our requests
                if "id" in message and message["id"] in self.pending_requests:
                    future = self.pending_requests.pop(message["id"])

                    if "error" in message:
                        error_msg = message["error"].get("message", "Request failed")
                        future.set_exception(Exception(error_msg))
                    else:
                        future.set_result(message.get("result"))

                # Handle notifications from server (method without id)
                elif "method" in message and "id" not in message:
                    logger.debug(f"[{self.name}] Notification: {message['method']}")

            except json.JSONDecodeError as e:
                logger.error(f"[{self.name}] Failed to parse message: {line[:100]}... Error: {e}")

    async def send_request(self, method: str, params: Optional[Dict] = None) -> Any:
        """Send a JSON-RPC request to the external server"""
        if not self.process or self.process.poll() is not None:
            raise Exception("Server not connected")

        request_id = self.request_id_counter
        self.request_id_counter += 1

        request = {
            "jsonrpc": "2.0",
            "id": request_id,
            "method": method,
            "params": params or {},
        }

        # Create future for this request
        future: asyncio.Future = asyncio.Future()
        self.pending_requests[request_id] = future

        # Send request
        try:
            request_json = json.dumps(request) + "\n"
            self.process.stdin.write(request_json)
            self.process.stdin.flush()

            logger.debug(f"[{self.name}] Sent request: {method} (id={request_id})")

            # Wait for response with timeout
            result = await asyncio.wait_for(future, timeout=30.0)
            return result

        except asyncio.TimeoutError:
            # Clean up pending request
            self.pending_requests.pop(request_id, None)
            raise Exception(f"Request timeout: {method}")
        except Exception as e:
            # Clean up pending request
            self.pending_requests.pop(request_id, None)
            raise

    async def send_initialize(self) -> Dict:
        """Send initialize request to external server"""
        return await self.send_request(
            "initialize",
            {
                "protocolVersion": "2024-11-05",
                "capabilities": {},
                "clientInfo": {"name": "agenticwerx-mcp-client", "version": "2.0.0"},
            },
        )

    async def list_tools(self) -> List[Dict]:
        """Get list of tools from external server"""
        result = await self.send_request("tools/list")
        return result.get("tools", [])

    async def call_tool(self, tool_name: str, args: Optional[Dict] = None) -> Any:
        """Call a tool on the external server"""
        result = await self.send_request(
            "tools/call", {"name": tool_name, "arguments": args or {}}
        )
        return result

    def disconnect(self):
        """Disconnect from the external server"""
        self._should_stop = True

        if self._read_task:
            self._read_task.cancel()
            self._read_task = None

        if self.process:
            try:
                self.process.terminate()
                # Give it a moment to terminate gracefully
                try:
                    self.process.wait(timeout=2)
                except subprocess.TimeoutExpired:
                    self.process.kill()
            except Exception as e:
                logger.error(f"[{self.name}] Error terminating process: {e}")

            self.process = None

        self.is_initialized = False
        logger.info(f"[{self.name}] Disconnected")


class MCPConnectionManager:
    """Manages multiple external MCP server connections with persistence"""

    def __init__(self, options: Optional[Dict[str, Any]] = None):
        options = options or {}
        self.connections: Dict[str, MCPServerConnection] = {}

        # Default persistence path: ~/.agenticwerx/mcp-connections.json
        default_path = Path.home() / ".agenticwerx" / "mcp-connections.json"
        persistence_path = options.get("persistencePath")
        # Use default if None or not provided
        self.persistence_path = Path(persistence_path) if persistence_path else default_path

        self.auto_reconnect = options.get("autoReconnect", True)
        self.is_initialized = False

    async def initialize(self):
        """Initialize the connection manager and restore saved connections"""
        if self.is_initialized:
            return

        try:
            # Ensure persistence directory exists
            self.persistence_path.parent.mkdir(parents=True, exist_ok=True)

            # Load and restore saved connections
            if self.auto_reconnect:
                await self.restore_connections()

            self.is_initialized = True
            logger.info("Connection manager initialized")

        except Exception as e:
            logger.error(f"Failed to initialize connection manager: {e}")
            self.is_initialized = True  # Continue even if restore fails

    async def save_connections(self):
        """Save current connections to disk"""
        try:
            connections_data = []

            for name, connection in self.connections.items():
                if connection.is_initialized:
                    connections_data.append(
                        {
                            "name": connection.name,
                            "command": connection.command,
                            "args": connection.args,
                            "env": connection.env,
                            "connectedAt": datetime.utcnow().isoformat() + "Z",
                        }
                    )

            with open(self.persistence_path, "w") as f:
                json.dump(connections_data, f, indent=2)

            logger.info(
                f"Saved {len(connections_data)} connection(s) to {self.persistence_path}"
            )

        except Exception as e:
            logger.error(f"Failed to save connections: {e}")

    async def restore_connections(self) -> List[Dict]:
        """Restore saved connections from disk"""
        try:
            if not self.persistence_path.exists():
                logger.info("No saved connections found")
                return []

            with open(self.persistence_path, "r") as f:
                saved_connections = json.load(f)

            logger.info(f"Restoring {len(saved_connections)} saved connection(s)...")

            results = []
            for config in saved_connections:
                try:
                    logger.info(f"Reconnecting to {config['name']}...")
                    result = await self.connect(config)
                    results.append(
                        {"name": config["name"], "success": True, "result": result}
                    )
                except Exception as e:
                    logger.error(f"Failed to reconnect to {config['name']}: {e}")
                    results.append(
                        {"name": config["name"], "success": False, "error": str(e)}
                    )

            successful = sum(1 for r in results if r["success"])
            logger.info(
                f"Successfully restored {successful}/{len(saved_connections)} connection(s)"
            )

            return results

        except Exception as e:
            logger.error(f"Failed to restore connections: {e}")
            return []

    async def clear_saved_connections(self):
        """Clear saved connections from disk"""
        try:
            if self.persistence_path.exists():
                with open(self.persistence_path, "w") as f:
                    json.dump([], f, indent=2)
                logger.info("Cleared saved connections")
        except Exception as e:
            logger.error(f"Failed to clear saved connections: {e}")

    async def connect(self, server_config: Dict[str, Any]) -> Dict[str, Any]:
        """
        Connect to an external MCP server

        Args:
            server_config: Server configuration with keys:
                - name: Server name
                - command: Command to run (e.g., 'uvx', 'npx')
                - args: Command arguments
                - env: Environment variables
                - persist: Whether to persist this connection (default: True)

        Returns:
            Connection result with server info and tools
        """
        # Ensure manager is initialized
        if not self.is_initialized:
            await self.initialize()

        # Check if already connected
        if server_config["name"] in self.connections:
            raise Exception(f"Already connected to {server_config['name']}")

        connection = MCPServerConnection(server_config)

        try:
            result = await connection.connect()
            self.connections[server_config["name"]] = connection

            # Save connections to disk (unless explicitly disabled)
            if server_config.get("persist", True):
                await self.save_connections()

            return {
                "success": True,
                "serverName": server_config["name"],
                "serverInfo": result["serverInfo"],
                "capabilities": result["capabilities"],
                "tools": result["tools"],
            }

        except Exception as e:
            connection.disconnect()
            raise

    async def disconnect(
        self, server_name: str, remove_persistence: bool = True
    ) -> Dict:
        """Disconnect from an external MCP server"""
        connection = self.connections.get(server_name)
        if connection:
            connection.disconnect()
            del self.connections[server_name]

            # Update saved connections
            if remove_persistence:
                await self.save_connections()

            return {"success": True}

        return {"success": False, "error": "Server not connected"}

    def list_connections(self) -> List[Dict]:
        """Get list of connected servers"""
        connections = []
        for name, connection in self.connections.items():
            connections.append(
                {
                    "name": name,
                    "serverInfo": connection.server_info,
                    "tools": [t["name"] for t in connection.tools],
                    "isConnected": connection.is_initialized,
                }
            )
        return connections

    async def call_tool(
        self, server_name: str, tool_name: str, args: Optional[Dict] = None
    ) -> Any:
        """Call a tool on a connected external server"""
        connection = self.connections.get(server_name)
        if not connection:
            raise Exception(f"Not connected to {server_name}")

        if not connection.is_initialized:
            raise Exception(f"{server_name} is not initialized")

        return await connection.call_tool(tool_name, args)

    def get_connection(self, server_name: str) -> Optional[MCPServerConnection]:
        """Get connection by server name"""
        return self.connections.get(server_name)

    def is_connected(self, server_name: str) -> bool:
        """Check if connected to a server"""
        connection = self.connections.get(server_name)
        return connection is not None and connection.is_initialized

    async def disconnect_all(self, clear_persistence: bool = False):
        """Disconnect all servers"""
        for connection in self.connections.values():
            connection.disconnect()
        self.connections.clear()

        if clear_persistence:
            await self.clear_saved_connections()
        else:
            await self.save_connections()

    def get_persistence_path(self) -> str:
        """Get persistence file path"""
        return str(self.persistence_path)

    async def get_saved_connections(self) -> List[Dict]:
        """Get saved connections without connecting"""
        try:
            if not self.persistence_path.exists():
                return []

            with open(self.persistence_path, "r") as f:
                return json.load(f)

        except Exception as e:
            logger.error(f"Failed to read saved connections: {e}")
            return []
