//! Model registry for multi-model serving

//!

//! Provides a central registry to manage multiple models in production environments.

//! Supports dynamic model loading/unloading, caching, and thread-safe concurrent access.

//!

//! ## Features

//!

//! - Thread-safe model registry with concurrent access

//! - Integration with `ModelCache` for efficient memory management

//! - Support for multiple model formats (GGUF, Safetensors)

//! - Model metadata and configuration

//! - Graceful loading/unloading with error handling

//!

//! ## Example

//!

//! ```rust,ignore

//! use realizar::registry::{ModelRegistry, ModelConfig};

//!

//! let registry = ModelRegistry::new(5); // Max 5 models cached

//! registry.register("llama-7b", model, tokenizer)?;

//! let model = registry.get("llama-7b")?;

//! ```

use std::{

    collections::HashMap,

    sync::{Arc, Mutex},

};

use arc_swap::ArcSwap;

use crate::{

    cache::ModelCache,

    error::{RealizarError, Result},

    layers::Model,

    tokenizer::BPETokenizer,

};

/// Information about a registered model

#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]

pub struct ModelInfo {

    /// Unique model identifier

    pub id: String,

    /// Human-readable model name

    pub name: String,

    /// Model description

    pub description: String,

    /// Model format (GGUF, Safetensors, etc.)

    pub format: String,

    /// Whether the model is currently loaded

    pub loaded: bool,

}

/// Entry in the model registry

#[derive(Clone)]

struct ModelEntry {

    /// Model instance

    model: Arc<Model>,

    /// Tokenizer for this model

    tokenizer: Arc<BPETokenizer>,

    /// Model metadata

    info: ModelInfo,

}

/// Type alias for the models map (immutable snapshot)

type ModelsMap = HashMap<String, ModelEntry>;

/// Type alias for model and tokenizer tuple

type ModelTuple = (Arc<Model>, Arc<BPETokenizer>);

/// Central registry for managing multiple models

///

/// The `ModelRegistry` provides thread-safe access to multiple models,

/// with automatic caching and lifecycle management.

///

/// Uses `ArcSwap` for lock-free reads (per `McKenney` 2011).

pub struct ModelRegistry {

    /// Registry of loaded models (lock-free reads via `ArcSwap`)

    models: ArcSwap<ModelsMap>,

    /// Write lock to serialize modifications

    write_lock: Mutex<()>,

    /// Model cache for memory management (reserved for future use)

    #[allow(dead_code)]

    cache: Arc<ModelCache>,

}

impl ModelRegistry {

    /// Create a new model registry

    ///

    /// # Arguments

    ///

    /// * `cache_capacity` - Maximum number of models to keep in cache

    #[must_use]

    pub fn new(cache_capacity: usize) -> Self {

        Self {

            models: ArcSwap::from_pointee(HashMap::new()),

            write_lock: Mutex::new(()),

            cache: Arc::new(ModelCache::new(cache_capacity)),

        }

    }

    /// Register a new model

    ///

    /// # Arguments

    ///

    /// * `id` - Unique identifier for the model

    /// * `model` - Model instance

    /// * `tokenizer` - Tokenizer for the model

    ///

    /// # Errors

    ///

    /// Returns error if model ID already exists

    pub fn register(&self, id: &str, model: Model, tokenizer: BPETokenizer) -> Result<()> {

        let _guard = self.write_lock.lock().map_err(|_| 
{0

            RealizarError::RegistryError("Failed to acquire write lock".to_string())

        })?;

        let current = self.models.load();

        if current.contains_key(id) {

            return Err(RealizarError::ModelAlreadyExists(id.to_string()));

        }

        let entry = ModelEntry {

            model: Arc::new(model),

            tokenizer: Arc::new(tokenizer),

            info: ModelInfo {

                id: id.to_string(),

                name: id.to_string(),

                description: String::new(),

                format: "unknown".to_string(),

                loaded: true,

            },

        };

        let mut new_map: ModelsMap = (**current).clone();

        new_map.insert(id.to_string(), entry);

        self.models.store(Arc::new(new_map));

        Ok(())

    }

    /// Register a model with full metadata

    ///

    /// # Arguments

    ///

    /// * `info` - Model metadata

