deepvoice_client/local_plugins/open_ai_service/ios/Classes/MCPClient.swift

import Foundation

/// 工具类型枚举
enum ToolType {
    case localFunction  // 本地函数
    case mcpTool        // MCP工具
}

/// 函数处理器协议
protocol FunctionHandler {
    /// 处理函数调用
    /// - Parameters:
    ///   - arguments: 函数参数，Dictionary格式
    /// - Returns: 函数执行结果，字符串
    func handle(arguments: [String: Any]) async throws -> String
}

/// MCPClient - iOS版本，负责连接到Model Context Protocol服务器和管理本地函数
public class MCPClient {
    private let TAG = "MCPClient"
    
    // 本地函数Map，函数名 -> 处理器
    private var localFunctions = [String: FunctionHandler]()
    
    // 工具定义Map
    private var toolDefinitions = [[String: Any]]()
    
    // 是否已连接
    private var isConnected = false
    
    // 系统功能处理器
    private var systemFunctionHandler: SystemFunctionHandler?
    
    // SSE会话
    private var sseSession: URLSession?
    private var sseTask: URLSessionDataTask?
    
    /// 初始化MCPClient
    /// - Parameter context: 可选的应用上下文
    init(_ context: Any? = nil) {
        if let appContext = context {
            systemFunctionHandler = SystemFunctionHandler(context: appContext)
            
            // 初始化语音功能
            initializeVoiceFunctions()
        }
    }
    
    /// 初始化语音功能处理器
    private func initializeVoiceFunctions() {
        do {
            // 使用SystemFunctionHandler注册所有功能
            systemFunctionHandler?.registerAllFunctions(client: self)
            
            print("\(TAG) 语音功能初始化完成")
        } catch {
            print("\(TAG) 语音功能初始化失败: \(error.localizedDescription)")
        }
    }
    
    /// 注册本地函数
    /// - Parameters:
    ///   - name: 函数名称
    ///   - description: 函数描述
    ///   - parameters: 函数参数定义
    ///   - handler: 函数处理器
    /// - Returns: 是否注册成功
    func registerLocalFunction(name: String, description: String, parameters: [String: Any], handler: FunctionHandler) -> Bool {
        do {
            // 存储函数处理器
            localFunctions[name] = handler
            
            // 创建函数定义
            let functionDef: [String: Any] = [
                "name": name,
                "description": description,
                "parameters": parameters
            ]
            
            // 添加到工具定义列表
            let toolDef: [String: Any] = [
                "type": "function",
                "function": functionDef
            ]
            
            // 更新或添加工具定义
            if let index = toolDefinitions.firstIndex(where: { 
                guard let function = $0["function"] as? [String: Any],
                      let functionName = function["name"] as? String else {
                    return false
                }
                return functionName == name
            }) {
                toolDefinitions[index] = toolDef
            } else {
                toolDefinitions.append(toolDef)
            }
            
            print("\(TAG) 成功注册本地函数: \(name)")
            return true
        } catch {
            print("\(TAG) 注册函数失败: \(error.localizedDescription)")
            return false
        }
    }
    
    /// 连接到MCP服务器(通过SSE)
    /// - Parameter mcpServer: 服务器URL
    /// - Returns: 是否连接成功
    func connectToSSE(_ mcpServer: String) -> Bool {
        // 关闭现有连接
        closeConnection()
        
        // 如果URL为空，返回失败
        if mcpServer.isEmpty {
            print("\(TAG) MCP服务器URL为空")
            return false
        }
        
        // 解析服务器URL
        guard let serverURL = URL(string: mcpServer) else {
            print("\(TAG) 无效的MCP服务器URL: \(mcpServer)")
            return false
        }
        
        print("\(TAG) 正在连接到MCP服务器: \(mcpServer)")
        
        // 创建会话配置
        let config = URLSessionConfiguration.default
        config.timeoutIntervalForRequest = 30
        config.timeoutIntervalForResource = 60
        
        // 创建会话
        let session = URLSession(configuration: config)
        self.sseSession = session
        
        // 创建请求
        var request = URLRequest(url: serverURL)
        request.httpMethod = "GET"
        request.addValue("text/event-stream", forHTTPHeaderField: "Accept")
        
        // 创建并启动SSE任务
        let task = session.dataTask(with: request) { [weak self] data, response, error in
            guard let self = self else { return }
            
            if let error = error {
                print("\(self.TAG) MCP服务器连接失败: \(error.localizedDescription)")
                self.isConnected = false
                return
            }
            
            if let httpResponse = response as? HTTPURLResponse {
                if httpResponse.statusCode == 200 {
                    print("\(self.TAG) 成功连接到MCP服务器")
                    self.isConnected = true
                    
                    // 处理服务器响应数据
                    if let data = data, let responseString = String(data: data, encoding: .utf8) {
                        print("\(self.TAG) 收到初始响应: \(responseString)")
                        
                        // 在这里解析可用工具列表，暂时模拟
                        // 实际情况下，需要实现SSE流解析
                    }
                } else {
                    print("\(self.TAG) MCP服务器连接失败，状态码: \(httpResponse.statusCode)")
                    self.isConnected = false
                }
            }
        }
        
        task.resume()
        self.sseTask = task
        
        // 在iOS中，连接是异步的，此处返回true表示连接已开始
        // 实际连接状态通过isConnected属性检查
        return true
    }
    
    /// 检查是否已连接
    /// - Returns: 是否已连接
    func checkIsConnected() -> Bool {
        return isConnected
    }
    
    /// 关闭连接
    func closeConnection() {
        sseTask?.cancel()
        sseTask = nil
        isConnected = false
    }
    
    /// 获取工具定义列表，用于OpenAI请求
    /// - Returns: 工具定义列表
    func getToolMaps() -> [[String: Any]] {
        return toolDefinitions
    }
    
    /// 解析JSON参数
    /// - Parameter argumentsJson: JSON字符串
    /// - Returns: 解析后的参数字典
    func parseJsonArguments(_ argumentsJson: String) -> [String: Any] {
        guard !argumentsJson.isEmpty,
              let data = argumentsJson.data(using: .utf8),
              let parsedArgs = try? JSONSerialization.jsonObject(with: data) as? [String: Any] else {
            return [:]
        }
        return parsedArgs
    }
    
    /// 调用工具
    /// - Parameters:
    ///   - name: 工具名称
    ///   - arguments: 工具参数
    /// - Returns: 工具调用结果
    func callTool(name: String, arguments: [String: Any]) async -> [String: Any] {
        do {
            // 检查是否为本地函数
            if let handler = localFunctions[name] {
                let result = try await handler.handle(arguments: arguments)
                return [
                    "name": name,
                    "context": result
                ]
            }
            
            // 如果不是本地函数，尝试作为MCP工具调用
            if isConnected {
                // 这里应该实现真正的MCP工具调用
                // 暂时返回一个模拟结果
                return [
                    "name": name,
                    "context": "模拟MCP工具调用结果: \(name)(\(arguments))"
                ]
            }
            
            return ["context": "无法处理工具调用：找不到工具 \(name) 或MCP未连接"]
        } catch {
            print("\(TAG) 调用工具失败: \(error.localizedDescription)")
            return ["context": "工具调用失败: \(error.localizedDescription)"]
        }
    }
    
    /// 关闭MCP客户端
    func close() {
        closeConnection()
        localFunctions.removeAll()
        toolDefinitions.removeAll()
    }
}