use std::collections::HashMap;
use std::sync::{Arc, RwLock};

/// Trait all inference models must implement
trait Model: Send + Sync {
    fn name(&self) -> &str;
    fn version(&self) -> &str;
    fn infer(&self, input: &[f32]) -> Vec<f32>;
}

/// Simple embedding model
struct EmbedModel { dim: usize }
impl Model for EmbedModel {
    fn name(&self) -> &str { "embed" }
    fn version(&self) -> &str { "1.0.0" }
    fn infer(&self, input: &[f32]) -> Vec<f32> {
        let mut v = input.to_vec();
        v.resize(self.dim, 0.0);
        let norm: f32 = v.iter().map(|x| x * x).sum::<f32>().sqrt().max(1e-8);
        v.iter().map(|x| x / norm).collect()
    }
}

/// Sentiment classifier
struct SentimentModel;
impl Model for SentimentModel {
    fn name(&self) -> &str { "sentiment" }
    fn version(&self) -> &str { "2.1.0" }
    fn infer(&self, input: &[f32]) -> Vec<f32> {
        let score = input.iter().sum::<f32>() / input.len() as f32;
        let pos = 1.0 / (1.0 + (-score).exp()); // sigmoid
        vec![1.0 - pos, pos] // [negative, positive]
    }
}

/// Thread-safe registry — read-heavy, write-rarely
struct ModelRegistry {
    inner: Arc<RwLock<HashMap<String, Arc<dyn Model>>>>,
}

impl ModelRegistry {
    fn new() -> Self {
        Self { inner: Arc::new(RwLock::new(HashMap::new())) }
    }

    fn register(&self, model: Arc<dyn Model>) -> &Self {
        let key = format!("{}-{}", model.name(), model.version());
        self.inner.write().unwrap().insert(key, model);
        self
    }

    fn infer(&self, model_name: &str, input: &[f32]) -> Option<Vec<f32>> {
        let map = self.inner.read().unwrap();
        // Find by name prefix (any version)
        map.iter()
            .find(|(k, _)| k.starts_with(model_name))
            .map(|(_, m)| m.infer(input))
    }

    fn list(&self) -> Vec<String> {
        self.inner.read().unwrap().keys().cloned().collect()
    }
}

fn main() {
    let registry = ModelRegistry::new();
    registry
        .register(Arc::new(EmbedModel { dim: 8 }))
        .register(Arc::new(SentimentModel));

    println!("Available models:");
    for m in registry.list() { println!("  {}", m); }

    let text_vec = vec![0.5f32, -0.2, 0.8, -0.1];

    if let Some(embed) = registry.infer("embed", &text_vec) {
        println!("\nEmbedding (dim=8): {:?}", embed);
    }

    if let Some(scores) = registry.infer("sentiment", &text_vec) {
        println!("Sentiment [neg, pos]: {:.3}, {:.3}", scores[0], scores[1]);
    }
}