# scripts/check_gpu.py

"""
Скрипт для точной диагностики окружения GPU: проверяет модель карты, объем и процент 
использования памяти, версию PyTorch, а также системную и встроенную версии CUDA.
"""
import logging
import subprocess
import re
import torch
from typing import List, Optional

# Настройка подробного логирования
logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s [%(levelname)s] %(message)s",
    datefmt="%H:%M:%S"
)
logger: logging.Logger = logging.getLogger(__name__)

def get_system_cuda_version() -> Optional[str]:
    """
    Извлекает системную версию CUDA из вывода утилиты nvidia-smi.
    
    Returns:
        Optional[str]: Строка с версией (например, '12.2') или None, если не найдено.
    """
    try:
        # Запрашиваем полный вывод nvidia-smi
        result: subprocess.CompletedProcess = subprocess.run(
            ["nvidia-smi"], stdout=subprocess.PIPE, text=True, check=True
        )
        # Ищем паттерн "CUDA Version: XX.X"
        match: Optional[re.Match] = re.search(r"CUDA Version:\s+(\d+\.\d+)", result.stdout)
        if match:
            return match.group(1)
    except Exception as e:
        logger.debug(f"Не удалось получить системную версию CUDA: {e}")
    return None

def main() -> None:
    """
    Главная функция для сбора и логирования подробной информации о GPU и окружении.
    """
    logger.info("🔍 Запуск точной диагностики окружения и видеокарты...")
    
    # 1. Информация о фреймворке
    logger.info(f"📦 Версия PyTorch: {torch.__version__}")
    logger.info(f"⚙️ Версия CUDA (сборка PyTorch): {torch.version.cuda}")
    
    # 2. Системная информация драйвера
    sys_cuda: Optional[str] = get_system_cuda_version()
    logger.info(f"🖥️ Версия CUDA (системный драйвер): {sys_cuda if sys_cuda else 'Неизвестно'}")
    
    if not torch.cuda.is_available():
        logger.error("❌ GPU не найден для PyTorch! Проверь настройки Colab.")
        return
        
    # 3. Физические метрики видеокарты
    try:
        cmd: List[str] = [
            "nvidia-smi",
            "--query-gpu=name,memory.total,memory.used",
            "--format=csv,noheader,nounits"
        ]
        
        result: subprocess.CompletedProcess = subprocess.run(
            cmd, stdout=subprocess.PIPE, text=True, check=True
        )
        
        lines: List[str] = result.stdout.strip().split('\n')
        
        for i, line in enumerate(lines):
            parts: List[str] = [p.strip() for p in line.split(',')]
            
            if len(parts) >= 3:
                gpu_name: str = parts[0]
                mem_total_mb: float = float(parts[1])
                mem_used_mb: float = float(parts[2])
                
                # Переводим в гигабайты и высчитываем процент занятости
                mem_total_gb: float = mem_total_mb / 1024
                mem_used_gb: float = mem_used_mb / 1024
                percent_used: float = (mem_used_mb / mem_total_mb) * 100
                
                logger.info(f"💻 Устройство {i}: {gpu_name}")
                logger.info(
                    f"🧠 Видеопамять: занято {mem_used_gb:.2f} GB из {mem_total_gb:.2f} GB "
                    f"({percent_used:.1f}%)"
                )
            else:
                logger.warning(f"⚠️ Неожиданный формат вывода nvidia-smi: {line}")
                
    except FileNotFoundError:
        logger.error("❌ Утилита nvidia-smi не найдена в системе.")
    except Exception as e:
        logger.error(f"❌ Ошибка при диагностике GPU: {e}", exc_info=True)
        
    logger.info("=" * 50)

if __name__ == "__main__":
    main()