#!/bin/bash
set -e

# 颜色定义
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
BLUE='\033[0;34m'
CYAN='\033[0;36m'
NC='\033[0m'

# 内核版本和基础URL
KERNEL_VERSION="7.1.5-lqx1-kvm-cloud"
KERNEL_UNAME="7.1.5-lqx1-kvm-cloud-defu"
BASE_URL="https://depot.gggu.eu.org/lqx1/7.1.5"

info() { echo -e "${BLUE}[INFO]${NC} $1"; }
success() { echo -e "${GREEN}[OK]${NC} $1"; }
warn() { echo -e "${YELLOW}[WARN]${NC} $1"; }
error() { echo -e "${RED}[ERROR]${NC} $1"; }
title() { echo -e "${CYAN}$1${NC}"; }

# ========== 主菜单 ==========
show_menu() {
    echo ""
    echo -e "${YELLOW}========================================${NC}"
    echo -e "${YELLOW}  德芙丝滑 LQX Cloud 内核管理脚本              ${NC}"
    echo -e "${YELLOW}========================================${NC}"
    echo ""
    echo "  [1] 安装内核"
    echo "  [2] 卸载内核"
    echo ""
    read -rp "请输入选项 [1/2]: " choice

    case "$choice" in
        1) install_kernel ;;
        2) uninstall_kernel ;;
        *) error "无效选项，退出。"; exit 1 ;;
    esac
}

