#!/usr/bin/env bash
# ==============================================================================
# FlyGuard: Удобный скрипт сборки и запуска Docker-контейнера
# Автоматически определяет наличие NVIDIA GPU и передает флаги --gpus all
# ==============================================================================
set -e

DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
IMAGE_NAME="flyguard:latest"
DATA_DIR="${FLYGUARD_DATA:-$DIR/data}"
ARTIFACTS_DIR="$DIR/artifacts"

mkdir -p "$DATA_DIR" "$ARTIFACTS_DIR"

# Определение наличия драйвера NVIDIA
GPU_FLAG=""
if command -v nvidia-smi &>/dev/null; then
    GPU_FLAG="--gpus all"
fi

usage() {
    echo "Использование: $0 [команда]"
    echo ""
    echo "Команды:"
    echo "  build         Собрать универсальный образ с поддержкой NVIDIA GPU ($IMAGE_NAME)"
    echo "  build-cpu     Собрать облегченный CPU-образ (flyguard:cpu)"
    echo "  run           Общий запуск конвейера (по умолчанию обработка всех бэгов)"
    echo "  pipeline      Запуск конвейера с дополнительными аргументами"
    echo "  test          Запуск всех 40+ тестов верификации"
    echo "  benchmark     Генерация бенчмарка с ускорением на GPU"
    echo "  evaluate      Оценка качества моделей и кросс-валидация"
    echo "  train-mbon    Обучение клеток Кеньона (50 000 KC) на GPU"
    echo "  train-track   Обучение классификатора треков"
    echo "  info          Вывод телеметрии и статуса GPU из контейнера"
    echo "  shell         Интерактивная командная строка (bash) внутри контейнера"
    echo ""
    echo "Примеры:"
    echo "  $0 build"
    echo "  $0 run"
    echo "  $0 pipeline --all --verbose"
    echo "  $0 train-mbon --device cuda --n-kc 50000"
}

CMD="${1:-run}"

case "$CMD" in
    build)
        echo "==> Сборка универсального образа FlyGuard ($IMAGE_NAME)..."
        docker build -t "$IMAGE_NAME" -f "$DIR/Dockerfile" "$DIR"
        ;;
    build-cpu)
        echo "==> Сборка облегченного CPU-образа FlyGuard (flyguard:cpu)..."
        docker build --target cpu -t flyguard:cpu -f "$DIR/Dockerfile" "$DIR"
        ;;
    run|default)
        shift 2>/dev/null || true
        echo "==> Запуск контейнера FlyGuard (GPU: ${GPU_FLAG:-CPU Fallback})..."
        docker run --rm -it $GPU_FLAG \
            -v "$DATA_DIR":/data:ro \
            -v "$ARTIFACTS_DIR":/app/artifacts \
            --shm-size=8g \
            "$IMAGE_NAME" "$@"
        ;;
    pipeline)
        shift
        docker run --rm -it $GPU_FLAG \
            -v "$DATA_DIR":/data:ro \
            -v "$ARTIFACTS_DIR":/app/artifacts \
            --shm-size=8g \
            "$IMAGE_NAME" pipeline "$@"
        ;;
    test|tests)
        shift
        docker run --rm -it $GPU_FLAG \
            -v "$DATA_DIR":/data:ro \
            -v "$ARTIFACTS_DIR":/app/artifacts \
            --shm-size=8g \
            "$IMAGE_NAME" test "$@"
        ;;
    benchmark)
        shift
        docker run --rm -it $GPU_FLAG \
            -v "$DATA_DIR":/data:ro \
            -v "$ARTIFACTS_DIR":/app/artifacts \
            --shm-size=8g \
            "$IMAGE_NAME" benchmark "$@"
        ;;
    evaluate)
        shift
        docker run --rm -it $GPU_FLAG \
            -v "$DATA_DIR":/data:ro \
            -v "$ARTIFACTS_DIR":/app/artifacts \
            --shm-size=8g \
            "$IMAGE_NAME" evaluate "$@"
        ;;
    train-mbon)
        shift
        docker run --rm -it $GPU_FLAG \
            -v "$DATA_DIR":/data \
            -v "$ARTIFACTS_DIR":/app/artifacts \
            --shm-size=8g \
            "$IMAGE_NAME" train-mbon "$@"
        ;;
    train-track)
        shift
        docker run --rm -it $GPU_FLAG \
            -v "$DATA_DIR":/data \
            -v "$ARTIFACTS_DIR":/app/artifacts \
            --shm-size=8g \
            "$IMAGE_NAME" train-track "$@"
        ;;
    info|gpu|device)
        docker run --rm $GPU_FLAG "$IMAGE_NAME" info
        ;;
    shell|bash)
        shift
        docker run --rm -it $GPU_FLAG \
            -v "$DATA_DIR":/data \
            -v "$ARTIFACTS_DIR":/app/artifacts \
            --shm-size=8g \
            "$IMAGE_NAME" bash "$@"
        ;;
    help|--help|-h)
        usage
        exit 0
        ;;
    *)
        docker run --rm -it $GPU_FLAG \
            -v "$DATA_DIR":/data:ro \
            -v "$ARTIFACTS_DIR":/app/artifacts \
            --shm-size=8g \
            "$IMAGE_NAME" "$@"
        ;;
esac