//
//  OpenAIProvider.swift
//  oAI
//
//  OpenAI API provider with SSE streaming and tool support
//

import Foundation
import os

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

    private let apiKey: String
    private let baseURL = "https://api.openai.com/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: - Models

    /// Known models with pricing, used as fallback and for enrichment
    private static let knownModels: [String: (name: String, desc: String?, ctx: Int, prompt: Double, completion: Double, vision: Bool)] = [
        "gpt-4o":            ("GPT-4o",            "Most capable GPT-4 model",           128_000, 2.50,  10.0,  true),
        "gpt-4o-mini":       ("GPT-4o Mini",       "Affordable and fast",                128_000, 0.15,   0.60, true),
        "gpt-4-turbo":       ("GPT-4 Turbo",       "GPT-4 Turbo with vision",            128_000, 10.0,  30.0,  true),
        "gpt-3.5-turbo":     ("GPT-3.5 Turbo",     "Fast and affordable",                 16_385, 0.50,   1.50, false),
        "o1":                ("o1",                 "Advanced reasoning model",           200_000, 15.0,  60.0,  true),
        "o1-mini":           ("o1 Mini",            "Fast reasoning model",               128_000, 3.0,   12.0,  false),
        "o3-mini":           ("o3 Mini",            "Latest fast reasoning model",        200_000, 1.10,   4.40, false),
    ]

    func listModels() async throws -> [ModelInfo] {
        Log.api.info("Fetching model list from OpenAI")
        let url = URL(string: "\(baseURL)/models")!
        var request = URLRequest(url: url)
        request.addValue("Bearer \(apiKey)", forHTTPHeaderField: "Authorization")

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

            guard let httpResponse = response as? HTTPURLResponse, httpResponse.statusCode == 200 else {
                Log.api.warning("OpenAI models endpoint failed, using fallback models")
                return fallbackModels()
            }

            guard let json = try? JSONSerialization.jsonObject(with: data) as? [String: Any],
                  let modelsArray = json["data"] as? [[String: Any]] else {
                return fallbackModels()
            }

            // Filter to chat models
            let chatModelIds = modelsArray
                .compactMap { $0["id"] as? String }
                .filter { id in
                    id.contains("gpt") || id.hasPrefix("o1") || id.hasPrefix("o3")
                }
                .sorted()

            var models: [ModelInfo] = []
            for id in chatModelIds {
                if let known = Self.knownModels[id] {
                    models.append(ModelInfo(
                        id: id,
                        name: known.name,
                        description: known.desc,
                        contextLength: known.ctx,
                        pricing: .init(prompt: known.prompt, completion: known.completion),
                        capabilities: .init(vision: known.vision, tools: true, online: false)
                    ))
                } else {
                    models.append(ModelInfo(
                        id: id,
                        name: id,
                        description: nil,
                        contextLength: 128_000,
                        pricing: .init(prompt: 0, completion: 0),
                        capabilities: .init(vision: false, tools: true, online: false)
                    ))
                }
            }

            Log.api.info("OpenAI loaded \(models.count) models")
            return models.isEmpty ? fallbackModels() : models
        } catch {
            Log.api.warning("OpenAI models fetch failed: \(error.localizedDescription), using fallback")
            return fallbackModels()
        }
    }

    private func fallbackModels() -> [ModelInfo] {
        Self.knownModels.map { id, info in
            ModelInfo(
                id: id,
                name: info.name,
                description: info.desc,
                contextLength: info.ctx,
                pricing: .init(prompt: info.prompt, completion: info.completion),
                capabilities: .init(vision: info.vision, tools: true, online: false)
            )
        }.sorted { $0.name < $1.name }
    }

    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("OpenAI chat request: model=\(request.model), messages=\(request.messages.count)")
        let urlRequest = try buildURLRequest(from: request, stream: false)
        let (data, response) = try await session.data(for: urlRequest)

        guard let httpResponse = response as? HTTPURLResponse else {
            Log.api.error("OpenAI 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("OpenAI chat HTTP \(httpResponse.statusCode): \(message)")
                throw ProviderError.unknown(message)
            }
            Log.api.error("OpenAI chat HTTP \(httpResponse.statusCode)")
            throw ProviderError.unknown("HTTP \(httpResponse.statusCode)")
        }

        // Reuse the OpenRouter response format — OpenAI is identical
        let apiResponse = try JSONDecoder().decode(OpenRouterChatResponse.self, from: data)
        return convertToChatResponse(apiResponse)
    }

    // 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("OpenAI tool chat: model=\(model), messages=\(messages.count)")
        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 }
        // o1/o3 models don't support temperature
        if let temperature = temperature, !model.hasPrefix("o1"), !model.hasPrefix("o3") {
            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.httpBody = try JSONSerialization.data(withJSONObject: body)

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

        guard let httpResponse = response as? HTTPURLResponse else {
            Log.api.error("OpenAI 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("OpenAI tool chat HTTP \(httpResponse.statusCode): \(message)")
                throw ProviderError.unknown(message)
            }
            Log.api.error("OpenAI tool chat HTTP \(httpResponse.statusCode)")
            throw ProviderError.unknown("HTTP \(httpResponse.statusCode)")
        }

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

    // MARK: - Streaming Chat

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

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

                    var buffer = ""

                    for try await line in bytes.lines {
                        guard line.hasPrefix("data: ") else { continue }
                        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)
                                guard let choice = chunk.choices.first else { continue }

                                continuation.yield(StreamChunk(
                                    id: chunk.id,
                                    model: chunk.model,
                                    delta: .init(content: choice.delta.content, role: choice.delta.role, images: nil),
                                    finishReason: choice.finishReason,
                                    usage: chunk.usage.map {
                                        ChatResponse.Usage(promptTokens: $0.promptTokens, completionTokens: $0.completionTokens, totalTokens: $0.totalTokens)
                                    }
                                ))
                                buffer = ""
                            } catch {
                                continue // Partial JSON, keep buffering
                            }
                        }
                    }

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

    // MARK: - Credits

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

    // MARK: - Helpers

    private func buildURLRequest(from request: ChatRequest, stream: Bool) throws -> URLRequest {
        let url = URL(string: "\(baseURL)/chat/completions")!

        var apiMessages: [[String: Any]] = []

        // Add system prompt if present
        if let systemPrompt = request.systemPrompt {
            apiMessages.append(["role": "system", "content": systemPrompt])
        }

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

            if hasAttachments, let attachments = msg.attachments {
                // Multi-part content (OpenAI vision format)
                var contentArray: [[String: Any]] = [
                    ["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()
                        let dataURL = "data:\(attachment.mimeType);base64,\(base64)"
                        contentArray.append([
                            "type": "image_url",
                            "image_url": ["url": dataURL]
                        ])
                    case .text:
                        let filename = (attachment.path as NSString).lastPathComponent
                        let textContent = String(data: data, encoding: .utf8) ?? ""
                        contentArray.append(["type": "text", "text": "File: \(filename)\n\n\(textContent)"])
                    }
                }
                apiMessages.append(["role": msg.role.rawValue, "content": contentArray])
            } else {
                apiMessages.append(["role": msg.role.rawValue, "content": msg.content])
            }
        }

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

        if let maxTokens = request.maxTokens { body["max_tokens"] = maxTokens }
        // o1/o3 reasoning models don't support temperature
        if let temperature = request.temperature, !request.model.hasPrefix("o1"), !request.model.hasPrefix("o3") {
            body["temperature"] = temperature
        }
        if let tools = request.tools {
            let toolsData = try JSONEncoder().encode(tools)
            body["tools"] = try JSONSerialization.jsonObject(with: toolsData)
            body["tool_choice"] = "auto"
        }
        if stream {
            // Request usage in streaming mode
            body["stream_options"] = ["include_usage": true]
        }

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

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

        return urlRequest
    }

    private func convertToChatResponse(_ apiResponse: OpenRouterChatResponse) -> ChatResponse {
        guard let choice = apiResponse.choices.first else {
            return ChatResponse(id: apiResponse.id, model: apiResponse.model, content: "", role: "assistant", finishReason: nil, usage: nil, created: Date())
        }

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

        return ChatResponse(
            id: apiResponse.id,
            model: apiResponse.model,
            content: choice.message.content ?? "",
            role: choice.message.role,
            finishReason: choice.finishReason,
            usage: apiResponse.usage.map {
                ChatResponse.Usage(promptTokens: $0.promptTokens, completionTokens: $0.completionTokens, totalTokens: $0.totalTokens)
            },
            created: Date(timeIntervalSince1970: TimeInterval(apiResponse.created)),
            toolCalls: toolCalls
        )
    }
}