    /// * `model` - Model instance

    /// * `tokenizer` - Tokenizer for the model

    ///

    /// # Errors

    ///

    /// Returns error if model ID already exists

    pub fn register_with_info(

        &self,

        mut info: ModelInfo,

        model: Model,

        tokenizer: BPETokenizer,

    ) -> Result<()> {

        let _guard = self.write_lock.lock().map_err(|_| 
{0

            RealizarError::RegistryError("Failed to acquire write lock".to_string())

        })?;

        let current = self.models.load();

        if current.contains_key(&info.id) {

            return Err(RealizarError::ModelAlreadyExists(info.id));

        }

        info.loaded = true;

        let entry = ModelEntry {

            model: Arc::new(model),

            tokenizer: Arc::new(tokenizer),

            info,

        };

        let id = entry.info.id.clone();

        let mut new_map: ModelsMap = (**current).clone();

        new_map.insert(id, entry);

        self.models.store(Arc::new(new_map));

        Ok(())

    }

    /// Get a model by ID

    ///

    /// # Arguments

    ///

    /// * `id` - Model identifier

    ///

    /// # Errors

    ///

    /// Returns error if model not found

    pub fn get(&self, id: &str) -> Result<ModelTuple> {

        // Lock-free read via ArcSwap::load()

        let models = self.models.load();

        let 
entry4
 = models

            .get(id)

            .ok_or_else(|| RealizarError::ModelNotFound(
id1
.
to_string1
()))
?1
;

        Ok((Arc::clone(&entry.model), Arc::clone(&entry.tokenizer)))

    }

    /// Get model info by ID

    ///

    /// # Arguments

    ///

    /// * `id` - Model identifier

    ///

    /// # Errors

    ///

    /// Returns error if model not found

    pub fn get_info(&self, id: &str) -> Result<ModelInfo> {

        let models = self.models.load();

        let 
entry1
 = models

            .get(id)

            .ok_or_else(|| RealizarError::ModelNotFound(
id1
.
to_string1
()))
?1
;

        Ok(entry.info.clone())

    }

    /// List all registered models (lock-free)

    #[must_use]

    pub fn list(&self) -> Vec<ModelInfo> {

        let models = self.models.load();

        
models.values()1
.
map1
(|entry| entry.info.clone()).
collect1
()

    }

    /// Unregister a model

    ///

    /// # Arguments

    ///

    /// * `id` - Model identifier

    ///

    /// # Errors

    ///

    /// Returns error if model not found

    pub fn unregister(&self, id: &str) -> Result<()> {

        let _guard = self.write_lock.lock().map_err(|_| 
{0

            RealizarError::RegistryError("Failed to acquire write lock".to_string())

        })?;

        let current = self.models.load();

        if !current.contains_key(id) {

            return Err(RealizarError::ModelNotFound(id.to_string()));

        }

        let mut new_map: ModelsMap = (**current).clone();

        new_map.remove(id);

        self.models.store(Arc::new(new_map));

        Ok(())

    }

    /// Atomically replace a model (for hot-reload)

    ///

    /// # Arguments

    ///

    /// * `id` - Model identifier to replace

    /// * `model` - New model instance

    /// * `tokenizer` - New tokenizer instance

    ///

    /// # Errors

    ///

    /// Returns error if model not found or lock acquisition fails

    pub fn replace(&self, id: &str, model: Model, tokenizer: BPETokenizer) -> Result<()> {

        let _guard = self.write_lock.lock().map_err(|_| 
{0

            RealizarError::RegistryError("Failed to acquire write lock".to_string())

        })?;

        let current = self.models.load();

        if !current.contains_key(id) {

            return Err(RealizarError::ModelNotFound(id.to_string()));

        }

        // Get existing info to preserve metadata

        let existing_info = current.get(id).map_or_else(

            || ModelInfo {

                id: id.to_string(),

                name: id.to_string(),

                description: String::new(),

                format: "unknown".to_string(),

                loaded: true,

            },

            |e| e.info.clone(),

        );

        let entry = ModelEntry {

            model: Arc::new(model),

            tokenizer: Arc::new(tokenizer),

            info: existing_info,

        };

        let mut new_map: ModelsMap = (**current).clone();

        new_map.insert(id.to_string(), entry);

        self.models.store(Arc::new(new_map));

        Ok(())

    }

