"""
NexAgent MCP 客户端 - 使用 fastmcp 库支持 SSE 和 Streamable HTTP
"""
import json
from typing import Dict, List, Optional
from threading import Lock
import asyncio
import traceback
import threading
import concurrent.futures

try:
    from fastmcp import Client
    from fastmcp.client.transports import SSETransport, StreamableHttpTransport
    FASTMCP_AVAILABLE = True
except ImportError:
    FASTMCP_AVAILABLE = False
    print("[MCP] 警告: fastmcp 未安装，MCP 功能将不可用")

from ._version import __version__


class MCPClient:
    """MCP 客户端，使用 fastmcp 库"""
    
    TYPE_SSE = "sse"
    TYPE_STREAMABLE_HTTP = "streamable_http"
    
    def __init__(self, server_id: str, name: str, url: str, server_type: str = "sse", headers: Dict[str, str] = None):
        if not FASTMCP_AVAILABLE:
            raise RuntimeError("fastmcp 库未安装")
        
        self.server_id = server_id
        self.name = name
        self.url = url.rstrip('/')
        self.server_type = server_type
        self.custom_headers = headers or {}
        self.tools: List[Dict] = []
        self.connected = False
        self._lock = Lock()
        self._last_error = None
        self._client: Optional[Client] = None
        self._client_context = None
        
        # 专用事件循环线程
        self._loop: Optional[asyncio.AbstractEventLoop] = None
        self._loop_thread: Optional[threading.Thread] = None
        self._loop_ready = threading.Event()
        self._start_event_loop_thread()
    
    def _start_event_loop_thread(self):
        """启动专用的事件循环线程"""
        def run_event_loop():
            self._loop = asyncio.new_event_loop()
            asyncio.set_event_loop(self._loop)
            self._loop_ready.set()
            self._loop.run_forever()
        
        self._loop_thread = threading.Thread(target=run_event_loop, daemon=True, name=f"MCP-Loop-{self.server_id}")
        self._loop_thread.start()
        self._loop_ready.wait()  # 等待事件循环启动
    
    def _stop_event_loop_thread(self):
        """停止事件循环线程"""
        if self._loop and self._loop.is_running():
            self._loop.call_soon_threadsafe(self._loop.stop)
        if self._loop_thread:
            self._loop_thread.join(timeout=2)
    
    def _run_coroutine(self, coro, timeout: float = 60):
        """在专用事件循环中运行协程"""
        if not self._loop or not self._loop.is_running():
            raise RuntimeError("事件循环未运行")
        
        future = concurrent.futures.Future()
        
        async def wrapper():
            try:
                result = await asyncio.wait_for(coro, timeout=timeout)
                future.set_result(result)
            except asyncio.TimeoutError:
                future.set_exception(TimeoutError(f"操作超时({timeout}秒)"))
            except Exception as e:
                future.set_exception(e)
        
        asyncio.run_coroutine_threadsafe(wrapper(), self._loop)
        return future.result(timeout=timeout + 1)
    
    def connect(self) -> bool:
        """连接到 MCP 服务器并获取工具列表"""
        self._last_error = None
        try:
            print(f"[MCP] 正在连接 {self.name} ({self.url}) - 类型: {self.server_type}...")
            result = self._run_coroutine(self._connect_async(), timeout=30)
            return result
        except Exception as e:
            self._last_error = str(e)
            print(f"[MCP] 连接 {self.name} 失败: {e}")
            traceback.print_exc()
            self.connected = False
            return False
    
    async def _connect_async(self) -> bool:
        """异步连接"""
        try:
            # 创建传输
            if self.server_type == self.TYPE_STREAMABLE_HTTP:
                transport = StreamableHttpTransport(
                    url=self.url,
                    headers=self.custom_headers
                )
            else:  # SSE
                transport = SSETransport(
                    url=self.url,
                    headers=self.custom_headers
                )
            
            # 创建客户端
            self._client = Client(transport)
            
            # 进入上下文（保持连接）
            self._client_context = self._client.__aenter__()
            await self._client_context
            
            # 获取工具列表
            tools_result = await self._client.list_tools()
            self.tools = []
            
            for tool in tools_result:
                # fastmcp 返回的工具对象
                tool_schema = {
                    "type": "function",
                    "function": {
                        "name": f"mcp_{self.server_id}_{tool.name}",
                        "description": f"[MCP:{self.name}] {tool.description or ''}",
                        "parameters": tool.inputSchema if hasattr(tool, 'inputSchema') else {"type": "object", "properties": {}}
                    }
                }
                self.tools.append(tool_schema)
            
            self.connected = True
            print(f"[MCP] 连接成功! 获取到 {len(self.tools)} 个工具")
            return True
                
        except Exception as e:
            self._last_error = str(e)
            print(f"[MCP] 连接失败: {e}")
            traceback.print_exc()
            self.connected = False
            return False
    
    def get_tools_for_openai(self) -> List[Dict]:
        """获取 OpenAI 格式的工具定义"""
        return self.tools
    
    def call_tool(self, tool_name: str, arguments: Dict) -> str:
        """调用 MCP 工具"""
        try:
            result = self._run_coroutine(self._call_tool_async(tool_name, arguments), timeout=60)
            return result
        except Exception as e:
            traceback.print_exc()
            return f"MCP调用失败: {e}"
    
    async def _call_tool_async(self, tool_name: str, arguments: Dict) -> str:
        """异步调用工具"""
        try:
            if not self._client:
                return "MCP客户端未初始化"
            
            # 使用 fastmcp 调用工具
            result = await self._client.call_tool(tool_name, arguments)
            
            # 解析结果
            if isinstance(result, list):
                # 结果是内容列表
                texts = []
                for item in result:
                    if hasattr(item, 'type') and item.type == 'text':
                        texts.append(item.text)
                    elif hasattr(item, 'text'):
                        texts.append(item.text)
                    else:
                        texts.append(str(item))
                return "\n".join(texts) if texts else "(无返回内容)"
            elif hasattr(result, 'content'):
                # 结果有 content 属性
                content = result.content
                if isinstance(content, list):
                    texts = []
                    for item in content:
                        if hasattr(item, 'type') and item.type == 'text':
                            texts.append(item.text)
                        elif hasattr(item, 'text'):
                            texts.append(item.text)
                        else:
                            texts.append(str(item))
                    return "\n".join(texts) if texts else "(无返回内容)"
                return str(content)
            
            return str(result)
            
        except Exception as e:
            return f"MCP调用失败: {e}"
    
    def disconnect(self):
        """断开连接"""
        if self._client_context:
            try:
                self._run_coroutine(self._client.__aexit__(None, None, None), timeout=5)
            except:
                pass
        self._client = None
        self._client_context = None
        self.connected = False
        
        # 停止事件循环线程
        self._stop_event_loop_thread()


