import psutil
from .logger import logger

# 尝试导入GPUtil，如果失败则设置has_gpu为False
try:
    import GPUtil
    GPUtil_available = True
except ImportError:
    logger.warning("GPUtil模块未找到，将禁用GPU监控")
    GPUtil_available = False

class SystemMonitor:
    def __init__(self):
        self.cpu_usage = 0.0
        self.memory_usage = 0.0
        self.gpu_memory_usage = 0.0
        self.has_gpu = self._check_gpu()
        
    def _check_gpu(self):
        """检查是否有可用的GPU"""
        if not GPUtil_available:
            return False
        try:
            gpus = GPUtil.getGPUs()
            return len(gpus) > 0
        except Exception as e:
            logger.warning(f"GPU检查失败: {e}")
            return False
    
    def update_metrics(self):
        """更新系统资源使用情况"""
        try:
            # 更新CPU使用率
            self.cpu_usage = psutil.cpu_percent(interval=0.1)
            
            # 更新内存使用率
            memory = psutil.virtual_memory()
            self.memory_usage = memory.percent
            
            # 更新GPU内存使用率（如果有GPU）
            if self.has_gpu:
                try:
                    gpus = GPUtil.getGPUs()
                    if gpus:
                        self.gpu_memory_usage = gpus[0].memoryUtil * 100
                except Exception as e:
                    logger.warning(f"GPU内存使用率获取失败: {e}")
        except Exception as e:
            logger.error(f"系统资源监控更新失败: {e}")
    
    def get_status_text(self):
        """获取状态栏显示文本"""
        self.update_metrics()
        status = f"CPU: {self.cpu_usage:.1f}% | 内存: {self.memory_usage:.1f}%"
        if self.has_gpu:
            status += f" | GPU内存: {self.gpu_memory_usage:.1f}%"
        return status