#!/bin/bash

# =================================================================
# 脚本名称: switch_kernel.sh
# 描述: 自动更换指定内核版本，并确保驱动程序（网卡、显卡）自动同步编译。
#       支持在线安装和离线安装（从本地 .deb 包目录安装）。
#
# 用法:
#   sudo ./switch_kernel.sh                           # 交互式（自动检测网络）
#   sudo ./switch_kernel.sh --offline /path/to/debs   # 强制离线，指定包目录
#   sudo ./switch_kernel.sh --download /path/to/debs  # 仅下载包（需要网络）
# =================================================================

DEFAULT_KERNEL_VERSION="6.8.0-52-generic"
KERNEL_VERSION=""
OFFLINE_MODE=false
DOWNLOAD_ONLY=false
DEB_DIR=""
SKIP_DKMS=false
DKMS_BACKUP_DIR="/tmp/dkms_backup_$$"

# 日志文件路径（当前目录）
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
LOG_FILE="${SCRIPT_DIR}/log-$(date +%Y%m%d)-dacs.log"

RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
BLUE='\033[0;34m'
NC='\033[0m'

# 日志函数：同时输出到终端和文件
log_info()  {
    local msg="[INFO]  $*"
    echo -e "${YELLOW}${msg}${NC}"
    echo "[$(date '+%Y-%m-%d %H:%M:%S')] ${msg}" >> "$LOG_FILE"
}

log_ok()    {
    local msg="[OK]    $*"
    echo -e "${GREEN}${msg}${NC}"
    echo "[$(date '+%Y-%m-%d %H:%M:%S')] ${msg}" >> "$LOG_FILE"
}

log_err()   {
    local msg="[ERROR] $*"
    echo -e "${RED}${msg}${NC}"
    echo "[$(date '+%Y-%m-%d %H:%M:%S')] ${msg}" >> "$LOG_FILE"
}

log_step()  {
    local msg="[STEP]  $*"
    echo -e "${BLUE}${msg}${NC}"
    echo "[$(date '+%Y-%m-%d %H:%M:%S')] ${msg}" >> "$LOG_FILE"
}

# 清理函数：确保退出时恢复 DKMS 钩子
cleanup() {
    restore_dkms_hook
}

trap cleanup EXIT INT TERM

# --- 参数解析 ---

parse_args() {
    while [[ $# -gt 0 ]]; do
        case "$1" in
            --offline)
                OFFLINE_MODE=true
                DEB_DIR="${2:?'--offline 需要指定包目录路径'}"
                shift 2
                ;;
            --download)
                DOWNLOAD_ONLY=true
                DEB_DIR="${2:?'--download 需要指定保存目录路径'}"
                shift 2
                ;;
            --version)
                KERNEL_VERSION="$2"
                shift 2
                ;;
            --skip-dkms)
                SKIP_DKMS=true
                shift
                ;;
            -h|--help)
                echo "用法:"
                echo "  sudo $0                              # 交互式（自动检测网络）"
                echo "  sudo $0 --offline /path/to/debs      # 强制离线，指定包目录"
                echo "  sudo $0 --download /path/to/debs     # 仅下载包（需要网络）"
                echo "  sudo $0 --version 6.8.0-52-generic   # 指定内核版本（可与上述组合）"
                echo "  sudo $0 --skip-dkms                  # 跳过 DKMS 自动编译（解决驱动兼容问题）"
                exit 0
                ;;
            *)
                log_err "未知参数: $1，使用 --help 查看用法。"
                exit 1
                ;;
        esac
    done
}

# --- 网络检测 ---

check_network() {
    curl -s --max-time 3 http://archive.ubuntu.com/ubuntu/ -o /dev/null 2>&1
}

# --- 内核检测 ---

get_kernel_list() {
    grep menuentry_id_option /boot/grub/grub.cfg
}

check_kernel_installed() {
    [[ $(get_kernel_list) == *"$KERNEL_VERSION"* ]]
}

