//
//  AnthropicProvider.swift
//  oAI
//
//  Anthropic Messages API provider with SSE streaming and tool support
//
// SPDX-License-Identifier: AGPL-3.0-or-later
// Copyright (C) 2026 Rune Olsen
//
// This file is part of oAI.
//
// oAI is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as
// published by the Free Software Foundation, either version 3 of the
// License, or (at your option) any later version.
//
// oAI is distributed in the hope that it will be useful, but WITHOUT
// ANY WARRANTY; without even the implied warranty of MERCHANTABILITY
// or FITNESS FOR A PARTICULAR PURPOSE. See the GNU Affero General
// Public License for more details.
//
// You should have received a copy of the GNU Affero General Public
// License along with oAI. If not, see <https://www.gnu.org/licenses/>.


import Foundation
import os

class AnthropicProvider: AIProvider {
    let name = "Anthropic"
    let capabilities = ProviderCapabilities(
        supportsStreaming: true,
        supportsVision: true,
        supportsTools: true,
        supportsOnlineSearch: false,
        maxContextLength: nil
    )

    enum AuthMode {
        case apiKey(String)
        case oauth
    }

    private let authMode: AuthMode
    private let baseURL = "https://api.anthropic.com/v1"
    private let apiVersion = "2023-06-01"
    private let session: URLSession

    /// Create with a standard API key
    init(apiKey: String) {
        self.authMode = .apiKey(apiKey)
        let config = URLSessionConfiguration.default
        config.timeoutIntervalForRequest = 180  // 3 minutes for initial response (tool use needs thinking time)
        config.timeoutIntervalForResource = 600  // 10 minutes total
        self.session = URLSession(configuration: config)
    }

    /// Create with OAuth (Pro/Max subscription)
    init(oauth: Bool) {
        self.authMode = .oauth
        let config = URLSessionConfiguration.default
        config.timeoutIntervalForRequest = 180  // 3 minutes for initial response (tool use needs thinking time)
        config.timeoutIntervalForResource = 600  // 10 minutes total
        self.session = URLSession(configuration: config)
    }

    /// Whether this provider is using OAuth authentication
    var isOAuth: Bool {
        if case .oauth = authMode { return true }
        return false
    }

    // MARK: - Models

