//
//  AnthropicClient.swift
//  Zyquo Agent
//
//  Author: Simon-Pierre Boucher
//  Mail: contact@spboucher.ai
//
//  Native Anthropic Messages API client (/v1/messages) — NOT OpenAI-compatible.
//  Auth: x-api-key + anthropic-version headers. System prompt is a top-level
//  param, content is block-structured, max_tokens is mandatory, streaming uses
//  named SSE events.
//
//  Ported from Zyquo Cloud; Zyquo Agent adds native tool calling
//  (docs/PROVIDER-REUSE.md §5.1): `tools` [{name, description, input_schema}]
//  + `tool_choice` on requests; streaming `content_block_start` (tool_use id +
//  name) → `input_json_delta` (partial_json fragments, accumulated per block
//  index) → `content_block_stop`; stop_reason "tool_use" vs "end_turn"; and
//  threading of assistant tool_use blocks + user tool_result blocks from
//  history.
//

import Foundation

struct AnthropicClient: ProviderClient {
    let providerID: ProviderID = .anthropic

    private static let apiVersion = "2023-06-01"
    private static let defaultMaxTokens = 8192

    // MARK: - Wire types (requests)

    private struct WireRequest: Encodable {
        var model: String
        var maxTokens: Int
        var messages: [WireMessage]
        var system: String?
        var stream: Bool?
        var temperature: Double?
        var topP: Double?
        var thinking: Thinking?
        var tools: [WireToolDef]?
        var toolChoice: WireToolChoice?

        enum CodingKeys: String, CodingKey {
            case model, messages, system, stream, temperature, thinking, tools
            case maxTokens = "max_tokens"
            case topP = "top_p"
            case toolChoice = "tool_choice"
        }
    }

    private struct Thinking: Encodable {
        var type: String
        var budgetTokens: Int?
        enum CodingKeys: String, CodingKey {
            case type
            case budgetTokens = "budget_tokens"
        }
    }

    /// `tools: [{name, description, input_schema: <JSON Schema>}]`.
    private struct WireToolDef: Encodable {
        var name: String
        var description: String
        var inputSchema: JSONValue

        enum CodingKeys: String, CodingKey {
            case name, description
            case inputSchema = "input_schema"
        }
    }

    /// `tool_choice: {"type":"auto"|"any"|"none"|"tool","name":…}`.
    private struct WireToolChoice: Encodable {
        var type: String
        var name: String?
    }

    private struct WireMessage: Encodable {
        var role: String
        var content: [WireBlock]
    }

    private enum WireBlock: Encodable {
        case text(String)
        case image(mediaType: String, base64: String)
        /// Assistant tool call echoed back into history.
        case toolUse(id: String, name: String, input: JSONValue)
        /// Tool result returned inside a **user** turn.
        case toolResult(toolUseID: String, content: String, isError: Bool)

        func encode(to encoder: Encoder) throws {
            var container = encoder.container(keyedBy: Key.self)
            switch self {
            case .text(let s):
                try container.encode("text", forKey: .type)
                try container.encode(s, forKey: .text)
            case .image(let mediaType, let base64):
                try container.encode("image", forKey: .type)
                var source = container.nestedContainer(keyedBy: Key.self, forKey: .source)
                try source.encode("base64", forKey: .type)
                try source.encode(mediaType, forKey: .mediaType)
                try source.encode(base64, forKey: .data)
            case .toolUse(let id, let name, let input):
                try container.encode("tool_use", forKey: .type)
                try container.encode(id, forKey: .id)
                try container.encode(name, forKey: .name)
                try container.encode(input, forKey: .input)
            case .toolResult(let toolUseID, let content, let isError):
                try container.encode("tool_result", forKey: .type)
                try container.encode(toolUseID, forKey: .toolUseID)
                try container.encode(content, forKey: .content)
                if isError {
                    try container.encode(true, forKey: .isError)
                }
            }
        }

        enum Key: String, CodingKey {
            case type, text, source, data, id, name, input, content
            case mediaType = "media_type"
            case toolUseID = "tool_use_id"
            case isError = "is_error"
        }
    }

    // MARK: - Wire types (responses)

    private struct StreamEvent: Decodable {
        var type: String?
        /// Content block index (content_block_start/delta/stop events).
        var index: Int?
        var delta: Delta?
        var usage: WireUsage?
        var message: MessageStart?
        var contentBlock: ContentBlock?
        var error: WireError?

        enum CodingKeys: String, CodingKey {
            case type, index, delta, usage, message, error
            case contentBlock = "content_block"
        }

        struct Delta: Decodable {
            var type: String?
            var text: String?
            var thinking: String?
            /// input_json_delta fragments for a streaming tool_use block.
            var partialJSON: String?
            var stopReason: String?
            enum CodingKeys: String, CodingKey {
                case type, text, thinking
                case partialJSON = "partial_json"
                case stopReason = "stop_reason"
            }
        }

        /// content_block_start payload: a tool_use block announces id + name.
        struct ContentBlock: Decodable {
            var type: String?
            var id: String?
            var name: String?
        }