is_kernel_package_installed() {
    dpkg -l | grep -q "linux-image-$KERNEL_VERSION"
}

check_kernel_version_format() {
    [[ $1 =~ ^[0-9]+\.[0-9]+\.[0-9]+-[0-9]+-generic$ ]]
}

# --- DKMS 管理 ---

list_dkms_modules() {
    dkms status 2>/dev/null | awk -F'[,:]' '{print $1"/"$2}' | sort -u
}

check_dkms_compatibility() {
    log_info "检查 DKMS 驱动兼容性..."
    local modules
    modules=$(list_dkms_modules)

    if [[ -z "$modules" ]]; then
        log_info "未检测到 DKMS 驱动。"
        return 0
    fi

    echo -e "${YELLOW}当前系统已安装的 DKMS 驱动：${NC}"
    echo "$modules"
    echo ""

    local problematic_modules=()
    while IFS= read -r mod; do
        # 已知问题驱动列表
        if [[ "$mod" =~ elkeid-driver|bbr|sch_cake ]]; then
            problematic_modules+=("$mod")
        fi
    done <<< "$modules"

    if [[ ${#problematic_modules[@]} -gt 0 ]]; then
        echo -e "${YELLOW}警告：以下驱动可能与新内核不兼容：${NC}"
        for m in "${problematic_modules[@]}"; do
            echo "  - $m"
        done
        echo ""
        return 1
    fi
    return 0
}

disable_dkms_hook() {
    if [[ -x /etc/kernel/postinst.d/dkms ]]; then
        log_info "临时禁用 DKMS 自动编译钩子..."
        mkdir -p "$DKMS_BACKUP_DIR"
        mv /etc/kernel/postinst.d/dkms "$DKMS_BACKUP_DIR/"
        log_ok "DKMS 钩子已禁用（备份到 $DKMS_BACKUP_DIR）"
    fi
}

restore_dkms_hook() {
    if [[ -f "$DKMS_BACKUP_DIR/dkms" ]]; then
        log_info "恢复 DKMS 钩子..."
        mv "$DKMS_BACKUP_DIR/dkms" /etc/kernel/postinst.d/
        chmod +x /etc/kernel/postinst.d/dkms
        rm -rf "$DKMS_BACKUP_DIR"
        log_ok "DKMS 钩子已恢复"
    fi
}

handle_dkms_failure() {
    log_err "DKMS 驱动编译失败！"
    echo ""
    echo "可能的原因："
    echo "  1. 某些第三方驱动与新内核不兼容（如 elkeid-driver、旧版 nvidia 等）"
    echo "  2. 缺少编译依赖"
    echo ""
    echo "解决方案："
    echo "  A. 临时移除问题驱动后重试"
    echo "  B. 强制完成安装（忽略 DKMS 错误）"
    echo "  C. 退出并手动排查"
    echo ""

    read -p "请选择 [A/B/C]: " choice
    case "$choice" in
        A|a)
            echo ""
            list_dkms_modules
            echo ""
            read -p "请输入要移除的驱动名称（格式: driver/version，如 elkeid-driver/1.7.0.22）: " module_to_remove
            if [[ -n "$module_to_remove" ]]; then
                log_info "移除驱动: $module_to_remove"
                dkms remove "$module_to_remove" --all
                log_ok "驱动已移除，重新尝试安装..."
                return 0
            else
                log_err "未输入驱动名称，退出。"
                exit 1
            fi
            ;;
        B|b)
            log_info "强制完成安装（忽略 DKMS 错误）..."
            dpkg --configure -a
            apt-get install -f -y
            if is_kernel_package_installed; then
                log_ok "内核已安装（部分驱动可能未编译）"
                return 0
            else
                log_err "强制安装失败。"
                exit 1
            fi
            ;;
        C|c|*)
            log_info "退出。请查看 TROUBLESHOOTING.md 获取详细排查步骤。"
            exit 1
            ;;
    esac
}

# --- 在线安装 ---