    /// Local metadata used to enrich API results (pricing, context length) and as offline fallback.
    /// Entries are matched by exact ID first; if no exact match is found, the enrichment step
    /// falls back to prefix matching so newly-released model variants (e.g. "claude-sonnet-4-6-20260301")
    /// still inherit the correct pricing tier.
    private static let knownModels: [ModelInfo] = [
        // Claude 4.x series
        ModelInfo(
            id: "claude-opus-4-6",
            name: "Claude Opus 4.6",
            description: "Most capable and intelligent model",
            contextLength: 200_000,
            pricing: .init(prompt: 15.0, completion: 75.0),
            capabilities: .init(vision: true, tools: true, online: true)
        ),
        ModelInfo(
            id: "claude-sonnet-4-6",
            name: "Claude Sonnet 4.6",
            description: "Best balance of speed and capability",
            contextLength: 200_000,
            pricing: .init(prompt: 3.0, completion: 15.0),
            capabilities: .init(vision: true, tools: true, online: true)
        ),
        ModelInfo(
            id: "claude-haiku-4-6",
            name: "Claude Haiku 4.6",
            description: "Fastest and most affordable",
            contextLength: 200_000,
            pricing: .init(prompt: 0.80, completion: 4.0),
            capabilities: .init(vision: true, tools: true, online: true)
        ),
        // Claude 4.5 series
        ModelInfo(
            id: "claude-opus-4-5",
            name: "Claude Opus 4.5",
            description: "Previous generation Opus",
            contextLength: 200_000,
            pricing: .init(prompt: 15.0, completion: 75.0),
            capabilities: .init(vision: true, tools: true, online: true)
        ),
        ModelInfo(
            id: "claude-opus-4-5-20251101",
            name: "Claude Opus 4.5",
            description: "Previous generation Opus",
            contextLength: 200_000,
            pricing: .init(prompt: 15.0, completion: 75.0),
            capabilities: .init(vision: true, tools: true, online: true)
        ),
        ModelInfo(
            id: "claude-sonnet-4-5",
            name: "Claude Sonnet 4.5",
            description: "Best balance of speed and capability",
            contextLength: 200_000,
            pricing: .init(prompt: 3.0, completion: 15.0),
            capabilities: .init(vision: true, tools: true, online: true)
        ),
        ModelInfo(
            id: "claude-sonnet-4-5-20250929",
            name: "Claude Sonnet 4.5",
            description: "Best balance of speed and capability",
            contextLength: 200_000,
            pricing: .init(prompt: 3.0, completion: 15.0),
            capabilities: .init(vision: true, tools: true, online: true)
        ),
        ModelInfo(
            id: "claude-haiku-4-5",
            name: "Claude Haiku 4.5",
            description: "Fastest and most affordable",
            contextLength: 200_000,
            pricing: .init(prompt: 0.80, completion: 4.0),
            capabilities: .init(vision: true, tools: true, online: true)
        ),
        ModelInfo(
            id: "claude-haiku-4-5-20251001",
            name: "Claude Haiku 4.5",
            description: "Fastest and most affordable",
            contextLength: 200_000,
            pricing: .init(prompt: 0.80, completion: 4.0),
            capabilities: .init(vision: true, tools: true, online: true)
        ),
        // Claude 3.x series
        ModelInfo(
            id: "claude-3-7-sonnet-20250219",
            name: "Claude 3.7 Sonnet",
            description: "Previous generation Sonnet",
            contextLength: 200_000,
            pricing: .init(prompt: 3.0, completion: 15.0),
            capabilities: .init(vision: true, tools: true, online: true)
        ),
        ModelInfo(
            id: "claude-3-haiku-20240307",
            name: "Claude 3 Haiku",
            description: "Previous generation Haiku",
            contextLength: 200_000,
            pricing: .init(prompt: 0.25, completion: 1.25),
            capabilities: .init(vision: true, tools: true, online: true)
        ),
    ]

    /// Pricing tiers used for fuzzy fallback matching on unknown model IDs.
    /// Keyed by model name prefix (longest match wins).
    private static let pricingFallback: [(prefix: String, prompt: Double, completion: Double)] = [
        ("claude-opus",   15.0, 75.0),
        ("claude-sonnet",  3.0, 15.0),
        ("claude-haiku",   0.80, 4.0),
    ]

    /// Fetch live model list from GET /v1/models, enriched with local pricing/context metadata.
    /// Falls back to knownModels if the request fails (no key, offline, etc.).
    func listModels() async throws -> [ModelInfo] {
        guard let url = URL(string: "\(baseURL)/models") else { return Self.knownModels }

        var request = URLRequest(url: url)
        request.httpMethod = "GET"
        do {
            try await applyAuth(to: &request)
        } catch {
            Log.api.warning("Anthropic listModels: auth failed, using fallback — \(error.localizedDescription)")
            return Self.knownModels
        }

        do {
            let (data, response) = try await session.data(for: request)
            guard let http = response as? HTTPURLResponse, http.statusCode == 200 else {
                let code = (response as? HTTPURLResponse)?.statusCode ?? 0
                Log.api.warning("Anthropic listModels: HTTP \(code), using fallback")
                return Self.knownModels
            }

            guard let json = try? JSONSerialization.jsonObject(with: data) as? [String: Any],
                  let items = json["data"] as? [[String: Any]] else {
                Log.api.warning("Anthropic listModels: unexpected JSON shape, using fallback")
                return Self.knownModels
            }

            let enrichment = Dictionary(uniqueKeysWithValues: Self.knownModels.map { ($0.id, $0) })

            let models: [ModelInfo] = items.compactMap { item in
                guard let id = item["id"] as? String,
                      id.hasPrefix("claude-") else { return nil }
                let displayName = item["display_name"] as? String ?? id
                // Exact match first
                if let known = enrichment[id] { return known }
                // Fuzzy fallback: find the longest prefix that matches
                let fallback = Self.pricingFallback
                    .filter { id.hasPrefix($0.prefix) }
                    .max(by: { $0.prefix.count < $1.prefix.count })
                let pricing = fallback.map { ModelInfo.Pricing(prompt: $0.prompt, completion: $0.completion) }
                    ?? ModelInfo.Pricing(prompt: 0, completion: 0)
                return ModelInfo(
                    id: id,
                    name: displayName,
                    description: item["description"] as? String ?? "",
                    contextLength: 200_000,
                    pricing: pricing,
                    capabilities: .init(vision: true, tools: true, online: false)
                )
            }

            Log.api.info("Anthropic listModels: fetched \(models.count) model(s) from API")
            return models.isEmpty ? Self.knownModels : models

        } catch {
            Log.api.warning("Anthropic listModels: network error (\(error.localizedDescription)), using fallback")
            return Self.knownModels
        }
    }

