//
//  AIService.swift
//  Zyquo Atlas
//
//  Author: Simon-Pierre Boucher
//  Mail: contact@spboucher.ai
//
//  Orchestrates browser AI actions against the ported provider clients: builds
//  the grounded, injection-safe request, streams tokens into published state,
//  and cancels the in-flight request on navigation / Stop (docs/
//  AI-BROWSER-RESEARCH.md §5 — dropping the stream cancels the URLSession task).
//  One AIService drives one AI surface (e.g. a tab's chat-with-page sidebar).
//

import Foundation
import Combine

@MainActor
final class AIService: ObservableObject {
    /// Streamed answer text for the current action.
    @Published private(set) var output: String = ""
    @Published private(set) var isStreaming: Bool = false
    @Published private(set) var errorText: String?
    /// Progress note during map-reduce ("Reading section 3/9…").
    @Published private(set) var statusNote: String?

    private var task: Task<Void, Never>?

    /// Cancels any in-flight request (call on navigation, Stop, tab close).
    func cancel() {
        task?.cancel()
        task = nil
        if isStreaming { isStreaming = false }
    }

    deinit { task?.cancel() }

    /// Runs an action over a page context, streaming the answer into `output`.
    func run(_ action: AIAction,
             on context: PageContext,
             model: AIModel,
             apiKey: String,
             extra: String? = nil) {
        cancel()
        reset()
        isStreaming = true

        task = Task { [weak self] in
            guard let self else { return }
            do {
                if action == .summarize {
                    try await self.summarize(context: context, model: model, apiKey: apiKey)
                } else {
                    let messages = [
                        Message(role: .user, text: action.userPrompt(for: context, extra: extra))
                    ]
                    try await self.stream(system: action.systemGuidance,
                                          messages: messages, model: model, apiKey: apiKey)
                }
                await MainActor.run { self.isStreaming = false }
            } catch is CancellationError {
                // Expected on navigation/Stop — leave partial output in place.
            } catch {
                await MainActor.run {
                    self.errorText = (error as? LocalizedError)?.errorDescription
                        ?? error.localizedDescription
                    self.isStreaming = false
                }
            }
        }
    }

    /// Streams an arbitrary prompt (used by omnibox ask, ask-about-history,
    /// multi-tab compare). The caller builds the grounded, hygiene-wrapped text.
    func runRawPrompt(system: String?, userText: String, model: AIModel, apiKey: String) {
        cancel()
        reset()
        isStreaming = true
        task = Task { [weak self] in
            guard let self else { return }
            do {
                try await self.stream(system: system,
                                      messages: [Message(role: .user, text: userText)],
                                      model: model, apiKey: apiKey)
                await MainActor.run { self.isStreaming = false }
            } catch is CancellationError {
            } catch {
                await MainActor.run {
                    self.errorText = (error as? LocalizedError)?.errorDescription ?? error.localizedDescription
                    self.isStreaming = false
                }
            }
        }
    }

    // MARK: - Summarize (stuff-first, map-reduce for long pages)

    private func summarize(context: PageContext, model: AIModel, apiKey: String) async throws {
        switch Summarizer.plan(for: context, model: model) {
        case .stuff:
            let messages = [Message(role: .user, text: AIAction.summarize.userPrompt(for: context))]
            try await stream(system: AIAction.summarize.systemGuidance,
                             messages: messages, model: model, apiKey: apiKey)

        case .mapReduce:
            let chunks = Summarizer.mapChunks(for: context)
            var partials: [String] = []
            for (i, chunk) in chunks.enumerated() {
                try Task.checkCancellation()
                await MainActor.run { self.statusNote = "Reading section \(i + 1)/\(chunks.count)…" }
                let partial = try await complete(
                    system: AIAction.summarize.systemGuidance,
                    userText: Summarizer.mapPrompt(chunk: chunk, title: context.title),
                    model: model, apiKey: apiKey
                )
                partials.append(partial)
            }
            await MainActor.run { self.statusNote = nil }
            let messages = [Message(role: .user,
                                    text: Summarizer.reducePrompt(partials: partials, title: context.title))]
            try await stream(system: AIAction.summarize.systemGuidance,
                             messages: messages, model: model, apiKey: apiKey)
        }
    }

    // MARK: - Streaming primitive

    /// Streams a chat request into `output`. Cancellation propagates to the
    /// provider stream by dropping the for-await loop.
    private func stream(system: String?, messages: [Message], model: AIModel, apiKey: String) async throws {
        let client = ProviderRegistry.client(for: model)
        // Reasoning models stream hidden thinking before the answer — budget for it.
        var params = ChatParameters(maxTokens: model.capabilities.reasoning ? 8192 : 1024)
        if model.parameterSupport.reasoningEffort { params.reasoningEffort = "low" }
        let request = ChatRequest(model: model, systemPrompt: system, messages: messages, parameters: params)

        for try await event in client.streamChat(request, apiKey: apiKey) {
            try Task.checkCancellation()
            switch event {
            case .textDelta(let delta):
                await MainActor.run { self.output += delta }
            case .reasoningDelta, .citations, .usage, .finished:
                break
            }
        }
    }

    /// Non-streaming completion used by the map phase.
    private func complete(system: String?, userText: String, model: AIModel, apiKey: String) async throws -> String {
        let client = ProviderRegistry.client(for: model)
        var params = ChatParameters(maxTokens: 512)
        if model.parameterSupport.reasoningEffort { params.reasoningEffort = "low" }
        let request = ChatRequest(model: model, systemPrompt: system,
                                  messages: [Message(role: .user, text: userText)],
                                  parameters: params, stream: false)
        let reply = try await client.complete(request, apiKey: apiKey)
        return reply.text
    }

    private func reset() {
        output = ""
        errorText = nil
        statusNote = nil
    }
}