markbaseengine/Sources/MarkBaseServer/ModelManager.swift

import Foundation
import MarkBase

// ─────────────────────────────────────────────────────────────
// Model Manager - Runtime Model Switching
// ─────────────────────────────────────────────────────────────

/// Model information
public struct ModelInfo: Codable, Sendable {
    public let id: String
    public let path: String
    public let name: String
    public var loaded: Bool
    public let parameters: [String: String]?
    
    public init(
        id: String,
        path: String,
        name: String,
        loaded: Bool = false,
        parameters: [String: String]? = nil
    ) {
        self.id = id
        self.path = path
        self.name = name
        self.loaded = loaded
        self.parameters = parameters
    }
}

/// Model errors
public enum ModelError: Error, LocalizedError {
    case modelNotFound(String)
    case noModelLoaded
    case loadFailed(String)
    
    public var errorDescription: String? {
        switch self {
        case .modelNotFound(let id):
            return "Model not found: \(id)"
        case .noModelLoaded:
            return "No model is currently loaded"
        case .loadFailed(let detail):
            return "Failed to load model: \(detail)"
        }
    }
}

/// Model manager for runtime model switching
public actor ModelManager {
    private var models: [String: ModelInfo] = [:]
    private var currentModelId: String?
    
    // Active model instances
    private var engine: MarkBaseEngine?
    private var model: E4BModel?
    private var tokenizer: Tokenizer?
    private var generator: StreamingGenerator?
    private var sampler: Sampler?
    
    public init() {}
    
    /// Register a model
    public func register(id: String, path: String, name: String) {
        models[id] = ModelInfo(id: id, path: path, name: name, loaded: false)
    }
    
    /// Get list of registered models
    public func listModels() -> [ModelInfo] {
        return Array(models.values)
    }
    
    /// Get current model info
    public func getCurrentModel() -> ModelInfo? {
        guard let id = currentModelId else { return nil }
        return models[id]
    }
    
    /// Load a model
    public func loadModel(id: String) async throws {
        guard let modelInfo = models[id] else {
            throw ModelError.modelNotFound(id)
        }
        
        // Validate path
        try Validator.validateModelPath(modelInfo.path)
        
        // Create engine
        let newEngine = try MarkBaseEngine(autoCompile: true)
        
        // Load model
        let newModel = try E4BModel(
            modelDir: modelInfo.path,
            engine: newEngine,
            maxContextLength: 512
        )
        
        // Load tokenizer
        let newTokenizer = try TokenizerFactory.load(modelDir: modelInfo.path)
        
        // Create generator
        let newGenerator = StreamingGenerator(
            model: newModel,
            tokenizer: newTokenizer,
            engine: newEngine
        )
        
        // Update active instances
        engine = newEngine
        model = newModel
        tokenizer = newTokenizer
        generator = newGenerator
        sampler = Sampler()
        currentModelId = id
        models[id]?.loaded = true
    }
    
    /// Unload current model
    public func unloadModel() {
        guard let id = currentModelId else { return }
        
        engine = nil
        model = nil
        tokenizer = nil
        generator = nil
        sampler = nil
        currentModelId = nil
        models[id]?.loaded = false
    }
    
    /// Switch to a different model
    public func switchModel(to id: String) async throws {
        // Unload current model if loaded
        if currentModelId != nil {
            unloadModel()
        }
        
        // Load new model
        try await loadModel(id: id)
    }
    
    /// Get active engine
    public func getEngine() throws -> MarkBaseEngine {
        guard let engine = engine else {
            throw ModelError.noModelLoaded
        }
        return engine
    }
    
    /// Get active model
    public func getModel() throws -> E4BModel {
        guard let model = model else {
            throw ModelError.noModelLoaded
        }
        return model
    }
    
    /// Get active tokenizer
    public func getTokenizer() throws -> Tokenizer {
        guard let tokenizer = tokenizer else {
            throw ModelError.noModelLoaded
        }
        return tokenizer
    }
    
    /// Get active generator
    public func getGenerator() throws -> StreamingGenerator {
        guard let generator = generator else {
            throw ModelError.noModelLoaded
        }
        return generator
    }
    
    /// Get active sampler
    public func getSampler() throws -> Sampler {
        guard let sampler = sampler else {
            throw ModelError.noModelLoaded
        }
        return sampler
    }
}