//
//  CompatAdjuster.swift
//  Zyquo Router
//
//  Author: Simon-Pierre Boucher
//  Mail: contact@spboucher.ai
//
//  Per-provider request/response adjustment for the OpenAI-compatible
//  upstreams (everything except Anthropic and Gemini, which translate fully).
//  Tables from docs/ROUTER-RESEARCH.md §3.3: strip what 400s, rename what
//  differs, clamp what has narrower ranges, pass unknown keys through (D5),
//  and normalize response quirks into the OpenAI contract.
//

import Foundation

enum CompatAdjuster {
    // MARK: - Request adjustment

    /// Rewrites an inbound OpenAI request body for a specific compat upstream.
    static func adjustRequest(
        _ original: [String: Any],
        model: AIModel,
        stream: Bool
    ) -> [String: Any] {
        var body = original
        let provider = model.provider

        // Bare upstream model id (namespace/alias resolution already happened).
        body["model"] = model.id

        // max_tokens naming: send whichever the model documents, never both.
        let maxTokens = (body["max_completion_tokens"] as? Int) ?? (body["max_tokens"] as? Int)
        body["max_tokens"] = nil
        body["max_completion_tokens"] = nil
        if let maxTokens {
            body[model.parameterSupport.usesMaxCompletionTokens ? "max_completion_tokens" : "max_tokens"] = maxTokens
        }

        // stream/stream_options are router-controlled.
        body["stream"] = stream
        body["stream_options"] = nil
        if stream, supportsStreamUsage(provider) {
            body["stream_options"] = ["include_usage": true]
        }

        // Parameters the catalog says this model rejects.
        if !model.parameterSupport.temperature { body["temperature"] = nil }
        if !model.parameterSupport.topP { body["top_p"] = nil }
        if !model.parameterSupport.frequencyPenalty { body["frequency_penalty"] = nil }
        if !model.parameterSupport.presencePenalty { body["presence_penalty"] = nil }
        // reasoning_effort goes upstream only when the model accepts the
        // param (some reasoning models 400 on it — Grok 4.20, grok-code).
        // Mistral/Qwen re-map it below before this strip matters.
        if !model.capabilities.reasoning { body["reasoning_effort"] = nil }
        if !model.parameterSupport.reasoningEffort, provider != .mistral, provider != .qwen {
            body["reasoning_effort"] = nil
        }

        // Echo-back hygiene: strip reasoning_content/reasoning_details from
        // incoming assistant messages (wastes tokens; DeepSeek 400s) — except
        // DeepSeek tool-call loops, which REQUIRE reasoning_content back.
        let keepReasoning = provider == .deepseek && body["tools"] != nil
        if !keepReasoning, var messages = body["messages"] as? [[String: Any]] {
            for index in messages.indices where messages[index]["role"] as? String == "assistant" {
                messages[index]["reasoning_content"] = nil
                messages[index]["reasoning_details"] = nil
            }
            body["messages"] = messages
        }

        switch provider {
        case .openai:
            // gpt-5.6 family: tools + (implicit) reasoning are rejected on
            // /v1/chat/completions; an explicit reasoning_effort "none"
            // unlocks tool calling (verified Phase 7).
            if model.id.hasPrefix("gpt-5.6"), body["tools"] != nil, body["reasoning_effort"] == nil {
                body["reasoning_effort"] = "none"
            }
        case .xai:
            if model.capabilities.reasoning {
                // Grok reasoning models 400 on these instead of ignoring them.
                body["presence_penalty"] = nil
                body["frequency_penalty"] = nil
                body["stop"] = nil
            }
        case .mistral:
            if let seed = body["seed"] { body["random_seed"] = seed; body["seed"] = nil }
            body["logit_bias"] = nil
            body["user"] = nil
            body["logprobs"] = nil
            // Mistral reasoning: models with native reasoning_effort accept
            // only none|high; Magistral-style models use prompt_mode instead.
            if let effort = body["reasoning_effort"] as? String {
                if model.parameterSupport.reasoningEffort {
                    body["reasoning_effort"] = effort == "none" ? "none" : "high"
                } else {
                    body["reasoning_effort"] = nil
                    if body["prompt_mode"] == nil { body["prompt_mode"] = "reasoning" }
                }
            }
        case .qwen:
            body["logit_bias"] = nil
            // Hybrid-thinking models think by default; honor the client's ask.
            if model.parameterSupport.thinkingToggle, body["enable_thinking"] == nil {
                body["enable_thinking"] = body["reasoning_effort"] != nil
                body["reasoning_effort"] = nil
            }
        case .deepseek:
            body["logprobs"] = nil
            body["top_logprobs"] = nil
        case .kimi:
            if let temperature = body["temperature"] as? Double {
                body["temperature"] = min(max(temperature, 0), 1)
            }
        case .deepinfra:
            body["logit_bias"] = nil
        case .perplexity, .together, .cerebras, .anthropic, .gemini, .custom:
            break
        }
        return body
    }

