"""Device and environment utilities for Xoron-Dev."""

import os
import gc
from typing import Dict, Optional, Tuple, List, Union
from dataclasses import dataclass

# Handle torch import gracefully
try:
    import torch
    import torch.nn as nn
    TORCH_AVAILABLE = True
except ImportError:
    TORCH_AVAILABLE = False
    torch = None
    nn = None


def setup_cuda_memory_management(
    enable_unified_memory: bool = True,
    max_split_size_mb: int = 512,
    garbage_collection_threshold: float = 0.8,
    reserve_memory_fraction: float = 0.05,
):
    """
    Setup CUDA memory management for efficient training with VRAM spillover to system RAM.
    
    This enables CUDA Unified Memory features on supported GPUs (Pascal+, including T4/Turing),
    allowing allocations to exceed VRAM by automatically paging to system RAM.
    
    Args:
        enable_unified_memory: Enable CUDA managed memory features for VRAM oversubscription
        max_split_size_mb: Max size for memory splits (smaller = less fragmentation, more overhead)
        garbage_collection_threshold: Trigger GC when this fraction of memory is used
        reserve_memory_fraction: Reserve this fraction of VRAM for PyTorch allocator headroom
        
    How Unified Memory Works on T4/Turing GPUs:
        - Allocations can exceed physical VRAM (~15GB usable)
        - When VRAM fills up, CUDA driver evicts pages to system RAM
        - Pages migrate back to VRAM on GPU access (demand paging)
        - Prevents OOM crashes for large audio/video tensors
        
    Note: Performance degrades when actively using data in system RAM due to PCIe bandwidth,
    but this is better than crashing. Best used for:
        - Large batch audio processing
        - Long audio sequences  
        - Multi-modal training with competing VRAM needs
    """
    if not TORCH_AVAILABLE or not torch.cuda.is_available():
        return False
    
    # Build CUDA allocator config string
    alloc_conf_parts = [
        'expandable_segments:True',  # Allow memory to be released back to system
        f'max_split_size_mb:{max_split_size_mb}',  # Reduce fragmentation
        f'garbage_collection_threshold:{garbage_collection_threshold}',  # Aggressive GC
    ]
    
    # Set the allocator config
    alloc_conf = ','.join(alloc_conf_parts)
    os.environ['PYTORCH_CUDA_ALLOC_CONF'] = alloc_conf
    
    if enable_unified_memory:
        # Enable CUDA managed memory features
        # This allows oversubscription - allocating more than VRAM with auto-paging
        os.environ['PYTORCH_CUDA_ALLOC_CONF'] += ',backend:cudaMallocAsync'
        
        # Enable memory-efficient cudnn algorithms  
        if hasattr(torch.backends, 'cudnn'):
            torch.backends.cudnn.benchmark = True
            torch.backends.cudnn.allow_tf32 = True
        
        # Enable TF32 for faster matmuls with minimal precision loss
        if hasattr(torch.backends, 'cuda'):
            torch.backends.cuda.matmul.allow_tf32 = True
    
    # Reserve some VRAM headroom for the allocator
    if reserve_memory_fraction > 0:
        try:
            for device_idx in range(torch.cuda.device_count()):
                torch.cuda.set_per_process_memory_fraction(
                    1.0 - reserve_memory_fraction, 
                    device=device_idx
                )
        except Exception:
            pass  # Some CUDA versions don't support this
    
    return True


