//
//  OpenRouterProvider.swift
//  Confab
//
//  OpenRouter AI provider implementation with SSE streaming
//
// SPDX-License-Identifier: PolyForm-Noncommercial-1.0.0
// Copyright (C) 2026 Rune Olsen
//
// This file is part of Confab.
//
// Confab is licensed under the PolyForm Noncommercial License 1.0.0.
// You may use, study, modify, and share it for any noncommercial
// purpose. Commercial use — including selling Confab or any part of
// it, standalone or bundled into another product or service —
// requires a separate commercial license from the copyright holder.
//
// See the LICENSE file or
// <https://polyformproject.org/licenses/noncommercial/1.0.0> for
// the full license text. For commercial licensing, contact Rune
// Olsen via <https://confab.no>.


import Foundation
import os

class OpenRouterProvider: AIProvider {
    let name = "OpenRouter"
    let capabilities = ProviderCapabilities(
        supportsStreaming: true,
        supportsVision: true,
        supportsTools: true,
        supportsOnlineSearch: true,
        maxContextLength: nil
    )
    
    private let apiKey: String
    private let baseURL = "https://openrouter.ai/api/v1"
    private let session: URLSession
    
    init(apiKey: String) {
        self.apiKey = apiKey
        
        let config = URLSessionConfiguration.default
        config.timeoutIntervalForRequest = 60
        config.timeoutIntervalForResource = 300
        self.session = URLSession(configuration: config)
    }
    
    // MARK: - List Models
    
    func listModels() async throws -> [ModelInfo] {
        Log.api.info("Fetching model list from OpenRouter")

        // Fetch chat models and image models in parallel
        async let chatData = fetchRaw(path: "/models")
        async let imageData = fetchRaw(path: "/images/models")
        let (chatRaw, imageRaw) = try await (chatData, imageData)

        let modelsResponse = try JSONDecoder().decode(OpenRouterModelsResponse.self, from: chatRaw)
        Log.api.info("OpenRouter loaded \(modelsResponse.data.count) chat models")

        var models = modelsResponse.data.map { modelData in
            let promptPrice = Double(modelData.pricing.prompt) ?? 0.0
            let completionPrice = Double(modelData.pricing.completion) ?? 0.0
            
            var info = ModelInfo(
                id: modelData.id,
                name: modelData.name,
                description: modelData.description,
                contextLength: modelData.contextLength,
                pricing: ModelInfo.Pricing(
                    prompt: promptPrice * 1_000_000, // Convert to per 1M tokens
                    completion: completionPrice * 1_000_000
                ),
                capabilities: ModelInfo.ModelCapabilities(
                    vision: {
                        let mod = modelData.architecture?.modality ?? ""
                        return mod == "multimodal" || mod.hasPrefix("text+image")
                    }(),
                    tools: modelData.supportedParameters?.contains("tools") ?? false,
                    online: {
                        // OpenRouter supports :online suffix for all text models
                        let mod = modelData.architecture?.modality ?? ""
                        if let arrow = mod.range(of: "->") {
                            return !mod[arrow.upperBound...].contains("image")
                        }
                        return true
                    }(),
                    imageGeneration: {
                        if let mod = modelData.architecture?.modality,
                           let arrow = mod.range(of: "->") {
                            let output = mod[arrow.upperBound...]
                            return output.contains("image")
                        }
                        return false
                    }(),
                    thinking: modelData.supportedParameters?.contains("reasoning") ?? false
                ),
                architecture: modelData.architecture.map { arch in
                    ModelInfo.Architecture(
                        tokenizer: arch.tokenizer,
                        instructType: arch.instructType,
                        modality: arch.modality
                    )
                },
                topProvider: modelData.id.components(separatedBy: "/").first
            )
            info.releaseDate = modelData.created.map { Date(timeIntervalSince1970: TimeInterval($0)) }
            info.categories = ModelCategory.infer(
                name: modelData.name,
                id: modelData.id,
                description: modelData.description
            )
            return info
        }

        // Merge dedicated image models (these don't appear in /models)
        if let imageModelsResponse = try? JSONDecoder().decode(OpenRouterImageModelsResponse.self, from: imageRaw) {
            Log.api.info("OpenRouter loaded \(imageModelsResponse.data.count) image models")
            let existingIds = Set(models.map { $0.id })
            let imageModels = imageModelsResponse.data.compactMap { m -> ModelInfo? in
                guard !existingIds.contains(m.id) else { return nil }
                let acceptsImageInput = m.architecture?.inputModalities?.contains("image") ?? false
                var info = ModelInfo(
                    id: m.id,
                    name: m.name,
                    description: m.description,
                    contextLength: 0,
                    pricing: ModelInfo.Pricing(prompt: 0, completion: 0),
                    capabilities: ModelInfo.ModelCapabilities(
                        vision: acceptsImageInput,
                        tools: false,
                        online: false,
                        imageGeneration: true,
                        thinking: false,
                        usesImagesAPI: true
                    ),
                    topProvider: m.id.components(separatedBy: "/").first
                )
                info.categories = ModelCategory.infer(name: m.name, id: m.id, description: m.description)
                return info
            }
            models.append(contentsOf: imageModels)
        }

        return models
    }

