//
//  ConversationStore.swift
//  Zyquo Cloud
//
//  Author: Simon-Pierre Boucher
//  Mail: contact@spboucher.ai
//
//  Central chat state: the conversation list, selection, and the chat engine
//  that streams provider responses into messages. All mutation happens on the
//  main actor; streaming work runs in tasks that hop back for UI updates.
//

import Foundation

@MainActor
final class ConversationStore: ObservableObject {
    @Published var conversations: [Conversation] = []
    @Published var selectedID: Conversation.ID?
    @Published var searchText: String = ""
    /// Live streaming task per conversation (supports parallel streams in compare mode).
    @Published private(set) var streamingConversationIDs: Set<Conversation.ID> = []

    let catalog: ModelCatalog
    let vault: KeyVaultStore
    private let persistence: PersistenceService
    private var streamTasks: [Conversation.ID: Task<Void, Never>] = [:]

    /// Global default system prompt (Settings → Advanced).
    @Published var defaultSystemPrompt: String {
        didSet { persistence.save(defaultSystemPrompt, to: "default-system-prompt.json") }
    }

    init(
        catalog: ModelCatalog,
        vault: KeyVaultStore,
        persistence: PersistenceService = .shared
    ) {
        self.catalog = catalog
        self.vault = vault
        self.persistence = persistence
        self.defaultSystemPrompt = persistence.load(String.self, from: "default-system-prompt.json") ?? ""
        conversations = persistence.loadConversations()
        selectedID = conversations.first?.id
    }

    // MARK: - Selection & lookup

    var selected: Conversation? {
        get { conversations.first { $0.id == selectedID } }
    }

    func binding(for id: Conversation.ID) -> Int? {
        conversations.firstIndex { $0.id == id }
    }

    func isStreaming(_ id: Conversation.ID) -> Bool {
        streamingConversationIDs.contains(id)
    }

    // MARK: - CRUD

    @discardableResult
    func newConversation(model: AIModel? = nil, persona: Persona? = nil) -> Conversation {
        let chosen = model ?? catalog.defaultModel
        var conversation = Conversation(
            modelID: chosen?.id ?? "",
            provider: chosen?.provider ?? .openai,
            systemPrompt: persona?.systemPrompt ?? (defaultSystemPrompt.isEmpty ? nil : defaultSystemPrompt)
        )
        if let persona {
            conversation.personaID = persona.id
            conversation.parameters = persona.parameters
            if let modelID = persona.modelID, let provider = persona.provider {
                conversation.modelID = modelID
                conversation.provider = provider
            }
        }
        conversations.insert(conversation, at: 0)
        selectedID = conversation.id
        persistence.save(conversation)
        return conversation
    }

    func delete(_ id: Conversation.ID) {
        stopStreaming(id)
        if let conversation = conversations.first(where: { $0.id == id }) {
            persistence.delete(conversation)
        }
        conversations.removeAll { $0.id == id }
        if selectedID == id { selectedID = conversations.first?.id }
    }

    func update(_ conversation: Conversation) {
        guard let index = conversations.firstIndex(where: { $0.id == conversation.id }) else { return }
        var updated = conversation
        updated.updatedAt = Date()
        conversations[index] = updated
        persistence.save(updated)
    }

    func togglePin(_ id: Conversation.ID) {
        guard var conversation = conversations.first(where: { $0.id == id }) else { return }
        conversation.isPinned.toggle()
        update(conversation)
    }

    func rename(_ id: Conversation.ID, to title: String) {
        guard var conversation = conversations.first(where: { $0.id == id }) else { return }
        conversation.title = title
        conversation.hasAutoTitle = false
        update(conversation)
    }

    // MARK: - Sidebar grouping

    struct SidebarGroup: Identifiable {
        let id: String
        let title: String
        let conversations: [Conversation]
    }

