//
//  OpenAICompatibleClient.swift
//  Zyquo Cloud
//
//  Author: Simon-Pierre Boucher
//  Mail: contact@spboucher.ai
//
//  One client for every provider speaking the OpenAI /chat/completions schema:
//  OpenAI, xAI, Mistral, Gemini (compat endpoint), Qwen/DashScope, DeepSeek,
//  Kimi, Perplexity, Together, DeepInfra, Cerebras, and custom endpoints.
//  All provider quirks live HERE — nothing leaks into ViewModels or Views.
//

import Foundation

struct OpenAICompatibleClient: ProviderClient {
    let providerID: ProviderID
    /// Custom endpoints override the provider's default base URL.
    var baseURLOverride: URL?

    init(provider: ProviderID, baseURLOverride: URL? = nil) {
        self.providerID = provider
        self.baseURLOverride = baseURLOverride
    }

    // MARK: - Wire types (requests)

    private struct WireRequest: Encodable {
        var model: String
        var messages: [WireMessage]
        var stream: Bool?
        var streamOptions: StreamOptions?
        var temperature: Double?
        var topP: Double?
        var maxTokens: Int?
        var maxCompletionTokens: Int?
        var frequencyPenalty: Double?
        var presencePenalty: Double?
        var reasoningEffort: String?
        var enableThinking: Bool?

        enum CodingKeys: String, CodingKey {
            case model, messages, stream, temperature
            case streamOptions = "stream_options"
            case topP = "top_p"
            case maxTokens = "max_tokens"
            case maxCompletionTokens = "max_completion_tokens"
            case frequencyPenalty = "frequency_penalty"
            case presencePenalty = "presence_penalty"
            case reasoningEffort = "reasoning_effort"
            case enableThinking = "enable_thinking"
        }
    }

    private struct StreamOptions: Encodable {
        var includeUsage: Bool
        enum CodingKeys: String, CodingKey { case includeUsage = "include_usage" }
    }

    private struct WireMessage: Encodable {
        var role: String
        var content: WireContent
    }

    /// Message content: plain string, or an array of text/image parts for vision.
    private enum WireContent: Encodable {
        case text(String)
        case parts([WirePart])

        func encode(to encoder: Encoder) throws {
            var container = encoder.singleValueContainer()
            switch self {
            case .text(let s): try container.encode(s)
            case .parts(let p): try container.encode(p)
            }
        }
    }

    private enum WirePart: Encodable {
        case text(String)
        case imageURL(String)

        func encode(to encoder: Encoder) throws {
            var container = encoder.container(keyedBy: DynamicKey.self)
            switch self {
            case .text(let s):
                try container.encode("text", forKey: DynamicKey("type"))
                try container.encode(s, forKey: DynamicKey("text"))
            case .imageURL(let url):
                try container.encode("image_url", forKey: DynamicKey("type"))
                var nested = container.nestedContainer(keyedBy: DynamicKey.self, forKey: DynamicKey("image_url"))
                try nested.encode(url, forKey: DynamicKey("url"))
            }
        }
    }

    private struct DynamicKey: CodingKey {
        var stringValue: String
        var intValue: Int? { nil }
        init(_ s: String) { stringValue = s }
        init?(stringValue: String) { self.stringValue = stringValue }
        init?(intValue: Int) { nil }
    }

    // MARK: - Wire types (responses)

    private struct WireChunk: Decodable {
        var choices: [WireChoice]?
        var usage: WireUsage?
        var citations: [String]?
        var searchResults: [WireSearchResult]?

        enum CodingKeys: String, CodingKey {
            case choices, usage, citations
            case searchResults = "search_results"
        }
    }

    private struct WireChoice: Decodable {
        var delta: WireDelta?
        var message: WireDelta?
        /// Together streams some models completions-style: the token text
        /// lives in `choices[].text` instead of `delta.content`.
        var text: String?
        var finishReason: String?

        enum CodingKeys: String, CodingKey {
            case delta, message, text
            case finishReason = "finish_reason"
        }
    }

    private struct WireDelta: Decodable {
        var content: String?
        var reasoningContent: String?
        var reasoning: String?

