import asyncio
import json
import logging
import sys
from typing import Dict, List, Any, Optional, Tuple
from dataclasses import dataclass, asdict
from contextlib import asynccontextmanager
import subprocess
import os
from enum import Enum
import anthropic
from anthropic import Anthropic
import aiohttp
from openai import AsyncOpenAI
from llama_api_client import LlamaAPIClient

# Configure logging
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)

class ServerType(Enum):
    LOCAL = "local"
    SSE = "sse"
    HTTP = "http"

@dataclass
class MCPServerConfig:
    """Configuration for an MCP server"""
    name: str
    server_type: ServerType
    command: List[str] = None
    args: List[str] = None
    env: Dict[str, str] = None
    description: str = ""
    url: str = None
    port: int = None
    host: str = None
    api_key: str = None
    auth_token: str = None
    auth_header: str = "Authorization"
    timeout: int = 30
    headers: Dict[str, str] = None
    ssl_verify: bool = False
    session_id: str = None

@dataclass
class MCPMessage:
    """MCP protocol message"""
    jsonrpc: str = "2.0"
    id: Optional[str] = None
    method: Optional[str] = None
    params: Optional[Dict[str, Any]] = None
    result: Optional[Any] = None
    error: Optional[Dict[str, Any]] = None

class MCPServerConnection:
    """Manages connection to a single MCP server"""
    
    def __init__(self, config: MCPServerConfig):
        self.config = config
        self.process = None
        self.initialized = False
        self.tools = {}
        self.resources = {}
        self.message_id = 0
        
    async def start(self):
        """Start the MCP server process"""
        try:
            cmd = self.config.command + (self.config.args or [])
            env = os.environ.copy()
            if self.config.env:
                env.update(self.config.env)
                
            self.process = await asyncio.create_subprocess_exec(
                *cmd,
                stdin=asyncio.subprocess.PIPE,
                stdout=asyncio.subprocess.PIPE,
                stderr=asyncio.subprocess.PIPE,
                env=env
            )
            
            # Initialize the server
            await self.initialize()
            logger.info(f"Started MCP server: {self.config.name}")
            
        except Exception as e:
            logger.error(f"Failed to start MCP server {self.config.name}: {e}")
            raise
    
    async def initialize(self):
        """Initialize the MCP server connection"""
        # Send initialization message
        init_msg = MCPMessage(
            id=str(self.message_id),
            method="initialize",
            params={
                "protocolVersion": "2024-11-05",
                "capabilities": {
                    "tools": {},
                    "resources": {}
                },
                "clientInfo": {
                    "name": "multi-mcp-client",
                    "version": "1.0.0"
                }
            }
        )
        
        response = await self.send_message(init_msg)
        if response and not response.error:
            self.initialized = True
            
            # Send initialized notification
            initialized_msg = MCPMessage(method="notifications/initialized")
            await self.send_message(initialized_msg)
            
            # List available tools and resources
            await self.list_tools()
            await self.list_resources()
        else:
            raise Exception(f"Failed to initialize server {self.config.name}")
    
    async def send_message(self, message: MCPMessage) -> Optional[MCPMessage]:
        """Send a message to the MCP server"""
        if not self.process:
            raise Exception("Server not started")
            
        if message.id is None and message.method != "notifications/initialized":
            self.message_id += 1
            message.id = str(self.message_id)
            
        # Send message
        msg_str = json.dumps(asdict(message), separators=(',', ':'))
        self.process.stdin.write(f"{msg_str}\n".encode())
        await self.process.stdin.drain()
        
        # Read response (if expecting one)
        if message.method and not message.method.startswith("notifications/"):
            try:
                response_line = await self.process.stdout.readline()
                if response_line:
                    response_data = json.loads(response_line.decode().strip())
                    return MCPMessage(**response_data)
            except Exception as e:
                logger.error(f"Error reading response from {self.config.name}: {e}")
                
        return None
    
    async def list_tools(self):
        """List available tools from the server"""
        msg = MCPMessage(method="tools/list")
        response = await self.send_message(msg)
        if response and response.result:
            self.tools = {tool['name']: tool for tool in response.result.get('tools', [])}
            logger.info(f"Server {self.config.name} has {len(self.tools)} tools")
    
    async def list_resources(self):
        """List available resources from the server"""
        msg = MCPMessage(method="resources/list")
        response = await self.send_message(msg)
        if response and response.result:
            self.resources = {res['uri']: res for res in response.result.get('resources', [])}
            logger.info(f"Server {self.config.name} has {len(self.resources)} resources")
    
    async def call_tool(self, tool_name: str, arguments: Dict[str, Any]) -> Any:
        """Call a tool on the server"""
        if tool_name not in self.tools:
            raise Exception(f"Tool {tool_name} not found on server {self.config.name}")
            
        msg = MCPMessage(
            method="tools/call",
            params={
                "name": tool_name,
                "arguments": arguments
            }
        )
        
        response = await self.send_message(msg)
        if response and response.result:
            return response.result
        elif response and response.error:
            raise Exception(f"Tool call failed: {response.error}")
        else:
            raise Exception("No response from tool call")
    
    async def call_tool_openai(self, tool_name: str, arguments: Dict[str, Any]) -> Any:
        """Call a tool on the server"""
        if tool_name not in self.tools:
            raise Exception(f"Tool {tool_name} not found on server {self.config.name}")
            
        msg = MCPMessage(
            method="tools/call",
            params={
                "id": self.message_id,
                "name": tool_name,
                "arguments": arguments
            }
        )
        self.message_id += 1
        
        response = await self.send_message(msg)
        if response and response.result:
            return response.result
        elif response and response.error:
            raise Exception(f"Tool call failed: {response.error}")
        else:
            raise Exception("No response from tool call")

    async def read_resource(self, uri: str) -> Any:
        """Read a resource from the server"""
        if uri not in self.resources:
            raise Exception(f"Resource {uri} not found on server {self.config.name}")
            
        msg = MCPMessage(
            method="resources/read",
            params={"uri": uri}
        )
        
        response = await self.send_message(msg)
        if response and response.result:
            return response.result
        elif response and response.error:
            raise Exception(f"Resource read failed: {response.error}")
        else:
            raise Exception("No response from resource read")
    
    async def stop(self):
        """Stop the MCP server"""
        if self.process:
            self.process.terminate()
            await self.process.wait()
            logger.info(f"Stopped MCP server: {self.config.name}")