    private func fetchRaw(path: String) async throws -> Data {
        let url = URL(string: "\(baseURL)\(path)")!
        var request = URLRequest(url: url)
        request.addValue("Bearer \(apiKey)", forHTTPHeaderField: "Authorization")
        request.addValue("application/json", forHTTPHeaderField: "Content-Type")
        let (data, response) = try await session.data(for: request)
        guard let httpResponse = response as? HTTPURLResponse else { throw ProviderError.invalidResponse }
        guard httpResponse.statusCode == 200 else {
            if let err = try? JSONDecoder().decode(OpenRouterErrorResponse.self, from: data) {
                throw ProviderError.unknown(err.error.message)
            }
            throw ProviderError.unknown("HTTP \(httpResponse.statusCode)")
        }
        return data
    }

    // MARK: - Images API

    func generateImage(model: String, prompt: String) async throws -> ChatResponse {
        Log.api.info("OpenRouter images API: model=\(model)")
        let url = URL(string: "\(baseURL)/images")!
        var urlRequest = URLRequest(url: url)
        urlRequest.httpMethod = "POST"
        urlRequest.addValue("Bearer \(apiKey)", forHTTPHeaderField: "Authorization")
        urlRequest.addValue("application/json", forHTTPHeaderField: "Content-Type")
        urlRequest.addValue("https://github.com/yourusername/Confab", forHTTPHeaderField: "HTTP-Referer")
        urlRequest.addValue("Confab-Swift", forHTTPHeaderField: "X-Title")
        urlRequest.httpBody = try JSONSerialization.data(withJSONObject: ["model": model, "prompt": prompt])

        let (data, response) = try await session.data(for: urlRequest)
        guard let httpResponse = response as? HTTPURLResponse else { throw ProviderError.invalidResponse }

        if httpResponse.statusCode != 200 {
            if let err = try? JSONDecoder().decode(OpenRouterErrorResponse.self, from: data) {
                Log.api.error("OpenRouter images HTTP \(httpResponse.statusCode): \(err.error.message)")
                throw ProviderError.unknown(err.error.message)
            }
            Log.api.error("OpenRouter images HTTP \(httpResponse.statusCode)")
            throw ProviderError.unknown("HTTP \(httpResponse.statusCode)")
        }

        if let rawStr = String(data: data, encoding: .utf8) {
            Log.api.debug("Images API raw response (first 200 chars): \(rawStr.prefix(200))")
        }

        let imageResponse = try JSONDecoder().decode(OpenRouterImageGenerationResponse.self, from: data)
        let images: [Data] = imageResponse.data.compactMap { item in
            Data(base64Encoded: item.b64Json)
        }

        let usage: ChatResponse.Usage? = imageResponse.usage.map { u in
            ChatResponse.Usage(
                promptTokens: u.promptTokens,
                completionTokens: u.completionTokens,
                totalTokens: u.totalTokens,
                rawCostUSD: u.cost
            )
        }

        return ChatResponse(
            id: UUID().uuidString,
            model: model,
            content: "",
            role: "assistant",
            finishReason: "stop",
            usage: usage,
            created: Date(),
            generatedImages: images.isEmpty ? nil : images
        )
    }

    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("OpenRouter chat request: model=\(request.model), messages=\(request.messages.count)")
        let apiRequest = try buildAPIRequest(from: request)
        let url = URL(string: "\(baseURL)/chat/completions")!
        