    /// Models whose non-streaming upstream endpoint is broken or times out
    /// while streaming works (Phase 7 findings) — the router transparently
    /// streams upstream and aggregates for buffered clients.
    static func requiresStreamingOverride(_ model: AIModel) -> Bool {
        switch (model.provider, model.id) {
        case (.deepinfra, "google/gemini-2.5-pro"),
             (.gemini, "gemini-3.1-pro-preview"),
             (.together, "Qwen/Qwen3.5-9B"):
            return true
        default:
            return false
        }
    }

    /// Providers that honor `stream_options.include_usage` (research §3.3).
    static func supportsStreamUsage(_ provider: ProviderID) -> Bool {
        switch provider {
        case .openai, .xai, .mistral, .qwen, .deepseek, .kimi, .together, .deepinfra, .cerebras:
            return true
        case .perplexity, .anthropic, .gemini, .custom:
            return false
        }
    }

    // MARK: - Response normalization (non-streaming)

    /// Normalizes a compat upstream's `chat.completion` into the router
    /// contract: namespaced model echo, canonical finish_reason, flattened
    /// reasoning, usage detail mapping. Unknown fields pass through.
    static func normalizeResponse(
        _ original: [String: Any],
        namespacedModel: String,
        provider: ProviderID
    ) -> [String: Any] {
        var body = original
        body["model"] = namespacedModel

        if var choices = body["choices"] as? [[String: Any]] {
            for index in choices.indices {
                if let finish = choices[index]["finish_reason"] as? String {
                    choices[index]["finish_reason"] = normalizeFinishReason(finish)
                }
                if var message = choices[index]["message"] as? [String: Any] {
                    normalizeMessage(&message, provider: provider)
                    choices[index]["message"] = message
                }
                // Perplexity leaks a `delta` field into non-streaming choices.
                if provider == .perplexity { choices[index]["delta"] = nil }
            }
            body["choices"] = choices
        }

        if var usage = body["usage"] as? [String: Any] {
            normalizeUsage(&usage, provider: provider)
            body["usage"] = usage
        }
        return body
    }