class HTTPMCPConnection:
    """HTTP MCP server connection"""
    
    def __init__(self, config: MCPServerConfig):
        self.config = config
        self.process = None
        self.initialized = False
        self.tools = {}
        self.resources = {}
        self.message_id = 0
        self.session = None
        self.base_url = self.config.url
        
    async def start(self):
        """Start the HTTP session"""
        try:
            headers = self.config.headers or {}
            if self.config.auth_token:
                headers[self.config.auth_header] = f"Bearer {self.config.auth_token}"
            elif self.config.api_key:
                headers["X-API-Key"] = self.config.api_key
            
            connector = aiohttp.TCPConnector(
                ssl=self.config.ssl_verify,
                limit=100,
                limit_per_host=30
            )
            
            timeout = aiohttp.ClientTimeout(total=self.config.timeout)
            self.session = aiohttp.ClientSession(
                headers=headers,
                connector=connector,
                timeout=timeout
            )
            
            self.connected = True
            await self.initialize()
            logger.info(f"Connected to HTTP MCP server: {self.config.name}")
            
        except Exception as e:
            logger.error(f"Failed to connect to HTTP MCP server {self.config.name}: {e}")
            raise
    
    async def initialize(self):
        """Initialize the MCP server connection"""
        # Send initialization message
        init_msg = MCPMessage(
            id=str(self.message_id),
            method="initialize",
            params={
                "protocolVersion": "2024-11-05",
                "capabilities": {
                    "tools": {},
                    "resources": {}
                },
                "clientInfo": {
                    "name": "multi-mcp-client",
                    "version": "1.0.0"
                }
            }
        )
        
        response = await self.send_message(init_msg)
        if response and not response.error:
            self.initialized = True
            
            # Send initialized notification
            initialized_msg = MCPMessage(method="notifications/initialized")
            await self.send_message(initialized_msg)
            
            # List available tools and resources
            await self.list_tools()
            await self.list_resources()
        else:
            raise Exception(f"Failed to initialize server {self.config.name}")

    async def send_message(self, message: MCPMessage) -> Optional[MCPMessage]:
        """Send a message via HTTP POST"""
        if not self.session or not self.connected:
            raise Exception("HTTP session not started")
            
        if message.id is None and message.method != "notifications/initialized":
            self.message_id += 1
            message.id = str(self.message_id)
        
        # Map MCP methods to HTTP endpoints
        endpoint_map = {
            "initialize": "/initialize",
            "tools/list": "/tools",
            "tools/call": "/tools/call",
            "resources/list": "/resources",
            "resources/read": "/resources/read",
            "notifications/initialized": "/notifications/initialized"
        }
        
        endpoint = endpoint_map.get(message.method, f"/{message.method}")
        url = f"{self.base_url.rstrip('/')}{endpoint}"
        
        try:
            async with self.session.post(url, json=asdict(message), headers=self.config.headers) as response:
                if response.status == 200 and response.content_type == "application/json":
                    data = await response.json()
                    # logger.info(f"data json: {data}")
                    return MCPMessage(**data)
                elif response.status == 200 and response.content_type == "text/event-stream":
                    data = await response.text()
                    # logger.info(f"response: {response}")
                    self.config.session_id = response.headers['mcp-session-id']
                    self.config.headers["mcp-session-id"] = self.config.session_id
                    # logger.info(f"data: {data}")
                    # response_data = json.loads(data.strip())
                    # logger.info(f"response: {response_data}")
                    return MCPMessage(data)
                elif response.status == 202:
                    data = await response.text()
                    return data
                else:
                    error_text = await response.text()
                    raise Exception(f"HTTP error {response.status}: {error_text}")
        except Exception as e:
            logger.error(f"HTTP request error for {self.config.name}: {e}")
            raise

    async def list_tools(self):
        """List available tools from the server"""
        msg = MCPMessage(method="tools/list")
        response = await self.send_message(msg)
        if response and response.result:
            self.tools = {tool['name']: tool for tool in response.result.get('tools', [])}
            logger.info(f"Server {self.config.name} has {len(self.tools)} tools")
    
    async def list_resources(self):
        """List available resources from the server"""
        msg = MCPMessage(method="resources/list")
        response = await self.send_message(msg)
        if response and response.result:
            self.resources = {res['uri']: res for res in response.result.get('resources', [])}
            logger.info(f"Server {self.config.name} has {len(self.resources)} resources")
    
    async def call_tool(self, tool_name: str, arguments: Dict[str, Any]) -> Any:
        """Call a tool on the server"""
        if tool_name not in self.tools:
            raise Exception(f"Tool {tool_name} not found on server {self.config.name}")
            
        msg = MCPMessage(
            method="tools/call",
            params={
                "name": tool_name,
                "arguments": arguments
            }
        )
        
        response = await self.send_message(msg)
        if response and response.result:
            return response.result
        elif response and response.error:
            raise Exception(f"Tool call failed: {response.error}")
        else:
            raise Exception("No response from tool call")


    async def call_tool_openai(self, tool_name: str, arguments: Dict[str, Any]) -> Any:
        """Call a tool on the server"""
        if tool_name not in self.tools:
            raise Exception(f"Tool {tool_name} not found on server {self.config.name}")
            
        msg = MCPMessage(
            method="tools/call",
            params={
                "id": self.message_id,
                "name": tool_name,
                "arguments": arguments
            }
        )
        self.message_id += 1
        
        response = await self.send_message(msg)
        if response and response.result:
            return response.result
        elif response and response.error:
            raise Exception(f"Tool call failed: {response.error}")
        else:
            raise Exception("No response from tool call")
    
    async def read_resource(self, uri: str) -> Any:
        """Read a resource from the server"""
        if uri not in self.resources:
            raise Exception(f"Resource {uri} not found on server {self.config.name}")
            
        msg = MCPMessage(
            method="resources/read",
            params={"uri": uri}
        )
        
        response = await self.send_message(msg)
        if response and response.result:
            return response.result
        elif response and response.error:
            raise Exception(f"Resource read failed: {response.error}")
        else:
            raise Exception("No response from resource read")
    
    async def stop(self):
        """Stop the MCP server"""
        if self.process:
            self.process.terminate()
            await self.process.wait()
            logger.info(f"Stopped MCP server: {self.config.name}")


