//
//  GeminiTranslator.swift
//  Zyquo Router
//
//  Author: Simon-Pierre Boucher
//  Mail: contact@spboucher.ai
//
//  Native Gemini generateContent / streamGenerateContent translation per
//  docs/ROUTER-RESEARCH.md §3.2 (decision D10): contents/parts with role
//  rename, systemInstruction, generationConfig, functionDeclarations, the
//  finishReason table (incl. the STOP+functionCall → tool_calls override),
//  and SSE chunks → OpenAI chunks (role synthesized, whole-argument tool
//  deltas, router-added [DONE]).
//

import Foundation

enum GeminiTranslator {
    // MARK: - Request (OpenAI → Gemini)

    static func buildRequest(_ request: ChatCompletionRequest, model: AIModel) -> [String: Any] {
        var body: [String: Any] = [:]

        let systemText = request.messages
            .filter { $0.role == "system" || $0.role == "developer" }
            .map(\.flattenedText)
            .filter { !$0.isEmpty }
            .joined(separator: "\n\n")
        if !systemText.isEmpty {
            body["systemInstruction"] = ["parts": [["text": systemText]]]
        }

        // tool_call_id → function name, resolved from prior assistant turns
        // (functionResponse requires the name).
        var callNames: [String: String] = [:]
        for message in request.messages {
            for call in message.toolCalls ?? [] {
                if let id = call["id"] as? String,
                   let name = (call["function"] as? [String: Any])?["name"] as? String {
                    callNames[id] = name
                }
            }
        }

        var contents: [[String: Any]] = []
        func append(role: String, parts: [[String: Any]]) {
            guard !parts.isEmpty else { return }
            if var last = contents.last, last["role"] as? String == role {
                var merged = last["parts"] as? [[String: Any]] ?? []
                merged.append(contentsOf: parts)
                last["parts"] = merged
                contents[contents.count - 1] = last
            } else {
                contents.append(["role": role, "parts": parts])
            }
        }

        for message in request.messages {
            switch message.role {
            case "system", "developer":
                continue
            case "user":
                append(role: "user", parts: parts(for: message))
            case "assistant":
                var assistantParts = parts(for: message)
                for call in message.toolCalls ?? [] {
                    guard let function = call["function"] as? [String: Any],
                          let name = function["name"] as? String else { continue }
                    let arguments = function["arguments"] as? String ?? "{}"
                    let args = (try? JSONSerialization.jsonObject(with: Data(arguments.utf8))) as? [String: Any] ?? [:]
                    var functionCall: [String: Any] = ["name": name, "args": args]
                    if let id = call["id"] as? String { functionCall["id"] = id }
                    assistantParts.append(["functionCall": functionCall])
                }
                append(role: "model", parts: assistantParts)
            case "tool":
                let callID = message.toolCallID ?? ""
                let text = message.flattenedText
                // functionResponse.response must be a JSON object.
                let response = (try? JSONSerialization.jsonObject(with: Data(text.utf8))) as? [String: Any]
                    ?? ["result": text]
                var functionResponse: [String: Any] = [
                    "name": callNames[callID] ?? callID,
                    "response": response,
                ]
                if !callID.isEmpty { functionResponse["id"] = callID }
                append(role: "user", parts: [["functionResponse": functionResponse]])
            default:
                continue
            }
        }
        body["contents"] = contents

        var generation: [String: Any] = [:]
        if let temperature = request.raw["temperature"] as? Double { generation["temperature"] = temperature }
        if let topP = request.raw["top_p"] as? Double { generation["topP"] = topP }
        if let topK = request.raw["top_k"] as? Int { generation["topK"] = topK }
        if let maxTokens = request.maxTokens { generation["maxOutputTokens"] = maxTokens }
        if let seed = request.raw["seed"] as? Int { generation["seed"] = seed }
        if let presence = request.raw["presence_penalty"] as? Double { generation["presencePenalty"] = presence }
        if let frequency = request.raw["frequency_penalty"] as? Double { generation["frequencyPenalty"] = frequency }
        if let stop = request.raw["stop"] as? String {
            generation["stopSequences"] = [stop]
        } else if let stop = request.raw["stop"] as? [String] {
            generation["stopSequences"] = stop
        }
        if let format = request.raw["response_format"] as? [String: Any] {
            switch format["type"] as? String {
            case "json_object":
                generation["responseMimeType"] = "application/json"
            case "json_schema":
                generation["responseMimeType"] = "application/json"
                if let schema = (format["json_schema"] as? [String: Any])?["schema"] {
                    generation["responseJsonSchema"] = schema
                }
            default:
                break
            }
        }
        if model.capabilities.reasoning, let effort = request.raw["reasoning_effort"] as? String {
            let budget: Int
            switch effort {
            case "minimal", "low": budget = 1024
            case "high", "xhigh", "max": budget = 24576
            default: budget = 8192
            }
            generation["thinkingConfig"] = ["thinkingBudget": budget, "includeThoughts": true]
        }
        if !generation.isEmpty { body["generationConfig"] = generation }

        if let tools = request.tools {
            let declarations = tools.compactMap { tool -> [String: Any]? in
                guard let function = tool["function"] as? [String: Any],
                      let name = function["name"] as? String else { return nil }
                var declaration: [String: Any] = ["name": name]
                if let description = function["description"] as? String {
                    declaration["description"] = description
                }
                if var parameters = function["parameters"] as? [String: Any] {
                    parameters["$schema"] = nil
                    declaration["parameters"] = parameters
                }
                return declaration
            }
            if !declarations.isEmpty {
                body["tools"] = [["functionDeclarations": declarations]]
            }
        }
        var callingConfig: [String: Any]?
        switch request.raw["tool_choice"] {
        case let choice as String:
            switch choice {
            case "auto": callingConfig = ["mode": "AUTO"]
            case "required": callingConfig = ["mode": "ANY"]
            case "none": callingConfig = ["mode": "NONE"]
            default: break
            }
        case let choice as [String: Any]:
            if let function = choice["function"] as? [String: Any], let name = function["name"] as? String {
                callingConfig = ["mode": "ANY", "allowedFunctionNames": [name]]
            }
        default:
            break
        }
        if let callingConfig {
            body["toolConfig"] = ["functionCallingConfig": callingConfig]
        }
        return body
    }

