"""
Model Manager - Centralized model loading and caching for ChatterBox Voice
Handles model discovery, loading, and caching across different sources
"""

import os
import warnings
import torch
import folder_paths
from typing import Optional, List, Tuple, Dict, Any
# Import extra paths support
from utils.models.extra_paths import get_all_tts_model_paths, find_model_in_paths

# Import ComfyUI model wrapper for integration
from utils.models.comfyui_model_wrapper import tts_model_manager

# --- LAZY ENGINE IMPORTS ---
# Engine imports are deferred to first use to avoid pulling in heavy dependencies
# (torch._dynamo, diffusers, transformers, librosa) at module load time.
# This reduces plugin startup from ~8s to <1s by not importing engine code
# until a model is actually loaded.
#
# The _lazy_* helpers below cache the result of the first import attempt
# so subsequent accesses are instant.

_engine_import_cache = {}

def _lazy_import_chatterbox_tts():
    """Lazily import ChatterboxTTS on first use, caching the result."""
    key = "chatterbox_tts"
    if key not in _engine_import_cache:
        from utils.system.import_manager import import_manager
        success, cls, source = import_manager.import_chatterbox_tts()
        _engine_import_cache[key] = (success, cls, source)
    return _engine_import_cache[key]

def _lazy_import_chatterbox_vc():
    """Lazily import ChatterboxVC on first use, caching the result."""
    key = "chatterbox_vc"
    if key not in _engine_import_cache:
        from utils.system.import_manager import import_manager
        success, cls, source = import_manager.import_chatterbox_vc()
        _engine_import_cache[key] = (success, cls, source)
    return _engine_import_cache[key]

def _lazy_import_f5tts():
    """Lazily import F5-TTS on first use, caching the result."""
    key = "f5tts"
    if key not in _engine_import_cache:
        from utils.system.import_manager import import_manager
        success, cls, source = import_manager.import_f5tts()
        _engine_import_cache[key] = (success, cls, source)
    return _engine_import_cache[key]

def _lazy_import_vibevoice():
    """Lazily import VibeVoice on first use, caching the result."""
    key = "vibevoice"
    if key not in _engine_import_cache:
        try:
            from vibevoice.modular.modeling_vibevoice_inference import VibeVoiceForConditionalGenerationInference
            from vibevoice.processor.vibevoice_processor import VibeVoiceProcessor
            from engines.vibevoice_engine import VibeVoiceDownloader
            _engine_import_cache[key] = (True, VibeVoiceForConditionalGenerationInference, VibeVoiceProcessor, VibeVoiceDownloader)
        except ImportError:
            _engine_import_cache[key] = (False, None, None, None)
    return _engine_import_cache[key]


# Availability flags are now properties that lazily probe on first access.
# Most code paths only check these when actually loading a model, so
# deferring is safe and dramatically speeds up startup.

def _get_chatterbox_tts_available():
    return _lazy_import_chatterbox_tts()[0]

def _get_chatterbox_vc_available():
    return _lazy_import_chatterbox_vc()[0]

def _get_f5tts_available():
    return _lazy_import_f5tts()[0]

def _get_vibevoice_available():
    return _lazy_import_vibevoice()[0]

def _get_using_bundled_chatterbox():
    _, _, tts_source = _lazy_import_chatterbox_tts()
    _, _, vc_source = _lazy_import_chatterbox_vc()
    return tts_source == "bundled" or vc_source == "bundled"

# Backwards-compatible module-level names.
# These are checked at model-load time, not at import time, so the lazy
# evaluation is transparent to callers.
# NOTE: We use a module-level __getattr__ to make these truly lazy.
def __getattr__(name):
    """Module-level __getattr__ for lazy availability flags.

    This avoids importing heavy engine code until someone actually checks
    whether an engine is available (typically at model-load time, not startup).
    """
    if name == "CHATTERBOX_TTS_AVAILABLE":
        return _get_chatterbox_tts_available()
    elif name == "CHATTERBOX_VC_AVAILABLE":
        return _get_chatterbox_vc_available()
    elif name == "F5TTS_AVAILABLE":
        return _get_f5tts_available()
    elif name == "VIBEVOICE_AVAILABLE":
        return _get_vibevoice_available()
    elif name == "USING_BUNDLED_CHATTERBOX":
        return _get_using_bundled_chatterbox()
    elif name == "ChatterboxTTS":
        return _lazy_import_chatterbox_tts()[1]
    elif name == "ChatterboxVC":
        return _lazy_import_chatterbox_vc()[1]
    elif name == "F5TTS":
        return _lazy_import_f5tts()[1]
    elif name == "VibeVoiceForConditionalGenerationInference":
        return _lazy_import_vibevoice()[1]
    elif name == "VibeVoiceProcessor":
        return _lazy_import_vibevoice()[2]
    elif name == "VibeVoiceDownloader":
        return _lazy_import_vibevoice()[3]
    raise AttributeError(f"module {__name__!r} has no attribute {name!r}")


