//
//  AppleFoundationProvider.swift
//  oAI
//
//  Apple Foundation Models provider (on-device Apple Intelligence)
//
// SPDX-License-Identifier: PolyForm-Noncommercial-1.0.0
// Copyright (C) 2026 Rune Olsen
//
// This file is part of oAI.
//
// oAI 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 oAI 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://oai.pm>.


import Foundation
import FoundationModels
import os

final class AppleFoundationProvider: AIProvider {
    let name = "Apple Intelligence"

    let capabilities = ProviderCapabilities(
        supportsStreaming: true,
        supportsVision: false,
        supportsTools: false,
        supportsOnlineSearch: false,
        maxContextLength: 4096
    )

    // MARK: - Models

    func listModels() async throws -> [ModelInfo] {
        [
            ModelInfo(
                id: "apple-on-device",
                name: "Apple On-Device (Beta)",
                description: "On-device Apple Intelligence model. Private, free, and works offline. 4K context window. Apple's Foundation Models framework is still in active beta (currently macOS 27 beta) — expect occasional generation errors and rough edges.",
                contextLength: 4096,
                pricing: ModelInfo.Pricing(prompt: 0, completion: 0),
                capabilities: ModelInfo.ModelCapabilities(
                    vision: false,
                    tools: false,
                    online: false
                )
            )
        ]
    }

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

    func getCredits() async throws -> Credits? { nil }

    // MARK: - Streaming chat

    func streamChat(request: ChatRequest) -> AsyncThrowingStream<StreamChunk, Error> {
        AsyncThrowingStream { continuation in
            Task {
                do {
                    let session = try self.makeSession(for: request)
                    let prompt = self.lastUserMessage(from: request)

                    // streamResponse(to: String) → ResponseStream<String>
                    // Each snapshot.content is the full accumulated text so far (snapshot model).
                    // We compute deltas by comparing each snapshot to the previous.
                    let stream = session.streamResponse(to: prompt)
                    var lastContent = ""

                    for try await snapshot in stream {
                        let current = snapshot.content
                        if current.count > lastContent.count {
                            let delta = String(current.dropFirst(lastContent.count))
                            continuation.yield(StreamChunk(
                                id: UUID().uuidString,
                                model: request.model,
                                delta: StreamChunk.Delta(content: delta, role: "assistant"),
                                finishReason: nil,
                                usage: nil
                            ))
                            lastContent = current
                        }
                    }

                    continuation.yield(StreamChunk(
                        id: UUID().uuidString,
                        model: request.model,
                        delta: StreamChunk.Delta(content: nil, role: nil),
                        finishReason: "stop",
                        usage: nil
                    ))
                    continuation.finish()

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

    // MARK: - Non-streaming chat

    func chat(request: ChatRequest) async throws -> ChatResponse {
        let session = try makeSession(for: request)
        let prompt = lastUserMessage(from: request)
        do {
            let response: LanguageModelSession.Response<String> = try await session.respond(to: prompt)
            return ChatResponse(
                id: UUID().uuidString,
                model: request.model,
                content: response.content,
                role: "assistant",
                finishReason: "stop",
                usage: nil,
                created: Date()
            )
        } catch {
            throw mapProviderError(error)
        }
    }

    // MARK: - Tool messages (not supported in Phase 1)

    func chatWithToolMessages(model: String, messages: [[String: Any]], tools: [Tool]?, maxTokens: Int?, temperature: Double?) async throws -> ChatResponse {
        throw ProviderError.unknown("Tool calling requires Apple Foundation Models Phase 3.")
    }

    // MARK: - Session construction

    private func makeSession(for request: ChatRequest) throws -> LanguageModelSession {
        guard case .available = SystemLanguageModel.default.availability else {
            throw availabilityError()
        }

        // Build instructions: system prompt + prior conversation turns as formatted text.
        // Foundation Models sessions don't accept a message array — we inject history inline.
        var instructions = request.systemPrompt ?? ""
        let priorMessages = request.messages.dropLast().filter { $0.role != .system }

        if !priorMessages.isEmpty {
            let history = priorMessages
                .map { m -> String in
                    let label = m.role == .user ? "User" : "Assistant"
                    return "\(label): \(m.content)"
                }
                .joined(separator: "\n")
            instructions += "\n\nConversation so far:\n\(history)\n\nContinue from here."
        }

        return instructions.isEmpty
            ? LanguageModelSession()
            : LanguageModelSession(instructions: instructions)
    }

    private func lastUserMessage(from request: ChatRequest) -> String {
        request.messages.last(where: { $0.role == .user })?.content ?? ""
    }

    // MARK: - Error mapping

    private func availabilityError() -> Error {
        switch SystemLanguageModel.default.availability {
        case .unavailable(.deviceNotEligible):
            return ProviderError.unknown("This Mac doesn't support Apple Intelligence. Apple Silicon is required.")
        case .unavailable(.appleIntelligenceNotEnabled):
            return ProviderError.unknown("Apple Intelligence is not enabled. Open System Settings → Apple Intelligence to turn it on.")
        case .unavailable(.modelNotReady):
            return ProviderError.unknown("Apple Intelligence model is still downloading. Please wait and try again.")
        default:
            return ProviderError.unknown("Apple Intelligence is not available on this device.")
        }
    }

    /// Dispatches to whichever generation-error type the running OS actually throws.
    /// `LanguageModelSession.GenerationError` was deprecated in macOS 27.0 in favor of the new
    /// top-level `LanguageModelError` — on a macOS 27+ runtime, generation failures come through
    /// as `LanguageModelError`, not `GenerationError`, even though the app's deployment target
    /// (26.2) still needs to handle macOS 26.x runtimes throwing the older type.
    private func mapProviderError(_ error: Error) -> Error {
        if #available(macOS 27.0, *), let lmError = error as? LanguageModelError {
            return mapLanguageModelError(lmError)
        }
        if let genError = error as? LanguageModelSession.GenerationError {
            return mapGenerationError(genError)
        }
        return error
    }

    @available(macOS 27.0, *)
    private func mapLanguageModelError(_ error: LanguageModelError) -> Error {
        switch error {
        case .contextSizeExceeded(let detail):
            return ProviderError.unknown("Apple Intelligence context limit exceeded (\(detail.tokenCount)/\(detail.contextSize) tokens). Start a new chat or enable Progressive Summarization in Settings → Advanced.")
        case .rateLimited:
            return ProviderError.rateLimitExceeded
        case .guardrailViolation:
            return ProviderError.unknown("Apple Intelligence declined to respond to this message.")
        case .timeout:
            return ProviderError.timeout
        default:
            return error
        }
    }

    /// macOS 26.x fallback — `GenerationError` is deprecated on macOS 27+ but still the type
    /// actually thrown on 26.x runtimes, which the 26.2 deployment target must keep supporting.
    private func mapGenerationError(_ error: LanguageModelSession.GenerationError) -> Error {
        switch error {
        case .exceededContextWindowSize:
            return ProviderError.unknown("Apple Intelligence context limit exceeded (4,096 tokens). Start a new chat or enable Progressive Summarization in Settings → Advanced.")
        case .rateLimited:
            return ProviderError.rateLimitExceeded
        case .guardrailViolation:
            return ProviderError.unknown("Apple Intelligence declined to respond to this message.")
        default:
            return error
        }
    }
}