# ========== 安装功能 ==========
install_kernel() {
    echo ""
    echo -e "${YELLOW}========================================${NC}"
    echo ""
    echo -e "  内核版本: ${BLUE}${KERNEL_VERSION}${NC}"
    echo -e "  实际内核: ${BLUE}${KERNEL_UNAME}${NC}"
    echo -e "  支持系统: ${BLUE}Debian 12/13, Ubuntu 24.04/26.04${NC}"
    echo -e "  支持架构: ${BLUE}amd64 (x86_64) KVM 环境${NC}"
    echo -e "  CPU要求:  ${BLUE}V2/V3/V4 (V1 不支持)${NC}"
    echo ""
    echo -e "${RED}警告: 安装新内核存在风险，请确保已备份重要数据！${NC}"
    echo ""
    echo -e "${YELLOW}按回车键开始安装，或按 Ctrl+C 取消...${NC}"
    read -r

    if [[ $EUID -ne 0 ]]; then
        error "请使用 root 权限运行此脚本 (sudo ./install_kernel.sh)"
        exit 1
    fi

    # 2. amd64
    info "检测系统架构..."
    ARCH=$(uname -m)
    if [[ "$ARCH" != "x86_64" ]]; then
        error "当前架构为 $ARCH，仅支持 amd64 (x86_64) 架构。"
        exit 1
    fi
    success "架构检测通过: amd64"

    # 3. KVM
    info "检测虚拟化环境..."
    VIRT_TYPE=""
    if command -v systemd-detect-virt &>/dev/null; then
        VIRT_TYPE=$(systemd-detect-virt 2>/dev/null || true)
    fi
    if [[ -z "$VIRT_TYPE" ]] && [[ -r /sys/class/dmi/id/product_name ]]; then
        DMI_PRODUCT=$(cat /sys/class/dmi/id/product_name 2>/dev/null | tr '[:upper:]' '[:lower:]')
        if [[ "$DMI_PRODUCT" == *"kvm"* ]] || [[ "$DMI_PRODUCT" == *"qemu"* ]]; then
            VIRT_TYPE="kvm"
        fi
    fi
    if [[ -z "$VIRT_TYPE" ]] && grep -q "^flags.*hypervisor" /proc/cpuinfo 2>/dev/null; then
        VIRT_TYPE="kvm"
    fi
    if [[ "$VIRT_TYPE" != "kvm" ]]; then
        error "当前虚拟化环境为: ${VIRT_TYPE:-未知}，该内核为 KVM Cloud 精简版，仅适用于 KVM 环境。"
        exit 1
    fi
    success "虚拟化检测通过: KVM"

    # 4. OS 版本
    info "检测操作系统版本..."
    if [[ ! -f /etc/os-release ]]; then
        error "无法读取 /etc/os-release。"
        exit 1
    fi
    source /etc/os-release
    OS_ID="$ID"
    OS_VERSION_ID="$VERSION_ID"
    SUPPORTED=0
    case "$OS_ID" in
        debian) [[ "$OS_VERSION_ID" == "12" ]] || [[ "$OS_VERSION_ID" == "13" ]] && SUPPORTED=1 ;;
        ubuntu) [[ "$OS_VERSION_ID" == "24.04" ]] || [[ "$OS_VERSION_ID" == "26.04" ]] && SUPPORTED=1 ;;
    esac
    if [[ "$SUPPORTED" -ne 1 ]]; then
        error "当前系统为 ${OS_ID} ${OS_VERSION_ID}，仅支持 Debian 12/13 或 Ubuntu 24.04/26.04。"
        exit 1
    fi
    success "系统检测通过: ${OS_ID} ${OS_VERSION_ID}"

    # 5. CPU V级别
    info "检测 CPU 微架构级别..."
    CPU_FLAGS=$(grep -m1 "^flags" /proc/cpuinfo | cut -d: -f2)
    
    # 辅助函数：检查 flags 中是否包含指定 flag（前后加空格确保精确匹配）
    has_flag() {
        local flags="$1"
        local flag="$2"
        [[ " ${flags} " == *" ${flag} "* ]]
    }

    check_v4() {
        local f="$1"
        has_flag "$f" "avx512f" && has_flag "$f" "avx512bw" && has_flag "$f" "avx512cd" && has_flag "$f" "avx512dq" && has_flag "$f" "avx512vl"
    }
    check_v3() {
        local f="$1"
        local has_bmi=0
        has_flag "$f" "bmi1" && has_bmi=1
        has_flag "$f" "bmi"  && has_bmi=1
        local has_lzcnt=0
        has_flag "$f" "lzcnt" && has_lzcnt=1
        has_flag "$f" "abm"   && has_lzcnt=1
        has_flag "$f" "avx" && has_flag "$f" "avx2" && [[ "$has_bmi" -eq 1 ]] && has_flag "$f" "bmi2" && has_flag "$f" "f16c" && has_flag "$f" "fma" && [[ "$has_lzcnt" -eq 1 ]] && has_flag "$f" "movbe" && has_flag "$f" "xsave"
    }
    check_v2() {
        local f="$1"
        has_flag "$f" "cx16" && has_flag "$f" "lahf_lm" && has_flag "$f" "popcnt" && has_flag "$f" "sse4_1" && has_flag "$f" "sse4_2" && has_flag "$f" "ssse3"
    }

    CPU_ARCH=""
    if check_v4 "$CPU_FLAGS"; then CPU_ARCH="V4"
    elif check_v3 "$CPU_FLAGS"; then CPU_ARCH="V3"
    elif check_v2 "$CPU_FLAGS"; then CPU_ARCH="V2"
    else CPU_ARCH="V1"
    fi

    info "检测到 CPU 微架构级别: $CPU_ARCH"
    if [[ "$CPU_ARCH" == "V1" ]]; then
        error "检测到 CPU 为 V1 架构，不满足安装要求（最低需要 V2）。"
        exit 1
    fi

    INSTALL_ARCH="$CPU_ARCH"
    success "将安装 ${INSTALL_ARCH} 版本的内核"

    # 6. 下载安装
    info "准备安装内核..."
    HEADERS_URL="${BASE_URL}/${INSTALL_ARCH}/linux-headers-${KERNEL_VERSION}-${INSTALL_ARCH}_amd64.deb"
    IMAGE_URL="${BASE_URL}/${INSTALL_ARCH}/linux-image-${KERNEL_VERSION}-${INSTALL_ARCH}_amd64.deb"

    TEMP_DIR=$(mktemp -d)
    trap "rm -rf $TEMP_DIR" EXIT

    info "下载 linux-headers..."
    if ! curl -fSL --progress-bar -o "${TEMP_DIR}/linux-headers.deb" "$HEADERS_URL"; then
        error "下载 headers 失败: $HEADERS_URL"; exit 1
    fi
    success "Headers 下载完成"

    info "下载 linux-image..."
    if ! curl -fSL --progress-bar -o "${TEMP_DIR}/linux-image.deb" "$IMAGE_URL"; then
        error "下载 image 失败: $IMAGE_URL"; exit 1
    fi
    success "Image 下载完成"

    if dpkg -l 2>/dev/null | grep -q "linux-image-${KERNEL_VERSION}-${INSTALL_ARCH}"; then
        warn "检测到已安装相同版本内核包"
        read -rp "是否强制重新安装? [y/N]: " confirm
        [[ "$confirm" == "y" || "$confirm" == "Y" ]] || { info "已取消安装。"; exit 0; }
    fi

    info "检查并安装必要的依赖..."
    apt-get update -qq
    apt-get install -y -qq initramfs-tools 2>/dev/null || true

    info "安装 linux-headers..."
    dpkg -i "${TEMP_DIR}/linux-headers.deb" || apt-get install -f -y -qq
    info "安装 linux-image..."
    dpkg -i "${TEMP_DIR}/linux-image.deb" || apt-get install -f -y -qq
    success "内核包安装完成"

    # 7. 获取实际内核版本并设置默认启动
    info "检测实际安装的内核版本..."
    KERNEL_FILE=$(ls -1 /boot/vmlinuz-* 2>/dev/null | grep "${KERNEL_UNAME}" | head -n1)
    if [[ -z "$KERNEL_FILE" ]]; then
        KERNEL_FILE=$(ls -1 /boot/vmlinuz-* 2>/dev/null | grep "${KERNEL_VERSION}" | sort -V | tail -n1)
    fi
    if [[ -z "$KERNEL_FILE" ]]; then
        error "无法找到新安装的内核镜像文件 /boot/vmlinuz-*${KERNEL_UNAME}* 或 /boot/vmlinuz-*${KERNEL_VERSION}*"; exit 1
    fi
    INSTALLED_KERNEL_VER=$(basename "$KERNEL_FILE" | sed 's/^vmlinuz-//')
    success "实际安装的内核版本: ${INSTALLED_KERNEL_VER}"

    # ========== 精确设置 grub 默认项 ==========
    info "设置新内核为默认启动项..."
    set_default_kernel "$INSTALLED_KERNEL_VER"

    echo ""
    echo -e "${GREEN}========================================${NC}"
    echo -e "${GREEN}    内核安装成功！                      ${NC}"
    echo -e "${GREEN}========================================${NC}"
    echo -e "  系统:       ${BLUE}${OS_ID} ${OS_VERSION_ID}${NC}"
    echo -e "  安装包:     ${BLUE}${KERNEL_VERSION}-${INSTALL_ARCH}${NC}"
    echo -e "  实际内核:   ${BLUE}${INSTALLED_KERNEL_VER}${NC}"
    echo -e "  CPU架构:    ${BLUE}${CPU_ARCH}${NC}"
    echo ""
    echo -e "${YELLOW}请按回车键重启系统以应用新内核...${NC}"
    read -r
    reboot
}