class MultiMCPClient:
    """Client that manages multiple MCP server connections"""
    
    def __init__(self, anthropic_api_key: str, openai_api_key: str):
        self.servers: Dict[str, MCPServerConnection] = {}
        self.anthropic_client = Anthropic(api_key=anthropic_api_key, timeout=60.0)
        self.openai_client = AsyncOpenAI(api_key=openai_api_key)
        self.deepseek_client = AsyncOpenAI(api_key=openai_api_key, base_url="https://api.deepseek.com/v1")
        self.qwen_client = AsyncOpenAI(api_key=openai_api_key, base_url="https://dashscope.aliyuncs.com/compatible-mode/v1")
        self.llama_client = LlamaAPIClient(api_key=openai_api_key)

        
    async def add_server(self, config: MCPServerConfig):
        """Add and start a new MCP server"""
        if config.name in self.servers:
            raise Exception(f"Server {config.name} already exists")
            
        if config.server_type == ServerType.LOCAL:
            server = MCPServerConnection(config)
        elif config.server_type == ServerType.HTTP:
            server = HTTPMCPConnection(config)
        await server.start()
        self.servers[config.name] = server
        logger.info(f"Added server: {config.name}")
    
    async def remove_server(self, name: str):
        """Remove and stop an MCP server"""
        if name in self.servers:
            await self.servers[name].stop()
            del self.servers[name]
            logger.info(f"Removed server: {name}")
    
    def get_all_tools(self) -> Dict[str, Dict[str, Any]]:
        """Get all available tools from all servers"""
        all_tools = {}
        for server_name, server in self.servers.items():
            for tool_name, tool_info in server.tools.items():
                # Prefix tool name with server name to avoid conflicts
                prefixed_name = f"{server_name}__{tool_name}"
                all_tools[prefixed_name] = {
                    **tool_info,
                    "server": server_name
                }
        return all_tools
    
    def get_all_resources(self) -> Dict[str, Dict[str, Any]]:
        """Get all available resources from all servers"""
        all_resources = {}
        for server_name, server in self.servers.items():
            for uri, resource_info in server.resources.items():
                all_resources[uri] = {
                    **resource_info,
                    "server": server_name
                }
        return all_resources
    
    async def call_tool(self, tool_name: str, arguments: Dict[str, Any]) -> Any:
        """Call a tool on the appropriate server"""
        if '__' in tool_name:
            server_name, actual_tool_name = tool_name.split('__', 1)
            if server_name in self.servers:
                return await self.servers[server_name].call_tool(actual_tool_name, arguments)
        elif '_' in tool_name:
            server_name, actual_tool_name = tool_name.split('_', 1)
            if server_name in self.servers:
                return await self.servers[server_name].call_tool_openai(actual_tool_name, arguments)
        
        # Try to find the tool on any server
        for server in self.servers.values():
            if tool_name in server.tools:
                return await server.call_tool(tool_name, arguments)
        
        raise Exception(f"Tool {tool_name} not found on any server")
    
    async def read_resource(self, uri: str) -> Any:
        """Read a resource from the appropriate server"""
        for server in self.servers.values():
            if uri in server.resources:
                return await server.read_resource(uri)
        
        raise Exception(f"Resource {uri} not found on any server")
    
    def format_tools_for_claude(self) -> List[Dict[str, Any]]:
        """Format tools for Claude's function calling format"""
        tools = []
        for tool_name, tool_info in self.get_all_tools().items():
            claude_tool = {
                "name": tool_name,
                "description": tool_info.get("description", ""),
                "input_schema": tool_info.get("inputSchema", {})
            }
            tools.append(claude_tool)
        return tools
    
    def format_tools_for_openai(self) -> List[Dict[str, Any]]:
        """Format tools for Claude's function calling format"""
        tools = []
        for tool_name, tool_info in self.get_all_tools().items():
            openai_tool = {
                "type": "function",
                "function": {
                    "name": tool_name.replace("__", "_"),
                    "description": tool_info.get("description", ""),
                    "parameters": tool_info.get("inputSchema", {})
                }
            }
            tools.append(openai_tool)
        return tools
    
    async def chat_with_openai(self, message: str, max_turns: int = 5) -> str:
        tools = self.format_tools_for_openai()
        messages = [{"role": "user", "content": message}]
        for turn in range(max_turns):
            try:
                response = await self.openai_client.chat.completions.create(
                    model="gpt-4.1",
                    messages=messages,
                    tools=tools if tools else None,
                    tool_choice="auto" if tools else None
                )
                
                message = response.choices[0].message
                messages.append({
                    "role": "assistant",
                    "content": message.content or "",
                    "tool_calls": [tc.model_dump() for tc in message.tool_calls] if message.tool_calls else None
                })
                if message.tool_calls:
                    for tool_call in message.tool_calls:
                        try:
                            function_name = tool_call.function.name
                            function_args = json.loads(tool_call.function.arguments)
                            
                            # Call the MCP tool
                            result = await self.call_tool(function_name, function_args)
                            
                            # Add tool result to conversation
                            messages.append({
                                "role": "tool",
                                "tool_call_id": tool_call.id,
                                "content": json.dumps(result) if result else "Tool executed successfully"
                            })
                            
                        except Exception as e:
                            logger.error(f"Tool call failed: {e}")
                            messages.append({
                                "role": "tool",
                                "tool_call_id": tool_call.id,
                                "content": f"Error: {str(e)}"
                            })
                    
                    # Continue conversation after tool calls
                    continue
                else:
                    # No tool calls, return the response
                    return message.content or ""
                    
            except Exception as e:
                logger.error(f"OpenAI API call failed: {e}")
                return f"Error: {str(e)}"
        
        return current_messages[-1].get("content", "Max turns reached")

    async def chat_with_llama(self, message: str, max_turns: int = 5) -> str:
        tools = self.format_tools_for_openai()
        messages = [{"role": "user", "content": message}]
        for turn in range(max_turns):
            try:
                response = await self.llama_client.chat.completions.create(
                    model="model",
                    messages=messages,
                    tools=tools if tools else None,
                    tool_choice="auto" if tools else None
                )
                
                message = response.choices[0].message
                messages.append({
                    "role": "assistant",
                    "content": message.content or "",
                    "tool_calls": [tc.model_dump() for tc in message.tool_calls] if message.tool_calls else None
                })
                if message.tool_calls:
                    for tool_call in message.tool_calls:
                        try:
                            function_name = tool_call.function.name
                            function_args = json.loads(tool_call.function.arguments)
                            
                            # Call the MCP tool
                            result = await self.call_tool(function_name, function_args)
                            
                            # Add tool result to conversation
                            messages.append({
                                "role": "tool",
                                "tool_call_id": tool_call.id,
                                "content": json.dumps(result) if result else "Tool executed successfully"
                            })
                            
                        except Exception as e:
                            logger.error(f"Tool call failed: {e}")
                            messages.append({
                                "role": "tool",
                                "tool_call_id": tool_call.id,
                                "content": f"Error: {str(e)}"
                            })
                    
                    # Continue conversation after tool calls
                    continue
                else:
                    # No tool calls, return the response
                    return message.content or ""
                    
            except Exception as e:
                logger.error(f"Llama API call failed: {e}")
                return f"Error: {str(e)}"
        
        return current_messages[-1].get("content", "Max turns reached")

    
    async def chat_with_deepseek(self, message: str, max_turns: int = 5) -> str:
        tools = self.format_tools_for_openai()
        messages = [{"role": "user", "content": message}]
        for turn in range(max_turns):
            try:
                response = await self.deepseek_client.chat.completions.create(
                    model="deepseek-chat",
                    messages=messages,
                    tools=tools if tools else None,
                    tool_choice="auto" if tools else None
                )
                
                message = response.choices[0].message
                messages.append({
                    "role": "assistant",
                    "content": message.content or "",
                    "tool_calls": [tc.model_dump() for tc in message.tool_calls] if message.tool_calls else None
                })
                if message.tool_calls:
                    for tool_call in message.tool_calls:
                        try:
                            function_name = tool_call.function.name
                            function_args = json.loads(tool_call.function.arguments)
                            
                            # Call the MCP tool
                            result = await self.call_tool(function_name, function_args)
                            
                            # Add tool result to conversation
                            messages.append({
                                "role": "tool",
                                "tool_call_id": tool_call.id,
                                "content": json.dumps(result) if result else "Tool executed successfully"
                            })
                            
                        except Exception as e:
                            logger.error(f"Tool call failed: {e}")
                            messages.append({
                                "role": "tool",
                                "tool_call_id": tool_call.id,
                                "content": f"Error: {str(e)}"
                            })
                    
                    # Continue conversation after tool calls
                    continue
                else:
                    # No tool calls, return the response
                    return message.content or ""
                    
            except Exception as e:
                logger.error(f"DeepSeek API call failed: {e}")
                return f"Error: {str(e)}"
        
        return current_messages[-1].get("content", "Max turns reached")


    async def chat_with_qwen(self, message: str, max_turns: int = 5) -> str:
        tools = self.format_tools_for_openai()
        messages = [{"role": "user", "content": message}]
        for turn in range(max_turns):
            try:
                response = await self.qwen_client.chat.completions.create(
                    model="qwen-plus",
                    messages=messages,
                    tools=tools if tools else None,
                    tool_choice="auto" if tools else None
                )
                
                message = response.choices[0].message
                messages.append({
                    "role": "assistant",
                    "content": message.content or "",
                    "tool_calls": [tc.model_dump() for tc in message.tool_calls] if message.tool_calls else None
                })
                if message.tool_calls:
                    for tool_call in message.tool_calls:
                        try:
                            function_name = tool_call.function.name
                            function_args = json.loads(tool_call.function.arguments)
                            
                            # Call the MCP tool
                            result = await self.call_tool(function_name, function_args)
                            
                            # Add tool result to conversation
                            messages.append({
                                "role": "tool",
                                "tool_call_id": tool_call.id,
                                "content": json.dumps(result) if result else "Tool executed successfully"
                            })
                            
                        except Exception as e:
                            logger.error(f"Tool call failed: {e}")
                            messages.append({
                                "role": "tool",
                                "tool_call_id": tool_call.id,
                                "content": f"Error: {str(e)}"
                            })
                    
                    # Continue conversation after tool calls
                    continue
                else:
                    # No tool calls, return the response
                    return message.content or ""
                    
            except Exception as e:
                logger.error(f"QWen API call failed: {e}")
                return f"Error: {str(e)}"
        
        return current_messages[-1].get("content", "Max turns reached")


    async def chat_with_claude(self, message: str, max_tokens: int = 1024) -> str:
        """Chat with Claude using available MCP tools"""
        tools = self.format_tools_for_claude()
        
        # logger.info(f"tools: {tools}")

        messages = [{"role": "user", "content": message}]
        
        # Initial request to Claude
        try:
            response = self.anthropic_client.messages.create(
                model="claude-3-5-sonnet-20241022",
                max_tokens=1000,
                messages=messages,
                tools=tools if tools else None
            )
        except anthropic.APIConnectionError as e:
            logger.error(f"Connection error: {e}")
            raise
        
        except anthropic.RateLimitError as e:
            logger.error(f"Rate limit error: {e}")
        
        except anthropic.APIError as e:
            logger.error(f"API error: {e}")
            raise
        
        # Handle tool calls
        while response.stop_reason == "tool_use":
            # logger.info(f"content: {response.content}")
            messages.append({"role": "assistant", "content": response.content})
            
            # Execute tool calls
            tool_results = []
            for content_block in response.content:
                if content_block.type == "tool_use":
                    tool_name = content_block.name
                    tool_args = content_block.input
                    
                    try:
                        result = await self.call_tool(tool_name, tool_args)
                        tool_results.append({
                            "type": "tool_result",
                            "tool_use_id": content_block.id,
                            "content": json.dumps(result, indent=2)
                        })
                    except Exception as e:
                        tool_results.append({
                            "type": "tool_result",
                            "tool_use_id": content_block.id,
                            "content": f"Error: {str(e)}",
                            "is_error": True
                        })
            
            # Send tool results back to Claude
            messages.append({"role": "user", "content": tool_results})
            
            response = self.anthropic_client.messages.create(
                model="claude-3-5-sonnet-20241022",
                max_tokens=1000,
                messages=messages,
                tools=tools if tools else None
            )
        
        # Extract text response
        text_content = ""
        for content_block in response.content:
            if content_block.type == "text":
                text_content += content_block.text
        
        return text_content
    
    async def interactive_chat(self, mode: int):
        """Start an interactive chat session"""
        print("Type 'quit' to exit, 'tools' to list available tools, 'resources' to list resources")
        print("-" * 50)
        
        while True:
            try:
                user_input = input("\nUser Input: ").strip()
                
                if user_input.lower() == 'quit':
                    break
                elif user_input.lower() == 'tools':
                    tools = self.get_all_tools()
                    print(f"\nAvailable tools ({len(tools)}):")
                    for name, info in tools.items():
                        print(f"  - {name}: {info.get('description', 'No description')}")
                    continue
                elif user_input.lower() == 'resources':
                    resources = self.get_all_resources()
                    print(f"\nAvailable resources ({len(resources)}):")
                    for uri, info in resources.items():
                        print(f"  - {uri}: {info.get('description', 'No description')}")
                    continue
                
                if not user_input:
                    continue
                
                if mode == 0:
                    print("\nClaude: ", end="", flush=True)
                    response = await self.chat_with_claude(user_input)
                elif mode == 1:
                    print("\nOpenAI: ", end="", flush=True)
                    response = await self.chat_with_openai(user_input)
                elif mode == 2:
                    print("\nDeepSeek: ", end="", flush=True)
                    response = await self.chat_with_deepseek(user_input)
                elif mode == 3:
                    print("\nQWen: ", end="", flush=True)
                    response = await self.chat_with_qwen(user_input)
                elif mode == 4:
                    print("\nLlama: ", end="", flush=True)
                    response = await self.chat_with_llama(user_input)
                print(response)
                
            except KeyboardInterrupt:
                print("\nExiting...")
                break
            except Exception as e:
                print(f"\nError: {e}")
    
    async def cleanup(self):
        """Clean up all server connections"""
        for server in self.servers.values():
            await server.stop()