def enable_audio_memory_optimization():
    """
    Enable memory optimizations specifically for audio training.
    
    Audio training is memory-intensive due to:
        - Long sequences (160k samples = 10 seconds at 16kHz)
        - Multiple intermediate representations (waveform, mel, features)
        - Voice cloning requires storing speaker embeddings
        
    This function configures PyTorch for optimal audio training memory usage.
    """
    if not TORCH_AVAILABLE or not torch.cuda.is_available():
        return False
    
    # Setup unified memory with audio-optimized settings
    setup_cuda_memory_management(
        enable_unified_memory=True,
        max_split_size_mb=256,  # Smaller splits for audio chunks
        garbage_collection_threshold=0.7,  # More aggressive GC for audio
        reserve_memory_fraction=0.1,  # More headroom for audio spikes
    )
    
    # Force synchronous CUDA operations for more predictable memory
    # (slightly slower but prevents memory spikes from async ops)
    os.environ['CUDA_LAUNCH_BLOCKING'] = '1'
    
    print("   🎧 Audio memory optimization enabled:")
    print("      • CUDA Unified Memory: spillover to system RAM on OOM")
    print("      • Expandable segments: dynamic memory allocation")
    print("      • Aggressive garbage collection: 70% threshold")
    
    return True


def get_available_memory() -> Dict[str, float]:
    """
    Get available memory on all devices (VRAM + system RAM).
    
    Returns:
        Dict with 'vram_free_gb', 'vram_total_gb', 'ram_free_gb', 'ram_total_gb'
    """
    result = {
        'vram_free_gb': 0.0,
        'vram_total_gb': 0.0,
        'ram_free_gb': 0.0,
        'ram_total_gb': 0.0,
    }
    
    # Get VRAM info
    if TORCH_AVAILABLE and torch.cuda.is_available():
        for i in range(torch.cuda.device_count()):
            free, total = torch.cuda.mem_get_info(i)
            result['vram_free_gb'] += free / (1024**3)
            result['vram_total_gb'] += total / (1024**3)
    
    # Get system RAM info
    try:
        import psutil
        mem = psutil.virtual_memory()
        result['ram_free_gb'] = mem.available / (1024**3)
        result['ram_total_gb'] = mem.total / (1024**3)
    except ImportError:
        pass
    
    return result


def smart_tensor_to_device(
    tensor: 'torch.Tensor',
    device: str,
    allow_cpu_fallback: bool = True,
    max_vram_usage: float = 0.9,
) -> 'torch.Tensor':
    """
    Move tensor to device with smart memory management.
    
    If VRAM is nearly full, either:
        1. Keep tensor on CPU (if allow_cpu_fallback=True)
        2. Let unified memory handle spillover (if enabled)
        
    Args:
        tensor: Tensor to move
        device: Target device ('cuda', 'cuda:0', etc.)
        allow_cpu_fallback: If True, keep on CPU when VRAM is low
        max_vram_usage: Max fraction of VRAM to use before fallback
        
    Returns:
        Tensor on appropriate device
    """
    if not TORCH_AVAILABLE:
        return tensor
        
    if not device.startswith('cuda') or not torch.cuda.is_available():
        return tensor.to(device)
    
    # Check VRAM usage
    device_idx = 0 if device == 'cuda' else int(device.split(':')[1])
    free, total = torch.cuda.mem_get_info(device_idx)
    usage = 1.0 - (free / total)
    
    tensor_size = tensor.numel() * tensor.element_size()
    
    # If VRAM is nearly full and tensor is large, consider fallback
    if usage > max_vram_usage and tensor_size > 100 * 1024 * 1024:  # >100MB
        if allow_cpu_fallback:
            # Keep on CPU, unified memory will handle if accessed on GPU
            return tensor.pin_memory() if tensor.device.type == 'cpu' else tensor.cpu().pin_memory()
    
    return tensor.to(device)


@dataclass
class EnvironmentInfo:
    """Information about the runtime environment."""
    name: str  # 'kaggle', 'colab', 'local'
    temp_dir: str
    datasets_dir: str
    output_dir: str
    model_dir: str
    final_model_dir: str


