kangkang/康康/AI/ModelStore.swift

import Foundation

nonisolated enum ModelKind: String, CaseIterable {
    /// 也是沙盒 Models/ 下的子目录名 / CDN 路径段。
    /// 主模型已从 Qwen3.5-2B → Gemma-3n → **Gemma 4 E2B(端侧 4bit,只走 MLX/GPU)**:
    /// - llm:**用户唯一下载的主模型**。MLX 文本引擎,经 LLMModelFactory 的 `gemma4` 路径加载
    ///        (mlx-swift-lm ≥3.31.4 已注册;`Gemma4Model` 只取 `language_model` 权重,跳过视觉/音频)。
    ///        e2b 检查点无视觉权重(视觉在 gemma4_unified 12B),故拍报告解读仍走 OCR + 文本 LLM。
    /// - vl:已废弃,保留枚举避免动一圈穷举 switch,不再下载/展示。
    /// - mnnLLM:已停用。本项目已放弃 MNN/SME2 路径;保留枚举只为兼容旧引擎判断,不再分发/计入下载。
    case llm    = "gemma-4-e2b-it-4bit"
    case vl     = "Qwen3-VL-4B-Instruct-4bit"
    case mnnLLM = "Qwen3.5-2B-MNN"

    var displayName: String {
        switch self {
        case .llm:    return "Gemma 4 E2B (MLX)"
        case .vl:     return "Qwen3-VL-4B"
        case .mnnLLM: return "Qwen3.5-2B (MNN/SME2)"
        }
    }

    /// HuggingFace 仓库 ID(org/name),用于下载
    var huggingFaceRepo: String { "mlx-community/\(rawValue)" }

    /// 用于判定该模型是否已就绪的最小标志文件
    var sentinelFilename: String { "config.json" }

    /// 面向用户的模型集合:模型管理页 / 下载全部 / 就绪计数对外只暴露统一的主模型
    /// Gemma 4 E2B(MLX,文本)。.vl/.mnnLLM 已停用,不展示、不计入「下载全部」与就绪计数。
    static let userFacing: [ModelKind] = [.llm]
}

/// `@unchecked Sendable`:rootURL 是 let,方法只读 filesystem(线程安全),
/// 可被任意 actor / Task 跨边界访问。
/// 实例方法显式 `nonisolated`:项目开了 `SWIFT_DEFAULT_ACTOR_ISOLATION = MainActor`,
/// 默认会把方法推到 MainActor,导致从 `AIRuntime` 等 actor 调用报错。
final class ModelStore: @unchecked Sendable {
    nonisolated static let shared: ModelStore = {
        do {
            let appSupport = try FileManager.default.url(
                for: .applicationSupportDirectory,
                in: .userDomainMask,
                appropriateFor: nil,
                create: true
            )
            let root = appSupport.appendingPathComponent("Models", isDirectory: true)
            return try ModelStore(rootURL: root)
        } catch {
            fatalError("ModelStore.shared init failed: \(error)")
        }
    }()

    let rootURL: URL

    init(rootURL: URL) throws {
        self.rootURL = rootURL
        try FileManager.default.createDirectory(at: rootURL, withIntermediateDirectories: true)
    }

    nonisolated func localURL(for kind: ModelKind) -> URL {
        rootURL.appendingPathComponent(kind.rawValue, isDirectory: true)
    }

    nonisolated func isReady(_ kind: ModelKind) -> Bool {
        let sentinel = localURL(for: kind).appendingPathComponent(kind.sentinelFilename)
        return FileManager.default.fileExists(atPath: sentinel.path)
    }

    nonisolated func totalBytes(for kind: ModelKind) -> Int {
        let folder = localURL(for: kind)
        guard let enumerator = FileManager.default.enumerator(
            at: folder,
            includingPropertiesForKeys: [.fileSizeKey]
        ) else { return 0 }
        var sum = 0
        for case let url as URL in enumerator {
            if let size = try? url.resourceValues(forKeys: [.fileSizeKey]).fileSize {
                sum += size
            }
        }
        return sum
    }

    /// Demo 现场旁路:从 Bundle 拷贝预装模型(W6 才真正使用,本周占位)
    nonisolated func seedFromBundle(_ kind: ModelKind) throws {
        guard let bundleURL = Bundle.main.url(forResource: kind.rawValue, withExtension: nil) else {
            #if DEBUG
            assertionFailure("Bundle 缺少 \(kind.rawValue),检查资源是否加入 target")
            #endif
            return
        }
        let target = localURL(for: kind)
        if FileManager.default.fileExists(atPath: target.path) {
            try FileManager.default.removeItem(at: target)
        }
        try FileManager.default.copyItem(at: bundleURL, to: target)
    }

    // MARK: - 下载 / 导入支撑

    /// 模型目录下某个相对路径文件的本地 URL。
    nonisolated func fileURL(for kind: ModelKind, relativePath: String) -> URL {
        localURL(for: kind).appendingPathComponent(relativePath)
    }

    /// 本地该文件当前字节数,不存在返回 0(用于断点续传偏移与跳过判断)。
    nonisolated func localBytes(for kind: ModelKind, relativePath: String) -> Int {
        let url = fileURL(for: kind, relativePath: relativePath)
        guard let size = try? url.resourceValues(forKeys: [.fileSizeKey]).fileSize else { return 0 }
        return size
    }

    /// 按清单校验模型是否完整:每个文件都存在且大小等于预期。
    /// `files` 默认取 `ModelManifest`;测试可注入小清单。
    nonisolated func isComplete(for kind: ModelKind, files: [ModelFile]? = nil) -> Bool {
        let manifest = files ?? ModelManifest.files(for: kind)
        guard !manifest.isEmpty else { return false }
        for file in manifest where localBytes(for: kind, relativePath: file.path) != file.bytes {
            return false
        }
        return true
    }

    /// 旁路导入:把一个含 config.json 的模型文件夹整体拷入沙盒(现场重装兜底)。
    nonisolated func importModel(_ kind: ModelKind, from sourceFolder: URL) throws {
        let configPath = sourceFolder.appendingPathComponent(kind.sentinelFilename).path
        guard FileManager.default.fileExists(atPath: configPath) else {
            throw ModelStoreError.missingConfig
        }
        let target = localURL(for: kind)
        if FileManager.default.fileExists(atPath: target.path) {
            try FileManager.default.removeItem(at: target)
        }
        try FileManager.default.createDirectory(
            at: target.deletingLastPathComponent(), withIntermediateDirectories: true)
        try FileManager.default.copyItem(at: sourceFolder, to: target)
    }
}

enum ModelStoreError: Error, LocalizedError {
    case missingConfig

    var errorDescription: String? {
        switch self {
        case .missingConfig:
            return String(appLoc: "所选文件夹缺少 config.json,不是有效的模型目录")
        }
    }
}