import Foundation

/// 通过 OpenAI 兼容 API 对转录文本进行纠错优化
final class LLMRefiner {

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

    /// 纠错结果
    struct RefinementResult {
        let refinedText: String
        let wasChanged: Bool
    }

    /// 对转录文本进行纠错
    /// - Parameters:
    ///   - text: 原始转录文本
    ///   - config: LLM 配置
    /// - Returns: 纠错结果
    func refine(text: String, config: Config) async throws -> RefinementResult {
        let baseURL = config.llmBaseURL
            .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.llmAPIKey)", forHTTPHeaderField: "Authorization")

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

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

        let body: [String: Any] = [
            "model": config.llmModel,
            "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 refined = message["content"] as? String else {
            throw NSError(domain: "LLMRefiner", code: -2,
                          userInfo: [NSLocalizedDescriptionKey: "Failed to parse response"])
        }

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

        print("[LLM] Refinement \(wasChanged ? "changed" : "unchanged"): \"\(trimmed.prefix(80))...\"")
        return RefinementResult(refinedText: trimmed, wasChanged: wasChanged)
    }

    /// 测试连接
    func testConnection(config: Config) async throws -> String {
        let testText = "测试连接"
        let result = try await refine(text: testText, config: config)
        return result.refinedText
    }
}