    /// Pinned / Today / Yesterday / Previous 7 Days / Older, filtered by search.
    var sidebarGroups: [SidebarGroup] {
        let filtered = searchText.isEmpty
            ? conversations
            : conversations.filter { conversation in
                conversation.title.localizedCaseInsensitiveContains(searchText)
                    || conversation.messages.contains {
                        $0.text.localizedCaseInsensitiveContains(searchText)
                    }
            }
        let calendar = Calendar.current
        let now = Date()
        var pinned: [Conversation] = []
        var today: [Conversation] = []
        var yesterday: [Conversation] = []
        var week: [Conversation] = []
        var older: [Conversation] = []
        for conversation in filtered {
            if conversation.isPinned { pinned.append(conversation); continue }
            if calendar.isDateInToday(conversation.updatedAt) { today.append(conversation) }
            else if calendar.isDateInYesterday(conversation.updatedAt) { yesterday.append(conversation) }
            else if conversation.updatedAt > now.addingTimeInterval(-7 * 86_400) { week.append(conversation) }
            else { older.append(conversation) }
        }
        return [
            SidebarGroup(id: "pinned", title: "Pinned", conversations: pinned),
            SidebarGroup(id: "today", title: "Today", conversations: today),
            SidebarGroup(id: "yesterday", title: "Yesterday", conversations: yesterday),
            SidebarGroup(id: "week", title: "Previous 7 Days", conversations: week),
            SidebarGroup(id: "older", title: "Older", conversations: older),
        ].filter { !$0.conversations.isEmpty }
    }

    // MARK: - Chat engine

    /// Sends the user's text (and attachments) in a conversation and streams
    /// the assistant reply. Optionally targets a different model for this
    /// message only (per-message model switch).
    func send(
        text: String,
        attachments: [Attachment] = [],
        in conversationID: Conversation.ID,
        overrideModel: AIModel? = nil
    ) {
        guard var conversation = conversations.first(where: { $0.id == conversationID }) else { return }
        guard let model = overrideModel
            ?? catalog.model(id: conversation.modelID, provider: conversation.provider)
            ?? catalog.defaultModel
        else { return }

        if let overrideModel {
            conversation.modelID = overrideModel.id
            conversation.provider = overrideModel.provider
        }

        var userMessage = Message(role: .user, text: text, attachments: attachments)
        userMessage.modelID = model.id
        userMessage.provider = model.provider
        conversation.messages.append(userMessage)
        update(conversation)

        generateReply(in: conversationID, model: model)
    }

    /// Regenerates the last assistant message (optionally with another model).
    func regenerate(in conversationID: Conversation.ID, with model: AIModel? = nil) {
        guard var conversation = conversations.first(where: { $0.id == conversationID }) else { return }
        if conversation.messages.last?.role == .assistant {
            conversation.messages.removeLast()
            update(conversation)
        }
        guard let target = model
            ?? catalog.model(id: conversation.modelID, provider: conversation.provider)
        else { return }
        generateReply(in: conversationID, model: target)
    }

    /// Replaces a user message's text and regenerates from that point.
    func editAndResend(messageID: Message.ID, newText: String, in conversationID: Conversation.ID) {
        guard var conversation = conversations.first(where: { $0.id == conversationID }),
              let index = conversation.messages.firstIndex(where: { $0.id == messageID })
        else { return }
        conversation.messages[index].text = newText
        conversation.messages.removeSubrange((index + 1)...)
        update(conversation)
        guard let model = catalog.model(id: conversation.modelID, provider: conversation.provider) else { return }
        generateReply(in: conversationID, model: model)
    }

    func deleteMessage(_ messageID: Message.ID, in conversationID: Conversation.ID) {
        guard var conversation = conversations.first(where: { $0.id == conversationID }) else { return }
        conversation.messages.removeAll { $0.id == messageID }
        update(conversation)
    }

    func stopStreaming(_ conversationID: Conversation.ID) {
        streamTasks[conversationID]?.cancel()
        streamTasks[conversationID] = nil
        streamingConversationIDs.remove(conversationID)
        finalizeStreamingMessage(in: conversationID)
    }