# ========== 精确设置 grub 默认内核 ==========
set_default_kernel() {
    local kernel_ver="$1"
    local grub_cfg="/boot/grub/grub.cfg"
    local grub_default="/etc/default/grub"

    info "在 grub.cfg 中查找内核 ${kernel_ver}..."

    # 找到目标内核 menuentry 的行号（排除恢复模式）
    local target_line
    target_line=$(grep -n "with Linux ${kernel_ver}" "$grub_cfg" 2>/dev/null | grep -iv "recovery" | head -n1 | cut -d: -f1)

    if [[ -z "$target_line" ]]; then
        warn "未找到内核 ${kernel_ver} 的 grub 菜单项，回退到 GRUB_DEFAULT=0"
        fallback_grub_default
        return
    fi

    # 向上查找该 menuentry 所在的 submenu
    local submenu_line
    submenu_line=$(head -n "$target_line" "$grub_cfg" | grep -n "submenu " | tail -n1 | cut -d: -f1)

    local submenu_name=""
    if [[ -n "$submenu_line" ]]; then
        submenu_name=$(sed -n "${submenu_line}p" "$grub_cfg" | sed -n "s/.*submenu '\([^']*\)'.*/\1/p")
    fi

    # 提取 menuentry 名称
    local menuentry_name
    menuentry_name=$(sed -n "${target_line}p" "$grub_cfg" | sed -n "s/.*menuentry '\([^']*\)'.*/\1/p")

    if [[ -z "$menuentry_name" ]]; then
        warn "无法提取菜单项名称，回退到 GRUB_DEFAULT=0"
        fallback_grub_default
        return
    fi

    # 构建 grub 路径
    local grub_entry
    if [[ -n "$submenu_name" ]]; then
        grub_entry="${submenu_name}>${menuentry_name}"
    else
        grub_entry="${menuentry_name}"
    fi

    info "找到 grub 菜单路径: ${grub_entry}"

    # 写入 /etc/default/grub（最可靠的方式）
    sed -i '/^GRUB_DEFAULT=/d' "$grub_default"
    echo "GRUB_DEFAULT='${grub_entry}'" >> "$grub_default"

    # 禁用 SAVEDEFAULT，防止用户之前的选择干扰
    sed -i 's/^GRUB_SAVEDEFAULT=.*/#GRUB_SAVEDEFAULT=true/' "$grub_default"

    # 同时尝试写入 grubenv（双重保险）
    grub-set-default "$grub_entry" 2>/dev/null || grub2-set-default "$grub_entry" 2>/dev/null || true

    update-grub

    # 验证：检查 grub.cfg 中第一个 linux /vmlinuz- 是否指向我们的内核
    local first_kernel
    first_kernel=$(awk '/^menuentry /{found=1} found && /linux \/vmlinuz-/{print; exit}' "$grub_cfg" | sed 's/.*linux \/vmlinuz-//' | awk '{print $1}')

    if [[ "$first_kernel" == "$kernel_ver" ]]; then
        success "验证通过: grub 默认项已精确指向 ${kernel_ver}"
    else
        warn "grub 首个条目内核为 ${first_kernel:-未知}，但已精确设置默认项为 ${grub_entry}"
        warn "重启后将启动 ${kernel_ver}"
    fi
}

