# -*- coding: utf-8 -*-
import json
import threading
import time
from typing import Any
from typing import Callable
from typing import Dict
from typing import List
from typing import Optional
from typing import cast
from urllib.parse import parse_qs
from urllib.parse import urlencode
from urllib.parse import urljoin

import requests  # 导入第三方库requests

from jarvis.jarvis_mcp import McpClient
from jarvis.jarvis_utils.output import PrettyOutput


class SSEMcpClient(McpClient):
    """远程MCP客户端实现

    参数:
        config: 配置字典，包含以下字段：
            - base_url: str - MCP服务器的基础URL
            - auth_token: str - 认证令牌（可选）
            - headers: Dict[str, str] - 额外的HTTP头（可选）
    """

    def __init__(self, config: Dict[str, Any]):
        self.config = config
        self.base_url = config.get("base_url", "")
        if not self.base_url:
            raise ValueError("No base_url specified in config")

        # 设置HTTP客户端
        self.session = requests.Session()
        self.session.headers.update(
            {
                "Content-Type": "application/json",
                "Accept": "application/json",
            }
        )

        # 添加认证令牌（如果提供）
        auth_token = config.get("auth_token")
        if auth_token:
            self.session.headers["Authorization"] = f"Bearer {auth_token}"

        # 添加额外的HTTP头
        extra_headers = config.get("headers", {})
        self.session.headers.update(extra_headers)

        # SSE相关属性
        self.sse_response: Optional[requests.Response] = None
        self.sse_thread: Optional[threading.Thread] = None
        self.messages_endpoint: Optional[str] = None
        self.session_id: Optional[str] = None
        self.pending_requests: Dict[
            str, threading.Event
        ] = {}  # 存储等待响应的请求 {id: Event}
        self.request_results: Dict[
            str, Dict[str, Any]
        ] = {}  # 存储请求结果 {id: result}
        self.notification_handlers: Dict[str, List[Callable]] = {}
        self.event_lock = threading.Lock()
        self.request_id_counter = 0

        # 初始化连接
        self._initialize()

    def _initialize(self) -> None:
        """初始化MCP连接"""
        try:
            # 启动SSE连接
            self._start_sse_connection()

            # 等待获取消息端点和会话ID
            start_time = time.time()
            while (
                not self.messages_endpoint or not self.session_id
            ) and time.time() - start_time < 5:
                time.sleep(0.1)

            if not self.messages_endpoint:
                self.messages_endpoint = "/messages"  # 默认端点
                PrettyOutput.auto_print(
                    f"⚠️ 未获取到消息端点，使用默认值: {self.messages_endpoint}"
                )

            if not self.session_id:
                PrettyOutput.auto_print("⚠️ 未获取到会话ID")

            # 发送初始化请求
            response = self._send_request(
                "initialize",
                {
                    "processId": None,  # 远程客户端不需要进程ID
                    "clientInfo": {"name": "jarvis", "version": "1.0.0"},
                    "capabilities": {},
                    "protocolVersion": "2025-03-26",
                },
            )

            # 验证服务器响应
            if "result" not in response:
                raise RuntimeError(
                    f"初始化失败: {response.get('error', 'Unknown error')}"
                )

            # 发送initialized通知
            self._send_notification("notifications/initialized", {})

        except Exception as e:
            PrettyOutput.auto_print(f"❌ MCP初始化失败: {str(e)}")
            raise

    def _start_sse_connection(self) -> None:
        """建立SSE连接并启动处理线程"""
        try:
            # 设置SSE请求头
            sse_headers = dict(self.session.headers)
            sse_headers.update(
                {
                    "Accept": "text/event-stream",
                    "Cache-Control": "no-cache",
                }
            )

            # 建立SSE连接
            sse_url = urljoin(self.base_url, "sse")
            self.sse_response = self.session.get(
                sse_url, stream=True, headers=sse_headers, timeout=30
            )
            if self.sse_response:
                self.sse_response.raise_for_status()

            # 启动事件处理线程
            self.sse_thread = threading.Thread(
                target=self._process_sse_events, daemon=True
            )
            if self.sse_thread:
                self.sse_thread.start()

        except Exception as e:
            PrettyOutput.auto_print(f"❌ SSE连接失败: {str(e)}")
            raise

    def _process_sse_events(self) -> None:
        """处理SSE事件流"""
        if not self.sse_response:
            return

        buffer = ""
        for line in self.sse_response.iter_lines(decode_unicode=True):
            if line:
                if line.startswith("data:"):
                    data = line[5:].strip()
                    # 检查是否包含消息端点信息
                    if data.startswith("/"):
                        # 这是消息端点信息，例如 "/messages/?session_id=xyz"
                        try:
                            # 提取消息端点路径和会话ID
                            url_parts = data.split("?")
                            self.messages_endpoint = url_parts[0]

                            # 如果有查询参数，尝试提取session_id
                            if len(url_parts) > 1:
                                query_string = url_parts[1]
                                query_params = parse_qs(query_string)
                                if "session_id" in query_params:
                                    self.session_id = query_params["session_id"][0]
                        except Exception as e:
                            PrettyOutput.auto_print(f"⚠️ 解析消息端点或会话ID失败: {e}")
                    else:
                        buffer += data
                elif line.startswith(":"):  # 忽略注释行
                    continue
                elif line.startswith("event:"):  # 事件类型
                    continue  # 我们不使用事件类型
                elif line.startswith("id:"):  # 事件ID
                    continue  # 我们不使用事件ID
                elif line.startswith("retry:"):  # 重连时间
                    continue  # 我们自己管理重连
            else:  # 空行表示事件结束
                if buffer:
                    try:
                        self._handle_sse_event(buffer)
                    except Exception as e:
                        PrettyOutput.auto_print(f"❌ 处理SSE事件出错: {e}")
                    buffer = ""

        PrettyOutput.auto_print("⚠️ SSE连接已关闭")

    def _handle_sse_event(self, data: str) -> None:
        """处理单个SSE事件数据"""
        try:
            event_data = json.loads(data)

            # 检查是请求响应还是通知
            if "id" in event_data:
                # 这是一个请求的响应
                req_id = event_data["id"]
                with self.event_lock:
                    self.request_results[req_id] = event_data
                    if req_id in self.pending_requests:
                        # 通知等待线程响应已到达
                        self.pending_requests[req_id].set()
            elif "method" in event_data:
                # 这是一个通知
                method = event_data.get("method", "")
                params = event_data.get("params", {})

                # 调用已注册的处理器
                if method in self.notification_handlers:
                    error_lines: List[str] = []
                    for handler in self.notification_handlers[method]:
                        try:
                            handler(params)
                        except Exception as e:
                            error_lines.append(f"处理通知时出错 ({method}): {e}")
                    if error_lines:
                        joined_errors = "\n".join(error_lines)
                        PrettyOutput.auto_print(f"❌ {joined_errors}")
        except Exception:
            PrettyOutput.auto_print(f"⚠️ 无法解析SSE事件: {data}")

    def register_notification_handler(self, method: str, handler: Callable) -> None:
        """注册通知处理器

        参数:
            method: 通知方法名
            handler: 处理通知的回调函数，接收params参数
        """
        with self.event_lock:
            if method not in self.notification_handlers:
                self.notification_handlers[method] = []
            self.notification_handlers[method].append(handler)

    def unregister_notification_handler(self, method: str, handler: Callable) -> None:
        """注销通知处理器

        参数:
            method: 通知方法名
            handler: 要注销的处理器函数
        """
        with self.event_lock:
            if method in self.notification_handlers:
                if handler in self.notification_handlers[method]:
                    self.notification_handlers[method].remove(handler)
                if not self.notification_handlers[method]:
                    del self.notification_handlers[method]

    def _get_next_request_id(self) -> str:
        """获取下一个请求ID"""
        with self.event_lock:
            self.request_id_counter += 1
            return str(self.request_id_counter)

    def _send_request(self, method: str, params: Dict[str, Any]) -> Dict[str, Any]:
        """发送请求到MCP服务器

        参数:
            method: 请求方法
            params: 请求参数

        返回:
            Dict[str, Any]: 响应结果
        """
        # 生成唯一请求ID
        req_id = self._get_next_request_id()

        # 创建事件标志，用于等待响应
        event = threading.Event()

        with self.event_lock:
            self.pending_requests[req_id] = event

        try:
            # 构建请求
            request = {
                "jsonrpc": "2.0",
                "method": method,
                "params": params,
                "id": req_id,
            }

            # 尝试不同的请求发送方式
            if self.session_id:
                # 方法1: 使用查询参数中的session_id
                query_params = {"session_id": self.session_id}
                messages_url = urljoin(self.base_url, self.messages_endpoint)

                # 尝试直接使用原始URL（不追加查询参数）
                try:
                    post_response = self.session.post(messages_url, json=request)
                    post_response.raise_for_status()
                except requests.HTTPError:
                    # 如果失败，尝试添加会话ID到查询参数
                    messages_url_with_session = (
                        f"{messages_url}?{urlencode(query_params)}"
                    )
                    post_response = self.session.post(
                        messages_url_with_session, json=request
                    )
                    post_response.raise_for_status()
            else:
                # 方法2: 不使用session_id
                if not self.messages_endpoint:
                    self.messages_endpoint = "/messages"

                messages_url = urljoin(self.base_url, self.messages_endpoint)

                # 尝试直接使用messages端点而不带任何查询参数
                try:
                    # 尝试1: 标准JSON-RPC格式
                    post_response = self.session.post(messages_url, json=request)
                    post_response.raise_for_status()
                except requests.HTTPError:
                    # 尝试2: JSON字符串作为请求参数
                    post_response = self.session.post(
                        messages_url, params={"request": json.dumps(request)}
                    )
                    post_response.raise_for_status()

            # 等待SSE通道返回响应（最多30秒）
            if not event.wait(timeout=30):
                raise TimeoutError(f"等待响应超时: {method}")

            # 获取响应结果
            with self.event_lock:
                result = self.request_results.pop(req_id, None)
                self.pending_requests.pop(req_id, None)

            if result is None:
                raise RuntimeError(f"未收到响应: {method}")

            return result

        except Exception as e:
            # 清理请求状态
            with self.event_lock:
                self.pending_requests.pop(req_id, None)
                self.request_results.pop(req_id, None)

            PrettyOutput.auto_print(f"❌ 发送请求失败: {str(e)}")
            raise

    def _send_notification(self, method: str, params: Dict[str, Any]) -> None:
        """发送通知到MCP服务器（不需要响应）

        参数:
            method: 通知方法
            params: 通知参数
        """
        try:
            # 构建通知
            notification = {"jsonrpc": "2.0", "method": method, "params": params}

            # 尝试不同的请求发送方式，与_send_request保持一致
            if self.session_id:
                # 方法1: 使用查询参数中的session_id
                query_params = {"session_id": self.session_id}
                messages_url = urljoin(
                    self.base_url, self.messages_endpoint or "/messages"
                )

                # 尝试直接使用原始URL（不追加查询参数）
                try:
                    post_response = self.session.post(messages_url, json=notification)
                    post_response.raise_for_status()
                except requests.HTTPError:
                    # 如果失败，尝试添加会话ID到查询参数
                    messages_url_with_session = (
                        f"{messages_url}?{urlencode(query_params)}"
                    )
                    post_response = self.session.post(
                        messages_url_with_session, json=notification
                    )
                    post_response.raise_for_status()
            else:
                # 方法2: 不使用session_id
                if not self.messages_endpoint:
                    self.messages_endpoint = "/messages"

                messages_url = urljoin(self.base_url, self.messages_endpoint)

                # 尝试直接使用messages端点而不带任何查询参数
                try:
                    # 尝试1: 标准JSON-RPC格式
                    post_response = self.session.post(messages_url, json=notification)
                    post_response.raise_for_status()
                except requests.HTTPError:
                    # 尝试2: JSON字符串作为请求参数
                    post_response = self.session.post(
                        messages_url, params={"request": json.dumps(notification)}
                    )
                    post_response.raise_for_status()

        except Exception as e:
            PrettyOutput.auto_print(f"❌ 发送通知失败: {str(e)}")
            raise

    def get_tool_list(self) -> List[Dict[str, Any]]:
        """获取工具列表

        返回:
            List[Dict[str, Any]]: 工具列表，每个工具包含以下字段：
                - name: str - 工具名称
                - description: str - 工具描述
                - parameters: Dict - 工具参数
        """
        try:
            response = self._send_request("tools/list", {})
            if "result" in response and "tools" in response["result"]:
                # 注意这里: 响应结构是 response['result']['tools']
                tools = response["result"]["tools"]
                # 将MCP协议字段转换为内部格式
                formatted_tools = []
                for tool in tools:
                    # 从inputSchema中提取参数定义
                    input_schema = tool.get("inputSchema", {})
                    parameters = {}
                    if "properties" in input_schema:
                        parameters = input_schema["properties"]

                    formatted_tools.append(
                        {
                            "name": tool.get("name", ""),
                            "description": tool.get("description", ""),
                            "parameters": parameters,
                        }
                    )
                return formatted_tools
            else:
                error_msg = "获取工具列表失败"
                if "error" in response:
                    error_msg += f": {response['error']}"
                elif "result" in response:
                    error_msg += f": 响应格式不正确 - {response['result']}"
                else:
                    error_msg += ": 未知错误"

                PrettyOutput.auto_print(f"❌ {error_msg}")
                return []
        except Exception as e:
            PrettyOutput.auto_print(f"❌ 获取工具列表失败: {str(e)}")
            return []

    def execute(self, tool_name: str, arguments: Dict[str, Any]) -> Dict[str, Any]:
        """执行工具

        参数:
            tool_name: 工具名称
            arguments: 参数字典，包含工具执行所需的参数

        返回:
            Dict[str, Any]: 执行结果，包含以下字段：
                - success: bool - 是否执行成功
                - stdout: str - 标准输出
                - stderr: str - 标准错误
        """
        try:
            response = self._send_request(
                "tools/call", {"name": tool_name, "arguments": arguments}
            )
            if "result" in response:
                result = response["result"]
                # 从content中提取输出信息
                stdout = ""
                stderr = ""
                for content in result.get("content", []):
                    if content.get("type") == "text":
                        stdout += content.get("text", "")
                    elif content.get("type") == "error":
                        stderr += content.get("text", "")

                return {"success": True, "stdout": stdout, "stderr": stderr}
            else:
                return {
                    "success": False,
                    "stdout": "",
                    "stderr": response.get("error", "Unknown error"),
                }
        except Exception as e:
            PrettyOutput.auto_print(f"❌ 执行工具失败: {str(e)}")
            return {"success": False, "stdout": "", "stderr": str(e)}

    def get_resource_list(self) -> List[Dict[str, Any]]:
        """获取资源列表

        返回:
            List[Dict[str, Any]]: 资源列表，每个资源包含以下字段：
                - uri: str - 资源的唯一标识符
                - name: str - 资源的名称
                - description: str - 资源的描述（可选）
                - mimeType: str - 资源的MIME类型（可选）
        """
        try:
            response = self._send_request("resources/list", {})
            if "result" in response and "resources" in response["result"]:
                return cast(List[Dict[str, Any]], response["result"]["resources"])
            else:
                error_msg = "获取资源列表失败"
                if "error" in response:
                    error_msg += f": {response['error']}"
                else:
                    error_msg += ": 未知错误"
                PrettyOutput.auto_print(f"❌ {error_msg}")
                return []
        except Exception as e:
            PrettyOutput.auto_print(f"❌ 获取资源列表失败: {str(e)}")
            return []

    def get_resource(self, uri: str) -> Dict[str, Any]:
        """获取指定资源的内容

        参数:
            uri: str - 资源的URI标识符

        返回:
            Dict[str, Any]: 执行结果，包含以下字段：
                - success: bool - 是否执行成功
                - stdout: str - 资源内容（文本或base64编码的二进制内容）
                - stderr: str - 错误信息
        """
        try:
            response = self._send_request("resources/read", {"uri": uri})
            if "result" in response and "contents" in response["result"]:
                contents = response["result"]["contents"]
                if contents:
                    content = contents[0]  # 获取第一个资源内容
                    # 根据资源类型返回内容
                    if "text" in content:
                        return {
                            "success": True,
                            "stdout": content["text"],
                            "stderr": "",
                        }
                    elif "blob" in content:
                        return {
                            "success": True,
                            "stdout": content["blob"],
                            "stderr": "",
                        }
                return {"success": False, "stdout": "", "stderr": "资源内容为空"}
            else:
                error_msg = "获取资源内容失败"
                if "error" in response:
                    error_msg += f": {response['error']}"
                else:
                    error_msg += ": 未知错误"
                PrettyOutput.auto_print(f"❌ {error_msg}")
                return {"success": False, "stdout": "", "stderr": error_msg}
        except Exception as e:
            error_msg = f"获取资源内容失败: {str(e)}"
            PrettyOutput.auto_print(f"❌ {error_msg}")
            return {"success": False, "stdout": "", "stderr": error_msg}

    def __del__(self):
        """清理资源"""
        # 清理请求状态
        with self.event_lock:
            for event in self.pending_requests.values():
                event.set()  # 释放所有等待的请求
            self.pending_requests.clear()
            self.request_results.clear()

        # 关闭SSE响应
        if self.sse_response:
            try:
                self.sse_response.close()
            except Exception:
                pass

        # 关闭HTTP会话
        if self.session:
            self.session.close()