    private func generateReply(in conversationID: Conversation.ID, model: AIModel) {
        guard let index = conversations.firstIndex(where: { $0.id == conversationID }) else { return }

        var placeholder = Message(role: .assistant, text: "")
        placeholder.modelID = model.id
        placeholder.provider = model.provider
        placeholder.isStreaming = true
        conversations[index].messages.append(placeholder)
        let messageID = placeholder.id

        let request = ChatRequest(
            model: model,
            systemPrompt: conversations[index].systemPrompt,
            messages: conversations[index].messages.filter { $0.id != messageID && $0.errorText == nil },
            parameters: conversations[index].parameters
        )

        streamingConversationIDs.insert(conversationID)
        let task = Task { [weak self] in
            guard let self else { return }
            do {
                let apiKey = try self.vault.apiKey(for: model.provider)
                let client = ProviderRegistry.client(for: model)
                var usage: TokenUsage?
                for try await event in client.streamChat(request, apiKey: apiKey) {
                    if Task.isCancelled { break }
                    switch event {
                    case .textDelta(let delta):
                        self.mutateMessage(messageID, in: conversationID) { $0.text += delta }
                    case .reasoningDelta(let delta):
                        self.mutateMessage(messageID, in: conversationID) {
                            $0.reasoning = ($0.reasoning ?? "") + delta
                        }
                    case .citations(let citations):
                        self.mutateMessage(messageID, in: conversationID) { $0.citations = citations }
                    case .usage(let u):
                        usage = u
                    case .finished:
                        break
                    }
                }
                let finalUsage = usage
                self.mutateMessage(messageID, in: conversationID) { message in
                    message.isStreaming = false
                    if let u = finalUsage {
                        message.usage = u
                        message.estimatedCost = model.pricing?.cost(
                            inputTokens: u.inputTokens, outputTokens: u.outputTokens
                        )
                    }
                }
            } catch let error as ProviderError {
                if case .cancelled = error {
                    self.mutateMessage(messageID, in: conversationID) { $0.isStreaming = false }
                } else {
                    self.mutateMessage(messageID, in: conversationID) {
                        $0.isStreaming = false
                        $0.errorText = error.localizedDescription
                    }
                }
            } catch {
                self.mutateMessage(messageID, in: conversationID) {
                    $0.isStreaming = false
                    $0.errorText = error.localizedDescription
                }
            }
            self.streamingConversationIDs.remove(conversationID)
            self.streamTasks[conversationID] = nil
            self.persistSnapshot(of: conversationID)
            self.autoTitleIfNeeded(conversationID)
        }
        streamTasks[conversationID] = task
    }

    private func mutateMessage(
        _ messageID: Message.ID,
        in conversationID: Conversation.ID,
        _ mutate: (inout Message) -> Void
    ) {
        guard let ci = conversations.firstIndex(where: { $0.id == conversationID }),
              let mi = conversations[ci].messages.firstIndex(where: { $0.id == messageID })
        else { return }
        mutate(&conversations[ci].messages[mi])
    }

    private func finalizeStreamingMessage(in conversationID: Conversation.ID) {
        guard let ci = conversations.firstIndex(where: { $0.id == conversationID }) else { return }
        for mi in conversations[ci].messages.indices where conversations[ci].messages[mi].isStreaming {
            conversations[ci].messages[mi].isStreaming = false
        }
        persistSnapshot(of: conversationID)
    }

    private func persistSnapshot(of conversationID: Conversation.ID) {
        guard var conversation = conversations.first(where: { $0.id == conversationID }) else { return }
        conversation.updatedAt = Date()
        if let index = conversations.firstIndex(where: { $0.id == conversationID }) {
            conversations[index] = conversation
        }
        persistence.save(conversation)
    }

    // MARK: - Auto titles

    /// After the first completed exchange, asks the provider's cheapest model
    /// for a short title.
    private func autoTitleIfNeeded(_ conversationID: Conversation.ID) {
        guard let conversation = conversations.first(where: { $0.id == conversationID }),
              conversation.hasAutoTitle,
              conversation.messages.filter({ $0.role == .assistant && !$0.text.isEmpty }).count == 1,
              let cheap = catalog.cheapestModel(for: conversation.provider),
              let firstUser = conversation.messages.first(where: { $0.role == .user })
        else { return }

        let assistantText = conversation.messages.last { $0.role == .assistant }?.text ?? ""
        let prompt = """
        Write a title of at most 5 words for this conversation. Reply with the title only, no quotes.

        User: \(firstUser.text.prefix(500))
        Assistant: \(assistantText.prefix(500))
        """
        Task { [weak self] in
            guard let self else { return }
            do {
                let apiKey = try self.vault.apiKey(for: cheap.provider)
                let client = ProviderRegistry.client(for: cheap)
                let request = ChatRequest(
                    model: cheap,
                    systemPrompt: nil,
                    messages: [Message(role: .user, text: prompt)],
                    parameters: ChatParameters(maxTokens: 24),
                    stream: false
                )
                let reply = try await client.complete(request, apiKey: apiKey)
                let title = reply.text
                    .trimmingCharacters(in: .whitespacesAndNewlines)
                    .trimmingCharacters(in: CharacterSet(charactersIn: "\"“”"))
                guard !title.isEmpty,
                      var current = self.conversations.first(where: { $0.id == conversationID }),
                      current.hasAutoTitle
                else { return }
                current.title = String(title.prefix(60))
                self.update(current)
            } catch {
                // Title generation is best-effort; keep "New Chat" on failure.
            }
        }
    }
}