def detect_environment() -> str:
    """
    Detect the runtime environment.
    
    Returns:
        'kaggle' - Running on Kaggle
        'colab' - Running on Google Colab
        'lightning' - Running on Lightning AI
        'local' - Running locally or other environment
    """
    # Check for Kaggle
    if os.path.exists('/kaggle') or 'KAGGLE_KERNEL_RUN_TYPE' in os.environ:
        return 'kaggle'
    
    # Check for Google Colab
    try:
        import google.colab
        return 'colab'
    except ImportError:
        pass
    
    # Check for Colab environment variable
    if 'COLAB_GPU' in os.environ or 'COLAB_RELEASE_TAG' in os.environ:
        return 'colab'
    
    # Check for Lightning AI
    # Lightning AI uses /teamspace/studios for studio environments
    # and sets LIGHTNING_* environment variables
    if (os.path.exists('/teamspace') or 
        'LIGHTNING_CLOUDSPACE_HOST' in os.environ or
        'LIGHTNING_CLOUD_PROJECT_ID' in os.environ or
        'LIGHTNING_STUDIO_ID' in os.environ or
        os.environ.get('LIGHTNING_CLOUD_URL')):
        return 'lightning'
    
    return 'local'


def get_environment_paths(env: Optional[str] = None) -> EnvironmentInfo:
    """
    Get appropriate paths based on the detected environment.
    
    Args:
        env: Override environment detection ('kaggle', 'colab', 'lightning', 'local')
        
    Returns:
        EnvironmentInfo with appropriate paths for the environment
    """
    if env is None:
        env = detect_environment()
    
    if env == 'kaggle':
        # Kaggle: /kaggle/tmp for checkpoints (to save disk space in /working)
        # Only final model goes to /kaggle/working for persistence
        return EnvironmentInfo(
            name='kaggle',
            temp_dir='/kaggle/tmp/xoron',
            datasets_dir='/kaggle/tmp/xoron/datasets',
            output_dir='/kaggle/tmp/xoron/checkpoints',  # Checkpoints in /tmp to save disk space
            model_dir='/kaggle/tmp/xoron/model',  # Built model in /tmp
            final_model_dir='/kaggle/working/xoron-final',  # Only final model in /working
        )
    
    elif env == 'colab':
        # Colab: /content for everything, can mount Google Drive
        return EnvironmentInfo(
            name='colab',
            temp_dir='/content/tmp',
            datasets_dir='/content/tmp/datasets',
            output_dir='/content/xoron-checkpoints',
            model_dir='/content/xoron-dev-model',
            final_model_dir='/content/xoron-final',
        )
    
    elif env == 'lightning':
        # Lightning AI: /teamspace/studios/this_studio for persistent storage
        # Use /tmp for temporary files, /teamspace for persistent outputs
        studio_path = '/teamspace/studios/this_studio'
        if not os.path.exists(studio_path):
            studio_path = '/teamspace'
        return EnvironmentInfo(
            name='lightning',
            temp_dir='/tmp/xoron',
            datasets_dir='/tmp/xoron/datasets',
            output_dir=f'{studio_path}/xoron-checkpoints',
            model_dir=f'{studio_path}/xoron-dev-model',
            final_model_dir=f'{studio_path}/xoron-final',
        )
    
    else:
        # Local: use current directory
        return EnvironmentInfo(
            name='local',
            temp_dir='./tmp',
            datasets_dir='./tmp/datasets',
            output_dir='./xoron-checkpoints',
            model_dir='./xoron-dev-model',
            final_model_dir='./xoron-final',
        )


def print_environment_info():
    """Print detected environment information."""
    env = detect_environment()
    paths = get_environment_paths(env)
    
    env_icons = {
        'kaggle': '🏆',
        'colab': '🔬',
        'lightning': '⚡️',
        'local': '💻',
    }
    
    print(f"\n{env_icons.get(env, '🖥️')} Environment: {env.upper()}")
    print(f"   Temp directory: {paths.temp_dir}")
    print(f"   Datasets directory: {paths.datasets_dir}")
    print(f"   Output directory: {paths.output_dir}")
    print(f"   Model directory: {paths.model_dir}")
    print(f"   Final model directory: {paths.final_model_dir}")