        var urlRequest = URLRequest(url: url)
        urlRequest.httpMethod = "POST"
        urlRequest.addValue("Bearer \(apiKey)", forHTTPHeaderField: "Authorization")
        urlRequest.addValue("application/json", forHTTPHeaderField: "Content-Type")
        urlRequest.addValue("https://github.com/yourusername/Confab", forHTTPHeaderField: "HTTP-Referer")
        urlRequest.addValue("Confab-Swift", forHTTPHeaderField: "X-Title")
        urlRequest.httpBody = try JSONEncoder().encode(apiRequest)
        
        let (data, response) = try await session.data(for: urlRequest)
        
        guard let httpResponse = response as? HTTPURLResponse else {
            throw ProviderError.invalidResponse
        }
        
        guard httpResponse.statusCode == 200 else {
            if let errorResponse = try? JSONDecoder().decode(OpenRouterErrorResponse.self, from: data) {
                Log.api.error("OpenRouter chat HTTP \(httpResponse.statusCode): \(errorResponse.error.message)")
                throw ProviderError.unknown(errorResponse.error.message)
            }
            Log.api.error("OpenRouter chat HTTP \(httpResponse.statusCode)")
            throw ProviderError.unknown("HTTP \(httpResponse.statusCode)")
        }

        // Debug: log raw response for image gen models
        if request.imageGeneration, let rawStr = String(data: data, encoding: .utf8) {
            Log.api.debug("Image gen raw response (first 3000 chars): \(rawStr.prefix(3000))")
        }

