#!/bin/bash

# 获取版本号从环境变量
CONTAINERD_VERSION=${CONTAINERD_VERSION:-1.6.8}
NERDCTL_VERSION=${NERDCTL_VERSION:-1.7.6}
CNI_VERSION=${CNI_VERSION:-1.1.1}
BUILDKIT_VERSION=${BUILDKIT_VERSION:-0.14.1}

# Docker Compose 安装路径
DOCKER_COMPOSE_PATH="/usr/local/bin/docker-compose"
LOGFILE="$(pwd)/install_log.txt"

install_docker() {
    echo "正在检查 Docker 是否已安装..." | tee -a $LOGFILE
    if ! [ -x "$(command -v docker)" ]; then
        echo "Docker 未安装。正在安装 Docker..." | tee -a $LOGFILE
        curl -fsSL https://get.docker.com -o get-docker.sh | tee -a $LOGFILE
        sh get-docker.sh | tee -a $LOGFILE
        echo "Docker 已安装。" | tee -a $LOGFILE
    else
        echo "Docker 已安装。" | tee -a $LOGFILE
    fi
}

install_docker_compose() {
    install_docker
    echo "正在安装 Docker Compose..." | tee -a $LOGFILE
    
    # 检查并安装 jq
    if ! [ -x "$(command -v jq)" ]; then
        echo "安装 jq..." | tee -a $LOGFILE
        sudo apt-get update && sudo apt-get install -y jq | tee -a $LOGFILE
    fi

    # 检查并安装 curl
    if ! [ -x "$(command -v curl)" ]; then
        echo "安装 curl..." | tee -a $LOGFILE
        sudo apt-get update && sudo apt-get install -y curl | tee -a $LOGFILE
    fi

    # 使用 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 的最新版本。" | tee -a $LOGFILE
        exit 1
    fi

    # 更新 Docker Compose 安装 URL 以使用最新版本
    DOCKER_COMPOSE_URL="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 | tee -a $LOGFILE
    sudo chmod +x $DOCKER_COMPOSE_PATH | tee -a $LOGFILE
    echo "Docker Compose ${LATEST_RELEASE} 已安装。" | tee -a $LOGFILE
}

uninstall_docker_compose() {
    echo "正在卸载 Docker Compose..." | tee -a $LOGFILE
    sudo rm -f $DOCKER_COMPOSE_PATH | tee -a $LOGFILE
    echo "Docker Compose 已卸载。" | tee -a $LOGFILE
}

install_containerd() {
    echo "正在安装 containerd..." | tee -a $LOGFILE
    sudo apt-get update | tee -a $LOGFILE
    sudo apt-get install -y containerd | tee -a $LOGFILE
    sudo mkdir -p /etc/containerd
    containerd config default | sudo tee /etc/containerd/config.toml | tee -a $LOGFILE
    sudo systemctl restart containerd | tee -a $LOGFILE
}

configure_containerd() {
    echo "正在配置 containerd 使用 nvidia-container-runtime..." | tee -a $LOGFILE
    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 | tee -a $LOGFILE
}

install_nerdctl() {
    echo "正在安装 nerdctl..." | tee -a $LOGFILE
    curl -LO https://github.com/containerd/nerdctl/releases/download/v${NERDCTL_VERSION}/nerdctl-${NERDCTL_VERSION}-linux-amd64.tar.gz | tee -a $LOGFILE
    sudo tar xzvf nerdctl-${NERDCTL_VERSION}-linux-amd64.tar.gz -C /usr/local/bin | tee -a $LOGFILE
    nerdctl version | tee -a $LOGFILE
}

install_cni_plugins() {
    echo "正在安装 CNI 插件..." | tee -a $LOGFILE
    mkdir -p /opt/cni/bin
    curl -L "https://github.com/containernetworking/plugins/releases/download/v${CNI_VERSION}/cni-plugins-linux-amd64-${CNI_VERSION}.tgz" | sudo tar -xz -C /opt/cni/bin | tee -a $LOGFILE
    ls /opt/cni/bin | tee -a $LOGFILE
    nerdctl network ls | tee -a $LOGFILE
}

install_buildkit() {
    echo "正在安装 BuildKit..." | tee -a $LOGFILE
    curl -LO https://github.com/moby/buildkit/releases/download/v${BUILDKIT_VERSION}/buildkit-${BUILDKIT_VERSION}.linux-amd64.tar.gz | tee -a $LOGFILE
    sudo tar -C /usr/local/bin --strip-components=1 -xzf buildkit-${BUILDKIT_VERSION}.linux-amd64.tar.gz | tee -a $LOGFILE
}

configure_buildkit_service() {
    echo "正在配置 BuildKit 守护进程..." | tee -a $LOGFILE
    sudo tee /etc/systemd/system/buildkitd.service > /dev/null <<EOT
[Unit]
Description=BuildKit Daemon
Documentation=https://github.com/moby/buildkit
After=network.target

[Service]
ExecStart=/usr/local/bin/buildkitd
Restart=always
RestartSec=10s
StartLimitInterval=0

[Install]
WantedBy=multi-user.target
EOT
    sudo systemctl daemon-reload | tee -a $LOGFILE
    sudo systemctl start buildkitd | tee -a $LOGFILE
    sudo systemctl enable buildkitd | tee -a $LOGFILE
    sudo systemctl status buildkitd | tee -a $LOGFILE
}

install_all() {
    install_docker_compose
    install_containerd
    configure_containerd
    install_nerdctl
    install_cni_plugins
    install_buildkit
    configure_buildkit_service
}

echo "选择操作：" | tee -a $LOGFILE
echo "1) 安装 Docker Compose" | tee -a $LOGFILE
echo "2) 卸载 Docker Compose" | tee -a $LOGFILE
echo "3) 安装并配置 containerd" | tee -a $LOGFILE
echo "4) 安装 nerdctl" | tee -a $LOGFILE
echo "5) 安装 CNI 插件" | tee -a $LOGFILE
echo "6) 安装 BuildKit 并配置守护进程" | tee -a $LOGFILE
echo "7) 全部安装" | tee -a $LOGFILE
echo "8) 退出" | tee -a $LOGFILE
read -p "输入选择（1-8）：" choice

case $choice in
    1) install_docker_compose ;;
    2) uninstall_docker_compose ;;
    3) install_containerd; configure_containerd ;;
    4) install_nerdctl ;;
    5) install_cni_plugins ;;
    6) install_buildkit; configure_buildkit_service ;;
    7) install_all ;;
    8) exit 0 ;;
    *) echo "无效输入。" | tee -a $LOGFILE ;;
esac