"""ModelManager singleton for managing transcription models.

This module provides a singleton ModelManager class that handles loading, caching,
and efficient management of transcription models used in the multi-pass pipeline.
Supports 8-bit quantization, memory management, and thread safety.
"""

import gc
import logging
import threading
from typing import Any, Dict, Optional, Union, List
from pathlib import Path

import psutil
import torch
from faster_whisper import WhisperModel
from transformers import AutoModelForSpeechSeq2Seq, AutoProcessor

logger = logging.getLogger(__name__)


class ModelManager:
    """Singleton class for managing transcription models.
    
    Provides efficient model loading, caching, and memory management with
    support for 8-bit quantization and thread safety.
    """
    
    _instance = None
    _lock = threading.Lock()
    
    def __new__(cls):
        """Ensure only one instance exists."""
        if cls._instance is None:
            with cls._lock:
                if cls._instance is None:
                    cls._instance = super(ModelManager, cls).__new__(cls)
                    cls._instance._initialize()
        return cls._instance
    
    def _initialize(self):
        """Initialize the ModelManager instance."""
        self.models: Dict[str, Any] = {}
        self.model_configs = {
            "fast_pass": {
                "model_id": "distil-small.en",
                "quantize": True,
                "compute_type": "int8",
                "device": "auto"
            },
            "refinement_pass": {
                "model_id": "distil-large-v3", 
                "quantize": True,
                "compute_type": "int8",
                "device": "auto"
            },
            "enhancement_pass": {
                "model_id": "distil-large-v3",
                "quantize": False,  # Full precision for enhancement
                "compute_type": "float16",
                "device": "auto"
            }
        }
        self._memory_threshold_mb = 6000  # 6GB threshold
        self._initialized = True
        logger.info("ModelManager singleton initialized")
    
    def load_model(self, model_key: str) -> Any:
        """Load a model by key with caching.
        
        Args:
            model_key: Key identifying the model to load
            
        Returns:
            Loaded model instance
            
        Raises:
            ValueError: If model_key is not recognized
        """
        if model_key not in self.model_configs:
            raise ValueError(f"Unknown model key: {model_key}")
        
        # Return cached model if available
        if model_key in self.models:
            logger.debug(f"Returning cached model: {model_key}")
            return self.models[model_key]
        
        # Check memory before loading
        self._check_memory_before_loading()
        
        config = self.model_configs[model_key]
        model_id = config["model_id"]
        quantize = config.get("quantize", False)
        compute_type = config.get("compute_type", "int8")
        device = config.get("device", "auto")
        
        logger.info(f"Loading model {model_id} with quantization={quantize}")
        
        try:
            if quantize:
                model = self._load_quantized_model(model_id, compute_type, device)
            else:
                model = self._load_full_precision_model(model_id, compute_type, device)
            
            self.models[model_key] = model
            logger.info(f"Successfully loaded model: {model_key}")
            return model
            
        except Exception as e:
            logger.error(f"Failed to load model {model_key}: {e}")
            raise
    
    def _load_quantized_model(self, model_id: str, compute_type: str, device: str) -> WhisperModel:
        """Load a quantized model using faster-whisper.
        
        Args:
            model_id: Model identifier
            compute_type: Compute type (int8, float16, etc.)
            device: Device to load on (auto, cpu, cuda)
            
        Returns:
            Loaded WhisperModel instance
        """
        return WhisperModel(
            model_id,
            device=device,
            compute_type=compute_type,
            download_root=None,  # Use default cache directory
            local_files_only=False
        )
    
    def _load_full_precision_model(self, model_id: str, compute_type: str, device: str) -> WhisperModel:
        """Load a full precision model using faster-whisper.
        
        Args:
            model_id: Model identifier
            compute_type: Compute type (float16, float32)
            device: Device to load on (auto, cpu, cuda)
            
        Returns:
            Loaded WhisperModel instance
        """
        return WhisperModel(
            model_id,
            device=device,
            compute_type=compute_type,
            download_root=None,
            local_files_only=False
        )
    
    def unload_model(self, model_key: str) -> None:
        """Unload a model and free memory.
        
        Args:
            model_key: Key of the model to unload
        """
        if model_key in self.models:
            logger.info(f"Unloading model: {model_key}")
            
            # Properly release model resources
            model = self.models[model_key]
            if hasattr(model, 'model'):
                del model.model
            if hasattr(model, 'processor'):
                del model.processor
            
            del self.models[model_key]
            
            # Force garbage collection
            gc.collect()
            
            # Clear CUDA cache if available
            if torch.cuda.is_available():
                torch.cuda.empty_cache()
            
            logger.info(f"Successfully unloaded model: {model_key}")
    
    def unload_all_models(self) -> None:
        """Unload all models and free memory."""
        logger.info("Unloading all models")
        model_keys = list(self.models.keys())
        for model_key in model_keys:
            self.unload_model(model_key)
        logger.info("All models unloaded")
    
    def get_memory_usage(self) -> Dict[str, float]:
        """Get current memory usage statistics.
        
        Returns:
            Dictionary with memory usage in MB
        """
        process = psutil.Process()
        memory_info = process.memory_info()
        
        memory_stats = {
            "rss_mb": memory_info.rss / (1024 * 1024),  # Resident Set Size
            "vms_mb": memory_info.vms / (1024 * 1024),  # Virtual Memory Size
            "percent": process.memory_percent()
        }
        
        # Add CUDA memory if available
        if torch.cuda.is_available():
            memory_stats["cuda_allocated_mb"] = torch.cuda.memory_allocated() / (1024 * 1024)
            memory_stats["cuda_reserved_mb"] = torch.cuda.memory_reserved() / (1024 * 1024)
        
        return memory_stats
    
    def _check_memory_before_loading(self) -> None:
        """Check memory usage before loading a new model."""
        memory_stats = self.get_memory_usage()
        rss_mb = memory_stats["rss_mb"]
        
        if rss_mb > self._memory_threshold_mb:
            logger.warning(f"High memory usage detected: {rss_mb:.1f}MB")
            
            # Try to free memory by unloading least recently used models
            if len(self.models) > 1:
                # For now, unload the first model (simple LRU)
                first_model_key = next(iter(self.models))
                logger.info(f"Unloading model {first_model_key} to free memory")
                self.unload_model(first_model_key)
    
    def set_model_config(self, model_key: str, config: Dict[str, Any]) -> None:
        """Update model configuration.
        
        Args:
            model_key: Key of the model to update
            config: New configuration dictionary
        """
        if model_key not in self.model_configs:
            raise ValueError(f"Unknown model key: {model_key}")
        
        logger.info(f"Updating config for model: {model_key}")
        
        # Update configuration
        self.model_configs[model_key].update(config)
        
        # If model is already loaded, reload with new config
        if model_key in self.models:
            logger.info(f"Reloading model {model_key} with new config")
            self.unload_model(model_key)
            self.load_model(model_key)
    
    def get_loaded_models(self) -> List[str]:
        """Get list of currently loaded model keys.
        
        Returns:
            List of loaded model keys
        """
        return list(self.models.keys())
    
    def is_model_loaded(self, model_key: str) -> bool:
        """Check if a model is currently loaded.
        
        Args:
            model_key: Key of the model to check
            
        Returns:
            True if model is loaded, False otherwise
        """
        return model_key in self.models
    
    def get_model_info(self, model_key: str) -> Optional[Dict[str, Any]]:
        """Get information about a model.
        
        Args:
            model_key: Key of the model
            
        Returns:
            Model information dictionary or None if not found
        """
        if model_key not in self.model_configs:
            return None
        
        info = {
            "config": self.model_configs[model_key].copy(),
            "loaded": model_key in self.models,
            "memory_usage": self.get_memory_usage()
        }
        
        return info

    def list_available_models(self) -> Dict[str, List[str]]:
        """List all available models by category.
        
        Returns:
            Dictionary mapping model categories to lists of available models
        """
        return {
            "whisper": ["distil-small.en", "distil-large-v3", "base", "small", "medium", "large"],
            "diarization": ["pyannote/speaker-diarization@2.1"],
            "domain": ["general", "technical", "medical", "academic"]
        }

    def download_model(self, model_type: str, model_name: str) -> None:
        """Download a specific model.
        
        Args:
            model_type: Type of model (whisper, diarization, domain)
            model_name: Name of the model to download
        """
        if model_type == "whisper":
            # For Whisper models, we just need to load them (they'll be downloaded automatically)
            if model_name in ["distil-small.en", "distil-large-v3", "base", "small", "medium", "large"]:
                logger.info(f"Whisper model {model_name} will be downloaded automatically when first used")
            else:
                raise ValueError(f"Unknown Whisper model: {model_name}")
        elif model_type == "diarization":
            if model_name == "pyannote/speaker-diarization@2.1":
                logger.info(f"Diarization model {model_name} will be downloaded automatically when first used")
            else:
                raise ValueError(f"Unknown diarization model: {model_name}")
        elif model_type == "domain":
            if model_name in ["general", "technical", "medical", "academic"]:
                logger.info(f"Domain adapter {model_name} will be downloaded automatically when first used")
            else:
                raise ValueError(f"Unknown domain: {model_name}")
        else:
            raise ValueError(f"Unknown model type: {model_type}")
    
    def __repr__(self) -> str:
        """String representation of ModelManager."""
        loaded_models = list(self.models.keys())
        memory_stats = self.get_memory_usage()
        
        return (
            f"ModelManager(loaded_models={loaded_models}, "
            f"memory_mb={memory_stats['rss_mb']:.1f})"
        )


# Global instance for easy access
model_manager = ModelManager()