def get_device_info() -> Dict[str, any]:
    """Get information about available devices."""
    if not TORCH_AVAILABLE:
        return {
            'cuda_available': False,
            'device': 'cpu',
            'num_gpus': 0,
            'gpus': [],
            'total_memory_gb': 0,
        }

    info = {
        'cuda_available': torch.cuda.is_available(),
        'device': 'cuda' if torch.cuda.is_available() else 'cpu',
        'num_gpus': 0,
        'gpus': [],
        'total_memory_gb': 0,
    }

    if torch.cuda.is_available():
        info['num_gpus'] = torch.cuda.device_count()
        total_memory = 0
        for i in range(info['num_gpus']):
            props = torch.cuda.get_device_properties(i)
            gpu_info = {
                'index': i,
                'name': torch.cuda.get_device_name(i),
                'memory_gb': props.total_memory / 1e9,
                'compute_capability': f"{props.major}.{props.minor}",
            }
            info['gpus'].append(gpu_info)
            total_memory += props.total_memory
        info['total_memory_gb'] = total_memory / 1e9

    return info


def print_device_info():
    """Print device information."""
    info = get_device_info()

    print("\n💻 Device Configuration:")
    print(f"   CUDA Available: {info['cuda_available']}")
    print(f"   Device: {info['device']}")

    if info['cuda_available']:
        print(f"   Number of GPUs: {info['num_gpus']}")
        for gpu in info['gpus']:
            print(f"   GPU {gpu['index']}: {gpu['name']} - {gpu['memory_gb']:.1f} GB")
        print(f"   Total GPU Memory: {info['total_memory_gb']:.1f} GB")


def clear_cuda_cache():
    """Clear CUDA cache to free memory."""
    if TORCH_AVAILABLE and torch.cuda.is_available():
        torch.cuda.synchronize()
        torch.cuda.empty_cache()


def get_optimal_device() -> str:
    """Get the optimal device for training."""
    if not TORCH_AVAILABLE:
        return 'cpu'
    if torch.cuda.is_available():
        return 'cuda'
    elif hasattr(torch.backends, 'mps') and torch.backends.mps.is_available():
        return 'mps'
    else:
        return 'cpu'


def move_to_device(tensor_or_model, device: str):
    """Move tensor or model to specified device."""
    if hasattr(tensor_or_model, 'to'):
        return tensor_or_model.to(device)
    return tensor_or_model


