#!/bin/bash
# ============================================================
#  Установка окружения MLflow Practice
#  Поддержка: Linux (GPU-сервер, RTX 3090) и macOS (Intel/Apple Silicon)
#  Запускать: bash setup_server.sh
# ============================================================
set -e

echo "============================================"
echo "  MLflow Practice — установка окружения"
echo "============================================"

# --- Определение ОС ---
OS="$(uname -s)"
case "$OS" in
    Darwin*) PLATFORM="macos" ;;
    Linux*)  PLATFORM="linux" ;;
    *)       echo "❌ Неподдерживаемая ОС: $OS"; exit 1 ;;
esac
ARCH="$(uname -m)"
echo "📍 Платформа: $PLATFORM ($ARCH)"

# --- Выбор интерпретатора Python для создания venv ---
# Предпочитаем python3.13, при отсутствии — любой python3.
if command -v python3.13 &> /dev/null; then
    PYBIN="python3.13"
elif command -v python3 &> /dev/null; then
    PYBIN="python3"
else
    echo "❌ python3 не найден в PATH"
    exit 1
fi
echo "🐍 Интерпретатор: $PYBIN ($($PYBIN --version 2>&1))"

# --- Создание venv (с обработкой отсутствия ensurepip/venv-пакета) ---
create_venv() {
    if [ -d ".venv" ]; then
        echo "   .venv уже существует — пропускаю создание"
        return 0
    fi
    if "$PYBIN" -m venv .venv; then
        return 0
    fi
    echo "⚠️  '$PYBIN -m venv .venv' не удался (частая причина — не установлен"
    echo "   пакет python3-venv / ensurepip в дистрибутиве)."
    if [ "$PLATFORM" = "linux" ]; then
        echo "   Debian/Ubuntu:  sudo apt-get install -y python3.13-venv python3.13-dev"
        echo "                  (или python3-venv, если версия другая)"
        echo "   Fedora/RHEL:    sudo dnf install -y python3-devel"
    fi
    echo "   Пробую fallback: venv без pip + bootstrap через get-pip.py..."
    if ! "$PYBIN" -m venv --without-pip .venv; then
        echo "❌ Не удалось создать venv даже без pip. Установите python3-venv и повторите."
        return 1
    fi
    if ! curl -fsSL https://bootstrap.pypa.io/get-pip.py -o /tmp/get-pip.py; then
        echo "❌ Не удалось скачать get-pip.py (нет сети?). Установите pip вручную."
        return 1
    fi
    if ! .venv/bin/python /tmp/get-pip.py; then
        echo "❌ bootstrap pip через get-pip.py не удался."
        return 1
    fi
    rm -f /tmp/get-pip.py
    return 0
}

# ============================================
#  macOS — установка без CUDA (есть MPS на Apple Silicon)
# ============================================
if [ "$PLATFORM" = "macos" ]; then

    echo ""
    echo "🍎 macOS: подготовка системы..."

    # --- Homebrew ---
    if ! command -v brew &> /dev/null; then
        echo "⚠️  Homebrew не найден. Установите его вручную:"
        echo "    /bin/bash -c \"\$(curl -fsSL https://raw.githubusercontent.com/Homebrew/install/HEAD/install.sh)\""
        echo "    и запустите скрипт снова."
        exit 1
    fi
    echo "   Homebrew: $(brew --version | head -1)"

    # --- Python через brew (если нет python3.13) ---
    if ! command -v python3.13 &> /dev/null; then
        echo "📥 Устанавливаю python@3.13 через brew..."
        brew install python@3.13
        # brew ставит команду python3.13
    fi

    # --- Создание venv ---
    echo ""
    echo "📦 Создание виртуального окружения venv..."
    create_venv || exit 1
    source .venv/bin/activate
    echo "   Активировано: $(which python)"

    # --- pip ---
    echo ""
    echo "⬆️  Обновление pip..."
    pip install --upgrade pip wheel setuptools

    # --- MLflow + базовые библиотеки ---
    echo ""
    echo "📥 Установка MLflow + scikit-learn + matplotlib..."
    pip install mlflow scikit-learn pandas numpy matplotlib

    # --- PyTorch (CPU на Intel, MPS на Apple Silicon) ---
    echo ""
    echo "🔥 Установка PyTorch (PyPI — без CUDA)..."
    pip install torch torchvision

    # --- Проверка PyTorch + устройства ---
    echo ""
    echo "✅ Проверка PyTorch..."
    python -c "