async def main():
    """Main function"""
    if len(sys.argv) < 2:
        print("Usage: uv run client.py <mode>")
        sys.exit(1)
    mode = int(sys.argv[1])
    if mode == 0:
    # Configuration
        ANTHROPIC_API_KEY = os.getenv("ANTHROPIC_API_KEY")
        if not ANTHROPIC_API_KEY:
            print("Error: ANTHROPIC_API_KEY environment variable is required")
            sys.exit(1)
    elif mode == 1:
        OPENAI_API_KEY = os.getenv("OPENAI_API_KEY")
        if not OPENAI_API_KEY:
            print("Error: OPENAI_API_KEY environment variable is required")
            sys.exit(1)
    elif mode == 2:
        DEEPSEEK_API_KEY = os.getenv("DEEPSEEK_API_KEY")
        if not DEEPSEEK_API_KEY:
            print("Error: DEEPSEEK_API_KEY environment variable is required")
            sys.exit(1)
    elif mode == 3:
        DASHSCOPE_API_KEY = os.getenv("DASHSCOPE_API_KEY")
        if not DASHSCOPE_API_KEY:
            print("Error: DASHSCOPE_API_KEY environment variable is required")
            sys.exit(1)
    elif mode == 4:
        LLAMA_API_KEY = os.getenv("LLAMA_API_KEY")
        if not LLAMA_API_KEY:
            print("Error: LLAMA_API_KEY environment variable is required")
            sys.exit(1)
    
    # Example MCP server configurations
    server_configs = [
        MCPServerConfig( 
            name="signature-checkers",
            server_type=ServerType.LOCAL,
            command=["uv"],
            args=["--directory", "/usr/games/mcp/MCPSecCode/mcpbench", "run", "./squatting.py"],
            description="Check file's signature"
        ),
        MCPServerConfig(
            name="compute-helper",
            server_type=ServerType.LOCAL,
            command=["uv"],
            args=["--directory", "/usr/games/mcp/MCPSecCode/mcpbench", "run", "addserver.py"],
            description="Basic computer"
        ),
        MCPServerConfig(
            name="compute-assistant",
            server_type=ServerType.LOCAL,
            command=["uv"],
            args=["--directory", "/usr/games/mcp/MCPSecCode/mcpbench", "run", "./maliciousadd.py"],
            description="Enhanced computer"
        ),
        MCPServerConfig(
            name="signature-checker",
            server_type=ServerType.HTTP,
            url="http://127.0.0.1:9001/mcp",
            headers={"accept": "application/json, text/event-stream"},
            api_key="",
            description="Check file's signature"
        ),
        MCPServerConfig( 
            name="schema",
            server_type=ServerType.LOCAL,
            command=["uv"],
            args=["--directory", "/usr/games/mcp/MCPSecCode/mcpbench", "run", "mcp", "run", "./squatting.py"],
            description="Enhanced computer"
        ),
        # Add more servers as needed
    ]
    
    if mode == 0:
        client = MultiMCPClient(ANTHROPIC_API_KEY, "")
    elif mode == 1:
        client = MultiMCPClient("", OPENAI_API_KEY)
    elif mode == 2:
        client = MultiMCPClient("", DEEPSEEK_API_KEY)
    elif mode == 3:
        client = MultiMCPClient("", DASHSCOPE_API_KEY)
    elif mode == 4:
        client = MultiMCPClient("", LLAMA_API_KEY)
    
    
    
    try:
        # Start all servers
        for config in server_configs:
            try:
                await client.add_server(config)
            except Exception as e:
                logger.error(f"Failed to add server {config.name}: {e}")
        
        # Start interactive chat
        await client.interactive_chat(mode)
        
    finally:
        # Clean up
        await client.cleanup()

if __name__ == "__main__":
    asyncio.run(main())