//
//  AnthropicTranslator.swift
//  Zyquo Router
//
//  Author: Simon-Pierre Boucher
//  Mail: contact@spboucher.ai
//
//  Bidirectional translation between OpenAI chat/completions and the
//  Anthropic Messages API, per docs/ROUTER-RESEARCH.md §3.1: system
//  extraction, turn merging, tools/tool_choice, required max_tokens,
//  stop_reason/usage mapping, and the SSE event state machine that emits
//  byte-exact OpenAI chunks.
//

import Foundation

enum AnthropicTranslator {
    // MARK: - Request (OpenAI → Anthropic)

    static func buildRequest(_ request: ChatCompletionRequest, model: AIModel) -> [String: Any] {
        var body: [String: Any] = [
            "model": model.id,
            // max_tokens is REQUIRED: synthesize from the catalog when omitted.
            "max_tokens": request.maxTokens ?? model.maxOutputTokens ?? 4096,
        ]

        // System/developer messages → top-level system.
        let systemText = request.messages
            .filter { $0.role == "system" || $0.role == "developer" }
            .map(\.flattenedText)
            .filter { !$0.isEmpty }
            .joined(separator: "\n\n")
        var system = systemText

        // Conversation messages with Anthropic's constraints: only
        // user/assistant roles, alternating, first must be user; tool results
        // become user tool_result blocks; consecutive same-role turns merge.
        var messages: [[String: Any]] = []
        func append(role: String, blocks: [[String: Any]]) {
            guard !blocks.isEmpty else { return }
            if var last = messages.last, last["role"] as? String == role {
                var content = last["content"] as? [[String: Any]] ?? []
                content.append(contentsOf: blocks)
                last["content"] = content
                messages[messages.count - 1] = last
            } else {
                messages.append(["role": role, "content": blocks])
            }
        }

        for message in request.messages {
            switch message.role {
            case "system", "developer":
                continue
            case "user":
                append(role: "user", blocks: contentBlocks(message))
            case "assistant":
                var blocks = contentBlocks(message)
                for call in message.toolCalls ?? [] {
                    guard let function = call["function"] as? [String: Any] else { continue }
                    let arguments = function["arguments"] as? String ?? "{}"
                    let input = (try? JSONSerialization.jsonObject(with: Data(arguments.utf8))) as? [String: Any] ?? [:]
                    blocks.append([
                        "type": "tool_use",
                        "id": call["id"] as? String ?? "call_\(UUID().uuidString.prefix(12))",
                        "name": function["name"] as? String ?? "unknown",
                        "input": input,
                    ])
                }
                append(role: "assistant", blocks: blocks)
            case "tool":
                append(role: "user", blocks: [[
                    "type": "tool_result",
                    "tool_use_id": message.toolCallID ?? "",
                    "content": message.flattenedText,
                ]])
            default:
                continue
            }
        }
        if messages.first?["role"] as? String != "user" {
            messages.insert(["role": "user", "content": [["type": "text", "text": "(continue)"]]], at: 0)
        }
        body["messages"] = messages

        // Sampling params (temperature clamped to Anthropic's 0–1).
        if let temperature = request.raw["temperature"] as? Double {
            body["temperature"] = min(temperature, 1.0)
        }
        if let topP = request.raw["top_p"] as? Double { body["top_p"] = topP }
        if let topK = request.raw["top_k"] as? Int { body["top_k"] = topK }
        if let stop = request.raw["stop"] as? String {
            body["stop_sequences"] = [stop]
        } else if let stop = request.raw["stop"] as? [String] {
            body["stop_sequences"] = stop
        }
        if let user = request.raw["user"] as? String {
            body["metadata"] = ["user_id": user]
        }

        // Tools + tool_choice.
        if let tools = request.tools {
            body["tools"] = tools.compactMap { tool -> [String: Any]? in
                guard let function = tool["function"] as? [String: Any],
                      let name = function["name"] as? String else { return nil }
                var entry: [String: Any] = [
                    "name": name,
                    "input_schema": function["parameters"] as? [String: Any]
                        ?? ["type": "object", "properties": [String: Any]()],
                ]
                if let description = function["description"] as? String {
                    entry["description"] = description
                }
                return entry
            }
        }
        var toolChoice: [String: Any]?
        switch request.raw["tool_choice"] {
        case let choice as String:
            switch choice {
            case "auto": toolChoice = ["type": "auto"]
            case "required": toolChoice = ["type": "any"]
            case "none": toolChoice = ["type": "none"]
            default: break
            }
        case let choice as [String: Any]:
            if let function = choice["function"] as? [String: Any], let name = function["name"] as? String {
                toolChoice = ["type": "tool", "name": name]
            }
        default:
            break
        }
        if request.raw["parallel_tool_calls"] as? Bool == false, request.hasTools {
            var choice = toolChoice ?? ["type": "auto"]
            choice["disable_parallel_tool_use"] = true
            toolChoice = choice
        }
        if let toolChoice { body["tool_choice"] = toolChoice }

        // JSON mode (best-effort system steering; json_schema via the current
        // structured-output surface is provider-verified in Phase 7).
        if let format = request.raw["response_format"] as? [String: Any],
           let type = format["type"] as? String, type == "json_object" || type == "json_schema" {
            let instruction = "You must respond with valid JSON only — no prose, no markdown fences."
            system = system.isEmpty ? instruction : system + "\n\n" + instruction
        }
        if !system.isEmpty { body["system"] = system }

        // Reasoning: standard reasoning_effort → thinking config; raw
        // `thinking` extra-body always wins. The Claude 4.7+/5 family rejects
        // {"type":"enabled"} and wants adaptive thinking (verified Phase 7).
        if let thinking = request.raw["thinking"] as? [String: Any] {
            body["thinking"] = thinking
        } else if model.capabilities.reasoning,
                  let effort = request.raw["reasoning_effort"] as? String {
            if usesAdaptiveThinking(model.id) {
                body["thinking"] = ["type": "adaptive"]
            } else {
                let budget: Int
                switch effort {
                case "minimal", "low": budget = 1024
                case "high", "xhigh", "max": budget = 24576
                default: budget = 8192
                }
                if let maxTokens = body["max_tokens"] as? Int, maxTokens <= budget {
                    body["max_tokens"] = budget + 4096
                }
                body["thinking"] = ["type": "enabled", "budget_tokens": budget]
            }
        }

        body["stream"] = request.stream
        return body
    }