        enum CodingKeys: String, CodingKey {
            case content, reasoning
            case reasoningContent = "reasoning_content"
        }

        init(from decoder: Decoder) throws {
            let container = try decoder.container(keyedBy: CodingKeys.self)
            reasoning = try? container.decodeIfPresent(String.self, forKey: .reasoning)
            reasoningContent = try? container.decodeIfPresent(String.self, forKey: .reasoningContent)
            // `content` is normally a string, but Mistral's reasoning models
            // return an array of chunks ({type: "thinking"|"text", …}).
            if let text = try? container.decodeIfPresent(String.self, forKey: .content) {
                content = text
            } else if let chunks = try? container.decodeIfPresent([ContentChunk].self, forKey: .content) {
                var textParts: [String] = []
                var thinkingParts: [String] = []
                for chunk in chunks {
                    if chunk.type == "thinking" {
                        thinkingParts.append(chunk.flattenedText)
                    } else {
                        textParts.append(chunk.flattenedText)
                    }
                }
                content = textParts.joined()
                let thinking = thinkingParts.joined()
                if !thinking.isEmpty, reasoningContent == nil {
                    reasoningContent = thinking
                }
            }
        }

        /// Mistral ThinkChunk/TextChunk: {"type":"text","text":…} or
        /// {"type":"thinking","thinking":[{"type":"text","text":…}]}.
        struct ContentChunk: Decodable {
            var type: String?
            var text: String?
            var thinking: [ContentChunkPart]?

            var flattenedText: String {
                if let text { return text }
                return (thinking ?? []).compactMap(\.text).joined()
            }
        }

        struct ContentChunkPart: Decodable {
            var text: String?
        }
    }

    private struct WireUsage: Decodable {
        var promptTokens: Int?
        var completionTokens: Int?
        var completionTokensDetails: Details?

        struct Details: Decodable {
            var reasoningTokens: Int?
            enum CodingKeys: String, CodingKey { case reasoningTokens = "reasoning_tokens" }
        }

        enum CodingKeys: String, CodingKey {
            case promptTokens = "prompt_tokens"
            case completionTokens = "completion_tokens"
            case completionTokensDetails = "completion_tokens_details"
        }

        var usage: TokenUsage {
            TokenUsage(
                inputTokens: promptTokens ?? 0,
                outputTokens: completionTokens ?? 0,
                reasoningTokens: completionTokensDetails?.reasoningTokens
            )
        }
    }

    private struct WireSearchResult: Decodable {
        var title: String?
        var url: String?
    }

    private struct WireModelList: Decodable {
        var data: [WireModelEntry]
    }

    private struct WireModelEntry: Decodable {
        var id: String
    }

    // MARK: - Request construction

    private var baseURL: URL? { baseURLOverride ?? providerID.defaultBaseURL }

    private func urlRequest(path: String, apiKey: String, method: String = "POST") throws -> URLRequest {
        guard let base = baseURL else {
            throw ProviderError.invalidResponse(providerID, detail: "no base URL configured")
        }
        // Preserve base path components ("…/v1", "…/compatible-mode/v1", "…/v1beta/openai").
        var request = URLRequest(url: base.appendingPathComponent(path))
        request.httpMethod = method
        request.setValue("Bearer \(apiKey)", forHTTPHeaderField: "Authorization")
        if method == "POST" {
            request.setValue("application/json", forHTTPHeaderField: "Content-Type")
        }
        return request
    }

    /// Providers whose final streamed chunk carries usage only when asked.
    private var wantsStreamOptions: Bool {
        switch providerID {
        case .openai, .xai, .gemini, .deepseek, .kimi, .together, .cerebras, .custom:
            return true
        // Qwen, DeepInfra, Perplexity include usage automatically; Mistral
        // rejects unknown params less gracefully — omit there.
        case .mistral, .qwen, .deepinfra, .perplexity:
            return false
        case .anthropic:
            return false // never routed here
        }
    }