    func getModel(_ id: String) async throws -> ModelInfo? {
        let models = try await listModels()
        return models.first { $0.id == id }
    }

    // MARK: - Chat Completion

    func chat(request: ChatRequest) async throws -> ChatResponse {
        Log.api.info("Anthropic chat request: model=\(request.model), messages=\(request.messages.count)")
        var (urlRequest, _) = try buildURLRequest(from: request, stream: false)
        try await applyAuth(to: &urlRequest)
        let (data, response) = try await session.data(for: urlRequest)

        guard let httpResponse = response as? HTTPURLResponse else {
            Log.api.error("Anthropic chat: invalid response (not HTTP)")
            throw ProviderError.invalidResponse
        }
        guard httpResponse.statusCode == 200 else {
            if let errorObj = try? JSONSerialization.jsonObject(with: data) as? [String: Any],
               let error = errorObj["error"] as? [String: Any],
               let message = error["message"] as? String {
                Log.api.error("Anthropic chat HTTP \(httpResponse.statusCode): \(message)")
                throw ProviderError.unknown(message)
            }
            Log.api.error("Anthropic chat HTTP \(httpResponse.statusCode)")
            throw ProviderError.unknown("HTTP \(httpResponse.statusCode)")
        }

        return try parseResponse(data: data)
    }

    // MARK: - Chat with raw tool messages

