import asyncio
import json
import sys
from typing import List, Dict, Any
from pathlib import Path

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

class MCPToolAdapter:
    def __init__(self, config_path: str = "mcp_server_config.json"):
        self.config_path = config_path
        self.servers: Dict[str, Dict] = {}
        #self.exit_stack: Dict[str, Any] = {}
        self.tool_registry: List[Dict[str, Any]] = []

    def _load_config(self) -> Dict[str, Any]:
        """Lädt die Server-Konfiguration aus der JSON-Datei."""
        path = Path(__file__).parent / self.config_path
        if not path.exists():
            print(f"Config file not found: {path}")
            return {}
        
        try:
            with open(path, 'r') as f:
                return json.load(f)
        except json.JSONDecodeError as e:
            print(f"Error decoding JSON config: {e}")
            return {}

    async def initialize_all_servers(self):
        """Lädt die Konfiguration und fragt alle Server ab, um die Tools zu registrieren."""
        print("Initializing MCP sessions...")
        config = self._load_config()
        print(f"Loaded config for servers: {list(config.keys())}")

        for server_name, params in config.items():
            print(f"Testing connection to {server_name}...")

            self.servers[server_name] = params
            server_script = str(Path(__file__).parent / params["args"][0])
            if params.get("command") in ["py", "python", "python3"]:
                server_command = sys.executable
            else:
                server_command = params["command"]
            
            server_params = StdioServerParameters(
                command=server_command,
                args=[server_script],
            )
            
            try:
               # Verbindung aufbauen
                async with stdio_client(server_params) as (read_stream, write_stream):
                    print(f"Connected to {server_name}. Initializing session...")
                    async with ClientSession(read_stream, write_stream) as session:
                        await session.initialize()
                        print(f"Session initialized for {server_name}. Requesting tools...")
                        result = await session.list_tools()
                        print(f"Tools received from {server_name}: {result}")
                        tools = result.tools
                print(f"Tools received from {server_name}: {result}")
                #tools = getattr(result, 'tools', [])

                for tool in tools:
                    t_params = tool.inputSchema.get("properties", {})
                    if t_params:
                        param_lines = []
                        for pname, pinfo in t_params.items():
                            ptype = pinfo.get("type", "any")
                            pdesc = pinfo.get("description", "")
                            param_lines.append(f"    - {pname} ({ptype}): {pdesc}")
                        param_str = "\n".join(param_lines)
                    else:
                        param_str = "    (none)"

                    t_definition = f"- {tool.name}: {tool.description}\nParameters:\n{param_str}"


                    self.tool_registry.append({
                        "server": server_name,
                        "tool_name": tool.name,
                        "tool_description": t_definition
                    })

                    print(f"Registered tool '{tool.name}' from {server_name}.")
                
                print(f"Session for {server_name} ready. {len(tools)} tools found.")


            except Exception as e:
                print(f"Failed to initialize {server_name}: {e}")

    def get_all_tools(self) -> List[Dict[str, Any]]:
        """Gibt alle gesammelten Tools zurück."""
        return self.tool_registry

    async def call_tool(self, tool_name: str, arguments: Dict[str, Any]):
        """Findet den richtigen Server für ein Tool und führt es aus."""
        # Suche in der Registry nach dem passenden Server
        tool_entry = next((t for t in self.tool_registry if t["tool_name"] == tool_name), None)
        
        if not tool_entry:
            print(f"Tool '{tool_name}' not found in MCP adapter registry.")
            return f"Error: Tool '{tool_name}' not found in registry."

        server_name = tool_entry["server"]
        s_params = self.servers.get(server_name)

        if s_params:
            server_script = str(Path(__file__).parent / s_params["args"][0])
            if s_params.get("command") in ["py", "python", "python3"]:
                server_command = sys.executable
            else:
                server_command = s_params["command"]
            
            server_params = StdioServerParameters(
                command=server_command,
                args=[server_script],
            )

            try:
                async with stdio_client(server_params) as (read_stream, write_stream):
                    async with ClientSession(read_stream, write_stream) as session:
                        await session.initialize()
                        result = await session.call_tool(tool_name, arguments)
                return result
            except Exception as e:
                return f"Error calling tool '{tool_name}' on server '{server_name}': {str(e)}"
        
        return f"Error: Session for server '{server_name}' not active."

    async def shutdown_all_sessions(self):
        """Schließt alle offenen Verbindungen sauber."""
        for server_name, (transport_gen, session) in self.exit_stack.items():
            try:
                await session.__aexit__(None, None, None)
                await transport_gen.__aexit__(None, None, None)
                print(f"Session for {server_name} shut down.")
            except Exception as e:
                print(f"Error during shutdown of {server_name}: {e}")

def main():
    adapter = MCPToolAdapter()
    asyncio.run(adapter.initialize_all_servers())
    print("All servers initialized. Registered tools:")
    for tool in adapter.get_all_tools():
        print(f"- {tool['tool_name']} (from {tool['server']})")


if __name__ == "__main__":
    main()