    /// Normalizes one streamed chunk in place. Returns nil for chunks the
    /// router should swallow (e.g. usage chunk when the client didn't ask).
    static func normalizeChunk(
        _ original: [String: Any],
        namespacedModel: String,
        provider: ProviderID,
        clientWantsUsage: Bool
    ) -> [String: Any]? {
        var chunk = original
        chunk["model"] = namespacedModel

        // Perplexity terminates streams with a non-spec summary event
        // (object: "chat.completion.done") — swallow it, keeping its usage.
        if let object = chunk["object"] as? String, object != "chat.completion.chunk" {
            return nil
        }
        // Some hosted models omit the object field entirely (DeepInfra-hosted
        // Gemini) — strict SDKs require it.
        if chunk["object"] == nil {
            chunk["object"] = "chat.completion.chunk"
        }

        if var choices = chunk["choices"] as? [[String: Any]] {
            for index in choices.indices {
                if let finish = choices[index]["finish_reason"] as? String {
                    choices[index]["finish_reason"] = normalizeFinishReason(finish)
                }
                if var delta = choices[index]["delta"] as? [String: Any] {
                    normalizeMessage(&delta, provider: provider)
                    choices[index]["delta"] = delta
                }
            }
            chunk["choices"] = choices

            // Usage-only chunk (empty choices): swallow unless requested.
            if choices.isEmpty, chunk["usage"] != nil, !clientWantsUsage {
                return nil
            }
        }

        if var usage = chunk["usage"] as? [String: Any] {
            normalizeUsage(&usage, provider: provider)
            chunk["usage"] = usage
        }
        return chunk
    }

    // MARK: - Shared pieces

    /// Everything lands in the OpenAI closed set (research §3.3 rule 4).
    static func normalizeFinishReason(_ raw: String) -> String {
        switch raw {
        case "stop", "length", "tool_calls", "content_filter", "function_call":
            return raw
        case "eos":
            return "stop"
        case "max_tokens", "model_length":
            return "length"
        case "safety", "recitation":
            return "content_filter"
        default:
            return "stop"
        }
    }

    /// Message/delta-level quirks: Mistral thinking arrays, Perplexity
    /// <think> tags, Together text field.
    private static func normalizeMessage(_ message: inout [String: Any], provider: ProviderID) {
        switch provider {
        case .mistral:
            // Magistral: content is an ARRAY of {type:"thinking"|"text"} chunks.
            if let parts = message["content"] as? [[String: Any]] {
                var text = ""
                var reasoning = message["reasoning_content"] as? String ?? ""
                for part in parts {
                    switch part["type"] as? String {
                    case "text":
                        text += part["text"] as? String ?? ""
                    case "thinking":
                        for inner in part["thinking"] as? [[String: Any]] ?? [] {
                            reasoning += inner["text"] as? String ?? ""
                        }
                    default:
                        break
                    }
                }
                message["content"] = text
                if !reasoning.isEmpty { message["reasoning_content"] = reasoning }
            }
        case .perplexity:
            // sonar-reasoning embeds <think>…</think> in content.
            if let content = message["content"] as? String,
               content.hasPrefix("<think>"),
               let closeRange = content.range(of: "</think>") {
                let reasoning = String(content[content.index(content.startIndex, offsetBy: 7)..<closeRange.lowerBound])
                message["reasoning_content"] = reasoning
                message["content"] = String(content[closeRange.upperBound...])
                    .trimmingCharacters(in: .whitespacesAndNewlines)
            }
        case .together:
            // Some Together models put streamed text in choices[].text.
            if message["content"] == nil, let text = message["text"] as? String {
                message["content"] = text
                message["text"] = nil
            }
        default:
            break
        }
    }

    /// Provider usage extras → OpenAI detail objects.
    private static func normalizeUsage(_ usage: inout [String: Any], provider: ProviderID) {
        if provider == .deepseek, let cacheHit = usage["prompt_cache_hit_tokens"] as? Int, cacheHit > 0 {
            var details = usage["prompt_tokens_details"] as? [String: Any] ?? [:]
            if details["cached_tokens"] == nil { details["cached_tokens"] = cacheHit }
            usage["prompt_tokens_details"] = details
        }
        if let reasoningTokens = usage["reasoning_tokens"] as? Int, reasoningTokens > 0 {
            var details = usage["completion_tokens_details"] as? [String: Any] ?? [:]
            if details["reasoning_tokens"] == nil { details["reasoning_tokens"] = reasoningTokens }
            usage["completion_tokens_details"] = details
            usage["reasoning_tokens"] = nil
        }
    }
}