#!/bin/bash
echo "开始安装，请稍等（约20分钟）...."
# 更新源
apt remove needrestart -y > /dev/null 2>&1
apt update > /dev/null 2>&1

# 安装依赖
apt-get install -y ca-certificates curl gnupg lsb-release > /dev/null 2>&1

# 添加 Docker GPG 密钥
curl -fsSL http://mirrors.aliyun.com/docker-ce/linux/ubuntu/gpg | sudo apt-key add - 

# 更新源列表
printf '\n' | add-apt-repository "deb [arch=amd64] http://mirrors.aliyun.com/docker-ce/linux/ubuntu $(lsb_release -cs) stable" 
# 安装 Docker 和 Docker Compose
apt-get install -y docker-ce docker-ce-cli containerd.io docker-compose 
# 检查 Docker 和 Docker Compose 是否安装成功
docker_version=$(docker --version | grep -oP 'Docker version \K[0-9.]+')
if [ -n "$docker_version" ]; then
    echo "Docker版本：$docker_version"
else
    echo "docker安装失败"
fi

# 获取 Docker Compose 版本
docker_compose_version=$(docker-compose --version | grep -oP 'docker-compose version \K[0-9.]+')
if [ -n "$docker_compose_version" ]; then
    echo "docker-compose版本：$docker_compose_version"
else
    echo "docker-compose安装失败。"
fi

# 导入 NVIDIA Docker 源到 /etc/apt/sources.list.d/nvidia-docker.list > /dev/null 2>&1
echo "deb https://nvidia.github.io/libnvidia-container/stable/deb/\$(ARCH) /" | tee /etc/apt/sources.list.d/nvidia-docker.list > /dev/null 2>&1
echo "deb https://nvidia.github.io/libnvidia-container/experimental/deb/\$(ARCH) /" | tee -a /etc/apt/sources.list.d/nvidia-docker.list > /dev/null 2>&1
curl -s -L https://nvidia.github.io/libnvidia-container/gpgkey | sudo apt-key add -  > /dev/null 2>&1
apt update > /dev/null 2>&1

# 安装 NVIDIA Docker 2
apt install -y nvidia-docker2 > /dev/null 2>&1

# 完成
echo "Docker 和 NVIDIA Docker 安装完成！"
read -p "是否要修改 /etc/docker/daemon.json 文件以使用 NVIDIA 容器运行时？(y/n): " modify_daemon_json

if [[ "$modify_daemon_json" == "y" || "$modify_daemon_json" == "Y" ]]; then
    # 创建或修改 /etc/docker/daemon.json 文件
    echo '{
        "default-runtime": "nvidia",
        "runtimes": {
            "nvidia": {
                "path": "nvidia-container-runtime",
                "runtimeArgs": []
            }
        }
    }' | sudo tee /etc/docker/daemon.json > /dev/null

    # 重新启动 Docker 服务以应用更改
    sudo systemctl restart docker
    echo "/etc/docker/daemon.json 已修改并已重新启动 Docker 服务。"
else
    echo "未修改 /etc/docker/daemon.json 文件。"
fi