    private func buildBody(_ request: ChatRequest) throws -> Data {
        var messages: [WireMessage] = []
        if let system = request.systemPrompt, !system.isEmpty {
            messages.append(WireMessage(role: "system", content: .text(system)))
        }
        for message in request.messages where message.role != .system {
            messages.append(wireMessage(from: message, vision: request.model.capabilities.vision))
        }

        let support = request.model.parameterSupport
        let params = request.parameters
        var wire = WireRequest(model: request.model.id, messages: messages)
        if request.stream {
            wire.stream = true
            if wantsStreamOptions {
                wire.streamOptions = StreamOptions(includeUsage: true)
            }
        }
        if support.temperature { wire.temperature = params.temperature }
        if support.topP { wire.topP = params.topP }
        if let max = params.maxTokens {
            if support.usesMaxCompletionTokens {
                wire.maxCompletionTokens = max
            } else {
                wire.maxTokens = max
            }
        }
        if support.frequencyPenalty { wire.frequencyPenalty = params.frequencyPenalty }
        if support.presencePenalty { wire.presencePenalty = params.presencePenalty }
        if support.reasoningEffort {
            // Mistral only accepts "high"/"none": map medium→high, low→none.
            if providerID == .mistral, let effort = params.reasoningEffort {
                wire.reasoningEffort = effort == "low" ? "none" : "high"
            } else {
                wire.reasoningEffort = params.reasoningEffort
            }
        }
        if support.thinkingToggle, providerID == .qwen {
            // DashScope: enable_thinking is only legal on streaming requests.
            if request.stream { wire.enableThinking = params.thinkingEnabled }
        }
        let encoder = JSONEncoder()
        return try encoder.encode(wire)
    }

    private func wireMessage(from message: Message, vision: Bool) -> WireMessage {
        let role = message.role == .assistant ? "assistant" : "user"
        var text = message.text
        // Text-file attachments are injected inline, fenced with the file name.
        for attachment in message.attachments where attachment.kind == .textFile {
            let contents = String(data: attachment.data, encoding: .utf8) ?? ""
            text += "\n\n```\(attachment.fileName)\n\(contents)\n```"
        }
        let images = message.attachments.filter { $0.kind == .image }
        guard vision, !images.isEmpty, message.role == .user else {
            return WireMessage(role: role, content: .text(text))
        }
        var parts: [WirePart] = [.text(text)]
        for image in images {
            let dataURI = "data:\(image.mimeType);base64,\(image.data.base64EncodedString())"
            parts.append(.imageURL(dataURI))
        }
        return WireMessage(role: role, content: .parts(parts))
    }

    // MARK: - ProviderClient

    func streamChat(_ request: ChatRequest, apiKey: String) -> AsyncThrowingStream<ChatEvent, Error> {
        AsyncThrowingStream { continuation in
            let task = Task {
                do {
                    var urlReq = try urlRequest(path: "chat/completions", apiKey: apiKey)
                    var streamRequest = request
                    streamRequest.stream = true
                    urlReq.httpBody = try buildBody(streamRequest)

                    var citationsSent = false
                    var finishReason: String?
                    let decoder = JSONDecoder()

                    for try await event in StreamingService.sseEvents(for: urlReq, provider: providerID) {
                        if event.data == "[DONE]" { break }
                        guard let data = event.data.data(using: .utf8),
                              let chunk = try? decoder.decode(WireChunk.self, from: data) else {
                            continue // tolerate unknown/malformed keep-alive chunks
                        }
                        if let choice = chunk.choices?.first {
                            if let reasoning = choice.delta?.reasoningContent ?? choice.delta?.reasoning,
                               !reasoning.isEmpty {
                                continuation.yield(.reasoningDelta(reasoning))
                            }
                            let deltaText = choice.delta?.content ?? choice.text
                            if let deltaText, !deltaText.isEmpty {
                                continuation.yield(.textDelta(deltaText))
                            }
                            if let reason = choice.finishReason {
                                finishReason = reason
                            }
                        }
                        if !citationsSent, let citations = Self.citations(from: chunk), !citations.isEmpty {
                            citationsSent = true
                            continuation.yield(.citations(citations))
                        }
                        if let usage = chunk.usage {
                            continuation.yield(.usage(usage.usage))
                        }
                    }
                    continuation.yield(.finished(reason: finishReason))
                    continuation.finish()
                } catch {
                    continuation.finish(throwing: error)
                }
            }
            continuation.onTermination = { _ in task.cancel() }
        }
    }

