#!/bin/bash

# Docker Compose 安装路径
DOCKER_COMPOSE_PATH="/usr/local/bin/docker-compose"

install_runc() {
    echo "正在安装 runc..."
    sudo apt-get update
    sudo apt-get install -y runc
    echo "runc 安装完成。"
}

install_containerd_and_cni() {
    install_runc
    echo "正在安装 containerd 和 CNI 插件..."
    sudo apt-get update
    sudo apt-get install -y containerd uidmap jq
    sudo mkdir -p /etc/containerd
    containerd config default | sudo tee /etc/containerd/config.toml

    sudo mkdir -p /opt/cni/bin
    curl -L https://github.com/containernetworking/plugins/releases/download/v1.1.1/cni-plugins-linux-amd64-v1.1.1.tgz | sudo tar -C /opt/cni/bin -xz

    sudo tee -a /etc/containerd/config.toml > /dev/null <<EOT
[plugins."io.containerd.grpc.v1.cri".containerd.runtimes.nvidia]
  runtime_type = "io.containerd.runc.v2"
  [plugins."io.containerd.grpc.v1.cri".containerd.runtimes.nvidia.options]
    BinaryName = "/usr/bin/nvidia-container-runtime"
EOT

    sudo systemctl restart containerd
    echo "containerd 和 CNI 插件已安装并配置。"

    echo "正在安装和配置 containerd-rootless..."
    
    # 获取最新的 containerd 版本
    LATEST_CONTAINERD_VERSION=$(curl -s https://api.github.com/repos/containerd/containerd/releases/latest | jq -r .tag_name)
    curl -LO https://github.com/containerd/containerd/releases/download/${LATEST_CONTAINERD_VERSION}/containerd-${LATEST_CONTAINERD_VERSION}-linux-amd64.tar.gz
    sudo tar -xvf containerd-${LATEST_CONTAINERD_VERSION}-linux-amd64.tar.gz -C /usr/local/bin/

    # 切换到普通用户执行安装脚本
    sudo -u $(getent passwd 1000 | cut -d: -f1) bash << EOF
containerd-rootless-setuptool.sh install
echo "设置环境变量..."
echo 'export XDG_RUNTIME_DIR=/run/user/$(id -u)' >> ~/.bashrc
EOF

    echo "containerd-rootless 安装和配置完成。"
}

install_buildkit() {
    echo "正在安装 buildkit..."
    
    # 获取最新的 buildkit 版本
    LATEST_BUILDKIT_VERSION=$(curl -s https://api.github.com/repos/moby/buildkit/releases/latest | jq -r .tag_name)
    curl -LO https://github.com/moby/buildkit/releases/download/${LATEST_BUILDKIT_VERSION}/buildkit-${LATEST_BUILDKIT_VERSION}.linux-amd64.tar.gz
    sudo tar -xvf buildkit-${LATEST_BUILDKIT_VERSION}.linux-amd64.tar.gz -C /usr/local/
    sudo rm buildkit-${LATEST_BUILDKIT_VERSION}.linux-amd64.tar.gz

    # 配置 systemd 服务
    sudo tee /etc/systemd/system/buildkit.service > /dev/null <<EOT
[Unit]
Description=BuildKit
Documentation=https://github.com/moby/buildkit

[Service]
Type=notify
NotifyAccess=all
ExecStart=/usr/local/bin/buildkitd --addr unix:///run/buildkit/buildkitd.sock
Restart=always
RestartSec=3

[Install]
WantedBy=multi-user.target
EOT

    sudo systemctl enable buildkit
    sudo systemctl start buildkit
    echo "buildkit 安装并配置完成。"
}

install_docker() {
    echo "正在安装或更新 Docker..."
    sudo apt update
    sudo apt install -y apt-transport-https ca-certificates curl software-properties-common
    curl -fsSL https://download.docker.com/linux/$(. /etc/os-release;echo $ID)/gpg | sudo gpg --dearmor -o /usr/share/keyrings/docker-archive-keyring.gpg
    echo "deb [arch=$(dpkg --print-architecture) signed-by=/usr/share/keyrings/docker-archive-keyring.gpg] https://download.docker.com/linux/$(. /etc/os-release;echo $ID) $(lsb_release -cs) stable" | sudo tee /etc/apt/sources.list.d/docker.list > /dev/null
    sudo apt update
    sudo apt install -y docker-ce docker-ce-cli --no-install-recommends
    echo "Docker 安装或更新完成。"
}

reinstall_docker() {
    echo "正在重新安装 Docker..."
    sudo apt-get remove -y docker docker-engine docker.io runc
    sudo apt update
    sudo apt install -y apt-transport-https ca-certificates curl software-properties-common
    curl -fsSL https://download.docker.com/linux/$(. /etc/os-release;echo $ID)/gpg | sudo gpg --dearmor -o /usr/share/keyrings/docker-archive-keyring.gpg
    echo "deb [arch=$(dpkg --print-architecture) signed-by=/usr/share/keyrings/docker-archive-keyring.gpg] https://download.docker.com/linux/$(. /etc/os-release;echo $ID) $(lsb_release -cs) stable" | sudo tee /etc/apt/sources.list.d/docker.list > /dev/null
    sudo apt update
    sudo apt install -y docker-ce docker-ce-cli containerd.io --no-install-recommends
    sudo systemctl enable docker
    sudo systemctl start docker
    echo "Docker 已重新安装。"
}

install_jq() {
    echo "正在安装 jq..."
    sudo apt-get update
    sudo apt-get install -y jq
    echo "jq 安装完成。"
}

install_docker_compose() {
    install_jq
    echo "正在安装 Docker Compose..."

    # 使用 GitHub API 获取最新版本
    LATEST_RELEASE=$(curl -s https://api.github.com/repos/docker/compose/releases/latest | jq -r .tag_name)
    if [ -z "$LATEST_RELEASE" ]; then
        echo "无法获取 Docker Compose 的最新版本。"
        exit 1
    fi

    # 更新 Docker Compose 安装 URL 以使用最新版本
    DOCKER_COMPOSE_URL="https://ghproxy.dockless.eu.org/https://github.com/docker/compose/releases/download/${LATEST_RELEASE}/docker-compose-$(uname -s)-$(uname -m)"

    sudo curl -L $DOCKER_COMPOSE_URL -o $DOCKER_COMPOSE_PATH
    sudo chmod +x $DOCKER_COMPOSE_PATH
    echo "Docker Compose ${LATEST_RELEASE} 已安装。"
}

install_nvidia_docker() {
    echo "正在安装 NVIDIA Docker..."
    sudo apt-get update
    sudo apt-get install -y curl gnupg lsb-release

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

    sudo apt-get update
    sudo apt-get install -y nvidia-container-toolkit
    sudo nvidia-ctk runtime configure --runtime=docker
    sudo systemctl restart docker
    echo "NVIDIA Docker 已安装。"
}

configure_docker() {
    echo "正在配置 Docker..."
    mkdir -p /opt/docker
    sudo tee /etc/docker/daemon.json > /dev/null <<EOT
{
    "debug": true,
    "log-level": "info",
    "runtimes": {
        "nvidia": {
            "path": "nvidia-container-runtime",
            "runtimeArgs": []
        },
        "containerd": {
            "path": "containerd",
            "runtimeArgs": []
        }
    },
    "default-runtime": "containerd",
    "live-restore": true,
    "data-root": "/opt/docker",
    "features": {
      "buildkit": true,
      "containerd": true
    },
    "experimental": true
}
EOT
    sudo systemctl restart docker
    echo "Docker 已配置。"
}

install_nerdctl() {
    install_jq
    echo "正在安装 nerdctl..."
    
    # 获取最新版本
    LATEST_VERSION=$(curl -s https://api.github.com/repos/containerd/nerdctl/releases/latest | jq -r .tag_name)
    if [ -z "$LATEST_VERSION" ]; then
        echo "无法获取 nerdctl 的最新版本。"
        exit 1
    fi

    # 下载并安装 nerdctl
    curl -LO https://github.com/containerd/nerdctl/releases/download/${LATEST_VERSION}/nerdctl-${LATEST_VERSION}-linux-amd64.tar.gz
    sudo tar xzvf nerdctl-${LATEST_VERSION}-linux-amd64.tar.gz -C /usr/local/bin
    sudo rm nerdctl-${LATEST_VERSION}-linux-amd64.tar.gz
    nerdctl version
    echo "nerdctl ${LATEST_VERSION} 已安装。"
}

uninstall_docker() {
    echo "正在卸载 Docker..."
    sudo systemctl stop docker
    sudo apt-get purge -y docker-ce docker-ce-cli containerd.io
    sudo rm -rf /var/lib/docker
    sudo rm -rf /etc/docker
    sudo rm /etc/systemd/system/docker.service
    sudo rm /etc/systemd/system/docker.socket
    sudo systemctl daemon-reload
    echo "Docker 已卸载。"
}

uninstall_all() {
    echo "正在卸载所有组件..."
    uninstall_docker
    sudo apt-get purge -y containerd runc
    sudo rm -rf /etc/containerd
    sudo rm -rf /opt/cni/bin
    sudo rm -rf /usr/local/bin/nerdctl
    sudo systemctl daemon-reload
    echo "所有组件已卸载。"
}

install_all() {
    install_containerd_and_cni
    install_docker
    install_docker_compose
    install_nerdctl
    install_nvidia_docker
    install_buildkit
    configure_docker
}

echo "选择操作："
echo "1) 安装 Docker Compose"
echo "2) 卸载 Docker Compose"
echo "3) 重新安装 Docker"
echo "4) 卸载 Docker"
echo "5) 安装 NVIDIA Docker"
echo "6) 安装 containerd 和 CNI 插件"
echo "7) 配置 Docker"
echo "8) 全部安装"
echo "9) 卸载所有组件"
echo "10) 安装 nerdctl"
echo "11) 安装 buildkit"
echo "12) 退出"
read -p "输入选择（1-12）：" choice

case $choice in
    1) install_docker_compose ;;
    2) uninstall_docker_compose ;;
    3) reinstall_docker ;;
    4) uninstall_docker ;;
    5) install_nvidia_docker ;;
    6) install_containerd_and_cni ;;
    7) configure_docker ;;
    8) install_all ;;
    9) uninstall_all ;;
    10) install_nerdctl ;;
    11) install_buildkit ;;
    12) exit 0 ;;
    *) echo "无效输入。" ;;
esac