#!/usr/bin/env bash
# ══════════════════════════════════════════════════
# 00_install_docker_ubuntu.sh
# Ubuntu 22.04 LTS 安装 Docker CE + nvidia-container-toolkit
# 用法：bash scripts/00_install_docker_ubuntu.sh
# ══════════════════════════════════════════════════
set -euo pipefail

RED='\033[0;31m'; GREEN='\033[0;32m'; YELLOW='\033[1;33m'; BLUE='\033[0;34m'; NC='\033[0m'
info()  { echo -e "${BLUE}[INFO]${NC} $*"; }
ok()    { echo -e "${GREEN}[OK]${NC} $*"; }
warn()  { echo -e "${YELLOW}[WARN]${NC} $*"; }
error() { echo -e "${RED}[ERROR]${NC} $*"; exit 1; }

# ── 检查 root 权限 ──────────────────────────────
if [[ $EUID -ne 0 ]]; then
    error "请以 root 或 sudo 运行：sudo bash scripts/00_install_docker_ubuntu.sh"
fi

# ── 检测 Ubuntu 版本 ────────────────────────────
. /etc/os-release
info "检测到 OS：$NAME $VERSION_ID"
if [[ "$ID" != "ubuntu" ]]; then
    warn "非 Ubuntu 系统，脚本可能不适用。继续（y/n）？"
    read -r ans; [[ "$ans" != "y" ]] && exit 0
fi

# ── Step 1：换国内源（可选）──────────────────────
info "Step 1/5：配置 APT 源..."
if [[ "${USE_MIRROR:-false}" == "true" ]]; then
    sed -i 's/archive.ubuntu.com/mirrors.aliyun.com/g' /etc/apt/sources.list
    sed -i 's/security.ubuntu.com/mirrors.aliyun.com/g' /etc/apt/sources.list
    ok "已切换到阿里云镜像"
fi
apt-get update -qq

# ── Step 2：安装依赖 ────────────────────────────
info "Step 2/5：安装依赖包..."
apt-get install -y -qq \
    ca-certificates \
    curl \
    gnupg \
    lsb-release \
    apt-transport-https

# ── Step 3：安装 Docker CE ──────────────────────
info "Step 3/5：安装 Docker CE..."
if command -v docker &>/dev/null; then
    DOCKER_VER=$(docker --version)
    warn "Docker 已安装：$DOCKER_VER"
    warn "跳过 Docker 安装。如需重装，请先运行：apt-get remove docker docker-engine docker.io containerd"
else
    # 添加 Docker 官方 GPG 密钥
    install -m 0755 -d /etc/apt/keyrings
    curl -fsSL https://download.docker.com/linux/ubuntu/gpg | \
        gpg --dearmor -o /etc/apt/keyrings/docker.gpg
    chmod a+r /etc/apt/keyrings/docker.gpg

    # 添加 Docker 仓库
    echo "deb [arch=$(dpkg --print-architecture) signed-by=/etc/apt/keyrings/docker.gpg] \
        https://download.docker.com/linux/ubuntu $(lsb_release -cs) stable" | \
        tee /etc/apt/sources.list.d/docker.list > /dev/null

    apt-get update -qq
    apt-get install -y -qq docker-ce docker-ce-cli containerd.io docker-buildx-plugin docker-compose-plugin

    # 启动并设置开机自启
    systemctl enable docker
    systemctl start docker
    ok "Docker CE 安装完成"
fi

# 验证
docker --version
docker compose version

# ── Step 4：将当前用户加入 docker 组 ────────────
info "Step 4/5：配置 Docker 用户组..."
CURRENT_USER=${SUDO_USER:-$USER}
if [[ -n "$CURRENT_USER" && "$CURRENT_USER" != "root" ]]; then
    usermod -aG docker "$CURRENT_USER"
    ok "用户 $CURRENT_USER 已加入 docker 组（重新登录后生效）"
fi

# ── Step 5：安装 nvidia-container-toolkit（可选）─
info "Step 5/5：检查 NVIDIA GPU..."
if command -v nvidia-smi &>/dev/null; then
    info "检测到 NVIDIA GPU，安装 nvidia-container-toolkit..."
    nvidia-smi --query-gpu=name --format=csv,noheader

    # 添加 NVIDIA 仓库
    curl -fsSL https://nvidia.github.io/libnvidia-container/gpgkey | \
        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' | \
        tee /etc/apt/sources.list.d/nvidia-container-toolkit.list

    apt-get update -qq
    apt-get install -y -qq nvidia-container-toolkit
    nvidia-ctk runtime configure --runtime=docker
    systemctl restart docker
    ok "nvidia-container-toolkit 安装完成"
else
    warn "未检测到 NVIDIA GPU，跳过 nvidia-container-toolkit 安装"
    warn "如有 GPU 请手动安装驱动后重新运行本脚本"
fi

echo ""
echo -e "${GREEN}══════════════════════════════════════════${NC}"
echo -e "${GREEN}  Docker 安装完成！${NC}"
echo -e "${GREEN}══════════════════════════════════════════${NC}"
echo ""
echo "  Docker 版本：$(docker --version)"
echo "  Compose 版本：$(docker compose version)"
echo ""
echo -e "${YELLOW}  注意：${NC}请重新登录以使 docker 组权限生效"
echo "  验证命令：docker run hello-world"