import Flutter
import UIKit

// 修改OpenAIService类，添加转发方法
extension OpenAIService {
    // 为插件提供的公共发送消息方法
    public func publicSendMessageStream(messages: [[String: Any]], systemPrompt: String, onToken: @escaping (String) -> Void, onComplete: @escaping () -> Void, onError: @escaping (Error) -> Void, onFunctionCall: @escaping ([String: Any]) -> Void, onFunctionCallResult: @escaping ([String: Any], [String: Any]) -> Void) {
        
        let callback = StreamCallback(
            onToken: onToken,
            onComplete: onComplete,
            onError: onError,
            onFunctionCall: onFunctionCall,
            onFunctionCallResult: onFunctionCallResult
        )
        
        sendMessageStream(messages: messages, systemPrompt: systemPrompt, callback: callback)
    }
    
    // 为插件提供的公共发送函数调用结果方法
    public func publicSendFunctionCallResult(messages: [[String: Any]], systemPrompt: String, functionCall: [String: Any], functionResult: String, onToken: @escaping (String) -> Void, onComplete: @escaping () -> Void, onError: @escaping (Error) -> Void, onFunctionCall: @escaping ([String: Any]) -> Void, onFunctionCallResult: @escaping ([String: Any], [String: Any]) -> Void) {
        
        let callback = StreamCallback(
            onToken: onToken,
            onComplete: onComplete,
            onError: onError,
            onFunctionCall: onFunctionCall,
            onFunctionCallResult: onFunctionCallResult
        )
        
        sendFunctionCallResult(messages: messages, systemPrompt: systemPrompt, functionCall: functionCall, functionResult: functionResult, callback: callback)
    }
}

public class OpenAIServicePlugin: NSObject, FlutterPlugin, FlutterStreamHandler {
    // OpenAI服务实例
    private let openAIService = OpenAIService()
    
    // 事件接收器
    private var eventSink: FlutterEventSink?
    
    // 注册插件
    public static func register(with registrar: FlutterPluginRegistrar) {
        let methodChannel = FlutterMethodChannel(name: "com.yunqiinnovation.open_ai_service/methods", binaryMessenger: registrar.messenger())
        let eventChannel = FlutterEventChannel(name: "com.yunqiinnovation.open_ai_service/events", binaryMessenger: registrar.messenger())
        
        let instance = OpenAIServicePlugin()
        registrar.addMethodCallDelegate(instance, channel: methodChannel)
        eventChannel.setStreamHandler(instance)
    }
    