    private static func parts(for message: OAIMessage) -> [[String: Any]] {
        if let text = message.contentString {
            return text.isEmpty ? [] : [["text": text]]
        }
        guard let contentParts = message.contentParts else { return [] }
        return contentParts.compactMap { part in
            switch part["type"] as? String {
            case "text":
                let text = part["text"] as? String ?? ""
                return text.isEmpty ? nil : ["text": text]
            case "image_url":
                guard let image = part["image_url"] as? [String: Any],
                      let url = image["url"] as? String,
                      url.hasPrefix("data:"),
                      let comma = url.firstIndex(of: ",") else { return nil }
                let header = url[url.index(url.startIndex, offsetBy: 5)..<comma]
                let mimeType = header.split(separator: ";").first.map(String.init) ?? "image/png"
                return ["inlineData": [
                    "mimeType": mimeType,
                    "data": String(url[url.index(after: comma)...]),
                ]]
            default:
                return nil
            }
        }
    }

    /// Remote (non data-URI) image URLs need inlining for Gemini; the route
    /// rejects them with a clear error instead of fetching arbitrary URLs.
    static func hasRemoteImageURL(_ request: ChatCompletionRequest) -> Bool {
        request.messages.contains { message in
            message.contentParts?.contains { part in
                guard (part["type"] as? String) == "image_url",
                      let url = (part["image_url"] as? [String: Any])?["url"] as? String else { return false }
                return !url.hasPrefix("data:")
            } ?? false
        }
    }

    // MARK: - Shared mapping

    static func finishReason(from raw: String?, hasFunctionCall: Bool) -> String {
        if hasFunctionCall { return "tool_calls" }
        switch raw {
        case "STOP", .none: return "stop"
        case "MAX_TOKENS": return "length"
        case "SAFETY", "PROHIBITED_CONTENT", "BLOCKLIST", "SPII", "IMAGE_SAFETY", "RECITATION":
            return "content_filter"
        default: return "stop"
        }
    }

    static func normalizedUsage(_ metadata: [String: Any]) -> [String: Any] {
        let prompt = metadata["promptTokenCount"] as? Int ?? 0
        let candidates = metadata["candidatesTokenCount"] as? Int ?? 0
        let thoughts = metadata["thoughtsTokenCount"] as? Int ?? 0
        return UsageBuilder.build(
            promptTokens: prompt,
            completionTokens: candidates + thoughts, // OpenAI counts reasoning inside completion
            cachedTokens: metadata["cachedContentTokenCount"] as? Int,
            reasoningTokens: thoughts > 0 ? thoughts : nil
        )
    }

    /// Empty-candidates + promptFeedback.blockReason → clear 400.
    static func blockReason(_ upstream: [String: Any]) -> String? {
        guard (upstream["candidates"] as? [[String: Any]] ?? []).isEmpty else { return nil }
        return (upstream["promptFeedback"] as? [String: Any])?["blockReason"] as? String
    }

    // MARK: - Response (Gemini → OpenAI), non-streaming

