//
//  OpenRouterProvider.swift
//  oAI
//
//  OpenRouter AI provider implementation with SSE streaming
//
// 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 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")
        let url = URL(string: "\(baseURL)/models")!
        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 {
            Log.api.error("OpenRouter models: invalid response (not HTTP)")
            throw ProviderError.invalidResponse
        }

        guard httpResponse.statusCode == 200 else {
            if let errorResponse = try? JSONDecoder().decode(OpenRouterErrorResponse.self, from: data) {
                Log.api.error("OpenRouter models HTTP \(httpResponse.statusCode): \(errorResponse.error.message)")
                throw ProviderError.unknown(errorResponse.error.message)
            }
            Log.api.error("OpenRouter models HTTP \(httpResponse.statusCode)")
            throw ProviderError.unknown("HTTP \(httpResponse.statusCode)")
        }
        
        let modelsResponse = try JSONDecoder().decode(OpenRouterModelsResponse.self, from: data)
        Log.api.info("OpenRouter loaded \(modelsResponse.data.count) models")
        return modelsResponse.data.map { modelData in
            let promptPrice = Double(modelData.pricing.prompt) ?? 0.0
            let completionPrice = Double(modelData.pricing.completion) ?? 0.0
            
            return 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
                    }()
                ),
                architecture: modelData.architecture.map { arch in
                    ModelInfo.Architecture(
                        tokenizer: arch.tokenizer,
                        instructType: arch.instructType,
                        modality: arch.modality
                    )
                },
                topProvider: modelData.id.components(separatedBy: "/").first
            )
        }
    }
    
    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/oAI", forHTTPHeaderField: "HTTP-Referer")
        urlRequest.addValue("oAI-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)")
        }

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

    // 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 }

        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/oAI", forHTTPHeaderField: "HTTP-Referer")
        urlRequest.addValue("oAI-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/oAI", forHTTPHeaderField: "HTTP-Referer")
                    urlRequest.addValue("oAI-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
    
    private 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
        }

        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
        )
    }
    
    private 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 images = choice.message.images.flatMap { decodeImageOutputs($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
                )
            },
            created: Date(timeIntervalSince1970: TimeInterval(apiResponse.created)),
            toolCalls: toolCalls,
            generatedImages: images
        )
    }

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

        let images = choice.delta.images.flatMap { decodeImageOutputs($0) }

        return StreamChunk(
            id: apiChunk.id,
            model: apiChunk.model,
            delta: StreamChunk.Delta(
                content: choice.delta.content,
                role: choice.delta.role,
                images: images
            ),
            finishReason: choice.finishReason,
            usage: apiChunk.usage.map { usage in
                ChatResponse.Usage(
                    promptTokens: usage.promptTokens,
                    completionTokens: usage.completionTokens,
                    totalTokens: usage.totalTokens
                )
            }
        )
    }

    /// Decode base64 data URL images from API response
    private 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
    }
}