import torch, platform
print(f'PyTorch: {torch.__version__}')
print(f'MPS available: {torch.backends.mps.is_available()}')
if torch.backends.mps.is_available():
    print('Устройство: Apple Silicon GPU (MPS)')
else:
    print('Устройство: CPU (Intel Mac)')
print('ℹ️  На macOS CUDA недоступна — train_gpu.py будет использовать CPU/MPS')
"

# ============================================
#  Linux — GPU-сервер с CUDA (RTX 3090)
# ============================================
else

    # --- Проверка GPU ---
    echo ""
    echo "📡 Проверка GPU..."
    if command -v nvidia-smi &> /dev/null; then
        nvidia-smi --query-gpu=name,memory.total,driver_version --format=csv,noheader
    else
        echo "⚠️  nvidia-smi не найден — CUDA может быть не установлена"
    fi

    # --- Создание venv ---
    echo ""
    echo "🐍 Проверка Python..."
    "$PYBIN" --version

    echo ""
    echo "📦 Создание виртуального окружения venv..."
    create_venv || exit 1
    source .venv/bin/activate
    echo "   Активировано: $(which python)"

    # --- pip ---
    echo ""
    echo "⬆️  Обновление pip..."
    pip install --upgrade pip wheel setuptools

    # --- MLflow + базовые библиотеки ---
    echo ""
    echo "📥 Установка MLflow + scikit-learn + matplotlib..."
    pip install mlflow scikit-learn pandas numpy matplotlib

    # --- PyTorch под CUDA 12.1 ---
    echo ""
    echo "🔥 Установка PyTorch (CUDA 12.1)..."
    pip install torch torchvision --index-url https://download.pytorch.org/whl/cu121

    # --- Проверка PyTorch + CUDA ---
    echo ""
    echo "✅ Проверка PyTorch + CUDA..."
    python -c "
import torch
print(f'PyTorch: {torch.__version__}')
print(f'CUDA available: {torch.cuda.is_available()}')
if torch.cuda.is_available():
    print(f'GPU: {torch.cuda.get_device_name(0)}')
    print(f'VRAM: {torch.cuda.get_device_properties(0).total_memory / 1e9:.1f} GB')
"
fi

# ============================================
#  Общая проверка MLflow (обе платформы)
# ============================================
echo ""
echo "✅ Проверка MLflow..."
python -c "import mlflow; print(f'MLflow: {mlflow.__version__}')"

echo ""
echo "============================================"
echo "  ✅ Установка завершена! ($PLATFORM)"
echo "============================================"
echo ""
echo "Следующие шаги:"
echo "  1. Активируйте окружение:  source .venv/bin/activate"
echo "  2. Запустите MLflow UI:    ./start_ui.sh"
echo "  3. В другом терминале:"
echo "     source .venv/bin/activate"
if [ "$PLATFORM" = "macos" ]; then
    echo "     python src/train_simple.py --n-estimators 100 --max-depth 8"
    echo "     python src/train_gpu.py --epochs 3 --batch-size 64 --lr 0.01  # CPU/MPS, меньше эпох"
else
    echo "     python src/train_simple.py --n-estimators 100 --max-depth 8"
    echo "     python src/train_gpu.py --epochs 10 --batch-size 256 --lr 0.001"
fi
echo ""