        let apiResponse = try JSONDecoder().decode(OpenRouterChatResponse.self, from: data)
        let chatResponse = try convertToChatResponse(apiResponse)
        if request.imageGeneration {
            Log.api.debug("Image gen decoded: content='\(chatResponse.content)', generatedImages=\(chatResponse.generatedImages?.count ?? 0)")
        }
        return chatResponse
    }

    // MARK: - Chat with raw tool messages

    /// Chat completion that accepts pre-encoded messages (for the tool call loop where
    /// message shapes vary: user, assistant+tool_calls, tool results).
    func chatWithToolMessages(model: String, messages: [[String: Any]], tools: [Tool]?, maxTokens: Int?, temperature: Double?) async throws -> ChatResponse {
        let url = URL(string: "\(baseURL)/chat/completions")!

        var body: [String: Any] = [
            "model": model,
            "messages": messages,
            "stream": false
        ]
        if let tools = tools {
            let toolsData = try JSONEncoder().encode(tools)
            body["tools"] = try JSONSerialization.jsonObject(with: toolsData)
            body["tool_choice"] = "auto"
        }
        if let maxTokens = maxTokens { body["max_tokens"] = maxTokens }
        if let temperature = temperature { body["temperature"] = temperature }
        // Anthropic models require an explicit cache_control opt-in on OpenRouter;
        // other providers cache automatically.
        if model.hasPrefix("anthropic/") {
            body["cache_control"] = ["type": "ephemeral"]
        }

        var urlRequest = URLRequest(url: url)
        urlRequest.httpMethod = "POST"
        urlRequest.addValue("Bearer \(apiKey)", forHTTPHeaderField: "Authorization")
        urlRequest.addValue("application/json", forHTTPHeaderField: "Content-Type")
        urlRequest.addValue("https://github.com/yourusername/Confab", forHTTPHeaderField: "HTTP-Referer")
        urlRequest.addValue("Confab-Swift", forHTTPHeaderField: "X-Title")
        urlRequest.httpBody = try JSONSerialization.data(withJSONObject: body)

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

        guard let httpResponse = response as? HTTPURLResponse else {
            throw ProviderError.invalidResponse
        }

        guard httpResponse.statusCode == 200 else {
            if let errorResponse = try? JSONDecoder().decode(OpenRouterErrorResponse.self, from: data) {
                Log.api.error("OpenRouter tool chat HTTP \(httpResponse.statusCode): \(errorResponse.error.message)")
                throw ProviderError.unknown(errorResponse.error.message)
            }
            Log.api.error("OpenRouter tool chat HTTP \(httpResponse.statusCode)")
            throw ProviderError.unknown("HTTP \(httpResponse.statusCode)")
        }

        let apiResponse = try JSONDecoder().decode(OpenRouterChatResponse.self, from: data)
        return try convertToChatResponse(apiResponse)
    }

    // MARK: - Streaming Chat
    
    func streamChat(request: ChatRequest) -> AsyncThrowingStream<StreamChunk, Error> {
        Log.api.info("OpenRouter stream request: model=\(request.model), messages=\(request.messages.count)")
        return AsyncThrowingStream { continuation in
            Task {
                do {
                    var apiRequest = try buildAPIRequest(from: request)
                    apiRequest.stream = true
                    
                    let url = URL(string: "\(baseURL)/chat/completions")!
                    var urlRequest = URLRequest(url: url)
                    urlRequest.httpMethod = "POST"
                    urlRequest.addValue("Bearer \(apiKey)", forHTTPHeaderField: "Authorization")
                    urlRequest.addValue("application/json", forHTTPHeaderField: "Content-Type")
                    urlRequest.addValue("text/event-stream", forHTTPHeaderField: "Accept")
                    urlRequest.addValue("https://github.com/yourusername/Confab", forHTTPHeaderField: "HTTP-Referer")
                    urlRequest.addValue("Confab-Swift", forHTTPHeaderField: "X-Title")
                    urlRequest.httpBody = try JSONEncoder().encode(apiRequest)
                    
                    let (bytes, response) = try await session.bytes(for: urlRequest)
                    
                    guard let httpResponse = response as? HTTPURLResponse else {
                        Log.api.error("OpenRouter stream: invalid response (not HTTP)")
                        continuation.finish(throwing: ProviderError.invalidResponse)
                        return
                    }

                    guard httpResponse.statusCode == 200 else {
                        Log.api.error("OpenRouter stream HTTP \(httpResponse.statusCode)")
                        continuation.finish(throwing: ProviderError.unknown("HTTP \(httpResponse.statusCode)"))
                        return
                    }

                    var buffer = ""
                    
                    for try await line in bytes.lines {
                        if line.hasPrefix("data: ") {
                            let jsonString = String(line.dropFirst(6))
                            
                            if jsonString == "[DONE]" {
                                continuation.finish()
                                return
                            }
                            
                            buffer += jsonString
                            
                            if let jsonData = buffer.data(using: .utf8) {
                                do {
                                    let chunk = try JSONDecoder().decode(OpenRouterStreamChunk.self, from: jsonData)
                                    let streamChunk = try convertToStreamChunk(chunk)
                                    continuation.yield(streamChunk)
                                    buffer = ""
                                } catch {
                                    // Partial JSON, keep buffering
                                    continue
                                }
                            }
                        }
                    }
                    
                    continuation.finish()
                } catch {
                    continuation.finish(throwing: error)
                }
            }
        }
    }
    
    // MARK: - Credits
    
    func getCredits() async throws -> Credits? {
        Log.api.info("Fetching OpenRouter credits")
        let url = URL(string: "\(baseURL)/credits")!
        var request = URLRequest(url: url)
        request.addValue("Bearer \(apiKey)", forHTTPHeaderField: "Authorization")

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

        guard let httpResponse = response as? HTTPURLResponse,
              httpResponse.statusCode == 200 else {
            return nil
        }

        let creditsResponse = try JSONDecoder().decode(OpenRouterCreditsResponse.self, from: data)
        let totalCredits = creditsResponse.data.totalCredits ?? 0
        let totalUsage = creditsResponse.data.totalUsage ?? 0
        let remaining = totalCredits - totalUsage

        return Credits(
            balance: remaining,
            currency: "USD",
            usage: totalUsage,
            limit: totalCredits
        )
    }
    
    // MARK: - Helper Methods
    
    func buildAPIRequest(from request: ChatRequest) throws -> OpenRouterChatRequest {
        let apiMessages = request.messages.map { message -> OpenRouterChatRequest.APIMessage in

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

            let content: OpenRouterChatRequest.APIMessage.MessageContent

            if hasAttachments {
                // Use array format for messages with attachments
                var contentArray: [OpenRouterChatRequest.APIMessage.ContentItem] = []

                // Add main text content
                contentArray.append(.text(message.content))

                // Add attachments
                if let attachments = message.attachments {
                    for attachment in attachments {
                        guard let data = attachment.data else { continue }

                        switch attachment.type {
                        case .image, .pdf:
                            // Send as base64 data URL with correct MIME type
                            let base64String = data.base64EncodedString()
                            let dataURL = "data:\(attachment.mimeType);base64,\(base64String)"
                            let imageContent = OpenRouterChatRequest.APIMessage.ContentItem.ImageContent(
                                type: "image_url",
                                imageUrl: .init(url: dataURL)
                            )
                            contentArray.append(.image(imageContent))

                        case .text:
                            // Inline text file content
                            let filename = (attachment.path as NSString).lastPathComponent
                            let textContent = String(data: data, encoding: .utf8) ?? ""
                            contentArray.append(.text("File: \(filename)\n\n\(textContent)"))
                        }
                    }
                }

                content = .array(contentArray)
            } else {
                // Use simple string format for text-only messages
                content = .string(message.content)
            }

            return OpenRouterChatRequest.APIMessage(
                role: message.role.rawValue,
                content: content
            )
        }
        
        // Append :online suffix for web search when online mode is enabled
        let effectiveModel: String
        if request.onlineMode && !request.imageGeneration && !request.model.hasSuffix(":online") {
            effectiveModel = request.model + ":online"
        } else {
            effectiveModel = request.model
        }

        let reasoningConfig: ReasoningAPIConfig? = request.reasoning.map {
            ReasoningAPIConfig(effort: $0.effort, exclude: $0.exclude ? true : nil)
        }

        // Anthropic models require an explicit cache_control opt-in on OpenRouter;
        // other providers (OpenAI, DeepSeek, Gemini, Grok, etc.) cache automatically.
        let cacheControl: OpenRouterChatRequest.CacheControl? = effectiveModel.hasPrefix("anthropic/")
            ? .init(type: "ephemeral")
            : nil

        return OpenRouterChatRequest(
            model: effectiveModel,
            messages: apiMessages,
            stream: request.stream,
            maxTokens: request.maxTokens,
            temperature: request.temperature,
            topP: request.topP,
            tools: request.tools,
            toolChoice: request.tools != nil ? "auto" : nil,
            modalities: request.imageGeneration ? ["text", "image"] : nil,
            reasoning: reasoningConfig,
            cacheControl: cacheControl
        )
    }
    
    func convertToChatResponse(_ apiResponse: OpenRouterChatResponse) throws -> ChatResponse {
        guard let choice = apiResponse.choices.first else {
            throw ProviderError.invalidResponse
        }

        let toolCalls: [ToolCallInfo]? = choice.message.toolCalls?.map { tc in
            ToolCallInfo(id: tc.id, type: tc.type, functionName: tc.function.name, arguments: tc.function.arguments)
        }

        let topLevelImages = choice.message.images.flatMap { decodeImageOutputs($0) } ?? []
        let blockImages = decodeImageOutputs(choice.message.contentBlockImages) ?? []
        let allImages = topLevelImages + blockImages
        let images: [Data]? = allImages.isEmpty ? nil : allImages

        if let details = apiResponse.usage?.promptTokensDetails,
           details.cachedTokens != nil || details.cacheWriteTokens != nil {
            Log.api.info("OpenRouter cache usage: model=\(apiResponse.model), created=\(details.cacheWriteTokens ?? 0), read=\(details.cachedTokens ?? 0)")
        }

        if choice.finishReason == "length" {
            Log.api.warning("OpenRouter response truncated: model=\(apiResponse.model), finishReason=length, completionTokens=\(apiResponse.usage?.completionTokens ?? 0)")
        }

        return ChatResponse(
            id: apiResponse.id,
            model: apiResponse.model,
            content: choice.message.content ?? "",
            role: choice.message.role,
            finishReason: choice.finishReason,
            usage: apiResponse.usage.map { usage in
                ChatResponse.Usage(
                    promptTokens: usage.promptTokens,
                    completionTokens: usage.completionTokens,
                    totalTokens: usage.totalTokens,
                    cacheCreationInputTokens: usage.promptTokensDetails?.cacheWriteTokens,
                    cacheReadInputTokens: usage.promptTokensDetails?.cachedTokens
                )
            },
            created: Date(timeIntervalSince1970: TimeInterval(apiResponse.created)),
            toolCalls: toolCalls,
            generatedImages: images
        )
    }

    func convertToStreamChunk(_ apiChunk: OpenRouterStreamChunk) throws -> StreamChunk {
        guard let choice = apiChunk.choices.first else {
            throw ProviderError.invalidResponse
        }

        // Merge images from both sources: top-level `images` field and content-block images
        let topLevelImages = choice.delta.images.flatMap { decodeImageOutputs($0) } ?? []
        let blockImages = decodeImageOutputs(choice.delta.contentBlockImages) ?? []
        let allImages = topLevelImages + blockImages
        let images: [Data]? = allImages.isEmpty ? nil : allImages

        if let details = apiChunk.usage?.promptTokensDetails,
           details.cachedTokens != nil || details.cacheWriteTokens != nil {
            Log.api.info("OpenRouter stream cache usage: model=\(apiChunk.model), created=\(details.cacheWriteTokens ?? 0), read=\(details.cachedTokens ?? 0)")
        }

        return StreamChunk(
            id: apiChunk.id,
            model: apiChunk.model,
            delta: StreamChunk.Delta(
                content: choice.delta.content,
                role: choice.delta.role,
                images: images,
                thinking: choice.delta.reasoning
            ),
            finishReason: choice.finishReason,
            usage: apiChunk.usage.map { usage in
                ChatResponse.Usage(
                    promptTokens: usage.promptTokens,
                    completionTokens: usage.completionTokens,
                    totalTokens: usage.totalTokens,
                    cacheCreationInputTokens: usage.promptTokensDetails?.cacheWriteTokens,
                    cacheReadInputTokens: usage.promptTokensDetails?.cachedTokens
                )
            }
        )
    }

    /// Decode base64 data URL images from API response
    func decodeImageOutputs(_ outputs: [OpenRouterChatResponse.ImageOutput]) -> [Data]? {
        let decoded = outputs.compactMap { output -> Data? in
            let url = output.imageUrl.url
            // Strip "data:image/...;base64," prefix
            guard let commaIndex = url.firstIndex(of: ",") else { return nil }
            let base64String = String(url[url.index(after: commaIndex)...])
            return Data(base64Encoded: base64String)
        }
        return decoded.isEmpty ? nil : decoded
    }
}