class MCPManager:
    """MCP 服务器管理器"""
    
    def __init__(self):
        self.clients: Dict[str, MCPClient] = {}
        self._lock = Lock()
    
    def add_server(self, server_id: str, name: str, url: str, server_type: str = "sse", headers: Dict[str, str] = None) -> bool:
        """添加并连接 MCP 服务器（异步，不阻塞）"""
        if not FASTMCP_AVAILABLE:
            print("[MCP] fastmcp 库未安装")
            return False
        
        client = MCPClient(server_id, name, url, server_type, headers)
        with self._lock:
            self.clients[server_id] = client
        
        # 在后台线程中连接
        from threading import Thread
        def connect_async():
            try:
                client.connect()
            except Exception as e:
                print(f"[MCP] 后台连接失败: {e}")
        
        thread = Thread(target=connect_async, daemon=True, name=f"MCP-Connect-{server_id}")
        thread.start()
        
        return False  # 返回 False 表示还在连接中
    
    def add_server_sync(self, server_id: str, name: str, url: str, server_type: str = "sse", headers: Dict[str, str] = None) -> bool:
        """添加并连接 MCP 服务器（同步）"""
        if not FASTMCP_AVAILABLE:
            print("[MCP] fastmcp 库未安装")
            return False
        
        client = MCPClient(server_id, name, url, server_type, headers)
        with self._lock:
            self.clients[server_id] = client
        
        try:
            return client.connect()
        except Exception as e:
            print(f"[MCP] 连接失败: {e}")
            return False
    
    def remove_server(self, server_id: str) -> bool:
        """移除 MCP 服务器"""
        with self._lock:
            if server_id in self.clients:
                client = self.clients[server_id]
                client.disconnect()
                del self.clients[server_id]
                return True
            return False
    
    def reconnect_server(self, server_id: str) -> bool:
        """重新连接 MCP 服务器"""
        with self._lock:
            if server_id not in self.clients:
                return False
            client = self.clients[server_id]
        
        try:
            client.disconnect()
            return client.connect()
        except Exception as e:
            print(f"[MCP] 重连失败: {e}")
            return False
    
    def update_server(self, server_id: str, name: str = None, url: str = None, server_type: str = None, headers: Dict[str, str] = None) -> bool:
        """更新服务器配置"""
        with self._lock:
            if server_id not in self.clients:
                return False
            client = self.clients[server_id]
            if name:
                client.name = name
            if url:
                client.url = url.rstrip('/')
            if server_type:
                client.server_type = server_type
            if headers is not None:
                client.custom_headers = headers
            return True
    
    def get_all_tools(self) -> List[Dict]:
        """获取所有已连接服务器的工具（OpenAI 格式）"""
        tools = []
        with self._lock:
            for client in self.clients.values():
                if client.connected:
                    tools.extend(client.get_tools_for_openai())
        return tools
    
    def call_tool(self, full_tool_name: str, arguments: Dict) -> str:
        """调用 MCP 工具"""
        if not full_tool_name.startswith("mcp_"):
            return f"无效的MCP工具名: {full_tool_name}"
        
        parts = full_tool_name[4:].split("_", 1)
        if len(parts) != 2:
            return f"无效的MCP工具名格式: {full_tool_name}"
        
        server_id, tool_name = parts
        
        with self._lock:
            if server_id not in self.clients:
                return f"MCP服务器未连接: {server_id}"
            client = self.clients[server_id]
        
        return client.call_tool(tool_name, arguments)
    
    def get_server_status(self) -> List[Dict]:
        """获取所有服务器状态"""
        result = []
        with self._lock:
            for server_id, client in self.clients.items():
                result.append({
                    "id": server_id,
                    "name": client.name,
                    "url": client.url,
                    "server_type": client.server_type,
                    "connected": client.connected,
                    "tool_count": len(client.tools),
                    "has_headers": bool(client.custom_headers),
                    "last_error": client._last_error
                })
        return result
    
    def is_mcp_tool(self, tool_name: str) -> bool:
        """判断是否是 MCP 工具"""
        return tool_name.startswith("mcp_")
