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

import Foundation

/// Errors surfaced when a directory is not a usable MLX model.
enum ModelValidationError: LocalizedError {
    case notADirectory(URL)
    case missingConfig(URL)
    case missingWeights(URL)
    case inconsistentShards(missing: [String])

    var errorDescription: String? {
        switch self {
        case .notADirectory(let url):
            "Not a folder: \(url.path)"
        case .missingConfig(let url):
            "No config.json found in \(url.lastPathComponent) — this is not an MLX model folder."
        case .missingWeights(let url):
            "No weight files (*.safetensors) found in \(url.lastPathComponent)."
        case .inconsistentShards(let missing):
            "The weight index references missing shards: \(missing.joined(separator: ", "))."
        }
    }
}

/// Scans, validates, and describes the local model library
/// (`~/Library/Application Support/ZyquoMLX/Models/`).
///
/// Validation follows the on-disk MLX format (docs/MLX-RESEARCH.md §5.3):
/// `config.json` (+ optional `quantization` dict), `model*.safetensors`
/// (single or sharded + index), tokenizer files.
actor ModelStore {

    static let shared = ModelStore()

    private let root: URL

    init(root: URL = PersistenceService.modelsDirectory) {
        self.root = root
    }

    /// All valid models in the library. Invalid directories are skipped
    /// (they surface through `validate` when the user targets them directly).
    func scan() throws -> [LocalModel] {
        let fm = FileManager.default
        guard fm.fileExists(atPath: root.path) else { return [] }
        let entries = try fm.contentsOfDirectory(
            at: root, includingPropertiesForKeys: [.isDirectoryKey],
            options: .skipsHiddenFiles)
        return entries.compactMap { try? describe(directory: $0) }
            .sorted { $0.installedAt > $1.installedAt }
    }

    /// Validate and describe one model directory (works for any path, not
    /// just the library root — Playground can open arbitrary folders).
    func describe(directory: URL) throws -> LocalModel {
        let fm = FileManager.default
        var isDir: ObjCBool = false
        guard fm.fileExists(atPath: directory.path, isDirectory: &isDir), isDir.boolValue else {
            throw ModelValidationError.notADirectory(directory)
        }

        let configURL = directory.appendingPathComponent("config.json")
        var config: [String: Any]?
        if let data = try? Data(contentsOf: configURL) {
            config = try JSONSerialization.jsonObject(with: data) as? [String: Any]
        }

        let type = ModelType.detect(directory: directory, config: config)

        // Image-generation layouts (FLUX) keep weights in component subdirs and
        // have no root config; everything else must have config.json.
        if config == nil && type != .imageGeneration {
            throw ModelValidationError.missingConfig(directory)
        }

        let allFiles = (try? fm.subpathsOfDirectory(atPath: directory.path)) ?? []
        let weightFiles = allFiles.filter { $0.hasSuffix(".safetensors") || $0.hasSuffix(".npz") }
        guard !weightFiles.isEmpty else {
            throw ModelValidationError.missingWeights(directory)
        }

        // Shard-index consistency (model.safetensors.index.json → weight_map).
        let indexURL = directory.appendingPathComponent("model.safetensors.index.json")
        var totalParameters: Int64?
        if let indexData = try? Data(contentsOf: indexURL),
            let index = try? JSONSerialization.jsonObject(with: indexData) as? [String: Any]
        {
            if let weightMap = index["weight_map"] as? [String: String] {
                let referenced = Set(weightMap.values)
                let present = Set(weightFiles.map { ($0 as NSString).lastPathComponent })
                let missing = referenced.subtracting(present).sorted()
                // Some mlx-community repos ship a consolidated
                // model.safetensors next to a stale sharded index (seen live:
                // Qwen3-VL-4B-Instruct-4bit) — treat that as valid.
                if !missing.isEmpty && !present.contains("model.safetensors") {
                    throw ModelValidationError.inconsistentShards(missing: missing)
                }
            }
            if let meta = index["metadata"] as? [String: Any],
                let params = meta["total_parameters"] as? Int64 ?? (meta["total_parameters"] as? Int).map(Int64.init)
            {
                totalParameters = params
            }
        }

        var weightsSize: Int64 = 0
        var diskSize: Int64 = 0
        for path in allFiles {
            let fileURL = directory.appendingPathComponent(path)
            guard let size = (try? fileURL.resourceValues(forKeys: [.fileSizeKey]))?.fileSize
            else { continue }
            diskSize += Int64(size)
            if path.hasSuffix(".safetensors") || path.hasSuffix(".npz") {
                weightsSize += Int64(size)
            }
        }

        var quantization: QuantizationInfo?
        if let quantDict = config?["quantization"] as? [String: Any],
            let data = try? JSONSerialization.data(withJSONObject: quantDict)
        {
            quantization = try? JSONDecoder().decode(QuantizationInfo.self, from: data)
        }

        let name = directory.lastPathComponent
        let installedAt =
            (try? directory.resourceValues(forKeys: [.creationDateKey]))?.creationDate ?? .now

        return LocalModel(
            id: name,
            name: name,
            repoID: Self.repoID(fromDirectoryName: name),
            type: type,
            directory: directory,
            architecture: config?["model_type"] as? String,
            parameterCount: totalParameters ?? Self.parameterCount(fromName: name, quantization: quantization, weightsSize: weightsSize),
            quantization: quantization,
            weightsSize: weightsSize,
            diskSize: diskSize,
            installedAt: installedAt
        )
    }

    /// Delete a model from the library.
    func delete(_ model: LocalModel) throws {
        try FileManager.default.removeItem(at: model.directory)
    }

    // MARK: - Name parsing

    /// Library directories are named "org--repo" (HF convention for local caches).
    private static func repoID(fromDirectoryName name: String) -> String? {
        guard name.contains("--") else { return nil }
        return name.replacingOccurrences(of: "--", with: "/")
    }

    /// Parse "…-4B-…" / "…-0.6B-…" / "…-135M-…" out of a repo name; fall back
    /// to estimating from weight bytes (docs/MODELS.md §3 bytes-per-param).
    static func parameterCount(fromName name: String, quantization: QuantizationInfo?, weightsSize: Int64) -> Int64? {
        if let range = name.range(of: #"(\d+(?:\.\d+)?)\s*[bB](?=[-_.]|$)"#, options: .regularExpression) {
            let token = name[range].dropLast()
            if let value = Double(token.trimmingCharacters(in: .whitespaces)) {
                return Int64(value * 1_000_000_000)
            }
        }
        if let range = name.range(of: #"(\d+(?:\.\d+)?)\s*[mM](?=[-_.]|$)"#, options: .regularExpression) {
            let token = name[range].dropLast()
            if let value = Double(token.trimmingCharacters(in: .whitespaces)) {
                return Int64(value * 1_000_000)
            }
        }
        guard weightsSize > 0 else { return nil }
        let bytesPerParam = MemoryAdvisor.bytesPerParameter(quantization: quantization)
        return Int64(Double(weightsSize) / bytesPerParam)
    }
}