import Foundation

/// 通过 OpenAI 兼容 API 进行 LLM 处理
final class LLMRefiner {

    private let session: URLSession = {
        let config = URLSessionConfiguration.default
        config.timeoutIntervalForRequest = 30
        config.timeoutIntervalForResource = 60
        return URLSession(configuration: config)
    }()

    struct RefinementResult {
        let text: String
        let wasChanged: Bool
    }

    // MARK: - Input 模式：语音输入纠错（Fn 长按触发）

    func refine(text: String, config: LLMConfig) async throws -> RefinementResult {
        let systemPrompt = """
        你是一个语音识别后处理工具。你唯一的任务是修正转录文本中的明显识别错误。规则：
        1. 修正中文同音字/词错误。语音识别经常把技术词汇识别为同音的日常词汇，例如：
           「变异」→「编译」、「程序圆」→「程序员」、「配森」→「Python」。
           请根据上下文判断修正方向，将同音错误还原为正确的词汇。
        2. 修正被错误转录为中文的英文技术术语。
        3. 修正极其明显的中文和英文拼写错误（包括错别字、漏字、多字）。
        4. 修正明显的乱码或无意义字符序列。
        5. 绝不重写、润色、改进或修改任何看起来正确的内容。
        6. 如果输入本身正确，请原样返回，不做任何修改。
        7. 保留所有标点符号、空格、大小写和格式。
        8. 只返回修正后的文本，不要添加任何解释、前缀、后缀或 Markdown 标记。
        """
        
        return try await callLLM(
            text: text,
            systemPrompt: systemPrompt,
            config: config
        )
    }

    // MARK: - Control 键模式：AI 助手对话

    func chat(text: String, config: LLMConfig) async throws -> RefinementResult {
        let systemPrompt = """
        你是一个 AI 助手，用户通过语音输入与你交流。用户的输入来自语音识别，可能存在同音字或转录错误，请结合上下文理解。

        任务：
        1. 回答用户的问题、提供信息或执行对话
        2. 当用户询问「xx 的英文是什么」「xx 的日文是什么」等翻译请求时，直接输出对应的翻译结果，不要添加任何前缀或解释

        要求：
        1. 保持回答简洁自然，适合口语化场景
        2. 用户输入模糊时，尝试猜测意图并给出最可能的回答
        3. 使用用户正在使用的语言，或用户明确要求的语言进行回复
        4. 不要包含任何解释性文字、前缀、后缀或 Markdown 标记
        5. 【最重要】你的输出将直接进入用户的输入框，只返回最终回答内容
        """
        
        return try await callLLM(
            text: text,
            systemPrompt: systemPrompt,
            config: config
        )
    }

    // MARK: - 通用调用

    private func callLLM(text: String, systemPrompt: String, config: LLMConfig) async throws -> RefinementResult {
        let baseURL = config.baseURL.trimmingCharacters(in: CharacterSet(charactersIn: "/"))
        let url = URL(string: "\(baseURL)/v1/chat/completions")!

        var urlRequest = URLRequest(url: url)
        urlRequest.httpMethod = "POST"
        urlRequest.setValue("application/json", forHTTPHeaderField: "Content-Type")
        urlRequest.setValue("Bearer \(config.apiKey)", forHTTPHeaderField: "Authorization")

        let messages: [[String: String]] = [
            ["role": "system", "content": systemPrompt],
            ["role": "user", "content": text]
        ]

        let body: [String: Any] = [
            "model": config.model,
            "messages": messages,
            "max_tokens": 4096,
            "thinking": ["type": "disabled"]
        ]

        urlRequest.httpBody = try JSONSerialization.data(withJSONObject: body)

        let (data, response) = try await session.data(for: urlRequest)

        guard let httpResponse = response as? HTTPURLResponse else {
            throw NSError(domain: "LLMRefiner", code: -1,
                          userInfo: [NSLocalizedDescriptionKey: "Invalid response"])
        }

        guard httpResponse.statusCode == 200 else {
            let body = String(data: data, encoding: .utf8) ?? "Unknown"
            throw NSError(domain: "LLMRefiner", code: httpResponse.statusCode,
                          userInfo: [NSLocalizedDescriptionKey: "HTTP \(httpResponse.statusCode): \(body)"])
        }

        guard let json = try JSONSerialization.jsonObject(with: data) as? [String: Any],
              let choices = json["choices"] as? [[String: Any]],
              let firstChoice = choices.first,
              let message = firstChoice["message"] as? [String: Any],
              let result = message["content"] as? String else {
            throw NSError(domain: "LLMRefiner", code: -2,
                          userInfo: [NSLocalizedDescriptionKey: "Failed to parse response"])
        }

        let trimmed = result.trimmingCharacters(in: .whitespacesAndNewlines)
        let wasChanged = trimmed != text.trimmingCharacters(in: .whitespacesAndNewlines)

        return RefinementResult(text: trimmed, wasChanged: wasChanged)
    }

    /// 测试连接
    func testConnection(config: LLMConfig) async throws -> String {
        let testText = "测试连接"
        let result = try await callLLM(
            text: testText,
            systemPrompt: "You are a helpful assistant. Reply with 'OK' only.",
            config: config
        )
        return result.text
    }
}