class ModelManager:
    """
    Centralized model loading and caching manager for ChatterBox Voice.
    Handles model discovery, loading from different sources, and caching.
    """
    
    # Class-level cache for shared model instances
    _model_cache: Dict[str, Any] = {}
    _model_sources: Dict[str, str] = {}
    
    def __init__(self, node_dir: Optional[str] = None):
        """
        Initialize ModelManager with optional node directory override.
        
        Args:
            node_dir: Optional override for the node directory path
        """
        self.node_dir = node_dir or os.path.dirname(os.path.dirname(__file__))
        self.bundled_chatterbox_dir = os.path.join(self.node_dir, "chatterbox")
        self.bundled_models_dir = os.path.join(self.node_dir, "models", "chatterbox")
        
        # Instance-level model references
        self.tts_model: Optional[Any] = None
        self.vc_model: Optional[Any] = None
        self.current_device: Optional[str] = None
    
    def find_chatterbox_models(self) -> List[Tuple[str, Optional[str]]]:
        """
        Find ChatterBox model files in order of priority.
        
        Returns:
            List of tuples containing (source_type, path) in priority order:
            - bundled: Models bundled with the extension
            - comfyui: Models in ComfyUI models directory
            - huggingface: Download from Hugging Face (path is None)
        """
        model_paths = []
        
        # 1. Check for bundled models in node folder
        bundled_model_path = os.path.join(self.bundled_models_dir, "s3gen.pt")
        if os.path.exists(bundled_model_path):
            model_paths.append(("bundled", self.bundled_models_dir))
            return model_paths  # Return immediately if bundled models found
        
        # 2. Check configured TTS model paths (extra_model_paths.yaml aware)
        tts_model_paths = get_all_tts_model_paths('TTS')
        
        for base_tts_path in tts_model_paths:
            comfyui_tts_dir = os.path.join(base_tts_path, "chatterbox")
            if not os.path.exists(comfyui_tts_dir):
                continue
            
            # Check for direct s3gen.pt in chatterbox folder
            direct_model = os.path.join(comfyui_tts_dir, "s3gen.pt")
            if os.path.exists(direct_model):
                model_paths.append(("comfyui", comfyui_tts_dir))
                return model_paths
            
            # Check for language subdirectories (English, German, etc.)
            try:
                for item in os.listdir(comfyui_tts_dir):
                    item_path = os.path.join(comfyui_tts_dir, item)
                    if os.path.isdir(item_path):
                        # Check if this subdirectory contains ChatterBox model files
                        required_files = ["s3gen.", "ve.", "t3_cfg.", "tokenizer.json"]
                        has_model = False
                        for file in os.listdir(item_path):
                            for required in required_files:
                                if file.startswith(required) and (file.endswith(".pt") or file.endswith(".safetensors")):
                                    has_model = True
                                    break
                            if has_model:
                                break
                        if has_model:
                            model_paths.append(("comfyui", comfyui_tts_dir))
                            return model_paths
            except OSError:
                pass
        
        # 3. Check legacy location (direct chatterbox) for backward compatibility
        comfyui_legacy_dir = os.path.join(folder_paths.models_dir, "chatterbox")
        if os.path.exists(comfyui_legacy_dir):
            # Check for direct s3gen.pt in chatterbox folder
            direct_model = os.path.join(comfyui_legacy_dir, "s3gen.pt")
            if os.path.exists(direct_model):
                model_paths.append(("comfyui", comfyui_legacy_dir))
                return model_paths
            
            # Check for language subdirectories
            try:
                for item in os.listdir(comfyui_legacy_dir):
                    item_path = os.path.join(comfyui_legacy_dir, item)
                    if os.path.isdir(item_path):
                        # Check if this subdirectory contains ChatterBox model files
                        required_files = ["s3gen.", "ve.", "t3_cfg.", "tokenizer.json"]
                        has_model = False
                        for file in os.listdir(item_path):
                            for required in required_files:
                                if file.startswith(required) and (file.endswith(".pt") or file.endswith(".safetensors")):
                                    has_model = True
                                    break
                            if has_model:
                                break
                        if has_model:
                            model_paths.append(("comfyui", comfyui_legacy_dir))
                            return model_paths
            except OSError:
                pass
        
        # 4. HuggingFace download as fallback
        model_paths.append(("huggingface", None))
        
        return model_paths
    
    def find_local_language_model(self, language: str) -> Optional[str]:
        """
        Find local ChatterBox model for a specific language.
        
        Args:
            language: Language to find model for
            
        Returns:
            Path to local model directory if found, None otherwise
        """
        # Import language models functionality
        try:
            from engines.chatterbox.language_models import find_local_model_path
            return find_local_model_path(language)
        except ImportError:
            # Fallback: check standard locations manually
            language_paths = [
                os.path.join(folder_paths.models_dir, "TTS", "chatterbox", language),
                os.path.join(folder_paths.models_dir, "TTS", "chatterbox", language.lower()),
                os.path.join(folder_paths.models_dir, "chatterbox", language),  # Legacy
                os.path.join(folder_paths.models_dir, "chatterbox", language.lower()),  # Legacy
                os.path.join(self.bundled_models_dir, language),
                os.path.join(self.bundled_models_dir, language.lower())
            ]
            
            for path in language_paths:
                if os.path.exists(path):
                    # Check if it contains the required model files
                    required_files = ["ve.", "t3_cfg.", "s3gen.", "tokenizer.json"]
                    has_all_files = True
                    
                    for required in required_files:
                        found = False
                        for ext in [".pt", ".safetensors"]:
                            if required == "tokenizer.json":
                                if os.path.exists(os.path.join(path, required)):
                                    found = True
                                    break
                            else:
                                if os.path.exists(os.path.join(path, required + ext.replace(".", ""))):
                                    found = True
                                    break
                        if not found:
                            has_all_files = False
                            break
                    
                    if has_all_files:
                        return path
            
            return None

    def get_model_cache_key(self, model_type: str, device: str, source: str, path: Optional[str] = None, language: str = "English") -> str:
        """
        Generate a cache key for model instances.
        
        Args:
            model_type: Type of model ('tts' or 'vc')
            device: Target device ('cuda', 'cpu')
            source: Model source ('bundled', 'comfyui', 'huggingface')
            path: Optional path for local models
            language: Language model (for multilingual support)
            
        Returns:
            Cache key string
        """
        path_component = path or "default"
        return f"{model_type}_{device}_{source}_{path_component}_{language}"
    
    def load_tts_model(self, device: str = "auto", language: str = "English", force_reload: bool = False) -> Any:
        """
        Load ChatterboxTTS model with ComfyUI-integrated caching and language support.
        
        This method now uses ComfyUI's model management system, enabling automatic
        memory management, "Clear VRAM" button functionality, and proper integration.
        
        Args:
            device: Target device ('auto', 'cuda', 'cpu')
            language: Language model to load ('English', 'German', 'Norwegian', etc.)
            force_reload: Force reload even if cached
            
        Returns:
            ChatterboxTTS model instance
            
        Raises:
            ImportError: If ChatterboxTTS is not available
            RuntimeError: If model loading fails
        """
        if not _get_chatterbox_tts_available():
            raise ImportError("ChatterboxTTS not available - check installation or add bundled version")
        
        # Resolve auto device
        if device == "auto":
            device = "cuda" if torch.cuda.is_available() else "cpu"
            
        # Use unified model interface for ComfyUI integration
        from utils.models.unified_model_interface import load_tts_model
        
        try:
            model = load_tts_model(
                engine_name="chatterbox",
                model_name=language,
                device=device,
                language=language,
                force_reload=force_reload
            )
            
            # Update instance variables for backward compatibility
            self.tts_model = model
            self.current_device = device
            
            return model
            
        except Exception as e:
            print(f"⚠️ Failed to load ChatterBox model via unified interface: {e}")
            # Fallback to original logic if unified interface fails
            pass
        
        # Check if we need to load/reload (including language check)
        if not force_reload and self.tts_model is not None and self.current_device == device:
            # Also check if the current model matches the requested language
            current_cache_key = getattr(self, '_current_tts_cache_key', None)
            if current_cache_key and language in current_cache_key:
                return self.tts_model
        
        # Include language in cache check for fallback logic
        cache_key_base = f"tts_{device}_{language}"
        
        # For English, also check the original model discovery paths
        if language == "English":
            # Check original model paths first for English (backward compatibility)
            model_paths = self.find_chatterbox_models()
            for source, path in model_paths:
                if source in ["bundled", "comfyui"] and path:
                    try:
                        cache_key = f"{cache_key_base}_local_{path}"
                        
                        # Check class-level cache first
                        if not force_reload and cache_key in self._model_cache:
                            self.tts_model = self._model_cache[cache_key]
                            self.current_device = device
                            self._current_tts_cache_key = cache_key
                            return self.tts_model
                        
                        print(f"📁 Loading local English ChatterBox model from: {path}")
                        with warnings.catch_warnings():
                            warnings.simplefilter("ignore")
                            model = ChatterboxTTS.from_local(path, device, "English")
                        
                        # Cache the loaded model
                        self._model_cache[cache_key] = model
                        self._model_sources[cache_key] = source
                        self.tts_model = model
                        self.current_device = device
                        self._current_tts_cache_key = cache_key
                        return self.tts_model
                        
                    except Exception as e:
                        print(f"⚠️ Failed to load local English model from {path}: {e}")
                        continue
        
        # Try to find local model for the specific language
        local_language_path = self.find_local_language_model(language)
        model_loaded = False
        last_error = None
        
        if local_language_path:
            # Load local language-specific model
            try:
                cache_key = f"{cache_key_base}_local_{local_language_path}"
                
                # Check class-level cache first
                if not force_reload and cache_key in self._model_cache:
                    cached_model = self._model_cache[cache_key]
                    self.tts_model = cached_model
                    self.current_device = device
                    self._current_tts_cache_key = cache_key
                    return self.tts_model
                
                print(f"📁 Loading local {language} ChatterBox model from: {local_language_path}")
                with warnings.catch_warnings():
                    warnings.simplefilter("ignore")
                    model = ChatterboxTTS.from_local(local_language_path, device, language)
                
                # Cache the loaded model
                self._model_cache[cache_key] = model
                self._model_sources[cache_key] = "local"
                self.tts_model = model
                self.current_device = device
                self._current_tts_cache_key = cache_key
                model_loaded = True
                
            except Exception as e:
                print(f"⚠️ Failed to load local {language} model: {e}")
                last_error = e
        
        # If local loading failed or no local model, try HuggingFace
        if not model_loaded:
            try:
                cache_key = f"{cache_key_base}_huggingface"
                
                # Check class-level cache first
                if not force_reload and cache_key in self._model_cache:
                    self.tts_model = self._model_cache[cache_key]
                    self.current_device = device
                    self._current_tts_cache_key = cache_key
                    return self.tts_model
                
                # Check if this language is supported before calling from_pretrained
                # This prevents the engine's internal fallback from bypassing our cache logic
                from engines.chatterbox.language_models import get_model_config
                if not get_model_config(language) and language != "English":
                    # Language not supported - skip to our fallback logic instead of engine's fallback
                    print(f"📦 Language '{language}' not supported by ChatterBox, skipping to fallback")
                    raise Exception(f"Language '{language}' not supported")
                
                print(f"📦 Loading {language} ChatterBox model from HuggingFace")
                model = ChatterboxTTS.from_pretrained(device, language=language)
                
                # Cache the loaded model
                self._model_cache[cache_key] = model
                self._model_sources[cache_key] = "huggingface"
                self.tts_model = model
                self.current_device = device
                self._current_tts_cache_key = cache_key
                model_loaded = True
                
            except Exception as e:
                print(f"⚠️ Failed to load {language} model from HuggingFace: {e}")
                last_error = e
        
        # Fallback: try English if requested language failed and it's not English
        if not model_loaded and language != "English":
            print(f"🔄 Falling back to English model...")
            
            # First check if English is already loaded in any form
            english_cache_keys = [key for key in self._model_cache.keys() if "English" in key and device in key]
            if english_cache_keys and not force_reload:
                existing_key = english_cache_keys[0]  # Use first found English model
                print(f"💾 Reusing already-loaded English model from cache")
                self.tts_model = self._model_cache[existing_key]
                self.current_device = device
                self._current_tts_cache_key = existing_key
                return self.tts_model
            
            # If no English model in cache, try loading English (check local first)
            try:
                cache_key = f"tts_{device}_English_fallback"
                
                if not force_reload and cache_key in self._model_cache:
                    self.tts_model = self._model_cache[cache_key]
                    self.current_device = device
                    self._current_tts_cache_key = cache_key
                    return self.tts_model
                
                # Try local English first before HuggingFace
                english_local_path = self.find_local_language_model("English")
                if english_local_path:
                    print(f"📁 Loading local English ChatterBox model as fallback from: {english_local_path}")
                    model = ChatterboxTTS.from_local(english_local_path, device, "English")
                else:
                    print(f"📦 Loading English ChatterBox model from HuggingFace as fallback")
                    model = ChatterboxTTS.from_pretrained(device, language="English")
                
                # Cache the loaded model
                self._model_cache[cache_key] = model
                self._model_sources[cache_key] = "huggingface"
                self.tts_model = model
                self.current_device = device
                self._current_tts_cache_key = cache_key
                model_loaded = True
                
            except Exception as e:
                print(f"❌ Even English fallback failed: {e}")
                last_error = e
        
        if not model_loaded:
            error_msg = f"Failed to load ChatterBox TTS model for language '{language}' from any source"
            if last_error:
                error_msg += f". Last error: {last_error}"
            raise RuntimeError(error_msg)
        
        return self.tts_model
    
    def load_vc_model(self, device: str = "auto", force_reload: bool = False, language: str = "English") -> Any:
        """
        Load ChatterboxVC model with ComfyUI-integrated caching and language support.
        
        This method now uses ComfyUI's model management system, enabling automatic
        memory management, "Clear VRAM" button functionality, and proper integration.
        
        Args:
            device: Target device ('auto', 'cuda', 'cpu')
            force_reload: Force reload even if cached
            language: Language model to use (English, German, Norwegian)
            
        Returns:
            ChatterboxVC model instance
            
        Raises:
            ImportError: If ChatterboxVC is not available
            RuntimeError: If model loading fails
        """
        if not _get_chatterbox_vc_available():
            raise ImportError("ChatterboxVC not available - check installation or add bundled version")
        
        # Resolve auto device
        if device == "auto":
            device = "cuda" if torch.cuda.is_available() else "cpu"
            
        # Use unified model interface for ComfyUI integration
        from utils.models.unified_model_interface import load_vc_model
        
        try:
            # Resolve local model path first (same logic as find_local_language_model)
            local_model_path = self.find_local_language_model(language)
            
            model = load_vc_model(
                engine_name="chatterbox",
                model_name=language,
                device=device,
                language=language,
                model_path=local_model_path,  # Pass resolved local path
                force_reload=force_reload
            )
            
            # Update instance variables for backward compatibility
            self.vc_model = model
            self.current_device = device
            self.current_vc_language = language
            
            return model
            
        except Exception as e:
            error_str = str(e)
            print(f"⚠️ Failed to load ChatterBox VC model via unified interface: {e}")
            
            # Re-raise clear errors about unsupported languages instead of falling back
            if "Voice conversion not supported" in error_str:
                raise e  # Re-raise the clear error message
            
            # Only fallback to legacy logic for other types of errors
            pass
        
        # Check if we need to load/reload (include language in check)
        current_language = getattr(self, 'current_vc_language', None)
        if (not force_reload and self.vc_model is not None and 
            self.current_device == device and current_language == language):
            return self.vc_model
        
        # For VC, try language-specific model first (like TTS does)
        local_language_path = self.find_local_language_model(language)
        if local_language_path:
            model_paths = [("comfyui", local_language_path), ("huggingface", None)]
        else:
            # Fall back to generic discovery (includes huggingface fallback)
            model_paths = self.find_chatterbox_models()
        
        model_loaded = False
        last_error = None
        
        for source, path in model_paths:
            try:
                cache_key = self.get_model_cache_key("vc", device, source, path, language)
                
                # Check class-level cache first
                if not force_reload and cache_key in self._model_cache:
                    self.vc_model = self._model_cache[cache_key]
                    self.current_device = device
                    self.current_vc_language = language
                    self._model_sources[cache_key] = source
                    model_loaded = True
                    break
                
                # Load model based on source
                if source in ["bundled", "comfyui"]:
                    # Ensure ChatterboxVC is available before attempting to use it
                    if ChatterboxVC is None:
                        raise RuntimeError(f"ChatterboxVC not available - cannot load local model from {path}")
                    
                    with warnings.catch_warnings():
                        warnings.simplefilter("ignore")
                        # Note: Local models currently don't support language selection
                        # They use whatever language they were trained on
                        if language != "English":
                            print(f"⚠️ Local VC model at {path} may not support {language} - using existing model")
                        print(f"🐛 Calling ChatterboxVC.from_local({path}, {device})")
                        model = ChatterboxVC.from_local(path, device)
                        print(f"🐛 from_local returned: {model} (type: {type(model)})")
                elif source == "huggingface":
                    # Ensure ChatterboxVC is available before attempting to use it
                    if ChatterboxVC is None:
                        raise RuntimeError(f"ChatterboxVC not available - cannot load model from HuggingFace")
                    
                    print(f"🐛 Calling ChatterboxVC.from_pretrained({device}, {language})")
                    model = ChatterboxVC.from_pretrained(device, language)
                    print(f"🐛 from_pretrained returned: {model} (type: {type(model)})")
                else:
                    continue
                
                # Cache the loaded model
                self._model_cache[cache_key] = model
                self._model_sources[cache_key] = source
                self.vc_model = model
                self.current_device = device
                self.current_vc_language = language
                model_loaded = True
                print(f"✅ Successfully loaded {language} ChatterBox VC model from {source}")
                break
                
            except Exception as e:
                last_error = e
                continue
        
        if not model_loaded:
            error_msg = f"Failed to load ChatterboxVC from any source"
            if last_error:
                error_msg += f". Last error: {last_error}"
            raise RuntimeError(error_msg)
        
        return self.vc_model
    
    def get_model_source(self, model_type: str) -> Optional[str]:
        """
        Get the source of the currently loaded model.
        
        Args:
            model_type: Type of model ('tts' or 'vc')
            
        Returns:
            Model source string or None if no model loaded
        """
        if model_type == "tts" and self.tts_model is not None:
            # Use the current cache key to determine source
            current_cache_key = getattr(self, '_current_tts_cache_key', None)
            if current_cache_key:
                # Extract source from cache key or _model_sources
                source = self._model_sources.get(current_cache_key)
                if source:
                    return source
                
                # Fallback: parse from cache key format
                if "_local_" in current_cache_key:
                    return "comfyui"
                elif "_huggingface" in current_cache_key:
                    return "huggingface"
                elif "_fallback" in current_cache_key:
                    return "huggingface (fallback)"
            
            # Legacy fallback
            device = self.current_device or "cpu"
            model_paths = self.find_chatterbox_models()
            if model_paths:
                source, path = model_paths[0]
                return source
        elif model_type == "vc" and self.vc_model is not None:
            device = self.current_device or "cpu"
            model_paths = self.find_chatterbox_models()
            if model_paths:
                source, path = model_paths[0]
                cache_key = self.get_model_cache_key("vc", device, source, path, getattr(self, 'current_vc_language', 'English'))
                return self._model_sources.get(cache_key)
        
        return None
    
    def clear_cache(self, model_type: Optional[str] = None):
        """
        Clear model cache.
        
        Args:
            model_type: Optional model type to clear ('tts', 'vc'), or None for all
        """
        if model_type is None:
            # Clear all
            self._model_cache.clear()
            self._model_sources.clear()
            self.tts_model = None
            self.vc_model = None
            self.current_device = None
        elif model_type == "tts":
            # Clear TTS models
            keys_to_remove = [k for k in self._model_cache.keys() if k.startswith("tts_")]
            for key in keys_to_remove:
                self._model_cache.pop(key, None)
                self._model_sources.pop(key, None)
            self.tts_model = None
        elif model_type == "vc":
            # Clear VC models
            keys_to_remove = [k for k in self._model_cache.keys() if k.startswith("vc_")]
            for key in keys_to_remove:
                self._model_cache.pop(key, None)
                self._model_sources.pop(key, None)
            self.vc_model = None
    
    @property
    def is_available(self) -> Dict[str, bool]:
        """
        Check availability of ChatterBox components.
        
        Returns:
            Dictionary with availability status
        """
        return {
            "tts": _get_chatterbox_tts_available(),
            "vc": _get_chatterbox_vc_available(),
            "bundled": _get_using_bundled_chatterbox(),
            "any": _get_chatterbox_tts_available() or _get_chatterbox_vc_available(),
            "vibevoice": _get_vibevoice_available()
        }
    
    def load_vibevoice_model(self, model_name: str = "vibevoice-1.5B", device: str = "auto", force_reload: bool = False):
        """
        Load VibeVoice model using ComfyUI model wrapper for consistency with other engines.
        
        Args:
            model_name: VibeVoice model name ("vibevoice-1.5B" or "vibevoice-7B")
            device: Target device ('cuda', 'cpu', 'auto')
            force_reload: Force reload even if cached
            
        Returns:
            Tuple of (model, processor) for VibeVoice
            
        Raises:
            ImportError: If VibeVoice is not available
            RuntimeError: If model loading fails
        """
        if not _get_vibevoice_available():
            raise ImportError("VibeVoice not available - check installation")
        
        # Use ComfyUI model manager for consistent behavior
        from utils.models.comfyui_model_wrapper import tts_model_manager
        
        # Resolve auto device
        if device == "auto":
            device = "cuda" if torch.cuda.is_available() else "cpu"
        
        # Create cache key for VibeVoice models
        cache_key = f"vibevoice_{model_name}_{device}"
        
        def vibevoice_factory(model_name, device):
            """Factory function to create VibeVoice model and processor"""
            # Get model path (downloads if necessary)
            downloader = VibeVoiceDownloader()
            model_path = downloader.get_model_path(model_name)
            if not model_path:
                raise RuntimeError(f"Failed to get VibeVoice model '{model_name}'")
            
            print(f"🔄 Creating VibeVoice model '{model_name}' on {device}...")
            
            # Load model and processor
            model = VibeVoiceForConditionalGenerationInference.from_pretrained(
                model_path,
                trust_remote_code=True,
                torch_dtype=torch.bfloat16,
                device_map=device if device != "auto" else None
            )
            processor = VibeVoiceProcessor.from_pretrained(model_path)
            
            # Move to device if needed
            if device == "cuda" and torch.cuda.is_available():
                model = model.cuda()
            
            # Create container to hold both model and processor
            class VibeVoiceModelContainer:
                def __init__(self, model, processor):
                    self.model = model
                    self.processor = processor
                    
                def to(self, device):
                    """Move both model and processor to device"""
                    if hasattr(self.model, 'to'):
                        self.model.to(device)
                    # Note: processor typically doesn't need device movement
                    return self
                    
                def parameters(self):
                    """Return model parameters for memory calculation"""
                    if hasattr(self.model, 'parameters'):
                        return self.model.parameters()
                    return iter([])
            
            return VibeVoiceModelContainer(model, processor)
        
        # Load using ComfyUI model manager
        wrapper = tts_model_manager.load_model(
            model_factory_func=vibevoice_factory,
            model_key=cache_key,
            model_type="tts",
            engine="vibevoice",
            device=device,
            force_reload=force_reload,
            model_name=model_name
        )
        
        # Extract model and processor from container
        container = wrapper.model
        return container.model, container.processor
    
    def unload_vibevoice_models(self):
        """Unload all VibeVoice models using ComfyUI model wrapper system."""
        from utils.models.comfyui_model_wrapper import tts_model_manager
        
        # Find all VibeVoice models in the ComfyUI model manager
        vibevoice_keys = []
        for key, wrapper in tts_model_manager._model_cache.items():
            if wrapper.model_info.engine == "vibevoice":
                vibevoice_keys.append(key)
        
        # Remove them using the ComfyUI model manager (handles unloading automatically)
        for key in vibevoice_keys:
            tts_model_manager.remove_model(key)
        
        if vibevoice_keys:
            print(f"🧹 Unloaded {len(vibevoice_keys)} VibeVoice models via ComfyUI model manager")


# Global model manager instance
model_manager = ModelManager()