    func chatWithToolMessages(model: String, messages: [[String: Any]], tools: [Tool]?, maxTokens: Int?, temperature: Double?) async throws -> ChatResponse {
        Log.api.info("Anthropic tool chat: model=\(model), messages=\(messages.count)")
        let url = messagesURL

        // Separate system message from conversation messages
        var systemText: String? = nil
        var conversationMessages: [[String: Any]] = []

        for msg in messages {
            let role = msg["role"] as? String ?? ""
            if role == "system" {
                systemText = msg["content"] as? String
            } else if role == "tool" {
                // Convert OpenAI tool result format to Anthropic tool_result format
                let toolCallId = msg["tool_call_id"] as? String ?? ""
                let content = msg["content"] as? String ?? ""
                conversationMessages.append([
                    "role": "user",
                    "content": [
                        ["type": "tool_result", "tool_use_id": toolCallId, "content": content]
                    ]
                ])
            } else if role == "assistant" {
                // Check for tool_calls — convert to Anthropic content blocks
                if let toolCalls = msg["tool_calls"] as? [[String: Any]], !toolCalls.isEmpty {
                    var contentBlocks: [[String: Any]] = []
                    if let text = msg["content"] as? String, !text.isEmpty {
                        contentBlocks.append(["type": "text", "text": text])
                    }
                    for tc in toolCalls {
                        let fn = tc["function"] as? [String: Any] ?? [:]
                        let name = fn["name"] as? String ?? ""
                        let argsStr = fn["arguments"] as? String ?? "{}"
                        let argsObj = (try? JSONSerialization.jsonObject(with: Data(argsStr.utf8))) ?? [:]
                        contentBlocks.append([
                            "type": "tool_use",
                            "id": tc["id"] as? String ?? UUID().uuidString,
                            "name": name,
                            "input": argsObj
                        ])
                    }
                    conversationMessages.append(["role": "assistant", "content": contentBlocks])
                } else {
                    conversationMessages.append(["role": "assistant", "content": msg["content"] as? String ?? ""])
                }
            } else {
                conversationMessages.append(["role": role, "content": msg["content"] as? String ?? ""])
            }
        }

        var body: [String: Any] = [
            "model": model,
            "messages": conversationMessages,
            "max_tokens": maxTokens ?? 16000,
            "stream": false
        ]
        if let systemText = systemText {
            body["system"] = systemText
        }
        if let temperature = temperature {
            body["temperature"] = temperature
        }
        if let tools = tools {
            body["tools"] = tools.map { tool -> [String: Any] in
                [
                    "name": tool.function.name,
                    "description": tool.function.description,
                    "input_schema": convertParametersToDict(tool.function.parameters)
                ]
            }
            body["tool_choice"] = ["type": "auto"]
        }

        var urlRequest = URLRequest(url: url)
        urlRequest.httpMethod = "POST"
        urlRequest.addValue("application/json", forHTTPHeaderField: "Content-Type")
        urlRequest.httpBody = try JSONSerialization.data(withJSONObject: body)
        try await applyAuth(to: &urlRequest)

        let (data, response) = try await session.data(for: urlRequest)

        guard let httpResponse = response as? HTTPURLResponse else {
            Log.api.error("Anthropic tool chat: invalid response (not HTTP)")
            throw ProviderError.invalidResponse
        }
        guard httpResponse.statusCode == 200 else {
            if let errorObj = try? JSONSerialization.jsonObject(with: data) as? [String: Any],
               let error = errorObj["error"] as? [String: Any],
               let message = error["message"] as? String {
                Log.api.error("Anthropic tool chat HTTP \(httpResponse.statusCode): \(message)")
                throw ProviderError.unknown(message)
            }
            Log.api.error("Anthropic tool chat HTTP \(httpResponse.statusCode)")
            throw ProviderError.unknown("HTTP \(httpResponse.statusCode)")
        }

        return try parseResponse(data: data)
    }

    // MARK: - Streaming Chat

    func streamChat(request: ChatRequest) -> AsyncThrowingStream<StreamChunk, Error> {
        Log.api.info("Anthropic stream request: model=\(request.model), messages=\(request.messages.count)")
        return AsyncThrowingStream { continuation in
            Task {
                do {
                    var (urlRequest, _) = try buildURLRequest(from: request, stream: true)
                    try await self.applyAuth(to: &urlRequest)
                    let (bytes, response) = try await session.bytes(for: urlRequest)

                    guard let httpResponse = response as? HTTPURLResponse else {
                        Log.api.error("Anthropic stream: invalid response (not HTTP)")
                        continuation.finish(throwing: ProviderError.invalidResponse)
                        return
                    }
                    guard httpResponse.statusCode == 200 else {
                        Log.api.error("Anthropic stream HTTP \(httpResponse.statusCode)")
                        continuation.finish(throwing: ProviderError.unknown("HTTP \(httpResponse.statusCode)"))
                        return
                    }

                    var currentId = ""
                    var currentModel = request.model
                    var inputTokens = 0

                    for try await line in bytes.lines {
                        // Anthropic SSE: "event: ..." and "data: {...}"
                        guard line.hasPrefix("data: ") else { continue }
                        let jsonString = String(line.dropFirst(6))

                        guard let jsonData = jsonString.data(using: .utf8),
                              let event = try? JSONSerialization.jsonObject(with: jsonData) as? [String: Any],
                              let eventType = event["type"] as? String else {
                            continue
                        }

                        switch eventType {
                        case "message_start":
                            if let message = event["message"] as? [String: Any] {
                                currentId = message["id"] as? String ?? ""
                                currentModel = message["model"] as? String ?? request.model
                                if let usageDict = message["usage"] as? [String: Any] {
                                    inputTokens = usageDict["input_tokens"] as? Int ?? 0
                                }
                            }

                        case "content_block_delta":
                            if let delta = event["delta"] as? [String: Any],
                               let deltaType = delta["type"] as? String,
                               deltaType == "text_delta",
                               let text = delta["text"] as? String {
                                continuation.yield(StreamChunk(
                                    id: currentId,
                                    model: currentModel,
                                    delta: .init(content: text, role: nil, images: nil),
                                    finishReason: nil,
                                    usage: nil
                                ))
                            }

                        case "message_delta":
                            let delta = event["delta"] as? [String: Any]
                            let stopReason = delta?["stop_reason"] as? String
                            var usage: ChatResponse.Usage? = nil
                            if let usageDict = event["usage"] as? [String: Any] {
                                let outputTokens = usageDict["output_tokens"] as? Int ?? 0
                                usage = ChatResponse.Usage(promptTokens: inputTokens, completionTokens: outputTokens, totalTokens: inputTokens + outputTokens)
                            }
                            continuation.yield(StreamChunk(
                                id: currentId,
                                model: currentModel,
                                delta: .init(content: nil, role: nil, images: nil),
                                finishReason: stopReason,
                                usage: usage
                            ))

                        case "message_stop":
                            continuation.finish()
                            return

                        default:
                            break
                        }
                    }

                    continuation.finish()
                } catch {
                    continuation.finish(throwing: error)
                }
            }
        }
    }