class DualGPUTrainingManager:
    """
    Manager for dual-GPU training of video/audio components.
    
    This class handles:
    1. Splitting model components across two GPUs for maximum VRAM utilization
    2. Coordinating forward/backward passes across GPUs
    3. Memory management and synchronization
    
    For video training:
    - GPU0: Encoders, first half of transformer blocks
    - GPU1: Second half of transformer blocks, decoders
    
    For audio training:
    - GPU0: Audio encoder, first half of processing
    - GPU1: Second half of processing, decoders
    
    Usage:
        manager = DualGPUTrainingManager(['cuda:0', 'cuda:1'], mode='video')
        manager.setup_model_parallel(video_generator)
        
        # In training loop:
        loss = manager.training_step(video_generator, video_frames, text_context)
    """
    
    def __init__(
        self,
        devices: List[str],
        mode: str = 'video',  # 'video' or 'audio'
        gradient_checkpointing: bool = True,
    ):
        if not TORCH_AVAILABLE:
            raise RuntimeError("PyTorch not available")
        
        self.devices = devices
        self.mode = mode
        self.gradient_checkpointing = gradient_checkpointing
        self.gpu0 = torch.device(devices[0])
        self.gpu1 = torch.device(devices[1]) if len(devices) > 1 else self.gpu0
        self._is_setup = False
        self._model_splits = {}
        
    def get_gpu_memory_info(self) -> Dict[str, Dict[str, float]]:
        """Get memory info for all GPUs."""
        info = {}
        for device in self.devices:
            if 'cuda' in device:
                idx = int(device.split(':')[-1])
                if idx < torch.cuda.device_count():
                    total = torch.cuda.get_device_properties(idx).total_memory / 1e9
                    allocated = torch.cuda.memory_allocated(idx) / 1e9
                    cached = torch.cuda.memory_reserved(idx) / 1e9
                    info[device] = {
                        'total_gb': total,
                        'allocated_gb': allocated,
                        'cached_gb': cached,
                        'free_gb': total - allocated,
                    }
        return info
    
    def print_memory_status(self, label: str = ""):
        """Print current GPU memory status."""
        info = self.get_gpu_memory_info()
        print(f"\n📊 GPU Memory Status {label}:")
        for device, mem in info.items():
            print(f"   {device}: {mem['allocated_gb']:.2f}GB / {mem['total_gb']:.2f}GB "
                  f"(free: {mem['free_gb']:.2f}GB)")
    
    def setup_video_generator_parallel(self, video_generator):
        """
        Split video generator across two GPUs for maximum VRAM utilization.
        
        Strategy:
        - GPU0: VAE encoder, first half of transformer blocks
        - GPU1: Second half of transformer blocks, VAE decoder
        
        This allows training with larger video resolutions and more frames.
        """
        if video_generator is None:
            print("   ⚠️ No video generator to setup")
            return
        
        # Check if already on a device
        try:
            gen_device = next(video_generator.parameters()).device
        except StopIteration:
            gen_device = self.gpu0
        
        # Move VAE encoder to GPU0
        if hasattr(video_generator, 'vae') and video_generator.vae is not None:
            if hasattr(video_generator.vae, 'encoder'):
                video_generator.vae.encoder = video_generator.vae.encoder.to(self.gpu0)
                self._model_splits['vae_encoder'] = self.gpu0
                print(f"   ✅ VAE encoder -> {self.gpu0}")
            
            # Move VAE decoder to GPU1
            if hasattr(video_generator.vae, 'decoder'):
                video_generator.vae.decoder = video_generator.vae.decoder.to(self.gpu1)
                self._model_splits['vae_decoder'] = self.gpu1
                print(f"   ✅ VAE decoder -> {self.gpu1}")
        
        # Split transformer blocks (UNet or DiT)
        if hasattr(video_generator, 'transformer') and video_generator.transformer is not None:
            transformer = video_generator.transformer
            if hasattr(transformer, 'blocks'):
                blocks = transformer.blocks
                n_blocks = len(blocks)
                mid = n_blocks // 2
                
                # First half to GPU0
                for i in range(mid):
                    blocks[i] = blocks[i].to(self.gpu0)
                    
                # Second half to GPU1
                for i in range(mid, n_blocks):
                    blocks[i] = blocks[i].to(self.gpu1)
                
                self._model_splits['transformer_blocks'] = {
                    'gpu0': list(range(mid)),
                    'gpu1': list(range(mid, n_blocks)),
                }
                print(f"   ✅ Transformer: blocks 0-{mid-1} -> {self.gpu0}, blocks {mid}-{n_blocks-1} -> {self.gpu1}")
        
        # If using UNet instead
        if hasattr(video_generator, 'unet') and video_generator.unet is not None:
            unet = video_generator.unet
            # Split down/up blocks
            if hasattr(unet, 'down_blocks'):
                for i, block in enumerate(unet.down_blocks):
                    device = self.gpu0 if i < len(unet.down_blocks) // 2 else self.gpu1
                    unet.down_blocks[i] = block.to(device)
            
            if hasattr(unet, 'mid_block'):
                unet.mid_block = unet.mid_block.to(self.gpu1)
            
            if hasattr(unet, 'up_blocks'):
                for i, block in enumerate(unet.up_blocks):
                    device = self.gpu1 if i < len(unet.up_blocks) // 2 else self.gpu0
                    unet.up_blocks[i] = block.to(device)
            
            print(f"   ✅ UNet split across {self.gpu0} and {self.gpu1}")
        
        # Enable gradient checkpointing if requested
        if self.gradient_checkpointing and hasattr(video_generator, 'gradient_checkpointing_enable'):
            video_generator.gradient_checkpointing_enable()
            print(f"   ✅ Gradient checkpointing enabled")
        
        self._is_setup = True
        self.print_memory_status("after video generator setup")
    
    def setup_audio_components_parallel(self, audio_encoder, audio_decoder, waveform_decoder=None):
        """
        Split audio components across two GPUs.
        
        Strategy:
        - GPU0: Audio encoder (input processing)
        - GPU1: Audio decoder, waveform decoder (output generation)
        """
        if audio_encoder is not None:
            audio_encoder = audio_encoder.to(self.gpu0)
            self._model_splits['audio_encoder'] = self.gpu0
            print(f"   ✅ Audio encoder -> {self.gpu0}")
        
        if audio_decoder is not None:
            audio_decoder = audio_decoder.to(self.gpu1)
            self._model_splits['audio_decoder'] = self.gpu1
            print(f"   ✅ Audio decoder -> {self.gpu1}")
        
        if waveform_decoder is not None:
            waveform_decoder = waveform_decoder.to(self.gpu1)
            self._model_splits['waveform_decoder'] = self.gpu1
            print(f"   ✅ Waveform decoder -> {self.gpu1}")
        
        self._is_setup = True
        self.print_memory_status("after audio setup")
    
    def sync_gpus(self):
        """Synchronize both GPUs."""
        if TORCH_AVAILABLE and torch.cuda.is_available():
            for device in self.devices:
                if 'cuda' in device:
                    idx = int(device.split(':')[-1])
                    if idx < torch.cuda.device_count():
                        torch.cuda.synchronize(idx)
    
    def clear_cache(self):
        """Clear CUDA cache on both GPUs."""
        if TORCH_AVAILABLE and torch.cuda.is_available():
            self.sync_gpus()
            gc.collect()
            torch.cuda.empty_cache()
    
    def transfer_tensor(self, tensor, target_device):
        """Transfer tensor between GPUs efficiently."""
        if tensor is None:
            return None
        current_device = tensor.device
        target = torch.device(target_device)
        if current_device != target:
            return tensor.to(target, non_blocking=True)
        return tensor