        struct MessageStart: Decodable {
            var usage: WireUsage?
        }

        struct WireError: Decodable {
            var message: String?
        }
    }

    private struct WireUsage: Decodable {
        var inputTokens: Int?
        var outputTokens: Int?
        enum CodingKeys: String, CodingKey {
            case inputTokens = "input_tokens"
            case outputTokens = "output_tokens"
        }
    }

    private struct WireResponse: Decodable {
        var content: [Block]?
        var usage: WireUsage?
        var stopReason: String?

        struct Block: Decodable {
            var type: String?
            var text: String?
            var thinking: String?
            /// tool_use block fields.
            var id: String?
            var name: String?
            var input: JSONValue?
        }

        enum CodingKeys: String, CodingKey {
            case content, usage
            case stopReason = "stop_reason"
        }
    }

    private struct WireModelList: Decodable {
        var data: [Entry]
        struct Entry: Decodable { var id: String }
    }

    // MARK: - Request construction

    private func urlRequest(path: String, apiKey: String, method: String = "POST") throws -> URLRequest {
        guard let base = providerID.defaultBaseURL else {
            throw ProviderError.invalidResponse(providerID, detail: "no base URL")
        }
        var request = URLRequest(url: base.appendingPathComponent(path))
        request.httpMethod = method
        request.setValue(apiKey, forHTTPHeaderField: "x-api-key")
        request.setValue(Self.apiVersion, forHTTPHeaderField: "anthropic-version")
        if method == "POST" {
            request.setValue("application/json", forHTTPHeaderField: "Content-Type")
        }
        return request
    }

    private func buildBody(_ request: ChatRequest) throws -> Data {
        var messages: [WireMessage] = []
        for message in request.messages where message.role != .system {
            messages.append(wireMessage(from: message, vision: request.model.capabilities.vision))
        }
        let params = request.parameters
        var wire = WireRequest(
            model: request.model.id,
            maxTokens: params.maxTokens ?? Self.defaultMaxTokens,
            messages: messages
        )
        if let system = request.systemPrompt, !system.isEmpty {
            wire.system = system
        }
        if request.stream { wire.stream = true }
        // Claude 4.7+ removed temperature/top_p; ParameterSupport encodes that per model.
        let support = request.model.parameterSupport
        if support.temperature { wire.temperature = params.temperature }
        if support.topP { wire.topP = params.topP }
        if support.thinkingToggle, let enabled = params.thinkingEnabled {
            wire.thinking = enabled
                ? Thinking(type: "enabled", budgetTokens: 8000)
                : Thinking(type: "disabled")
        }
        // Native tool calling.
        if !request.tools.isEmpty, request.model.capabilities.tools {
            wire.tools = request.tools.map { spec in
                WireToolDef(
                    name: spec.name,
                    description: spec.description,
                    inputSchema: JSONValue.parse(spec.parametersJSONSchema) ?? .emptyObject
                )
            }
            switch request.toolChoice {
            case .auto:
                break // Anthropic default — omit
            case .none:
                wire.toolChoice = WireToolChoice(type: "none")
            case .required:
                wire.toolChoice = WireToolChoice(type: "any")
            case .named(let name):
                wire.toolChoice = WireToolChoice(type: "tool", name: name)
            }
        }
        return try JSONEncoder().encode(wire)
    }

