#!/bin/bash
set -euo pipefail

# Containerd安装脚本（Ubuntu专用）
# 集成配置：默认配置生成、镜像源替换、Cgroup驱动、镜像加速
# 适用：Ubuntu 20.04/22.04/24.04 LTS

#######################################
# 配置参数
#######################################
DEFAULT_VERSION="1.7.28-1"
ACTION=""
CONTAINERD_VERSION="$DEFAULT_VERSION"
LOG_FILE="/var/log/containerd_manage.log"
DOCKER_REPO="https://download.docker.com/linux/ubuntu"
GPG_KEY_PATH="/etc/apt/keyrings/docker.gpg"
REPO_LIST_PATH="/etc/apt/sources.list.d/docker.list"
CONFIG_PATH="/etc/containerd/config.toml"
PACKAGE_NAME="containerd.io"
CERTS_D_PATH="/etc/containerd/certs.d"
DOCKER_IO_HOSTS="${CERTS_D_PATH}/docker.io/hosts.toml"

#######################################
# 日志函数
#######################################
log() {
    local level=$1
    local message=$2
    local timestamp=$(date +"%Y-%m-%d %H:%M:%S")
    echo "[$timestamp] [$level] $message" | tee -a "$LOG_FILE"
}

#######################################
# 帮助信息
#######################################
usage() {
    echo "Usage: $0 [OPTIONS]"
    echo "Options:"
    echo "  --install          安装containerd（必填）"
    echo "  --uninstall        卸载containerd（必填）"
    echo "  --version VERSION  指定版本（如1.7.28-1，默认：$DEFAULT_VERSION）"
    echo "  --help             显示帮助"
    exit 1
}

#######################################
# 解析参数
#######################################
parse_args() {
    while [[ $# -gt 0 ]]; do
        case "$1" in
            --install)
                if [[ "$ACTION" == "uninstall" ]]; then log "ERROR" "参数冲突"; usage; fi
                ACTION="install"
                shift
                ;;
            --uninstall)
                if [[ "$ACTION" == "install" ]]; then log "ERROR" "参数冲突"; usage; fi
                ACTION="uninstall"
                shift
                ;;
            --version)
                if [[ $# -lt 2 || "$2" == -* ]]; then log "ERROR" "需指定版本"; usage; fi
                CONTAINERD_VERSION="$2"
                shift 2
                ;;
            --help)
                usage
                ;;
            *)
                log "ERROR" "未知参数：$1"
                usage
                ;;
        esac
    done
    if [[ -z "$ACTION" ]]; then log "ERROR" "需指定--install或--uninstall"; usage; fi
}

#######################################
# 前置检查
#######################################
pre_check() {
    log "INFO" "前置环境检查"
    if [[ "$(id -u)" -ne 0 ]]; then log "ERROR" "需root权限"; exit 1; fi
    if ! grep -q "Ubuntu" /etc/os-release; then log "ERROR" "仅支持Ubuntu"; exit 1; fi
    local os_version=$(grep VERSION_ID /etc/os-release | cut -d= -f2 | tr -d '"')
    if ! [[ "$os_version" =~ ^20\.04$ || "$os_version" =~ ^22\.04$ || "$os_version" =~ ^24\.04$ ]]; then
        log "ERROR" "仅支持Ubuntu 20.04/22.04/24.04"
        exit 1
    fi
    local required_tools=("apt-get" "curl" "gpg" "lsb_release" "systemctl")
    for tool in "${required_tools[@]}"; do
        if ! command -v "$tool" &> /dev/null; then
            log "INFO" "安装缺失工具：$tool"
            apt-get update -qq >/dev/null
            apt-get install -qq -y "$tool" >/dev/null
        fi
    done
    log "INFO" "前置检查通过"
}

#######################################
# 配置Docker源
#######################################
configure_repo() {
    log "INFO" "配置Docker源"
    mkdir -p "$(dirname "$GPG_KEY_PATH")"
    if [[ ! -f "$GPG_KEY_PATH" ]]; then
        log "INFO" "下载Docker GPG密钥"
        curl -fsSL --retry 3 --retry-delay 5 "${DOCKER_REPO}/gpg" | gpg --dearmor -o "$GPG_KEY_PATH" >/dev/null 2>&1 || {
            log "ERROR" "GPG密钥下载失败"
            exit 1
        }
        chmod a+r "$GPG_KEY_PATH"
    fi
    local codename=$(lsb_release -cs)
    if [[ "$codename" == "noble" ]]; then codename="jammy"; fi  # 兼容24.04
    local repo_content="deb [arch=$(dpkg --print-architecture) signed-by=$GPG_KEY_PATH] $DOCKER_REPO $codename stable"
    if [[ ! -f "$REPO_LIST_PATH" || "$(cat "$REPO_LIST_PATH")" != "$repo_content" ]]; then
        echo "$repo_content" | tee "$REPO_LIST_PATH" >/dev/null
        apt-get update -qq >/dev/null
    fi
    log "INFO" "Docker源配置完成"
}

#######################################
# 验证版本可用性
#######################################
verify_version() {
    log "INFO" "验证版本${CONTAINERD_VERSION}可用性"
    local available_versions=$(apt-cache madison "$PACKAGE_NAME" | awk '{print $3}' | sed 's/~ubuntu.*//' | sort -u)
    if ! echo "$available_versions" | grep -qxF "$CONTAINERD_VERSION"; then
        log "ERROR" "版本不可用，可用版本：$available_versions"
        exit 1
    fi
    log "INFO" "版本验证通过"
}