    /// Claude 4.7+, 4.8 and the 5 family take adaptive thinking only;
    /// budget-based "enabled" thinking 400s (Phase 7 verified).
    static func usesAdaptiveThinking(_ modelID: String) -> Bool {
        if modelID.range(of: #"claude-(opus|sonnet|haiku|fable)-5"#, options: .regularExpression) != nil {
            return true
        }
        return modelID.range(of: #"claude-(opus|sonnet|haiku)-4-(7|8|9)"#, options: .regularExpression) != nil
    }

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

    // MARK: - Shared mapping

    static func finishReason(from stopReason: String?) -> String {
        switch stopReason {
        case "end_turn", "stop_sequence", "pause_turn", .none: return "stop"
        case "max_tokens", "model_context_window_exceeded": return "length"
        case "tool_use": return "tool_calls"
        case "refusal": return "content_filter"
        default: return "stop"
        }
    }

    /// prompt_tokens include cache reads/writes (Anthropic excludes them).
    static func normalizedUsage(_ usage: [String: Any]) -> [String: Any] {
        let input = usage["input_tokens"] as? Int ?? 0
        let cacheRead = usage["cache_read_input_tokens"] as? Int ?? 0
        let cacheCreation = usage["cache_creation_input_tokens"] as? Int ?? 0
        let output = usage["output_tokens"] as? Int ?? 0
        return UsageBuilder.build(
            promptTokens: input + cacheRead + cacheCreation,
            completionTokens: output,
            cachedTokens: cacheRead > 0 ? cacheRead : nil
        )
    }

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

    static func translateResponse(_ upstream: [String: Any], emitter: ChunkEmitter) -> [String: Any] {
        var text = ""
        var reasoning = ""
        var toolCalls: [[String: Any]] = []
        for block in upstream["content"] as? [[String: Any]] ?? [] {
            switch block["type"] as? String {
            case "text":
                text += block["text"] as? String ?? ""
            case "thinking":
                reasoning += block["thinking"] as? String ?? ""
            case "tool_use":
                let input = block["input"] as? [String: Any] ?? [:]
                let arguments = String(
                    data: (try? JSONSerialization.data(withJSONObject: input)) ?? Data("{}".utf8),
                    encoding: .utf8
                ) ?? "{}"
                toolCalls.append([
                    "id": block["id"] as? String ?? "call_\(UUID().uuidString.prefix(12))",
                    "type": "function",
                    "function": ["name": block["name"] as? String ?? "", "arguments": arguments],
                ])
            default:
                break
            }
        }

        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: upstream["stop_reason"] as? String),
            usage: normalizedUsage(upstream["usage"] as? [String: Any] ?? [:])
        )
    }

    // MARK: - Streaming state machine (Anthropic events → OpenAI chunks)

    /// Feed each upstream SSE event in; write the produced OpenAI chunks out.
    struct StreamMachine {
        let emitter: ChunkEmitter
        let includeUsage: Bool

        private var toolIndex = -1
        private var currentBlockIsTool = false
        /// Accumulated assistant text for the request-log inspector (capped).
        private(set) var textPreview = ""
        private(set) var promptTokens = 0
        private(set) var cachedTokens = 0
        private(set) var completionTokens = 0
        private(set) var finishReasonSent: String?
        private(set) var upstreamError: (status: Int, message: String)?

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

        /// Returns the OpenAI SSE `data:` payloads to emit for one event,
        /// and whether the stream is finished.
        mutating func consume(_ event: SSEEvent) -> (payloads: [Data], done: Bool) {
            guard let json = (try? JSONSerialization.jsonObject(with: Data(event.data.utf8))) as? [String: Any] else {
                return ([], false)
            }
            switch event.event ?? json["type"] as? String ?? "" {
            case "message_start":
                if let usage = (json["message"] as? [String: Any])?["usage"] as? [String: Any] {
                    promptTokens = (usage["input_tokens"] as? Int ?? 0)
                        + (usage["cache_read_input_tokens"] as? Int ?? 0)
                        + (usage["cache_creation_input_tokens"] as? Int ?? 0)
                    cachedTokens = usage["cache_read_input_tokens"] as? Int ?? 0
                }
                return ([emitter.roleChunk()], false)

            case "content_block_start":
                guard let block = json["content_block"] as? [String: Any] else { return ([], false) }
                if block["type"] as? String == "tool_use" {
                    toolIndex += 1
                    currentBlockIsTool = true
                    return ([emitter.toolCallStartChunk(
                        toolIndex: toolIndex,
                        callID: block["id"] as? String ?? "call_\(UUID().uuidString.prefix(12))",
                        name: block["name"] as? String ?? ""
                    )], false)
                }
                currentBlockIsTool = false
                return ([], false)

            case "content_block_delta":
                guard let delta = json["delta"] as? [String: Any] else { return ([], false) }
                switch delta["type"] as? String {
                case "text_delta":
                    let text = delta["text"] as? String ?? ""
                    if textPreview.count < 20_000 { textPreview += text }
                    return (text.isEmpty ? [] : [emitter.contentChunk(text)], false)
                case "input_json_delta":
                    let fragment = delta["partial_json"] as? String ?? ""
                    guard !fragment.isEmpty, currentBlockIsTool else { return ([], false) }
                    return ([emitter.toolCallArgumentsChunk(toolIndex: toolIndex, fragment: fragment)], false)
                case "thinking_delta":
                    let thinking = delta["thinking"] as? String ?? ""
                    return (thinking.isEmpty ? [] : [emitter.reasoningChunk(thinking)], false)
                default: // signature_delta and friends: log-only material
                    return ([], false)
                }

            case "content_block_stop":
                currentBlockIsTool = false
                return ([], false)

            case "message_delta":
                if let usage = json["usage"] as? [String: Any],
                   let output = usage["output_tokens"] as? Int {
                    completionTokens = output // cumulative
                }
                if let delta = json["delta"] as? [String: Any],
                   let stopReason = delta["stop_reason"] as? String {
                    let reason = AnthropicTranslator.finishReason(from: stopReason)
                    finishReasonSent = reason
                    return ([emitter.finishChunk(reason: reason)], false)
                }
                return ([], false)

            case "message_stop":
                var payloads: [Data] = []
                if includeUsage {
                    payloads.append(emitter.usageChunk(UsageBuilder.build(
                        promptTokens: promptTokens,
                        completionTokens: completionTokens,
                        cachedTokens: cachedTokens > 0 ? cachedTokens : nil
                    )))
                }
                return (payloads, true)

            case "error":
                let detail = (json["error"] as? [String: Any])?["message"] as? String ?? "upstream stream error"
                upstreamError = (502, detail)
                let frame = OpenAIError(error: .init(
                    message: detail, type: "api_error", param: nil,
                    code: (json["error"] as? [String: Any])?["type"] as? String
                ))
                return ([(try? JSONEncoder().encode(frame)) ?? Data()], true)

            default: // ping etc.
                return ([], false)
            }
        }
    }
}