    private func wireMessage(from message: Message, vision: Bool) -> WireMessage {
        // Tool results travel as tool_result blocks inside a USER turn — not a
        // special role. Parallel calls put multiple results in one turn.
        if let results = message.toolResults, !results.isEmpty {
            let blocks = results.map { result in
                WireBlock.toolResult(
                    toolUseID: result.toolCallID,
                    content: result.content,
                    isError: result.isError
                )
            }
            return WireMessage(role: "user", content: blocks)
        }
        let role = message.role == .assistant ? "assistant" : "user"
        var text = message.text
        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```"
        }
        var blocks: [WireBlock] = []
        if vision, message.role == .user {
            for image in message.attachments where image.kind == .image {
                blocks.append(.image(mediaType: image.mimeType, base64: image.data.base64EncodedString()))
            }
        }
        // Assistant turns that called tools: text block (when non-empty)
        // followed by the tool_use blocks, echoed back verbatim.
        if let calls = message.toolCalls, !calls.isEmpty, message.role == .assistant {
            if !text.isEmpty {
                blocks.append(.text(text))
            }
            for call in calls {
                blocks.append(.toolUse(
                    id: call.id,
                    name: call.name,
                    input: JSONValue.parse(call.argumentsJSON) ?? .emptyObject
                ))
            }
            return WireMessage(role: role, content: blocks)
        }
        blocks.append(.text(text.isEmpty ? " " : text))
        return WireMessage(role: role, content: blocks)
    }

    // MARK: - ProviderClient

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

                    var usage = TokenUsage()
                    var stopReason: String?
                    // Streaming tool_use blocks, keyed by content block index.
                    struct PartialToolUse {
                        var id: String
                        var name: String
                        var inputJSON = ""
                    }
                    var partials: [Int: PartialToolUse] = [:]
                    let decoder = JSONDecoder()

                    for try await sse in StreamingService.sseEvents(for: urlReq, provider: providerID) {
                        guard let data = sse.data.data(using: .utf8),
                              let event = try? decoder.decode(StreamEvent.self, from: data) else {
                            continue
                        }
                        let type = sse.event ?? event.type ?? ""
                        switch type {
                        case "message_start":
                            if let u = event.message?.usage {
                                usage.inputTokens = u.inputTokens ?? 0
                            }
                        case "content_block_start":
                            if let block = event.contentBlock, block.type == "tool_use" {
                                let index = event.index ?? (partials.keys.max().map { $0 + 1 } ?? 0)
                                let id = block.id ?? "toolu_\(index)_\(UUID().uuidString.prefix(8))"
                                let name = block.name ?? ""
                                partials[index] = PartialToolUse(id: id, name: name)
                                continuation.yield(.toolCallStarted(index: index, id: id, name: name))
                            }
                        case "content_block_delta":
                            if let text = event.delta?.text, !text.isEmpty {
                                continuation.yield(.textDelta(text))
                            }
                            if let thinking = event.delta?.thinking, !thinking.isEmpty {
                                continuation.yield(.reasoningDelta(thinking))
                            }
                            if let fragment = event.delta?.partialJSON, !fragment.isEmpty,
                               let index = event.index, partials[index] != nil {
                                partials[index]?.inputJSON += fragment
                                continuation.yield(.toolCallArgumentsDelta(index: index, delta: fragment))
                            }
                        case "content_block_stop":
                            break // block complete; finalized set is emitted at stream end
                        case "message_delta":
                            if let u = event.usage {
                                usage.outputTokens = u.outputTokens ?? usage.outputTokens
                            }
                            if let reason = event.delta?.stopReason {
                                stopReason = reason
                            }
                        case "error":
                            throw ProviderError.serverError(
                                providerID, status: 200, message: event.error?.message
                            )
                        case "message_stop":
                            break
                        default:
                            break // ping, unknown future events
                        }
                    }
                    let calls = partials.sorted { $0.key < $1.key }.map { _, partial in
                        ToolCall(
                            id: partial.id,
                            name: partial.name,
                            argumentsJSON: partial.inputJSON.isEmpty ? "{}" : partial.inputJSON
                        )
                    }
                    if !calls.isEmpty {
                        continuation.yield(.toolCalls(calls))
                    }
                    continuation.yield(.usage(usage))
                    continuation.yield(.finished(
                        reason: stopReason,
                        stop: StopReason.normalize(stopReason, hasToolCalls: !calls.isEmpty)
                    ))
                    continuation.finish()
                } catch {
                    continuation.finish(throwing: error)
                }
            }
            continuation.onTermination = { _ in task.cancel() }
        }
    }

    func complete(_ request: ChatRequest, apiKey: String) async throws -> Message {
        var urlReq = try urlRequest(path: "messages", apiKey: apiKey)
        var plainRequest = request
        plainRequest.stream = false
        urlReq.httpBody = try buildBody(plainRequest)
        let data = try await StreamingService.postJSON(urlReq, provider: providerID)
        guard let response = try? JSONDecoder().decode(WireResponse.self, from: data) else {
            throw ProviderError.invalidResponse(providerID, detail: "undecodable messages response")
        }
        let blocks = response.content ?? []
        let text = blocks.compactMap { $0.type == "text" ? $0.text : nil }.joined()
        let thinking = blocks.compactMap { $0.type == "thinking" ? $0.thinking : nil }.joined()
        let calls: [ToolCall] = blocks.enumerated().compactMap { index, block in
            guard block.type == "tool_use" else { return nil }
            return ToolCall(
                id: block.id ?? "toolu_\(index)_\(UUID().uuidString.prefix(8))",
                name: block.name ?? "",
                argumentsJSON: (block.input ?? .emptyObject).jsonString
            )
        }
        var message = Message(
            role: .assistant,
            text: text,
            reasoning: thinking.isEmpty ? nil : thinking,
            toolCalls: calls.isEmpty ? nil : calls,
            modelID: request.model.id,
            provider: providerID
        )
        if let u = response.usage {
            let usage = TokenUsage(inputTokens: u.inputTokens ?? 0, outputTokens: u.outputTokens ?? 0)
            message.usage = usage
            message.estimatedCost = request.model.pricing?.cost(
                inputTokens: usage.inputTokens, outputTokens: usage.outputTokens
            )
        }
        return message
    }

    func listModelIDs(apiKey: String) async throws -> [String] {
        var urlReq = try urlRequest(path: "models", apiKey: apiKey, method: "GET")
        urlReq.url = urlReq.url.flatMap {
            URL(string: $0.absoluteString + "?limit=100")
        }
        let data = try await StreamingService.getJSON(urlReq, provider: providerID)
        guard let list = try? JSONDecoder().decode(WireModelList.self, from: data) else {
            throw ProviderError.invalidResponse(providerID, detail: "unrecognized /models response shape")
        }
        return list.data.map(\.id)
    }
}