apt_update() {
    log_info "正在更新软件包列表..."
    if apt-get update 2>&1 | tee -a "$LOG_FILE"; then
        return 0
    else
        log_err "apt-get update 失败，请检查网络连接。"
        exit 1
    fi
}

install_kernel_online() {
    log_step "在线安装内核: ${KERNEL_VERSION}"

    # 如果启用了 --skip-dkms，临时禁用 DKMS 钩子
    if [[ "$SKIP_DKMS" == true ]]; then
        disable_dkms_hook
    fi

    if apt-get install -y \
        linux-image-"$KERNEL_VERSION" \
        linux-headers-"$KERNEL_VERSION" \
        linux-modules-extra-"$KERNEL_VERSION" \
        dkms \
        build-essential 2>&1 | tee -a "$LOG_FILE"; then
        log_ok "内核相关基础包安装成功！"

        # 恢复 DKMS 钩子
        if [[ "$SKIP_DKMS" == true ]]; then
            restore_dkms_hook
        fi
    else
        # 检查是否是 DKMS 导致的失败
        if grep -q "dkms.*failed" /var/log/apt/term.log 2>/dev/null || \
           grep -q "post-installation script subprocess returned error" /var/log/apt/term.log 2>/dev/null; then

            # 恢复钩子（如果之前禁用了）
            restore_dkms_hook

            # 提供交互式修复
            handle_dkms_failure

            # 用户选择修复后，重新尝试
            install_kernel_online
        else
            log_err "内核安装失败，请检查网络或版本号是否存在。"
            restore_dkms_hook
            exit 1
        fi
    fi
}

# --- 离线安装 ---