class CPUOffloadManager:
    """
    Manager for CPU offloading of frozen model components.
    
    This class handles:
    1. Moving frozen components to CPU to free GPU VRAM
    2. Efficient tensor transfer for inference passes
    3. Memory pinning for faster CPU-GPU transfers
    
    Usage:
        manager = CPUOffloadManager()
        manager.offload_component(model.vision_encoder, 'vision_encoder')
        
        # During forward pass (automatic if hooked):
        output = model.vision_encoder(input)  # Auto-transfers as needed
    """
    
    def __init__(self, pin_memory: bool = True):
        if not TORCH_AVAILABLE:
            raise RuntimeError("PyTorch not available")
        
        self.pin_memory = pin_memory
        self._offloaded_components = {}
        self._component_devices = {}  # Original devices for potential restore
        
    def offload_component(self, component, name: str, force: bool = False):
        """
        Offload a model component to CPU.
        
        Args:
            component: The nn.Module to offload
            name: Name for tracking
            force: Force offload even if already on CPU
        """
        if component is None:
            return
        
        try:
            current_device = next(component.parameters()).device
        except StopIteration:
            current_device = torch.device('cpu')
        
        if current_device.type == 'cpu' and not force:
            print(f"   ℹ️ {name} already on CPU")
            return
        
        # Store original device for potential restore
        self._component_devices[name] = current_device
        
        # Move to CPU
        component = component.cpu()
        
        # Pin memory for faster transfers (if enabled and not already pinned)
        if self.pin_memory:
            for param in component.parameters():
                if not param.is_pinned():
                    try:
                        param.data = param.data.pin_memory()
                    except Exception:
                        pass  # Pin memory not always available
        
        self._offloaded_components[name] = component
        print(f"   ✅ {name} offloaded to CPU (was on {current_device})")
        
    def restore_component(self, name: str):
        """Restore a component to its original GPU device."""
        if name not in self._offloaded_components:
            return None
        
        component = self._offloaded_components[name]
        original_device = self._component_devices.get(name, 'cuda:0')
        
        component = component.to(original_device)
        print(f"   ✅ {name} restored to {original_device}")
        
        del self._offloaded_components[name]
        return component
    
    def get_offloaded_names(self) -> List[str]:
        """Get list of offloaded component names."""
        return list(self._offloaded_components.keys())
    
    def estimate_cpu_memory_usage(self) -> float:
        """Estimate CPU memory used by offloaded components (in GB)."""
        total_bytes = 0
        for name, component in self._offloaded_components.items():
            for param in component.parameters():
                total_bytes += param.numel() * param.element_size()
        return total_bytes / 1e9


