"""
Unified MCP Client for connecting to MCP servers.
Works with any MCP server via stdio connection.
"""

import os
import sys
import json
from typing import Optional, Dict, Any, List
from pathlib import Path

from mcp.client.session import ClientSession
from mcp.client.stdio import stdio_client, StdioServerParameters


class MCPClient:
    """Generic client for interacting with any MCP server via stdio."""

    def __init__(self, server_script_path: str | Path, client_name: str = "MCP"):
        """
        Initialize MCP client.

        Args:
            server_script_path: Path to the MCP server Python script
            client_name: Name for logging/debugging (e.g., "OpenParlData", "BFS")
        """
        self.server_script = Path(server_script_path)
        self.client_name = client_name
        self.session: Optional[ClientSession] = None
        self.available_tools: List[Dict[str, Any]] = []

    async def connect(self):
        """Connect to the MCP server."""
        if not self.server_script.exists():
            raise FileNotFoundError(
                f"{self.client_name} MCP server script not found at {self.server_script}"
            )

        # Server parameters for stdio connection
        server_params = StdioServerParameters(
            command=sys.executable,  # Python interpreter
            args=[str(self.server_script)],
            env=None
        )

        # Create stdio client context
        self.stdio_context = stdio_client(server_params)
        read, write = await self.stdio_context.__aenter__()

        # Create session
        self.session = ClientSession(read, write)
        await self.session.__aenter__()

        # Initialize and get available tools
        await self.session.initialize()

        # List available tools
        tools_result = await self.session.list_tools()
        self.available_tools = [
            {
                "name": tool.name,
                "description": tool.description,
                "input_schema": tool.inputSchema
            }
            for tool in tools_result.tools
        ]

        return self.available_tools

    async def disconnect(self):
        """Disconnect from the MCP server."""
        if self.session:
            await self.session.__aexit__(None, None, None)
        if hasattr(self, 'stdio_context'):
            await self.stdio_context.__aexit__(None, None, None)

    async def call_tool(self, tool_name: str, arguments: Dict[str, Any]) -> str:
        """
        Call an MCP tool with given arguments.

        Args:
            tool_name: Name of the tool to call
            arguments: Dictionary of arguments for the tool

        Returns:
            Tool response as string
        """
        if not self.session:
            raise RuntimeError(
                f"Not connected to {self.client_name} MCP server. Call connect() first."
            )

        # Wrap arguments in 'params' key as expected by MCP server
        tool_arguments = {"params": arguments}

        # DEBUG: Log MCP payload before sending
        print(f"\n📤 [{self.client_name}Client] Sending to MCP server:")
        print(f"  Tool: {tool_name}")
        print(f"  Wrapped payload: {tool_arguments}")
        print(f"  Payload types: {dict((k, type(v).__name__) for k, v in tool_arguments.items())}")

        # Call the tool
        result = await self.session.call_tool(tool_name, arguments=tool_arguments)

        # Extract text content from result
        if result.content:
            # MCP returns list of content blocks
            text_parts = []
            for content in result.content:
                if hasattr(content, 'text'):
                    text_parts.append(content.text)
                elif isinstance(content, dict) and 'text' in content:
                    text_parts.append(content['text'])
            return "\n".join(text_parts)

        return "No response from tool"

    def get_tool_info(self) -> List[Dict[str, Any]]:
        """Get information about available tools."""
        return self.available_tools


# Convenience factory functions for backwards compatibility

def create_openparldata_client() -> MCPClient:
    """Create an MCP client for OpenParlData."""
    server_script = Path(__file__).parent.parent / "mcp_openparldata" / "openparldata_mcp.py"
    return MCPClient(server_script, client_name="OpenParlData")


def create_bfs_client() -> MCPClient:
    """Create an MCP client for BFS."""
    server_script = Path(__file__).parent.parent / "mcp_bfs" / "bfs_mcp_server.py"
    return MCPClient(server_script, client_name="BFS")