#######################################
# 安装containerd
#######################################
install_containerd() {
    log "INFO" "安装containerd"
    local full_version=$(apt-cache madison "$PACKAGE_NAME" | grep "$CONTAINERD_VERSION" | awk '{print $3}' | head -n 1)
    if [[ -z "$full_version" ]]; then log "ERROR" "未找到完整版本"; exit 1; fi
    if dpkg -l "$PACKAGE_NAME" &>/dev/null && [[ "$(dpkg -l "$PACKAGE_NAME" | awk '/ii/ {print $3}')" == "$full_version" ]]; then
        log "INFO" "已安装$full_version，跳过"
        return
    fi
    log "INFO" "安装$full_version"
    apt-get install -qq -y "${PACKAGE_NAME}=${full_version}" >/dev/null 2>&1 || {
        log "ERROR" "安装失败，手动执行：sudo apt-get install -y ${PACKAGE_NAME}=${full_version}"
        exit 1
    }
}

#######################################
# 核心配置（严格按你的要求执行）
#######################################
configure_containerd() {
    log "INFO" "开始配置containerd"

    # 1. 生成默认配置文件
    log "INFO" "生成默认配置文件"
    containerd config default > "$CONFIG_PATH" 2>/dev/null || {
        log "ERROR" "生成默认配置失败"
        exit 1
    }

    # 2. 替换镜像源（registry.k8s.io -> 华为云）
    log "INFO" "替换镜像源为华为云"
    sed -i 's#registry.k8s.io#swr.cn-north-4.myhuaweicloud.com/ddn-k8s/registry.k8s.io#' "$CONFIG_PATH" || {
        log "ERROR" "镜像源替换失败"
        exit 1
    }

    # 3. 配置Cgroup驱动器（SystemdCgroup = true）
    log "INFO" "配置Cgroup驱动器"
    sed -i 's/SystemdCgroup\ =\ false/SystemdCgroup\ =\ true/g' "$CONFIG_PATH" || {
        log "ERROR" "Cgroup配置失败"
        exit 1
    }

    # 4. 配置镜像加速
    log "INFO" "配置镜像加速"
    # 4.1 设置registry.config_path
    local registry_block="[plugins.\"io.containerd.grpc.v1.cri\".registry]"
    local config_path_line='  config_path = "/etc/containerd/certs.d"'
    if ! grep -qxF "$config_path_line" "$CONFIG_PATH"; then
        if ! grep -qxF "$registry_block" "$CONFIG_PATH"; then
            echo -e "\n$registry_block" >> "$CONFIG_PATH"
        fi
        sed -i "/$registry_block/a $config_path_line" "$CONFIG_PATH"
    fi

    # 4.2 创建目录并写入加速配置
    mkdir -p "$(dirname "$DOCKER_IO_HOSTS")"
    local hosts_content=$(cat << EOF
server = "https://docker.io"
[host."https://o2j0mc5x.mirror.aliyuncs.com"]
  capabilities = ["pull", "resolve"]
server = "https://k8s.gcr.io"
[host."https://gcr.mirrors.ustc.edu.cn/google-containers/"]
  capabilities = ["pull", "resolve"]
server = "https://quay.io"
[host."https://mirror.ccs.tencentyun.com"]
  capabilities = ["pull", "resolve"]
EOF
    )
    echo "$hosts_content" | tee "$DOCKER_IO_HOSTS" >/dev/null

    # 重启服务
    systemctl restart containerd >/dev/null 2>&1
    log "INFO" "配置完成，已重启服务"
}

#######################################
# 启动服务并设置开机启动
#######################################
start_service() {
    log "INFO" "设置开机自启并启动服务"
    systemctl enable --now containerd >/dev/null 2>&1 || {
        log "ERROR" "服务启动失败"
        exit 1
    }
}

#######################################
# 验证安装
#######################################
verify_installation() {
    log "INFO" "验证安装结果"
    if ! systemctl is-active --quiet containerd; then log "ERROR" "服务未运行"; exit 1; fi
    if ! grep -q "swr.cn-north-4.myhuaweicloud.com/ddn-k8s/registry.k8s.io" "$CONFIG_PATH"; then
        log "WARNING" "镜像源替换未生效"
    fi
    if ! grep -q "SystemdCgroup \= true" "$CONFIG_PATH"; then
        log "WARNING" "Cgroup配置未生效"
    fi
    if [[ ! -f "$DOCKER_IO_HOSTS" ]]; then
        log "WARNING" "镜像加速配置不存在"
    fi
    log "INFO" "验证完成"
}

#######################################
# 卸载containerd
#######################################
uninstall_containerd() {
    log "INFO" "卸载containerd"
    systemctl stop containerd >/dev/null 2>&1 || true
    if dpkg -l "$PACKAGE_NAME" &>/dev/null; then
        apt-get purge -qq -y "$PACKAGE_NAME" >/dev/null 2>&1
        apt-get autoremove -qq -y >/dev/null 2>&1
    fi
    rm -rf "$CONFIG_PATH" "$CERTS_D_PATH" /var/lib/containerd /var/log/containerd
    systemctl daemon-reload >/dev/null 2>&1
    log "INFO" "卸载完成"
}

#######################################
# 主流程
#######################################
main() {
    > "$LOG_FILE"
    log "INFO" "===== 脚本启动 ====="
    parse_args "$@"
    pre_check

    case "$ACTION" in
        install)
            configure_repo
            verify_version
            install_containerd
            configure_containerd  # 执行你的所有配置步骤
            start_service
            verify_installation
            ;;
        uninstall)
            uninstall_containerd
            ;;
    esac

    log "INFO" "===== 操作完成 ====="
    log "INFO" "日志：$LOG_FILE"
}

main "$@"