    static func translateResponse(_ upstream: [String: Any], emitter: ChunkEmitter) -> [String: Any] {
        let candidate = (upstream["candidates"] as? [[String: Any]])?.first ?? [:]
        let parts = (candidate["content"] as? [String: Any])?["parts"] as? [[String: Any]] ?? []

        var text = ""
        var reasoning = ""
        var toolCalls: [[String: Any]] = []
        for part in parts {
            if let partText = part["text"] as? String {
                if part["thought"] as? Bool == true {
                    reasoning += partText
                } else {
                    text += partText
                }
            }
            if let functionCall = part["functionCall"] as? [String: Any] {
                let args = functionCall["args"] as? [String: Any] ?? [:]
                let arguments = String(
                    data: (try? JSONSerialization.data(withJSONObject: args)) ?? Data("{}".utf8),
                    encoding: .utf8
                ) ?? "{}"
                toolCalls.append([
                    "id": functionCall["id"] as? String ?? "call_\(UUID().uuidString.prefix(12))",
                    "type": "function",
                    "function": ["name": functionCall["name"] as? String ?? "", "arguments": arguments],
                ])
            }
        }

        var message: [String: Any] = ["role": "assistant"]
        message["content"] = toolCalls.isEmpty || !text.isEmpty ? text : NSNull()
        if !reasoning.isEmpty { message["reasoning_content"] = reasoning }
        if !toolCalls.isEmpty { message["tool_calls"] = toolCalls }

        return emitter.completion(
            message: message,
            finishReason: finishReason(
                from: candidate["finishReason"] as? String,
                hasFunctionCall: !toolCalls.isEmpty
            ),
            usage: normalizedUsage(upstream["usageMetadata"] as? [String: Any] ?? [:])
        )
    }

    // MARK: - Streaming (Gemini SSE → OpenAI chunks)

    /// Each Gemini SSE data payload is a complete GenerateContentResponse
    /// carrying the increment; there is no upstream [DONE].
    struct StreamMachine {
        let emitter: ChunkEmitter
        let includeUsage: Bool

        private var roleSent = false
        private var toolIndex = -1
        private var finishSent = false
        /// Accumulated assistant text for the request-log inspector (capped).
        private(set) var textPreview = ""
        private(set) var lastUsage: [String: Any]?
        private(set) var promptTokens = 0
        private(set) var completionTokens = 0

        init(emitter: ChunkEmitter, includeUsage: Bool) {
            self.emitter = emitter
            self.includeUsage = includeUsage
        }

        mutating func consume(_ event: SSEEvent) -> [Data] {
            guard let json = (try? JSONSerialization.jsonObject(with: Data(event.data.utf8))) as? [String: Any] else {
                return []
            }
            var payloads: [Data] = []
            if !roleSent {
                roleSent = true
                payloads.append(emitter.roleChunk())
            }

            if let metadata = json["usageMetadata"] as? [String: Any] {
                lastUsage = metadata // cumulative; last one wins
                promptTokens = metadata["promptTokenCount"] as? Int ?? promptTokens
                completionTokens = (metadata["candidatesTokenCount"] as? Int ?? 0)
                    + (metadata["thoughtsTokenCount"] as? Int ?? 0)
            }

            let candidate = (json["candidates"] as? [[String: Any]])?.first ?? [:]
            var sawFunctionCall = false
            for part in (candidate["content"] as? [String: Any])?["parts"] as? [[String: Any]] ?? [] {
                if let text = part["text"] as? String, !text.isEmpty {
                    if part["thought"] as? Bool == true {
                        payloads.append(emitter.reasoningChunk(text))
                    } else {
                        if textPreview.count < 20_000 { textPreview += text }
                        payloads.append(emitter.contentChunk(text))
                    }
                }
                if let functionCall = part["functionCall"] as? [String: Any] {
                    // Tool calls arrive complete: announce, then one full-
                    // arguments delta (strict SDKs accept single-shot args).
                    sawFunctionCall = true
                    toolIndex += 1
                    let args = functionCall["args"] as? [String: Any] ?? [:]
                    let arguments = String(
                        data: (try? JSONSerialization.data(withJSONObject: args)) ?? Data("{}".utf8),
                        encoding: .utf8
                    ) ?? "{}"
                    payloads.append(emitter.toolCallStartChunk(
                        toolIndex: toolIndex,
                        callID: functionCall["id"] as? String ?? "call_\(UUID().uuidString.prefix(12))",
                        name: functionCall["name"] as? String ?? ""
                    ))
                    payloads.append(emitter.toolCallArgumentsChunk(toolIndex: toolIndex, fragment: arguments))
                }
            }

            if let rawFinish = candidate["finishReason"] as? String, !finishSent {
                finishSent = true
                payloads.append(emitter.finishChunk(reason: GeminiTranslator.finishReason(
                    from: rawFinish,
                    hasFunctionCall: sawFunctionCall || toolIndex >= 0
                )))
            }
            return payloads
        }

        /// Called when the upstream stream ends (no [DONE] from Gemini).
        mutating func finalPayloads() -> [Data] {
            var payloads: [Data] = []
            if !finishSent {
                finishSent = true
                payloads.append(emitter.finishChunk(reason: "stop"))
            }
            if includeUsage {
                payloads.append(emitter.usageChunk(
                    lastUsage.map(GeminiTranslator.normalizedUsage)
                        ?? UsageBuilder.build(promptTokens: promptTokens, completionTokens: completionTokens, estimated: true)
                ))
            }
            return payloads
        }
    }
}