check_offline_packages() {
    local dir="$1"
    local missing=()
    local required_patterns=(
        "linux-image-${KERNEL_VERSION}"
        "linux-headers-${KERNEL_VERSION}"
        "linux-modules-extra-${KERNEL_VERSION}"
    )

    for pattern in "${required_patterns[@]}"; do
        if ! ls "$dir"/${pattern}*.deb &>/dev/null; then
            missing+=("$pattern")
        fi
    done

    if [[ ${#missing[@]} -gt 0 ]]; then
        log_err "离线包目录 '$dir' 中缺少以下包："
        for m in "${missing[@]}"; do
            echo "  - ${m}*.deb"
        done
        echo ""
        echo "请先在有网络的机器上运行："
        echo "  sudo $0 --download $dir --version $KERNEL_VERSION"
        return 1
    fi
    return 0
}

install_kernel_offline() {
    local dir="$1"
    log_step "离线安装内核: ${KERNEL_VERSION}，包目录: $dir"

    if ! check_offline_packages "$dir"; then
        exit 1
    fi

    # 如果启用了 --skip-dkms，临时禁用 DKMS 钩子
    if [[ "$SKIP_DKMS" == true ]]; then
        disable_dkms_hook
    fi

    # 安装 dkms / build-essential（如果包目录中有的话）
    for pkg in dkms build-essential; do
        if ! dpkg -l | grep -q "^ii  $pkg"; then
            if ls "$dir"/${pkg}*.deb &>/dev/null; then
                log_info "从包目录安装 $pkg ..."
                dpkg -i "$dir"/${pkg}*.deb || true
            else
                log_info "包目录中无 $pkg，跳过（DKMS 驱动编译可能受影响）。"
            fi
        fi
    done

    log_info "安装内核包（按依赖顺序）..."
    # headers → modules → modules-extra → image
    local debs=()
    for pattern in \
        "linux-headers-${KERNEL_VERSION}" \
        "linux-modules-${KERNEL_VERSION}" \
        "linux-modules-extra-${KERNEL_VERSION}" \
        "linux-image-${KERNEL_VERSION}"; do
        while IFS= read -r -d '' f; do
            debs+=("$f")
        done < <(find "$dir" -name "${pattern}*.deb" -print0)
    done

    if [[ ${#debs[@]} -eq 0 ]]; then
        log_err "未找到任何匹配的 .deb 文件。"
        restore_dkms_hook
        exit 1
    fi

    if dpkg -i "${debs[@]}" 2>&1 | tee -a "$LOG_FILE"; then
        log_ok "内核离线安装成功！"
        restore_dkms_hook
    else
        # 检查是否是 DKMS 错误
        if dmesg | tail -50 | grep -q "dkms.*failed" || \
           dpkg -l | grep -q "^iU.*linux-image-${KERNEL_VERSION}"; then

            restore_dkms_hook
            handle_dkms_failure

            # 重新尝试
            install_kernel_offline "$dir"
        else
            log_info "尝试修复依赖关系（apt-get install -f）..."
            apt-get install -f -y 2>&1 | tee -a "$LOG_FILE"

            if is_kernel_package_installed; then
                log_ok "内核离线安装成功！"
                restore_dkms_hook
            else
                log_err "内核离线安装失败，请检查 .deb 包是否完整或版本是否匹配。"
                restore_dkms_hook
                exit 1
            fi
        fi
    fi
}

# --- 下载模式 ---

download_packages() {
    local dir="$1"
    mkdir -p "$dir"
    log_step "下载内核包到目录: $dir（内核版本: $KERNEL_VERSION）"

    apt_update

    log_info "下载内核相关包（不安装）..."
    apt-get install --download-only -y \
        linux-image-"$KERNEL_VERSION" \
        linux-headers-"$KERNEL_VERSION" \
        linux-modules-extra-"$KERNEL_VERSION" \
        dkms \
        build-essential 2>&1 | tee -a "$LOG_FILE"

    log_info "复制下载的包到 $dir ..."
    find /var/cache/apt/archives/ \( \
        -name "linux-image-${KERNEL_VERSION}*.deb" \
        -o -name "linux-headers-${KERNEL_VERSION}*.deb" \
        -o -name "linux-modules-${KERNEL_VERSION}*.deb" \
        -o -name "linux-modules-extra-${KERNEL_VERSION}*.deb" \
        -o -name "dkms*.deb" \
        -o -name "build-essential*.deb" \
    \) -exec cp {} "$dir/" \; 2>&1 | tee -a "$LOG_FILE"

    log_ok "下载完成，包列表："
    ls -lh "$dir"/*.deb 2>&1 | tee -a "$LOG_FILE" || log_err "目录中没有找到 .deb 文件，下载可能失败。"
    echo ""
    echo "将 $dir 目录复制到目标机器后，运行："
    echo "  sudo $0 --offline $dir --version $KERNEL_VERSION"
}

# --- GRUB 配置 ---

get_kernel_uuid() {
    local kernel_list
    kernel_list=$(get_kernel_list)
    local gnulinux="gnulinux-advanced"
    local kernel_advanced="gnulinux-$KERNEL_VERSION-advanced"
    local gnulinux_uuid=""
    local kernel_advanced_uuid=""

    while read -r l; do
        if [[ $l == *"$gnulinux"* ]]; then
            gnulinux_uuid=$(echo "$l" | sed -n "s/.*\($gnulinux[^']*\).*/\1/p")
        elif [[ $l == *"$kernel_advanced"* ]]; then
            kernel_advanced_uuid=$(echo "$l" | sed -n "s/.*\($kernel_advanced[^']*\).*/\1/p")
        fi
    done <<< "$kernel_list"
    echo "$gnulinux_uuid>$kernel_advanced_uuid"
}

set_boot_kernel() {
    local grub_path="/etc/default/grub"
    local newlines=""
    local uuid_list
    uuid_list=$(get_kernel_uuid)

    if [[ -z "$uuid_list" || "$uuid_list" == ">" ]]; then
        log_err "无法从 grub.cfg 中获取内核 UUID，请确认内核已安装并运行过 update-grub。"
        return 1
    fi

    # 重写 GRUB_DEFAULT，同时确保 GRUB_SAVEDEFAULT 不干扰
    while read -r l; do
        if [[ $l == GRUB_DEFAULT=* ]]; then
            newlines+="GRUB_DEFAULT=\"$uuid_list\"\n"
        elif [[ $l == GRUB_SAVEDEFAULT=* ]]; then
            newlines+="GRUB_SAVEDEFAULT=false\n"
        else
            newlines+="$l\n"
        fi
    done < "$grub_path"

    # 如果文件里没有 GRUB_SAVEDEFAULT 这行，追加一条
    if ! grep -q "^GRUB_SAVEDEFAULT=" "$grub_path"; then
        newlines+="GRUB_SAVEDEFAULT=false\n"
    fi

    log_info "写入 GRUB_DEFAULT: $uuid_list"
    echo -e "$newlines" > "$grub_path"
    return 0
}

set_default_boot_kernel() {
    if ! check_kernel_installed; then
        if [[ "$OFFLINE_MODE" == true ]]; then
            install_kernel_offline "$DEB_DIR"
        else
            apt_update
            install_kernel_online
        fi
    fi

    if set_boot_kernel; then
        log_info "正在更新 GRUB 引导程序..."
        update-grub 2>&1 | tee -a "$LOG_FILE"

        if [[ "$SKIP_DKMS" == false ]]; then
            log_info "正在为新内核 ($KERNEL_VERSION) 构建驱动程序 (DKMS)..."
            if dkms autoinstall -k "$KERNEL_VERSION" 2>&1 | tee -a "$LOG_FILE"; then
                log_ok "DKMS 驱动编译成功。"
            else
                log_err "DKMS 驱动编译失败（部分驱动可能不可用）。"
                echo "可以重启后手动运行: sudo dkms autoinstall -k $KERNEL_VERSION"
            fi
        else
            log_info "已跳过 DKMS 自动编译（可重启后手动运行: sudo dkms autoinstall -k $KERNEL_VERSION）。"
        fi
    else
        log_err "GRUB 配置写入失败，请手动检查 /etc/default/grub。"
        exit 1
    fi
}

hold_default_boot_kernel() {
    log_info "正在锁定内核版本，防止被系统更新替换..."
    # 锁定具体版本包
    apt-mark hold linux-image-"$KERNEL_VERSION"
    apt-mark hold linux-headers-"$KERNEL_VERSION"
    apt-mark hold linux-modules-extra-"$KERNEL_VERSION"
    # 锁定元包，防止 apt 升级时将默认内核切换到新版本
    apt-mark hold linux-image-generic
    apt-mark hold linux-headers-generic
    log_ok "内核版本已锁定。"
}

# --- 主程序 ---

parse_args "$@"

if [[ $EUID -ne 0 ]]; then
    log_err "请使用 root 权限运行此脚本！"
    exit 1
fi

# 记录脚本启动信息
{
    echo ""
    echo "=========================================="
    echo "脚本启动时间: $(date '+%Y-%m-%d %H:%M:%S')"
    echo "执行命令: $0 $*"
    echo "当前用户: $(whoami)"
    echo "主机名: $(hostname)"
    echo "当前内核: $(uname -r)"
    echo "=========================================="
    echo ""
} >> "$LOG_FILE"

log_info "日志文件: $LOG_FILE"

# 下载模式：只下载，不安装
if [[ "$DOWNLOAD_ONLY" == true ]]; then
    if [[ -z "$KERNEL_VERSION" ]]; then
        read -p "请输入要下载的内核版本号 (直接回车使用 $DEFAULT_KERNEL_VERSION): " user_input
        KERNEL_VERSION="${user_input:-$DEFAULT_KERNEL_VERSION}"
    fi
    if ! check_kernel_version_format "$KERNEL_VERSION"; then
        log_err "版本格式不正确（示例: 6.8.0-52-generic）"
        exit 1
    fi
    download_packages "$DEB_DIR"
    exit 0
fi

# 交互式选择内核版本（命令行未指定时）
if [[ -z "$KERNEL_VERSION" ]]; then
    read -p "请输入要安装的内核版本号 (直接回车使用 $DEFAULT_KERNEL_VERSION): " user_input
    if [[ -z "$user_input" ]]; then
        KERNEL_VERSION="$DEFAULT_KERNEL_VERSION"
    else
        KERNEL_VERSION="$user_input"
        if ! check_kernel_version_format "$KERNEL_VERSION"; then
            log_err "格式不正确，将使用默认版本 $DEFAULT_KERNEL_VERSION。"
            KERNEL_VERSION="$DEFAULT_KERNEL_VERSION"
        fi
    fi
fi

# 自动检测网络（未强制指定离线模式时）
if [[ "$OFFLINE_MODE" == false ]]; then
    if check_network; then
        log_info "网络可用，使用在线模式。"
    else
        log_info "网络不可用，自动切换到离线模式。"
        OFFLINE_MODE=true
        if [[ -z "$DEB_DIR" ]]; then
            read -p "请输入本地 .deb 包目录路径: " DEB_DIR
            if [[ -z "$DEB_DIR" || ! -d "$DEB_DIR" ]]; then
                log_err "目录不存在或未指定，退出。"
                exit 1
            fi
        fi
    fi
fi

# DKMS 兼容性检查（除非用户已指定 --skip-dkms）
if [[ "$SKIP_DKMS" == false ]]; then
    if ! check_dkms_compatibility; then
        echo ""
        read -p "检测到可能不兼容的 DKMS 驱动，是否跳过 DKMS 自动编译？(Y/N): " answer
        if [[ "$answer" == "Y" || "$answer" == "y" ]]; then
            SKIP_DKMS=true
            log_info "将跳过 DKMS 自动编译。"
        else
            log_info "继续安装（如遇 DKMS 错误会提供修复选项）。"
        fi
    fi
fi

# 安装阶段
if is_kernel_package_installed; then
    log_info "内核包 linux-image-$KERNEL_VERSION 已经安装。"
    # 确保 headers 存在以便编译驱动
    if [[ "$OFFLINE_MODE" == true ]]; then
        if ls "$DEB_DIR"/linux-headers-"${KERNEL_VERSION}"*.deb &>/dev/null; then
            dpkg -i "$DEB_DIR"/linux-headers-"${KERNEL_VERSION}"*.deb || true
            dpkg -i "$DEB_DIR"/dkms*.deb "$DEB_DIR"/build-essential*.deb 2>/dev/null || true
        fi
    else
        apt-get install -y linux-headers-"$KERNEL_VERSION" dkms build-essential
    fi
else
    if [[ "$OFFLINE_MODE" == true ]]; then
        install_kernel_offline "$DEB_DIR"
    else
        install_kernel_online
    fi
fi

# 安装成功后：自动设置默认启动内核 + 锁定版本（无需用户确认）
log_step "设置 $KERNEL_VERSION 为默认启动内核..."
set_default_boot_kernel

log_step "锁定内核版本，防止 apt 自动升级切换..."
hold_default_boot_kernel

log_ok "操作完成！"
echo -e "${YELLOW}注意：如果重启后显卡驱动不正常，请检查 Secure Boot 是否已关闭。${NC}"

# 记录脚本完成信息
{
    echo ""
    echo "=========================================="
    echo "脚本完成时间: $(date '+%Y-%m-%d %H:%M:%S')"
    echo "目标内核版本: $KERNEL_VERSION"
    echo "安装模式: $([ "$OFFLINE_MODE" == true ] && echo "离线" || echo "在线")"
    echo "DKMS 状态: $([ "$SKIP_DKMS" == true ] && echo "已跳过" || echo "已执行")"
    echo "=========================================="
    echo ""
} >> "$LOG_FILE"

read -p "是否立刻重启系统？(Y/N): " answer
if [[ "$answer" == "Y" || "$answer" == "y" ]]; then
    log_info "系统即将重启..."
    reboot
fi