fallback_grub_default() {
    local grub_default="/etc/default/grub"
    sed -i '/^GRUB_DEFAULT=/d' "$grub_default"
    echo "GRUB_DEFAULT=0" >> "$grub_default"
    sed -i 's/^GRUB_SAVEDEFAULT=.*/#GRUB_SAVEDEFAULT=true/' "$grub_default"
    update-grub
}

# ========== 卸载功能 ==========
uninstall_kernel() {
    echo ""
    echo -e "${YELLOW}========================================${NC}"
    echo -e "${YELLOW}            内核卸载                  ${NC}"
    echo -e "${YELLOW}========================================${NC}"
    echo ""

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

    info "扫描已安装的 LQX 内核..."
    mapfile -t INSTALLED_IMAGES < <(dpkg -l 2>/dev/null | grep '^ii' | grep 'linux-image-.*lqx1-kvm-cloud' | awk '{print $2}')
    mapfile -t INSTALLED_HEADERS < <(dpkg -l 2>/dev/null | grep '^ii' | grep 'linux-headers-.*lqx1-kvm-cloud' | awk '{print $2}')

    local TOTAL=${#INSTALLED_IMAGES[@]}
    if [[ $TOTAL -eq 0 ]]; then
        error "未检测到已安装的 LQX 内核。"
        exit 1
    fi

    echo ""
    title "已安装的 LQX 内核镜像:"
    for i in "${!INSTALLED_IMAGES[@]}"; do
        echo "  [$((i+1))] ${INSTALLED_IMAGES[$i]}"
    done
    echo "  [0] 全部卸载"
    echo ""
    read -rp "请选择要卸载的内核编号 [0-${TOTAL}]: " select_idx

    if ! [[ "$select_idx" =~ ^[0-9]+$ ]]; then
        error "无效输入"; exit 1
    fi

    local TARGET_IMAGES=()
    local TARGET_HEADERS=()

    if [[ "$select_idx" -eq 0 ]]; then
        TARGET_IMAGES=("${INSTALLED_IMAGES[@]}")
        TARGET_HEADERS=("${INSTALLED_HEADERS[@]}")
    elif [[ "$select_idx" -ge 1 && "$select_idx" -le "$TOTAL" ]]; then
        TARGET_IMAGES=("${INSTALLED_IMAGES[$((select_idx-1))]}")
        local img_ver
        img_ver=$(echo "${INSTALLED_IMAGES[$((select_idx-1))]}" | sed 's/linux-image-//')
        for h in "${INSTALLED_HEADERS[@]}"; do
            if [[ "$h" == *"$img_ver"* ]]; then
                TARGET_HEADERS+=("$h")
            fi
        done
    else
        error "无效编号"; exit 1
    fi

    info "检查系统是否还有其他可用内核..."
    mapfile -t ALL_IMAGES < <(dpkg -l 2>/dev/null | grep '^ii' | grep 'linux-image-' | awk '{print $2}')

    local OTHER_COUNT=0
    for pkg in "${ALL_IMAGES[@]}"; do
        local is_target=0
        for t in "${TARGET_IMAGES[@]}"; do
            [[ "$pkg" == "$t" ]] && is_target=1 && break
        done
        if [[ $is_target -eq 0 ]]; then
            OTHER_COUNT=$((OTHER_COUNT + 1))
        fi
    done

    if [[ $OTHER_COUNT -eq 0 ]]; then
        echo ""
        error "━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━"
        error "  危险！卸载后将没有可用内核！"
        error "  系统将无法启动，拒绝执行卸载。"
        error "━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━"
        error "请先安装其他内核（如系统自带内核）后再卸载。"
        exit 1
    fi

    success "安全检查通过: 还有 ${OTHER_COUNT} 个其他内核可用"

    echo ""
    title "即将卸载以下包:"
    for p in "${TARGET_IMAGES[@]}"; do echo "  - $p"; done
    for p in "${TARGET_HEADERS[@]}"; do echo "  - $p"; done
    echo ""
    read -rp "确认卸载? [y/N]: " confirm
    if [[ "$confirm" != "y" && "$confirm" != "Y" ]]; then
        info "已取消卸载。"
        exit 0
    fi

    info "正在卸载..."
    apt-get remove --purge -y "${TARGET_IMAGES[@]}" "${TARGET_HEADERS[@]}"

    info "清理残留..."
    apt-get autoremove -y 2>/dev/null || true

    info "更新 GRUB..."
    update-grub

    # ========== 卸载后清理失效的 grub 默认项 ==========
    info "检查 GRUB 默认启动项..."
    local grub_default="/etc/default/grub"
    local current_default
    current_default=$(grep '^GRUB_DEFAULT=' "$grub_default" 2>/dev/null | sed 's/^GRUB_DEFAULT=//' | tr -d "'" | tr -d '"' || true)

    if [[ -n "$current_default" ]] && [[ "$current_default" != "0" ]] && [[ "$current_default" == *"${KERNEL_UNAME}"* ]]; then
        warn "默认启动项指向已卸载的内核，重置为 GRUB_DEFAULT=0"
        sed -i '/^GRUB_DEFAULT=/d' "$grub_default"
        echo "GRUB_DEFAULT=0" >> "$grub_default"
        update-grub
    fi

    echo ""
    success "━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━"
    success "  LQX 内核卸载完成！"
    success "━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━"
    info "建议重启系统以完全清理旧内核模块。"
    echo ""
}

# ========== 入口 ==========
show_menu