    func complete(_ request: ChatRequest, apiKey: String) async throws -> Message {
        // Some models reject non-streaming calls — aggregate a stream instead.
        if request.model.parameterSupport.requiresStreaming {
            return try await completeViaStream(request, apiKey: apiKey)
        }
        var urlReq = try urlRequest(path: "chat/completions", apiKey: apiKey)
        var plainRequest = request
        plainRequest.stream = false
        urlReq.httpBody = try buildBody(plainRequest)
        let data = try await StreamingService.postJSON(urlReq, provider: providerID)
        let chunk = try decodeOrThrow(WireChunk.self, from: data)
        guard let choice = chunk.choices?.first, let content = choice.message ?? choice.delta else {
            throw ProviderError.invalidResponse(providerID, detail: "response contained no message")
        }
        var message = Message(
            role: .assistant,
            text: content.content ?? choice.text ?? "",
            reasoning: content.reasoningContent ?? content.reasoning,
            modelID: request.model.id,
            provider: providerID
        )
        if let citations = Self.citations(from: chunk) {
            message.citations = citations
        }
        if let usage = chunk.usage?.usage {
            message.usage = usage
            message.estimatedCost = request.model.pricing?.cost(
                inputTokens: usage.inputTokens, outputTokens: usage.outputTokens
            )
        }
        return message
    }

    func listModelIDs(apiKey: String) async throws -> [String] {
        let urlReq = try urlRequest(path: "models", apiKey: apiKey, method: "GET")
        let data = try await StreamingService.getJSON(urlReq, provider: providerID)
        // Together returns a bare array; everyone else wraps in {"data": […]}.
        // Gemini's compat endpoint prefixes IDs with "models/" — normalize.
        let ids: [String]
        if let list = try? JSONDecoder().decode(WireModelList.self, from: data) {
            ids = list.data.map(\.id)
        } else if let bare = try? JSONDecoder().decode([WireModelEntry].self, from: data) {
            ids = bare.map(\.id)
        } else {
            throw ProviderError.invalidResponse(providerID, detail: "unrecognized /models response shape")
        }
        return ids.map { $0.hasPrefix("models/") ? String($0.dropFirst(7)) : $0 }
    }

    /// Non-streaming result assembled from the streaming endpoint, for models
    /// that only support `stream: true`.
    private func completeViaStream(_ request: ChatRequest, apiKey: String) async throws -> Message {
        var text = ""
        var reasoning = ""
        var citations: [Citation] = []
        var usage: TokenUsage?
        for try await event in streamChat(request, apiKey: apiKey) {
            switch event {
            case .textDelta(let delta): text += delta
            case .reasoningDelta(let delta): reasoning += delta
            case .citations(let c): citations = c
            case .usage(let u): usage = u
            case .finished: break
            }
        }
        var message = Message(
            role: .assistant,
            text: text,
            reasoning: reasoning.isEmpty ? nil : reasoning,
            citations: citations,
            modelID: request.model.id,
            provider: providerID
        )
        if let usage {
            message.usage = usage
            message.estimatedCost = request.model.pricing?.cost(
                inputTokens: usage.inputTokens, outputTokens: usage.outputTokens
            )
        }
        return message
    }

    // MARK: - Helpers

    private func decodeOrThrow<T: Decodable>(_ type: T.Type, from data: Data) throws -> T {
        do {
            return try JSONDecoder().decode(type, from: data)
        } catch {
            throw ProviderError.invalidResponse(providerID, detail: "decode failed: \(error.localizedDescription)")
        }
    }

    /// Perplexity: `citations` is an array of URL strings; `search_results`
    /// adds titles. Merge both into numbered citations.
    private static func citations(from chunk: WireChunk) -> [Citation]? {
        guard let urls = chunk.citations, !urls.isEmpty else { return nil }
        let titles = chunk.searchResults ?? []
        return urls.enumerated().compactMap { index, urlString in
            guard let url = URL(string: urlString) else { return nil }
            let title = index < titles.count ? titles[index].title : nil
            return Citation(index: index + 1, url: url, title: title)
        }
    }
}