#!/bin/bash
# Скрипт для проверки подключения GPU в Docker контейнере Ollama

set -e

CONTAINER_NAME="ollama"

echo "=== Проверка GPU в Docker контейнере ==="
echo ""

# Определяем команду docker
DOCKER_CMD="docker"
if ! docker ps >/dev/null 2>&1; then
    DOCKER_CMD="sudo docker"
fi

echo "1. Проверка конфигурации в docker-compose.yml..."
if grep -q "driver: nvidia" docker-compose.yml 2>/dev/null; then
    echo "   ✓ GPU настроен в docker-compose.yml"
    echo "   Конфигурация:"
    grep -A 5 "driver: nvidia" docker-compose.yml | head -6
else
    echo "   ✗ GPU не настроен в docker-compose.yml"
fi

echo ""
echo "2. Проверка конфигурации контейнера..."
GPU_CONFIG=$($DOCKER_CMD inspect $CONTAINER_NAME 2>/dev/null | grep -i "nvidia\|gpu" | head -5)
if [ -n "$GPU_CONFIG" ]; then
    echo "   ✓ Найдена конфигурация GPU:"
    echo "$GPU_CONFIG"
else
    echo "   ⚠ Конфигурация GPU не найдена в inspect"
fi

echo ""
echo "3. Проверка доступности nvidia-smi в контейнере..."
if $DOCKER_CMD exec $CONTAINER_NAME which nvidia-smi >/dev/null 2>&1; then
    echo "   ✓ nvidia-smi доступен"
    echo ""
    echo "   Вывод nvidia-smi:"
    $DOCKER_CMD exec $CONTAINER_NAME nvidia-smi 2>&1 | head -15 || echo "   ⚠ nvidia-smi не может выполниться"
else
    echo "   ✗ nvidia-smi не найден в контейнере"
    echo "   Возможно, контейнер не имеет доступа к GPU"
fi

echo ""
echo "4. Проверка переменных окружения NVIDIA..."
NVIDIA_ENV=$($DOCKER_CMD exec $CONTAINER_NAME env | grep -i nvidia)
if [ -n "$NVIDIA_ENV" ]; then
    echo "   ✓ Переменные NVIDIA найдены:"
    echo "$NVIDIA_ENV"
else
    echo "   ⚠ Переменные NVIDIA не найдены"
fi

echo ""
echo "5. Проверка логов Ollama на использование GPU..."
RECENT_LOGS=$($DOCKER_CMD logs $CONTAINER_NAME --tail 50 2>&1)
if echo "$RECENT_LOGS" | grep -qi "gpu\|cuda\|nvidia"; then
    echo "   ✓ Найдены упоминания GPU в логах:"
    echo "$RECENT_LOGS" | grep -i "gpu\|cuda\|nvidia" | tail -5
else
    echo "   ⚠ Нет упоминаний GPU в последних логах"
fi

echo ""
echo "6. Проверка на хосте..."
if command -v nvidia-smi >/dev/null 2>&1; then
    echo "   GPU на хосте:"
    nvidia-smi --query-gpu=name,driver_version,memory.total --format=csv,noheader 2>/dev/null || echo "   ⚠ Не удалось получить информацию о GPU"
else
    echo "   ⚠ nvidia-smi не найден на хосте"
fi

echo ""
echo "=== Резюме ==="
echo ""
if $DOCKER_CMD exec $CONTAINER_NAME nvidia-smi >/dev/null 2>&1; then
    echo "✓ GPU подключен и доступен в контейнере"
    echo ""
    echo "Для детальной информации выполните:"
    echo "  sudo docker exec ollama nvidia-smi"
else
    echo "✗ GPU не доступен в контейнере"
    echo ""
    echo "Возможные причины:"
    echo "  1. Docker не имеет доступа к GPU (нужен nvidia-docker2 или nvidia-container-toolkit)"
    echo "  2. Контейнер не был перезапущен после изменения docker-compose.yml"
    echo "  3. Драйверы NVIDIA не установлены на хосте"
    echo ""
    echo "Решение:"
    echo "  1. Установите nvidia-container-toolkit:"
    echo "     sudo apt-get install -y nvidia-container-toolkit"
    echo "     sudo systemctl restart docker"
    echo ""
    echo "  2. Перезапустите контейнер:"
    echo "     docker compose restart ollama"
fi