#!/bin/bash

# Данный скрипт написан ИИ для быстрой подготовки окружения, установка драйверов и докера
# Остановка скрипта при возникновении любой ошибки
set -e

# Цвета для красивого вывода в консоль
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
RED='\033[0;31m'
BLUE='\033[0;34m'
NC='\033[0m'

echo -e "${BLUE}[INFO]${NC} Инициализация проверки окружения для проекта EmoM..."

# 1. ПРОВЕРКА DOCKER
if ! command -v docker &> /dev/null; then
    echo -e "${YELLOW}[SETUP]${NC} Docker не найден. Начинаем установку..."
    # Использование официального скрипта установки Docker
    curl -fsSL https://get.docker.com -o get-docker.sh
    sudo sh get-docker.sh
    rm get-docker.sh
    
    # Добавляем текущего пользователя в группу docker, чтобы не писать sudo docker
    sudo usermod -aG docker $USER
    echo -e "${GREEN}[OK]${NC} Docker успешно установлен."
    echo -e "${YELLOW}[ВНИМАНИЕ]${NC} Для применения прав группы docker потребуется перезайти в сессию SSH после завершения скрипта."
else
    echo -e "${GREEN}[OK]${NC} Docker установлен ($(docker --version))."
fi

# 2. ПРОВЕРКА DOCKER COMPOSE
if ! docker compose version &> /dev/null; then
    echo -e "${YELLOW}[SETUP]${NC} Плагин Docker Compose не найден. Устанавливаем..."
    sudo apt-get update && sudo apt-get install -y docker-compose-plugin
    echo -e "${GREEN}[OK]${NC} Docker Compose успешно установлен."
else
    echo -e "${GREEN}[OK]${NC} Плагин Docker Compose доступен."
fi

# 3. ПРОВЕРКА NVIDIA И ПРОБРОСА GPU В DOCKER
if command -v nvidia-smi &> /dev/null; then
    echo -e "${GREEN}[OK]${NC} Драйверы NVIDIA обнаружены."
    
    # Проверяем наличие NVIDIA Container Toolkit
    if ! dpkg -l | grep -q nvidia-container-toolkit; then
        echo -e "${YELLOW}[SETUP]${NC} NVIDIA Container Toolkit не найден. Выполняется установка..."
        
        # Настройка репозиториев NVIDIA
        curl -fsSL https://nvidia.github.io/libnvidia-container/gpgkey | sudo gpg --dearmor -o /usr/share/keyrings/nvidia-container-toolkit-keyring.gpg
        curl -s -L https://nvidia.github.io/libnvidia-container/stable/deb/nvidia-container-toolkit.list | \
            sed 's#deb https://#deb [signed-by=/usr/share/keyrings/nvidia-container-toolkit-keyring.gpg] https://#g' | \
            sudo tee /etc/apt/sources.list.d/nvidia-container-toolkit.list
            
        sudo apt-get update
        sudo apt-get install -y nvidia-container-toolkit
        
        # Конфигурация Docker для работы с NVIDIA
        echo -e "${YELLOW}[SETUP]${NC} Конфигурация runtime NVIDIA для Docker..."
        sudo nvidia-ctk runtime configure --runtime=docker
        sudo systemctl restart docker
        
        echo -e "${GREEN}[OK]${NC} NVIDIA Container Toolkit установлен и настроен."
    else
        echo -e "${GREEN}[OK]${NC} NVIDIA Container Toolkit уже установлен."
    fi
else
    echo -e "${RED}[WARN]${NC} Утилита nvidia-smi не найдена! Убедитесь, что драйверы видеокарты установлены, иначе Docker будет использовать только CPU."
fi

echo -e "\n${BLUE}[INFO]${NC} ========================================="
echo -e "${GREEN}[SUCCESS]${NC} Окружение готово к работе!"
echo -e "Теперь вы можете запустить проект командой: ${YELLOW}make up${NC}"