"""
Model Manager for preloading and managing AI models in memory
"""

import torch
import numpy as np
from transformers import pipeline
import time
from typing import Optional, Dict, Any
import gc

class ModelManager:
    """Singleton class to manage preloaded models"""
    _instance = None
    
    def __new__(cls):
        if cls._instance is None:
            cls._instance = super(ModelManager, cls).__new__(cls)
            cls._instance._initialized = False
        return cls._instance
    
    def __init__(self):
        if self._initialized:
            return
        
        self.depth_estimator = None
        self.upscaler = None
        self.models_loaded = False
        self.device = "cuda" if torch.cuda.is_available() else "cpu"
        self.model_config = {
            "depth_model": "Intel/dpt-hybrid-midas",  # Lighter model for faster loading
            "upscale_model": None,  # Placeholder for future upscaling models
        }
        self._initialized = True
    
    def load_models_sync(self):
        """Load models synchronously to avoid threading issues"""
        if self.models_loaded:
            return True
        
        try:
            print("🔄 Loading AI models...")
            start_time = time.time()
            
            # Load depth estimation model
            print("Loading depth estimation model...")
            self.depth_estimator = pipeline(
                "depth-estimation",
                model=self.model_config["depth_model"],
                device=self.device,
                torch_dtype=torch.float16 if self.device == "cuda" else torch.float32
            )
            
            # Load upscaler model (placeholder for now)
            self.upscaler = None
            
            self.models_loaded = True
            load_time = time.time() - start_time
            print(f"✅ Models loaded successfully in {load_time:.2f}s")
            return True
            
        except Exception as e:
            print(f"❌ Error loading models: {e}")
            self.models_loaded = False
            return False
    
    def wait_for_models(self, timeout: Optional[float] = None) -> bool:
        """Wait for models to be loaded"""
        if self.models_loaded:
            return True
        
        return self.load_models_sync()
    
    def get_depth_estimator(self):
        """Get the depth estimation model"""
        if not self.models_loaded:
            self.load_models_sync()
        return self.depth_estimator
    
    def get_upscaler(self):
        """Get the upscaler model"""
        if not self.models_loaded:
            self.load_models_sync()
        return self.upscaler
    
    def is_ready(self) -> bool:
        """Check if models are ready"""
        return self.models_loaded
    
    def get_model_info(self) -> Dict[str, Any]:
        """Get information about loaded models"""
        return {
            "models_loaded": self.models_loaded,
            "device": self.device,
            "depth_model": self.model_config["depth_model"],
            "upscale_model": self.model_config["upscale_model"]
        }
    
    def cleanup(self):
        """Clean up models and free memory"""
        if self.depth_estimator is not None:
            del self.depth_estimator
            self.depth_estimator = None
        
        if self.upscaler is not None:
            del self.upscaler
            self.upscaler = None
        
        self.models_loaded = False
        gc.collect()
        
        if torch.cuda.is_available():
            torch.cuda.empty_cache()

# Global instance
model_manager = ModelManager()