# modelManager.py - Model Cache Manager
# PRODUCTION HARDENED

import time
import gc
import psutil
from typing import Dict, Optional, Any

class ModelManager:
    """Memory-aware model management with LRU cache."""
    
    def __init__(self, max_models: int = 3, memory_limit_mb: int = 6000):
        self.models: Dict[str, Any] = {}
        self.load_times: Dict[str, float] = {}
        self.use_counts: Dict[str, int] = {}
        self.max_models = max_models
        self.memory_limit_mb = memory_limit_mb
    
    def get_model(self, name: str, loader_func) -> Optional[Any]:
        """Get or load a model with caching."""
        # Check if already loaded
        if name in self.models:
            self.use_counts[name] = self.use_counts.get(name, 0) + 1
            return self.models[name]
        
        # Check memory before loading
        try:
            mem = psutil.virtual_memory()
            if mem.percent > 85:
                self._evict_models()
        except:
            pass
        
        # Load model
        start = time.time()
        try:
            session = loader_func(name)
            load_time = time.time() - start
            self.models[name] = session
            self.load_times[name] = load_time
            self.use_counts[name] = 1
            
            # Evict if over limit
            if len(self.models) > self.max_models:
                self._evict_models()
            
            return session
        except Exception as e:
            return None
    
    def _evict_models(self):
        """Evict least recently used models."""
        if len(self.models) <= 1:
            return
        
        sorted_models = sorted(
            self.models.keys(),
            key=lambda x: self.use_counts.get(x, 0)
        )
        
        to_remove = sorted_models[0]
        del self.models[to_remove]
        if to_remove in self.use_counts:
            del self.use_counts[to_remove]
        gc.collect()
    
    def get_stats(self) -> Dict:
        return {
            'loaded_models': list(self.models.keys()),
            'model_count': len(self.models),
            'use_counts': self.use_counts,
            'load_times': self.load_times
        }