    // MARK: - Credits

    func getCredits() async throws -> Credits? {
        // Anthropic doesn't have a public credits API
        return nil
    }

    // MARK: - Auth Helpers

    /// Apply auth headers based on mode (API key or OAuth Bearer)
    private func applyAuth(to request: inout URLRequest) async throws {
        switch authMode {
        case .apiKey(let key):
            request.addValue(key, forHTTPHeaderField: "x-api-key")
            request.addValue(apiVersion, forHTTPHeaderField: "anthropic-version")

        case .oauth:
            let token = try await AnthropicOAuthService.shared.getValidAccessToken()
            request.addValue("Bearer \(token)", forHTTPHeaderField: "Authorization")
            request.addValue(apiVersion, forHTTPHeaderField: "anthropic-version")
            request.addValue("oauth-2025-04-20,interleaved-thinking-2025-05-14", forHTTPHeaderField: "anthropic-beta")
        }
    }

    /// Build the messages endpoint URL, appending ?beta=true for OAuth
    private var messagesURL: URL {
        switch authMode {
        case .apiKey:
            return URL(string: "\(baseURL)/messages")!
        case .oauth:
            return URL(string: "\(baseURL)/messages?beta=true")!
        }
    }

    // MARK: - Request Building

    private func buildURLRequest(from request: ChatRequest, stream: Bool) throws -> (URLRequest, Data) {
        let url = messagesURL

        // Separate system message
        var systemText: String? = request.systemPrompt
        var apiMessages: [[String: Any]] = []

        for msg in request.messages {
            if msg.role == .system {
                systemText = msg.content
                continue
            }

            let hasAttachments = msg.attachments?.contains(where: { $0.data != nil }) ?? false

            if hasAttachments, let attachments = msg.attachments {
                var contentBlocks: [[String: Any]] = []
                contentBlocks.append(["type": "text", "text": msg.content])

                for attachment in attachments {
                    guard let data = attachment.data else { continue }
                    switch attachment.type {
                    case .image, .pdf:
                        let base64 = data.base64EncodedString()
                        contentBlocks.append([
                            "type": "image",
                            "source": [
                                "type": "base64",
                                "media_type": attachment.mimeType,
                                "data": base64
                            ]
                        ])
                    case .text:
                        let filename = (attachment.path as NSString).lastPathComponent
                        let textContent = String(data: data, encoding: .utf8) ?? ""
                        contentBlocks.append(["type": "text", "text": "File: \(filename)\n\n\(textContent)"])
                    }
                }
                apiMessages.append(["role": msg.role.rawValue, "content": contentBlocks])
            } else {
                apiMessages.append(["role": msg.role.rawValue, "content": msg.content])
            }
        }

        var body: [String: Any] = [
            "model": request.model,
            "messages": apiMessages,
            "max_tokens": request.maxTokens ?? 16000,
            "stream": stream
        ]

        if let systemText = systemText {
            body["system"] = systemText
        }
        if let temperature = request.temperature {
            body["temperature"] = temperature
        }
        var toolsArray: [[String: Any]] = []
        if let tools = request.tools {
            toolsArray += tools.map { tool -> [String: Any] in
                [
                    "name": tool.function.name,
                    "description": tool.function.description,
                    "input_schema": convertParametersToDict(tool.function.parameters)
                ]
            }
        }
        if request.onlineMode {
            toolsArray.append([
                "type": "web_search_20250305",
                "name": "web_search",
                "max_uses": 5
            ])
        }
        if !toolsArray.isEmpty {
            body["tools"] = toolsArray
            if request.tools != nil {
                body["tool_choice"] = ["type": "auto"]
            }
        }

        let bodyData = try JSONSerialization.data(withJSONObject: body)

        var urlRequest = URLRequest(url: url)
        urlRequest.httpMethod = "POST"
        urlRequest.addValue("application/json", forHTTPHeaderField: "Content-Type")
        if stream {
            urlRequest.addValue("text/event-stream", forHTTPHeaderField: "Accept")
        }
        urlRequest.httpBody = bodyData

        // Auth is applied async in the caller (chat/streamChat)
        return (urlRequest, bodyData)
    }

