//
//  AppModel.swift
//  Zyquo Local
//
//  Author: Simon-Pierre Boucher
//  Mail: contact@spboucher.ai
//

import Foundation
import MLX
import Observation

/// Root coordinator: owns the engine, stores, downloads and conversations.
@MainActor
@Observable
final class AppModel {
    let settings: AppSettings
    let store: ModelStore
    let downloads: DownloadManager
    let engine = InferenceEngine()
    let chat: ChatController

    var conversations: [Conversation] = []
    var selectedConversationID: UUID?

    /// Mirrors the engine actor's state for the UI.
    var engineState: EngineState = .unloaded
    /// Live MLX active memory while a model is loaded (footer chip).
    var liveMemoryBytes: Int = 0
    /// Last load error to surface in the UI.
    var lastError: String?

    var personaStore: PersonaStore
    var promptLibrary: PromptLibrary

    /// Detail column routing.
    enum DetailRoute: Hashable {
        case chat
        case library
    }
    var route: DetailRoute = .chat

    var selectedConversation: Conversation? {
        get { conversations.first { $0.id == selectedConversationID } }
        set {
            guard let newValue, let i = conversations.firstIndex(where: { $0.id == newValue.id }) else { return }
            conversations[i] = newValue
        }
    }

    var loadedModelID: String? {
        switch engineState {
        case .ready(let id), .generating(let id), .loading(let id): id
        case .unloaded: nil
        }
    }

    init() {
        let settings = AppSettings()
        self.settings = settings
        let store = ModelStore()
        self.store = store
        self.downloads = DownloadManager(hub: HubService(token: settings.hfToken), store: store)
        self.personaStore = PersonaStore()
        self.promptLibrary = PromptLibrary()
        self.chat = ChatController()
        chat.bind(to: self)

        conversations = PersistenceService.loadConversations()
        if conversations.isEmpty {
            newConversation()
        } else {
            selectedConversationID = conversations.first?.id
        }

        applyGPUCacheLimit()

        // Load the default model on launch when configured.
        if let defaultID = settings.defaultModelID, store.model(for: defaultID) != nil {
            Task { await loadModel(repoID: defaultID) }
        }
        // Keep the memory readout fresh while a model is loaded.
        Task { await memoryTicker() }
    }

    func applyGPUCacheLimit() {
        if settings.gpuCacheLimitMB > 0 {
            MLX.Memory.cacheLimit = settings.gpuCacheLimitMB * 1_048_576
        }
    }

    /// Push the (possibly updated) HF token into services that need it.
    func refreshToken() {
        downloads.hub = HubService(token: settings.hfToken)
    }

    // MARK: - Conversations

    @discardableResult
    func newConversation(persona: Persona? = nil) -> Conversation {
        var conversation = Conversation(
            modelID: loadedModelID ?? settings.defaultModelID,
            systemPrompt: persona?.systemPrompt
                ?? (settings.defaultSystemPrompt.isEmpty ? nil : settings.defaultSystemPrompt),
            params: persona?.params ?? settings.defaultParams
        )
        if let persona {
            conversation.title = persona.name
            if let preferred = persona.preferredModelID { conversation.modelID = preferred }
        }
        conversations.insert(conversation, at: 0)
        selectedConversationID = conversation.id
        route = .chat
        PersistenceService.save(conversation)
        return conversation
    }

    func delete(conversationID: UUID) {
        conversations.removeAll { $0.id == conversationID }
        PersistenceService.delete(conversationID: conversationID)
        if selectedConversationID == conversationID {
            selectedConversationID = conversations.first?.id
        }
    }

    func update(_ conversation: Conversation, touch: Bool = true) {
        var conversation = conversation
        if touch { conversation.updatedAt = Date() }
        if let i = conversations.firstIndex(where: { $0.id == conversation.id }) {
            conversations[i] = conversation
        }
        PersistenceService.save(conversation)
    }

    func togglePin(conversationID: UUID) {
        guard var c = conversations.first(where: { $0.id == conversationID }) else { return }
        c.pinned.toggle()
        update(c, touch: false)
    }

    // MARK: - Model lifecycle

    func loadModel(repoID: String) async {
        guard let model = store.model(for: repoID) else {
            lastError = "\(repoID) is not downloaded."
            return
        }
        engineState = .loading(repoID: repoID)
        do {
            try await engine.load(model: model)
            engineState = await engine.state
            store.markUsed(repoID: repoID)
            // Bind the current conversation to the newly loaded model.
            if var c = selectedConversation {
                c.modelID = repoID
                update(c, touch: false)
            }
        } catch {
            engineState = .unloaded
            lastError = error.localizedDescription
        }
    }

    func unloadModel() async {
        await engine.unload()
        engineState = .unloaded
        liveMemoryBytes = 0
    }

    private func memoryTicker() async {
        while !Task.isCancelled {
            try? await Task.sleep(for: .seconds(2))
            if loadedModelID != nil {
                liveMemoryBytes = MemoryAdvisor.activeMemoryBytes
            }
        }
    }

    // MARK: - Sidebar grouping & search

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

    func sidebarGroups(query: String) -> [SidebarGroup] {
        let filtered = query.isEmpty
            ? conversations
            : conversations.filter { c in
                c.title.localizedCaseInsensitiveContains(query)
                    || c.messages.contains { $0.content.localizedCaseInsensitiveContains(query) }
            }
        var groups: [SidebarGroup] = []
        let pinned = filtered.filter(\.pinned)
        if !pinned.isEmpty {
            groups.append(SidebarGroup(title: "Pinned", conversations: pinned))
        }
        let rest = filtered.filter { !$0.pinned }
        let calendar = Calendar.current
        let now = Date()
        func bucket(_ date: Date) -> String {
            if calendar.isDateInToday(date) { return "Today" }
            if calendar.isDateInYesterday(date) { return "Yesterday" }
            if date > calendar.date(byAdding: .day, value: -7, to: now)! { return "Previous 7 Days" }
            return "Older"
        }
        for title in ["Today", "Yesterday", "Previous 7 Days", "Older"] {
            let matching = rest.filter { bucket($0.updatedAt) == title }
            if !matching.isEmpty {
                groups.append(SidebarGroup(title: title, conversations: matching))
            }
        }
        return groups
    }
}