    /// Check if a model is registered (lock-free)

    #[must_use]

    pub fn contains(&self, id: &str) -> bool {

        let models = self.models.load();

        models.contains_key(id)

    }

    /// Get the number of registered models (lock-free)

    #[must_use]

    pub fn len(&self) -> usize {

        let models = self.models.load();

        models.len()

    }

    /// Check if the registry is empty (lock-free)

    #[must_use]

    pub fn is_empty(&self) -> bool {

        self.len() == 0

    }

}

#[cfg(test)]

mod tests {

    use super::*;

    use crate::layers::ModelConfig;

    fn create_test_model() -> (Model, BPETokenizer) {

        let config = ModelConfig {

            vocab_size: 100,

            hidden_dim: 32,

            num_heads: 1,

            num_layers: 1,

            intermediate_dim: 64,

            eps: 1e-5,

        };

        let model = Model::new(config).expect("test");

        let vocab: Vec<String> = (0..100)

            .
map21
(|i| {

                if i == 0 {

                    "<unk>".to_string()

                } else {

                    format!("token{i}")

                }

            })

            .collect();

        let tokenizer = BPETokenizer::new(vocab, vec![], "<unk>").expect("test");

        (model, tokenizer)

    }

    #[test]

    fn test_registry_creation() {

        let registry = ModelRegistry::new(5);

        assert_eq!(registry.len(), 0);

        assert!(registry.is_empty());

    }

    #[test]

    fn test_register_model() {

        let registry = ModelRegistry::new(5);

        let (model, tokenizer) = create_test_model();

        registry

            .register("test-model", model, tokenizer)

            .expect("test");

        assert_eq!(registry.len(), 1);

        assert!(!registry.is_empty());

        assert!(registry.contains("test-model"));

    }

    #[test]

    fn test_register_duplicate_error() {

        let registry = ModelRegistry::new(5);

        let (model1, tokenizer1) = create_test_model();

        let (model2, tokenizer2) = create_test_model();

        registry

            .register("test-model", model1, tokenizer1)

            .expect("test");

        let result = registry.register("test-model", model2, tokenizer2);

        assert!(result.is_err());

        assert_eq!(registry.len(), 1);

    }

    #[test]

    fn test_get_model() {

        let registry = ModelRegistry::new(5);

        let (model, tokenizer) = create_test_model();

        registry

            .register("test-model", model, tokenizer)

            .expect("test");

        let (retrieved_model, retrieved_tokenizer) = registry.get("test-model").expect("test");

        assert!(Arc::strong_count(&retrieved_model) >= 2); // Registry + local

        assert!(Arc::strong_count(&retrieved_tokenizer) >= 2);

    }

    #[test]

    fn test_get_nonexistent_model() {

        let registry = ModelRegistry::new(5);

        let result = registry.get("nonexistent");

        assert!(result.is_err());

    }

    #[test]

    fn test_register_with_info() {

        let registry = ModelRegistry::new(5);

        let (model, tokenizer) = create_test_model();

        let info = ModelInfo {

            id: "llama-7b".to_string(),

            name: "Llama 7B".to_string(),

            description: "7B parameter Llama model".to_string(),

            format: "GGUF".to_string(),

            loaded: false,

        };

        registry

            .register_with_info(info.clone(), model, tokenizer)

            .expect("test");

        let retrieved_info = registry.get_info("llama-7b").expect("test");

        assert_eq!(retrieved_info.id, "llama-7b");

        assert_eq!(retrieved_info.name, "Llama 7B");

        assert_eq!(retrieved_info.description, "7B parameter Llama model");

        assert_eq!(retrieved_info.format, "GGUF");

        assert!(retrieved_info.loaded); // Should be set to true

    }

    #[test]

    fn test_list_models() {

        let registry = ModelRegistry::new(5);

        let (model1, tokenizer1) = create_test_model();

        let (model2, tokenizer2) = create_test_model();

        registry

            .register("model-1", model1, tokenizer1)

            .expect("test");

        registry

            .register("model-2", model2, tokenizer2)

            .expect("test");

        let model_list = registry.list();

        assert_eq!(model_list.len(), 2);

        let 
ids1
: 
Vec<String>1
 = 
model_list.iter()1
.
map1
(|m| m.id.clone()).
collect1
();

        assert!(ids.contains(&"model-1".to_string()));

        assert!(ids.contains(&"model-2".to_string()));

    }