    private func parseResponse(data: Data) throws -> ChatResponse {
        guard let json = try? JSONSerialization.jsonObject(with: data) as? [String: Any] else {
            throw ProviderError.invalidResponse
        }

        let id = json["id"] as? String ?? ""
        let model = json["model"] as? String ?? ""
        let contentBlocks = json["content"] as? [[String: Any]] ?? []

        var textContent = ""
        var toolCalls: [ToolCallInfo] = []

        for block in contentBlocks {
            let blockType = block["type"] as? String ?? ""
            switch blockType {
            case "text":
                textContent += block["text"] as? String ?? ""
            case "tool_use":
                let tcId = block["id"] as? String ?? UUID().uuidString
                let tcName = block["name"] as? String ?? ""
                let tcInput = block["input"] ?? [:]
                let argsData = try JSONSerialization.data(withJSONObject: tcInput)
                let argsStr = String(data: argsData, encoding: .utf8) ?? "{}"
                toolCalls.append(ToolCallInfo(id: tcId, type: "function", functionName: tcName, arguments: argsStr))
            default:
                break
            }
        }

        let usageDict = json["usage"] as? [String: Any]
        let inputTokens = usageDict?["input_tokens"] as? Int ?? 0
        let outputTokens = usageDict?["output_tokens"] as? Int ?? 0

        return ChatResponse(
            id: id,
            model: model,
            content: textContent,
            role: "assistant",
            finishReason: json["stop_reason"] as? String,
            usage: ChatResponse.Usage(
                promptTokens: inputTokens,
                completionTokens: outputTokens,
                totalTokens: inputTokens + outputTokens
            ),
            created: Date(),
            toolCalls: toolCalls.isEmpty ? nil : toolCalls
        )
    }

    private func convertParametersToDict(_ params: Tool.Function.Parameters) -> [String: Any] {
        var props: [String: Any] = [:]
        for (key, prop) in params.properties {
            var propDict: [String: Any] = [
                "type": prop.type,
                "description": prop.description
            ]
            if let enumVals = prop.enum {
                propDict["enum"] = enumVals
            }
            props[key] = propDict
        }
        var dict: [String: Any] = [
            "type": params.type,
            "properties": props
        ]
        if let required = params.required {
            dict["required"] = required
        }
        return dict
    }
}