def get_multi_gpu_info() -> Dict:
    """
    Get detailed information about all available GPUs.
    
    Returns:
        Dict with GPU info including memory, compute capability, and recommended usage.
    """
    if not TORCH_AVAILABLE or not torch.cuda.is_available():
        return {'available': False, 'num_gpus': 0, 'gpus': []}
    
    info = {
        'available': True,
        'num_gpus': torch.cuda.device_count(),
        'gpus': [],
        'total_vram_gb': 0,
        'recommended_training_mode': 'single_gpu',
    }
    
    for i in range(info['num_gpus']):
        props = torch.cuda.get_device_properties(i)
        gpu_info = {
            'index': i,
            'device': f'cuda:{i}',
            'name': props.name,
            'memory_gb': props.total_memory / 1e9,
            'compute_capability': (props.major, props.minor),
            'supports_bf16': props.major >= 8,
            'supports_flash_attn': (props.major, props.minor) >= (7, 5),
            'multi_processor_count': props.multi_processor_count,
        }
        info['gpus'].append(gpu_info)
        info['total_vram_gb'] += gpu_info['memory_gb']
    
    # Recommend training mode based on hardware
    if info['num_gpus'] >= 2:
        if info['total_vram_gb'] >= 40:
            info['recommended_training_mode'] = 'dual_gpu_training'
        elif info['total_vram_gb'] >= 24:
            info['recommended_training_mode'] = 'cpu_offload_with_model_parallel'
        else:
            info['recommended_training_mode'] = 'cpu_offload'
    elif info['num_gpus'] == 1:
        if info['gpus'][0]['memory_gb'] >= 24:
            info['recommended_training_mode'] = 'single_gpu_full'
        else:
            info['recommended_training_mode'] = 'single_gpu_with_offload'
    
    return info


def print_multi_gpu_info():
    """Print detailed multi-GPU information."""
    info = get_multi_gpu_info()
    
    print("\n🖥️  Multi-GPU Configuration:")
    print(f"   Available: {info['available']}")
    print(f"   Number of GPUs: {info['num_gpus']}")
    print(f"   Total VRAM: {info['total_vram_gb']:.1f} GB")
    print(f"   Recommended mode: {info['recommended_training_mode']}")
    
    for gpu in info['gpus']:
        print(f"\n   GPU {gpu['index']} ({gpu['device']}):")
        print(f"      Name: {gpu['name']}")
        print(f"      Memory: {gpu['memory_gb']:.1f} GB")
        print(f"      Compute: SM{gpu['compute_capability'][0]}.{gpu['compute_capability'][1]}")
        print(f"      BF16: {'✅' if gpu['supports_bf16'] else '❌'}")
        print(f"      Flash Attention: {'✅' if gpu['supports_flash_attn'] else '❌'}")