    // 处理方法调用
    public func handle(_ call: FlutterMethodCall, result: @escaping FlutterResult) {
        switch call.method {
        case "initialize":
            if let args = call.arguments as? [String: Any],
               let apiKey = args["apiKey"] as? String {
                let baseUrl = args["baseUrl"] as? String ?? ""
                let model = args["model"] as? String ?? ""
                let mcpServer = args["mcpServer"] as? String ?? ""
                let initialized = openAIService.initialize(apiKey: apiKey, baseUrl: baseUrl, model: model, mcpServer: mcpServer)
                result(initialized)
            } else {
                result(FlutterError(code: "INVALID_ARGUMENT", message: "初始化参数无效", details: nil))
            }
            
        case "registerFunction":
            if let args = call.arguments as? [String: Any],
               let name = args["name"] as? String,
               let description = args["description"] as? String,
               let parameters = args["parameters"] as? [String: Any] {
                
                let registered = openAIService.registerFunction(name: name, description: description, parameters: parameters)
                result(registered)
            } else {
                result(FlutterError(code: "INVALID_ARGUMENT", message: "函数注册参数无效", details: nil))
            }
            
        case "sendMessage":
            guard let args = call.arguments as? [String: Any],
                  let messagesRaw = args["messages"] as? [[String: Any]] else {
                result(FlutterError(code: "INVALID_ARGUMENT", message: "发送消息参数无效", details: nil))
                return
            }
            
            let systemPrompt = args["systemPrompt"] as? String ?? ""
            
            // 在后台线程执行
            DispatchQueue.global(qos: .userInitiated).async {
                do {
                    let response = try self.openAIService.sendMessage(messages: messagesRaw, systemPrompt: systemPrompt)
                    // 在主线程返回结果
                    DispatchQueue.main.async {
                        result(response)
                    }
                } catch {
                    // 在主线程返回错误
                    DispatchQueue.main.async {
                        result(FlutterError(code: "OPENAI_ERROR", message: error.localizedDescription, details: nil))
                    }
                }
            }
            
        case "sendMessageStream":
            guard let args = call.arguments as? [String: Any],
                  let messagesRaw = args["messages"] as? [[String: Any]] else {
                result(FlutterError(code: "INVALID_ARGUMENT", message: "发送消息参数无效", details: nil))
                return
            }
            
            let systemPrompt = args["systemPrompt"] as? String ?? ""
            
            // 检查事件接收器
            guard let eventSink = self.eventSink else {
                result(FlutterError(code: "NO_EVENT_SINK", message: "没有可用的事件流接收器", details: nil))
                return
            }
            
            // 在后台线程执行
            DispatchQueue.global(qos: .userInitiated).async {
                self.openAIService.publicSendMessageStream(
                    messages: messagesRaw, 
                    systemPrompt: systemPrompt,
                    onToken: { token in
                        DispatchQueue.main.async {
                            eventSink(["type": "token", "content": token])
                        }
                    },
                    onComplete: {
                        DispatchQueue.main.async {
                            eventSink(["type": "complete"])
                        }
                    },
                    onError: { error in
                        DispatchQueue.main.async {
                            eventSink(["type": "error", "content": error.localizedDescription])
                        }
                    },
                    onFunctionCall: { functionCall in
                        DispatchQueue.main.async {
                            do {
                                let data = try JSONSerialization.data(withJSONObject: functionCall)
                                if let jsonString = String(data: data, encoding: .utf8) {
                                    eventSink(["type": "function_call", "content": jsonString])
                                } else {
                                    eventSink(["type": "error", "content": "函数调用数据序列化失败"])
                                }
                            } catch {
                                eventSink(["type": "error", "content": "函数调用数据序列化失败: \(error.localizedDescription)"])
                            }
                        }
                    },
                    onFunctionCallResult: { functionCall, functionCallResult in
                        DispatchQueue.main.async {
                            do {
                                let functionCallData = try JSONSerialization.data(withJSONObject: functionCall)
                                let resultData = try JSONSerialization.data(withJSONObject: functionCallResult)
                                
                                if let functionCallString = String(data: functionCallData, encoding: .utf8),
                                   let resultString = String(data: resultData, encoding: .utf8) {
                                    eventSink(["type": "function_call_result", "function_call": functionCallString, "result": resultString])
                                } else {
                                    eventSink(["type": "error", "content": "函数调用结果数据序列化失败"])
                                }
                            } catch {
                                eventSink(["type": "error", "content": "函数调用结果数据序列化失败: \(error.localizedDescription)"])
                            }
                        }
                    }
                )
                
                // 请求已开始
                DispatchQueue.main.async {
                    result(true)
                }
            }
            
        case "sendFunctionCallResult":
            guard let args = call.arguments as? [String: Any],
                  let messagesRaw = args["messages"] as? [[String: Any]],
                  let functionCallRaw = args["functionCall"] as? [String: Any],
                  let functionResult = args["functionResult"] as? String else {
                result(FlutterError(code: "INVALID_ARGUMENT", message: "发送函数调用结果参数无效", details: nil))
                return
            }
            
            let systemPrompt = args["systemPrompt"] as? String ?? ""
            
            // 检查事件接收器
            guard let eventSink = self.eventSink else {
                result(FlutterError(code: "NO_EVENT_SINK", message: "没有可用的事件流接收器", details: nil))
                return
            }
            
            // 在后台线程执行
            DispatchQueue.global(qos: .userInitiated).async {
                self.openAIService.publicSendFunctionCallResult(
                    messages: messagesRaw,
                    systemPrompt: systemPrompt,
                    functionCall: functionCallRaw,
                    functionResult: functionResult,
                    onToken: { token in
                        DispatchQueue.main.async {
                            eventSink(["type": "token", "content": token])
                        }
                    },
                    onComplete: {
                        DispatchQueue.main.async {
                            eventSink(["type": "complete"])
                        }
                    },
                    onError: { error in
                        DispatchQueue.main.async {
                            eventSink(["type": "error", "content": error.localizedDescription])
                        }
                    },
                    onFunctionCall: { functionCall in
                        DispatchQueue.main.async {
                            do {
                                let data = try JSONSerialization.data(withJSONObject: functionCall)
                                if let jsonString = String(data: data, encoding: .utf8) {
                                    eventSink(["type": "function_call", "content": jsonString])
                                } else {
                                    eventSink(["type": "error", "content": "函数调用数据序列化失败"])
                                }
                            } catch {
                                eventSink(["type": "error", "content": "函数调用数据序列化失败: \(error.localizedDescription)"])
                            }
                        }
                    },
                    onFunctionCallResult: { functionCall, functionCallResult in
                        DispatchQueue.main.async {
                            do {
                                let functionCallData = try JSONSerialization.data(withJSONObject: functionCall)
                                let resultData = try JSONSerialization.data(withJSONObject: functionCallResult)
                                
                                if let functionCallString = String(data: functionCallData, encoding: .utf8),
                                   let resultString = String(data: resultData, encoding: .utf8) {
                                    eventSink(["type": "function_call_result", "function_call": functionCallString, "result": resultString])
                                } else {
                                    eventSink(["type": "error", "content": "函数调用结果数据序列化失败"])
                                }
                            } catch {
                                eventSink(["type": "error", "content": "函数调用结果数据序列化失败: \(error.localizedDescription)"])
                            }
                        }
                    }
                )
                
                // 请求已开始
                DispatchQueue.main.async {
                    result(true)
                }
            }
            
        case "createUserMessage":
            if let args = call.arguments as? [String: Any],
               let content = args["content"] as? String {
                let message = openAIService.createUserMessage(content: content)
                result(message)
            } else {
                result(FlutterError(code: "INVALID_ARGUMENT", message: "创建用户消息参数无效", details: nil))
            }
            
        case "createAssistantMessage":
            if let args = call.arguments as? [String: Any],
               let content = args["content"] as? String {
                let message = openAIService.createAssistantMessage(content: content)
                result(message)
            } else {
                result(FlutterError(code: "INVALID_ARGUMENT", message: "创建助手消息参数无效", details: nil))
            }
            
        case "createUserMessageWithImage":
            if let args = call.arguments as? [String: Any],
               let text = args["text"] as? String,
               let imageBase64 = args["imageBase64"] as? String {
                let message = openAIService.createUserMessageWithImage(text: text, imageBase64: imageBase64)
                result(message)
            } else {
                result(FlutterError(code: "INVALID_ARGUMENT", message: "创建图片消息参数无效", details: nil))
            }
            
        case "fileToBase64":
            if let args = call.arguments as? [String: Any],
               let filePath = args["filePath"] as? String {
                let maxSizeKB = args["maxSizeKB"] as? Int ?? 20480
                let base64 = openAIService.fileToBase64(filePath, maxSizeKB: maxSizeKB)
                if let base64 = base64 {
                    result(base64)
                } else {
                    result(FlutterError(code: "FILE_ERROR", message: "文件转换失败", details: nil))
                }
            } else {
                result(FlutterError(code: "INVALID_ARGUMENT", message: "文件转换参数无效", details: nil))
            }
            
        case "initializeMcpClient":
            if let args = call.arguments as? [String: Any],
               let serverUrl = args["serverUrl"] as? String {
                let success = openAIService.initializeMcpClient(serverUrl)
                result(success)
            } else {
                result(FlutterError(code: "INVALID_ARGUMENT", message: "初始化MCP客户端参数无效", details: nil))
            }
            
        case "isMcpInitialized":
            result(openAIService.checkMcpInitialized())
            
        case "closeMcpClient":
            result(openAIService.closeMcpClient())
            
        case "handleMcpToolCall":
            if let args = call.arguments as? [String: Any],
               let functionCallJson = args["functionCall"] as? String {
                
                // 解析函数调用JSON
                guard let data = functionCallJson.data(using: .utf8),
                      let functionCall = try? JSONSerialization.jsonObject(with: data) as? [String: Any] else {
                    result(FlutterError(code: "INVALID_ARGUMENT", message: "无效的函数调用JSON", details: nil))
                    return
                }
                
                // 在后台线程执行
                Task {
                    do {
                        let toolResult = await openAIService.handleMcpToolCall(functionCall)
                        
                        // 在主线程返回结果
                        DispatchQueue.main.async {
                            result(toolResult)
                        }
                    } catch {
                        // 在主线程返回错误
                        DispatchQueue.main.async {
                            result(FlutterError(code: "MCP_TOOL_ERROR", message: "处理MCP工具调用失败: \(error.localizedDescription)", details: nil))
                        }
                    }
                }
            } else {
                result(FlutterError(code: "INVALID_ARGUMENT", message: "处理MCP工具调用参数无效", details: nil))
            }
            
        default:
            result(FlutterMethodNotImplemented)
        }
    }
    
    // MARK: - FlutterStreamHandler
    
    public func onListen(withArguments arguments: Any?, eventSink events: @escaping FlutterEventSink) -> FlutterError? {
        self.eventSink = events
        return nil
    }
    
    public func onCancel(withArguments arguments: Any?) -> FlutterError? {
        self.eventSink = nil
        return nil
    }
}