import Foundation
import Hub
import MLXEmbedders
import MLXLMCommon

/// Loads Qwen3-Embedding model containers on demand and caches them per size.
/// Deduplicates concurrent loads of the same size.
actor ModelRegistry {
    private var loaded: [ModelSize: MLXEmbedders.ModelContainer] = [:]
    private var loading: [ModelSize: Task<MLXEmbedders.ModelContainer, Error>] = [:]

    private let hub: HubApi

    init(hub: HubApi = HubApi()) {
        self.hub = hub
    }

    /// Already-loaded sizes (read-only snapshot).
    func loadedSizes() -> Set<ModelSize> { Set(loaded.keys) }

    /// Get or lazily load the container for a given size.
    func container(for size: ModelSize) async throws -> MLXEmbedders.ModelContainer {
        if let c = loaded[size] { return c }
        if let inflight = loading[size] { return try await inflight.value }

        let hub = self.hub
        let task = Task<MLXEmbedders.ModelContainer, Error> {
            // Check for a local override env var first — lets us bypass the
            // Hub downloader entirely (useful when the Hub's metadata-file
            // management fails under sandboxing).
            let envKey: String
            switch size {
            case .s06: envKey = "MLX_MODEL_DIR_0_6B"
            case .s4:  envKey = "MLX_MODEL_DIR_4B"
            case .s8:  envKey = "MLX_MODEL_DIR_8B"
            }
            let dir: URL
            if let localPath = ProcessInfo.processInfo.environment[envKey],
               !localPath.isEmpty {
                dir = URL(fileURLWithPath: localPath)
            } else {
                let repoId = size.repoId
                // HubApi.snapshot downloads or returns an existing local snapshot directory.
                dir = try await hub.snapshot(
                    from: repoId,
                    matching: ["*.json", "*.safetensors", "*.txt", "tokenizer.*", "*.model"]
                )
            }
            let container = try await MLXEmbedders.loadModelContainer(
                from: dir, using: HFTokenizerLoader())
            return container
        }
        loading[size] = task
        do {
            let container = try await task.value
            loaded[size] = container
            loading[size] = nil
            return container
        } catch {
            loading[size] = nil
            throw error
        }
    }

    /// Eagerly load and warm a list of sizes. Failures are thrown after the
    /// first bad size (warmup is best-effort; callers decide policy).
    func preload(_ sizes: [ModelSize], maxBatch: Int, maxLen: Int) async throws {
        for size in sizes {
            let container = try await container(for: size)
            await WarmUp.warm(container: container, maxBatch: maxBatch, maxLen: maxLen)
            fputs("[embed] warmed \(size.rawValue)\n", stderr)
        }
    }
}