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

import AppKit
import Foundation
import Observation

/// The on-disk model library: scans the Models folder, validates model
/// directories, tracks sizes and last-used dates, deletes models.
@MainActor
@Observable
final class ModelStore {
    private(set) var models: [LocalModel] = []

    /// Root: ~/Library/Application Support/ZyquoLocal/Models
    let modelsRoot: URL

    /// Sidecar metadata (last used, per-model params) keyed by repoID.
    private var meta: [String: ModelMeta] = [:]

    struct ModelMeta: Codable {
        var lastUsed: Date?
        var defaultParams: GenerationParams?
    }

    var totalSizeBytes: Int64 {
        models.reduce(0) { $0 + $1.sizeBytes }
    }

    init(root: URL = PersistenceService.appSupportDirectory.appendingPathComponent("Models")) {
        self.modelsRoot = root
        try? FileManager.default.createDirectory(at: root, withIntermediateDirectories: true)
        loadMeta()
        rescan()
    }

    /// Directory a repo downloads into: Models/{org}/{repo}.
    func directory(for repoID: String) -> URL {
        modelsRoot.appendingPathComponent(repoID)
    }

    func model(for repoID: String) -> LocalModel? {
        models.first { $0.repoID == repoID }
    }

    /// Required files for a valid model directory (docs/MLX-RESEARCH.md §3).
    static func isValidModelDirectory(_ dir: URL) -> Bool {
        let fm = FileManager.default
        guard fm.fileExists(atPath: dir.appendingPathComponent("config.json").path) else { return false }
        guard fm.fileExists(atPath: dir.appendingPathComponent("tokenizer_config.json").path) else { return false }
        guard fm.fileExists(atPath: dir.appendingPathComponent("tokenizer.json").path) else { return false }
        guard let contents = try? fm.contentsOfDirectory(atPath: dir.path) else { return false }
        return contents.contains { $0.hasSuffix(".safetensors") && !$0.hasSuffix(".partial") }
    }

    /// Scans Models/{org}/{repo} two levels deep and rebuilds the library.
    func rescan() {
        let fm = FileManager.default
        var found: [LocalModel] = []
        let orgs = (try? fm.contentsOfDirectory(
            at: modelsRoot, includingPropertiesForKeys: nil,
            options: .skipsHiddenFiles)) ?? []
        for org in orgs where org.hasDirectoryPath || (try? org.resourceValues(forKeys: [.isDirectoryKey]).isDirectory) == true {
            let repos = (try? fm.contentsOfDirectory(
                at: org, includingPropertiesForKeys: nil,
                options: .skipsHiddenFiles)) ?? []
            for repo in repos {
                guard (try? repo.resourceValues(forKeys: [.isDirectoryKey]).isDirectory) == true else { continue }
                guard Self.isValidModelDirectory(repo) else { continue }
                let repoID = "\(org.lastPathComponent)/\(repo.lastPathComponent)"
                let size = (try? fm.allocatedSizeOfDirectory(at: repo)) ?? 0
                let config = Self.readConfig(directory: repo)
                found.append(LocalModel(
                    repoID: repoID,
                    directory: repo,
                    sizeBytes: size,
                    architecture: config.architecture,
                    quantization: Self.quantLabel(repoID: repoID, config: config),
                    contextWindow: InferenceEngine.readContextWindow(directory: repo),
                    lastUsed: meta[repoID]?.lastUsed,
                    defaultParams: meta[repoID]?.defaultParams
                ))
            }
        }
        models = found.sorted { ($0.lastUsed ?? .distantPast) > ($1.lastUsed ?? .distantPast) }
    }

    /// Deletes a model directory. Returns the reclaimed bytes.
    @discardableResult
    func delete(repoID: String) -> Int64 {
        guard let model = model(for: repoID) else { return 0 }
        let reclaimed = model.sizeBytes
        try? FileManager.default.removeItem(at: model.directory)
        // Prune the org folder when empty.
        let org = model.directory.deletingLastPathComponent()
        if let rest = try? FileManager.default.contentsOfDirectory(atPath: org.path), rest.isEmpty {
            try? FileManager.default.removeItem(at: org)
        }
        meta[repoID] = nil
        saveMeta()
        rescan()
        return reclaimed
    }

    func revealInFinder(repoID: String) {
        guard let model = model(for: repoID) else { return }
        NSWorkspace.shared.activateFileViewerSelecting([model.directory])
    }

    func markUsed(repoID: String) {
        var m = meta[repoID] ?? ModelMeta()
        m.lastUsed = Date()
        meta[repoID] = m
        saveMeta()
        if let i = models.firstIndex(where: { $0.repoID == repoID }) {
            models[i].lastUsed = m.lastUsed
        }
    }

    func setDefaultParams(_ params: GenerationParams?, for repoID: String) {
        var m = meta[repoID] ?? ModelMeta()
        m.defaultParams = params
        meta[repoID] = m
        saveMeta()
        if let i = models.firstIndex(where: { $0.repoID == repoID }) {
            models[i].defaultParams = params
        }
    }

    // MARK: - Config parsing

    struct ConfigInfo {
        var architecture: String?
        var quantBits: Int?
    }

    static func readConfig(directory: URL) -> ConfigInfo {
        let url = directory.appendingPathComponent("config.json")
        guard let data = try? Data(contentsOf: url),
            let json = try? JSONSerialization.jsonObject(with: data) as? [String: Any]
        else { return ConfigInfo() }
        var bits: Int?
        if let quant = json["quantization"] as? [String: Any] {
            bits = quant["bits"] as? Int
        }
        return ConfigInfo(architecture: json["model_type"] as? String, quantBits: bits)
    }

    /// "4bit", "8bit", "bf16"… parsed from repo name, falling back to config.
    static func quantLabel(repoID: String, config: ConfigInfo) -> String? {
        let name = repoID.lowercased()
        for candidate in ["4bit", "5bit", "6bit", "8bit", "bf16", "fp16", "mxfp4"] {
            if name.contains(candidate) { return candidate }
        }
        if let bits = config.quantBits { return "\(bits)bit" }
        return nil
    }

    // MARK: - Meta persistence

    private var metaURL: URL {
        PersistenceService.appSupportDirectory.appendingPathComponent("library-meta.json")
    }

    private func loadMeta() {
        guard let data = try? Data(contentsOf: metaURL),
            let decoded = try? JSONDecoder().decode([String: ModelMeta].self, from: data)
        else { return }
        meta = decoded
    }

    private func saveMeta() {
        try? JSONEncoder().encode(meta).write(to: metaURL, options: .atomic)
    }
}