import Foundation

enum AiChatClient {
    struct HistoryMessage {
        let role: String
        let text: String
    }

    static func complete(
        config: AiRuntimeConfig,
        system: String,
        user: String,
        imageJPEGBase64: String? = nil,
        history: [HistoryMessage] = [],
        capability: String = "chat"
    ) async throws -> String {
        var full = ""
        try await stream(
            config: config,
            system: system,
            user: user,
            imageJPEGBase64: imageJPEGBase64,
            history: history,
            capability: capability
        ) { delta in
            full = delta
        }
        return full
    }

    static func stream(
        config: AiRuntimeConfig,
        system: String,
        user: String,
        imageJPEGBase64: String? = nil,
        history: [HistoryMessage] = [],
        capability: String = "chat",
        onDelta: @escaping @Sendable (String) -> Void
    ) async throws {
        let want = imageJPEGBase64 == nil ? capability : "vision"
        guard let ep = config.endpoint(for: want) ?? config.endpoint(for: "chat") else {
            throw AiConfigResolveError.noCache
        }
        let base = ep.baseUrl.trimmingCharacters(in: CharacterSet(charactersIn: "/"))
        guard let url = URL(string: base + "/chat/completions") else {
            throw URLError(.badURL)
        }

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

        var apiMessages: [[String: Any]] = [["role": "system", "content": system]]
        for turn in history {
            apiMessages.append(["role": turn.role, "content": turn.text])
        }

        let userContent: Any
        if let b64 = imageJPEGBase64 {
            userContent = [
                ["type": "text", "text": user],
                ["type": "image_url", "image_url": ["url": "data:image/jpeg;base64,\(b64)"]],
            ]
        } else {
            userContent = user
        }
        apiMessages.append(["role": "user", "content": userContent])

        let body: [String: Any] = [
            "model": ep.model,
            "messages": apiMessages,
            "max_tokens": 800,
            "stream": true,
        ]
        req.httpBody = try JSONSerialization.data(withJSONObject: body)

        let (bytes, response) = try await URLSession.shared.bytes(for: req)
        guard let http = response as? HTTPURLResponse else { throw URLError(.badServerResponse) }
        guard (200..<300).contains(http.statusCode) else {
            var data = Data()
            for try await byte in bytes { data.append(byte) }
            let snippet = String(data: data, encoding: .utf8)?.prefix(200) ?? ""
            throw NSError(domain: "AiChat", code: http.statusCode, userInfo: [
                NSLocalizedDescriptionKey: "HTTP \(http.statusCode): \(snippet)",
            ])
        }

        var full = ""
        for try await line in bytes.lines {
            let trimmed = line.trimmingCharacters(in: .whitespaces)
            guard trimmed.hasPrefix("data:") else { continue }
            let data = trimmed.dropFirst(5).trimmingCharacters(in: .whitespaces)
            if data.isEmpty || data == "[DONE]" { continue }
            guard let jsonData = data.data(using: .utf8),
                  let json = try? JSONSerialization.jsonObject(with: jsonData) as? [String: Any],
                  let choices = json["choices"] as? [[String: Any]],
                  let first = choices.first else { continue }

            var piece = ""
            if let delta = first["delta"] as? [String: Any],
               let content = delta["content"] as? String {
                piece = content
            } else if let message = first["message"] as? [String: Any],
                      let content = message["content"] as? String {
                piece = content
            }
            if !piece.isEmpty {
                full += piece
                await MainActor.run { onDelta(full) }
            }
        }

        if full.trimmingCharacters(in: .whitespacesAndNewlines).isEmpty {
            let fallback = "(空回复)"
            await MainActor.run { onDelta(fallback) }
        }
    }
}