"""
Device utilities for VoxCPM nodes.

This module provides unified device handling using ComfyUI's model_management
functions for consistent device detection and management across all backends.
"""

import torch
from typing import List, Optional

# Lazy import to avoid circular dependencies
_model_management = None


def _get_model_management():
    """Lazy import of ComfyUI's model_management module."""
    global _model_management
    if _model_management is None:
        import comfy.model_management as mm
        _model_management = mm
    return _model_management


# Device type constants
DEVICE_CPU = "cpu"
DEVICE_CUDA = "cuda"
DEVICE_MPS = "mps"
DEVICE_XPU = "xpu"
DEVICE_NPU = "npu"
DEVICE_DIRECTML = "directml"
DEVICE_HIP = "hip"


def get_available_devices() -> List[str]:
    """Get list of available devices in order of preference.
    
    Uses ComfyUI's model_management for consistent device detection.
    
    Returns:
        List of device type strings (e.g., ["cuda", "cpu"])
    """
    devices = []
    mm = _get_model_management()
    
    # Check for GPU devices using ComfyUI's detection
    if mm.is_nvidia():
        devices.append(DEVICE_CUDA)
    elif mm.is_amd():
        # AMD uses HIP backend but we report as "cuda" for compatibility
        devices.append(DEVICE_CUDA)
    elif mm.mps_mode():
        devices.append(DEVICE_MPS)
    else:
        # Check for other GPU backends
        if hasattr(torch, 'xpu') and torch.xpu.is_available():
            devices.append(DEVICE_XPU)
        elif hasattr(torch, 'npu') and torch.npu.is_available():
            devices.append(DEVICE_NPU)
    
    # CPU is always available as fallback
    devices.append(DEVICE_CPU)
    
    return devices


def get_device_display_name(device_type: str) -> str:
    """Get human-readable display name for a device type.
    
    Args:
        device_type: Device type string (e.g., "cuda", "cpu")
        
    Returns:
        Human-readable device name
    """
    mm = _get_model_management()
    
    if device_type == DEVICE_CUDA:
        if mm.is_nvidia():
            try:
                gpu_name = torch.cuda.get_device_name(0)
                return f"CUDA ({gpu_name})"
            except:
                return "CUDA (NVIDIA GPU)"
        elif mm.is_amd():
            return "CUDA (AMD ROCm/HIP)"
        return "CUDA"
    elif device_type == DEVICE_MPS:
        return "MPS (Apple Silicon)"
    elif device_type == DEVICE_XPU:
        return "XPU (Intel Arc)"
    elif device_type == DEVICE_NPU:
        return "NPU (Huawei Ascend)"
    elif device_type == DEVICE_DIRECTML:
        return "DirectML (Windows)"
    elif device_type == DEVICE_CPU:
        return "CPU"
    
    return device_type.upper()


def get_torch_device(device_type: str = None) -> torch.device:
    """Get torch.device for the specified device type.
    
    Uses ComfyUI's get_torch_device() for consistent device handling.
    
    Args:
        device_type: Device type string. If None, uses ComfyUI's default device.
        
    Returns:
        torch.device instance
    """
    mm = _get_model_management()
    
    if device_type is None or device_type == "auto":
        return mm.get_torch_device()
    
    # Map our device types to torch devices
    if device_type == DEVICE_CPU:
        return torch.device(DEVICE_CPU)
    
    # For GPU devices, use ComfyUI's device selection
    return mm.get_torch_device()


def get_device_type(device: torch.device) -> str:
    """Convert torch.device to our device type string.
    
    Args:
        device: torch.device instance
        
    Returns:
        Device type string (e.g., "cuda", "cpu")
    """
    if device.type == DEVICE_CPU:
        return DEVICE_CPU
    elif device.type == DEVICE_MPS:
        return DEVICE_MPS
    elif device.type == DEVICE_XPU:
        return DEVICE_XPU
    elif device.type == DEVICE_NPU:
        return DEVICE_NPU
    elif device.type == DEVICE_CUDA:
        return DEVICE_CUDA
    else:
        return device.type


def get_offload_device() -> torch.device:
    """Get the offload device (typically CPU) for memory management.
    
    Uses ComfyUI's intermediate_device() for consistent offload handling.
    
    Returns:
        torch.device for offloading
    """
    mm = _get_model_management()
    return mm.intermediate_device()


def should_use_fp16(device: torch.device = None) -> bool:
    """Check if FP16 should be used for the given device.
    
    Uses ComfyUI's should_use_fp16() for consistent precision selection.
    
    Args:
        device: torch.device to check. If None, uses default device.
        
    Returns:
        True if FP16 should be used
    """
    mm = _get_model_management()
    if device is None:
        device = mm.get_torch_device()
    return mm.should_use_fp16(device)


def should_use_bf16(device: torch.device = None) -> bool:
    """Check if BF16 should be used for the given device.
    
    Uses ComfyUI's should_use_bf16() for consistent precision selection.
    
    Args:
        device: torch.device to check. If None, uses default device.
        
    Returns:
        True if BF16 should be used
    """
    mm = _get_model_management()
    if device is None:
        device = mm.get_torch_device()
    return mm.should_use_bf16(device)


def supports_fp8_compute(device: torch.device = None) -> bool:
    """Check if FP8 compute is supported on the given device.
    
    Args:
        device: torch.device to check. If None, uses default device.
        
    Returns:
        True if FP8 compute is supported
    """
    mm = _get_model_management()
    if device is None:
        device = mm.get_torch_device()
    return mm.supports_fp8_compute(device)


def get_autocast_device(device: torch.device) -> str:
    """Get the device type string for torch.autocast.
    
    Args:
        device: torch.device instance
        
    Returns:
        Device type string suitable for torch.autocast
    """
    mm = _get_model_management()
    return mm.get_autocast_device(device)


def is_gpu_device(device_type: str) -> bool:
    """Check if a device type is a GPU device.
    
    Args:
        device_type: Device type string
        
    Returns:
        True if the device is a GPU type
    """
    gpu_types = {DEVICE_CUDA, DEVICE_MPS, DEVICE_XPU, DEVICE_NPU, DEVICE_DIRECTML, DEVICE_HIP}
    return device_type.lower() in gpu_types


def get_device_memory_info(device_type: str = None) -> dict:
    """Get memory information for a device.
    
    Args:
        device_type: Device type string. If None, uses default device.
        
    Returns:
        Dictionary with 'total', 'used', 'free' memory in bytes
    """
    device = get_torch_device(device_type)
    
    if device.type == DEVICE_CUDA:
        try:
            total = torch.cuda.get_device_properties(device).total_memory
            reserved = torch.cuda.memory_reserved(device)
            allocated = torch.cuda.memory_allocated(device)
            free = total - reserved + (reserved - allocated)
            return {
                'total': total,
                'used': allocated,
                'free': free
            }
        except:
            pass
    
    # Return empty dict for non-CUDA devices or on error
    return {}


# Device options for node schema
def get_device_options() -> List[str]:
    """Get list of device options for node dropdown.
    
    Returns:
        List of device type strings for UI dropdown
    """
    return get_available_devices()


def get_device_option_labels() -> dict:
    """Get mapping of device options to display labels.
    
    Returns:
        Dictionary mapping device types to display names
    """
    return {dev: get_device_display_name(dev) for dev in get_available_devices()}