#!/bin/bash

VERSION="1.7.6"
FILENAME="nerdctl-full-${VERSION}-linux-amd64.tar.gz"
DOWNLOAD_URL="https://github.com/containerd/nerdctl/releases/download/v${VERSION}/${FILENAME}"
INSTALL_DIR="/usr/local"
NERDCTL_BIN="${INSTALL_DIR}/bin/nerdctl"

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

    if ! grep -q "^deb .\+nvidia-container-toolkit" /etc/apt/sources.list /etc/apt/sources.list.d/*; then
        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
    else
        echo "NVIDIA Docker 源已经配置，跳过此步骤。"
    fi

    sudo apt-get update
    sudo apt-get install -y nvidia-container-toolkit
}

configure_containerd() {
    echo "正在配置 containerd 使用 NVIDIA runtime..."

    # 创建 /etc/containerd 目录并生成默认配置文件
    sudo mkdir -p /etc/containerd
    containerd config default | sudo tee /etc/containerd/config.toml

    # 使用 sed 命令在特定位置添加配置项
    sudo sed -i '/\[plugins."io.containerd.grpc.v1.cri".registry.configs\]/a \\t[plugins."io.containerd.grpc.v1.cri".registry.configs."docker.io".tls]\n\t\tinsecure_skip_verify = true' /etc/containerd/config.toml

    # 检查语法
    sudo containerd config dump | sudo tee /etc/containerd/config.toml

    # 配置 NVIDIA runtime
    sudo nvidia-ctk runtime configure --runtime=containerd

    # 重启 containerd 服务
    sudo systemctl restart containerd
}

install_nerdctl() {
    if [ ! -f "$FILENAME" ]; then
        echo "Downloading $FILENAME..."
        wget $DOWNLOAD_URL
    else
        echo "$FILENAME already exists, skipping download."
    fi

    echo "Installing nerdctl..."
    sudo tar -xvf $FILENAME -C $INSTALL_DIR

    echo "Setting up nerdctl as docker..."
    cat << 'EOF' | sudo tee /usr/local/bin/docker
#!/bin/bash
/usr/local/bin/nerdctl $@
EOF
    sudo chmod +x /usr/local/bin/docker

    echo "Setting up bash completion for nerdctl and docker..."
    sudo apt install bash-completion -y
    sudo nerdctl completion bash | sudo tee /etc/bash_completion.d/nerdctl > /dev/null
    sudo nerdctl completion bash | sudo tee /etc/bash_completion.d/docker > /dev/null

    # Create containerd group and add current user
    # if ! getent group containerd > /dev/null; then
    #     echo "Creating containerd group..."
    #     sudo groupadd containerd
    # fi

    # # Use SUDO_USER to get the original username who invoked sudo
    # ORIGINAL_USER=${SUDO_USER:-$USER}

    # echo "Adding $ORIGINAL_USER to containerd group..."
    # sudo usermod -aG containerd $ORIGINAL_USER

    # Change ownership and permissions of completion files
    echo "Setting permissions for completion files..."
    sudo chown root:containerd /etc/bash_completion.d/nerdctl
    sudo chmod 755 /etc/bash_completion.d/nerdctl
    sudo chown root:containerd /etc/bash_completion.d/docker
    sudo chmod 755 /etc/bash_completion.d/docker

    # Load completion files
    echo "Loading bash completion files..."
    source /etc/bash_completion.d/nerdctl
    source /etc/bash_completion.d/docker

    echo "Nerdctl installation and setup completed."
}

install_nerdctl_rootless() {
    echo "Setting up nerdctl for rootless mode..."
    /usr/local/bin/containerd-rootless-setuptool.sh install
    /usr/local/bin/containerd-rootless-setuptool.sh install-buildkit
}

main() {
    install_nvidia_docker
    install_nerdctl
    install_nerdctl_rootless
    configure_containerd
    echo "Nerdctl installation completed."
}

main