    #[test]

    fn test_unregister_model() {

        let registry = ModelRegistry::new(5);

        let (model, tokenizer) = create_test_model();

        registry

            .register("test-model", model, tokenizer)

            .expect("test");

        assert_eq!(registry.len(), 1);

        registry.unregister("test-model").expect("test");

        assert_eq!(registry.len(), 0);

        assert!(!registry.contains("test-model"));

    }

    #[test]

    fn test_unregister_nonexistent() {

        let registry = ModelRegistry::new(5);

        let result = registry.unregister("nonexistent");

        assert!(result.is_err());

    }

    #[test]

    fn test_concurrent_access() {

        use std::thread;

        let registry = Arc::new(ModelRegistry::new(10));

        let mut handles = vec![];

        // Register models from multiple threads

        for 
i5
 in 0..5 {

            let registry_clone = Arc::clone(&registry);

            let handle = thread::spawn(move || {

                let (model, tokenizer) = create_test_model();

                registry_clone

                    .register(&format!("model-{i}"), model, tokenizer)

                    .expect("test");

            });

            handles.push(handle);

        }

        for 
handle5
 in handles {

            handle.join().expect("test");

        }

        assert_eq!(registry.len(), 5);

    }

    #[test]

    fn test_multiple_get_same_model() {

        let registry = ModelRegistry::new(5);

        let (model, tokenizer) = create_test_model();

        registry

            .register("test-model", model, tokenizer)

            .expect("test");

        // Get the same model multiple times

        let (model1, _) = registry.get("test-model").expect("test");

        let (model2, _) = registry.get("test-model").expect("test");

        // Both should point to the same underlying model

        assert!(Arc::ptr_eq(&model1, &model2));

    }

    #[test]

    fn test_replace_model() {

        let registry = ModelRegistry::new(5);

        let (model1, tokenizer1) = create_test_model();

        let (model2, tokenizer2) = create_test_model();

        // Register initial model

        registry

            .register("test-model", model1, tokenizer1)

            .expect("test");

        assert_eq!(registry.len(), 1);

        // Replace with new model

        registry

            .replace("test-model", model2, tokenizer2)

            .expect("test");

        assert_eq!(registry.len(), 1);

        // Verify replacement worked

        let (retrieved, _) = registry.get("test-model").expect("test");

        assert!(Arc::strong_count(&retrieved) >= 2);

    }

    #[test]

    fn test_replace_nonexistent_model() {

        let registry = ModelRegistry::new(5);

        let (model, tokenizer) = create_test_model();

        // Try to replace a model that doesn't exist

        let result = registry.replace("nonexistent", model, tokenizer);

        assert!(result.is_err());

    }

    #[test]

    fn test_contains_method() {

        let registry = ModelRegistry::new(5);

        let (model, tokenizer) = create_test_model();

        assert!(!registry.contains("test-model"));

        registry

            .register("test-model", model, tokenizer)

            .expect("test");

        assert!(registry.contains("test-model"));

    }

    #[test]

    fn test_is_empty_method() {

        let registry = ModelRegistry::new(5);

        assert!(registry.is_empty());

        let (model, tokenizer) = create_test_model();

        registry

            .register("test-model", model, tokenizer)

            .expect("test");

        assert!(!registry.is_empty());

    }

    #[test]

    fn test_get_info_nonexistent() {

        let registry = ModelRegistry::new(5);

        let result = registry.get_info("nonexistent");

        assert!(result.is_err());

    }

    #[test]

    fn test_len_method() {

        let registry = ModelRegistry::new(5);

        assert_eq!(registry.len(), 0);

        let (model1, tokenizer1) = create_test_model();

        registry

            .register("model-1", model1, tokenizer1)

            .expect("test");

        assert_eq!(registry.len(), 1);

        let (model2, tokenizer2) = create_test_model();

        registry

            .register("model-2", model2, tokenizer2)

            .expect("test");

        assert_eq!(registry.len(), 2);

    }

}