#!/bin/bash

# ==================== 公共函数 ====================
# ANSI颜色代码
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
RED='\033[0;31m'
NC='\033[0m'  # No Color
BOLD='\033[1m'
CYAN='\033[0;36m'
PURPLE='\033[0;35m'

# 显示标题函数
display_header() {
    echo -e "${BOLD}${YELLOW}======================================${NC}"
    echo -e "${BOLD}${YELLOW}|          Debian系统初始化工具       |${NC}"
    echo -e "${BOLD}${YELLOW}======================================${NC}"
}

# 权限检查函数
check_root() {
    if [ $EUID -ne 0 ]; then
        echo -e "${RED}请先切换到root用户：${NC}"
        echo "1. 使用以下命令切换："
        echo "   sudo -i"
        echo "2. 输入root密码"
        echo "3. 再次执行本脚本"
        exit 1
    fi
}

# 交互确认函数
confirm() {
    read -p "$1 [Y/N]: " yn
    yn=${yn:-y}  # 设置默认值为 y
    case $yn in
        [Yy]* ) return 0;;
        * ) return 1;;
    esac
}

# 检查并安装必要工具
install_dependencies() {
    local missing_tools=()

    # 检查 curl 是否存在
    if ! command -v curl &>/dev/null; then
        missing_tools+=("curl")
    fi

    # 检查 gpg 是否存在
    if ! command -v gpg &>/dev/null; then
        missing_tools+=("gnupg")  # gpg 属于 gnupg 包
    fi

    # 如果有缺失工具，则尝试安装
    if [ ${#missing_tools[@]} -gt 0 ]; then
        echo -e "${YELLOW}检测到以下工具缺失：${missing_tools[*]}${NC}"
        if confirm "是否自动安装缺失工具？"; then
            apt update -y > /dev/null 2>&1
            if ! apt install -y "${missing_tools[@]}" > /dev/null 2>&1; then
                echo -e "${RED}安装失败，请手动安装：sudo apt install ${missing_tools[*]}${NC}"
                return 1
            fi
            echo -e "${GREEN}缺失工具已成功安装${NC}"
        else
            echo -e "${RED}请手动安装缺失工具：sudo apt install ${missing_tools[*]}${NC}"
            return 1
        fi
    fi
}

# 等待用户按键继续
wait_for_key_press() {
    read -n 1 -s -r -p "按任意键返回主菜单..."
    echo ""
}

# 日志记录函数
log_message() {
    local message="$1"
    echo "$(date '+%Y-%m-%d %H:%M:%S') - $message" >> /var/log/debian_init.log
}

# 检查网络连接
check_network_connection() {
    if ! ping -c 1 google.com &>/dev/null; then
        echo -e "${RED}网络连接不可用，请检查网络设置${NC}"
        wait_for_key_press
        return 1
    fi
    return 0
}

# 环境变量配置模块
setup_environment() {
    log_message "开始环境变量配置"
    echo -e "${CYAN}\n=== 环境变量设置 ===${NC}"

    if ! cd ~ 2>/dev/null; then
        echo -e "${RED}无法切换到root用户的主目录${NC}"
        log_message "无法切换到root用户的主目录"
        wait_for_key_press
        return 1
    fi

    BASHRC_FILE=~/.bashrc
    [ ! -f "$BASHRC_FILE" ] && touch "$BASHRC_FILE"

    CONFIG_LINE="export PATH=\$PATH:/usr/sbin"
    if ! grep -qxF "$CONFIG_LINE" "$BASHRC_FILE"; then
        echo -e "\n# 系统工具路径\n$CONFIG_LINE" >> "$BASHRC_FILE"
        echo -e "${GREEN}PATH配置已添加至.bashrc${NC}"
        log_message "PATH配置已添加至.bashrc"

        # 确保当前 shell 会话加载新的配置
        if ! source "$BASHRC_FILE"; then
            echo -e "${RED}配置加载失败，请检查文件内容${NC}"
            log_message "配置加载失败，请检查文件内容"
            wait_for_key_press
            return 1
        fi
    else
        echo -e "${YELLOW}PATH配置已存在${NC}"
        log_message "PATH配置已存在"
    fi

    echo -e "${CYAN}当前环境PATH：${NC}"
    echo "$PATH"
    log_message "当前环境PATH: $PATH"
    wait_for_key_press
    return 0
}

# SSH服务配置模块
setup_ssh() {
    log_message "开始SSH服务配置"
    echo -e "${CYAN}\n=== SSH服务配置 ===${NC}"

    if confirm "是否允许root SSH登录并启用密码验证？"; then
        # 检查SSH服务是否安装
        if ! dpkg -l | grep -q 'openssh-server'; then
            echo -e "${YELLOW}检测到SSH服务未安装，正在尝试自动安装...${NC}"
            if confirm "是否自动安装SSH服务？"; then
                apt update -y > /dev/null 2>&1
                if ! apt install -y openssh-server > /dev/null 2>&1; then
                    echo -e "${RED}安装SSH服务失败，请手动安装：sudo apt install openssh-server${NC}"
                    log_message "SSH服务安装失败"
                    if ! confirm "是否继续执行后续配置？"; then
                        wait_for_key_press
                        return 1
                    fi
                else
                    echo -e "${GREEN}SSH服务安装成功${NC}"
                    log_message "SSH服务安装成功"
                fi
            else
                echo -e "${RED}请先手动安装SSH服务：sudo apt install openssh-server${NC}"
                log_message "用户取消SSH服务安装"
                wait_for_key_press
                return 1
            fi
        fi

        SSH_CONFIG="/etc/ssh/sshd_config"
        cp "$SSH_CONFIG" "${SSH_CONFIG}.bak" 2>/dev/null || {
            echo -e "${RED}无法备份sshd_config文件${NC}"
            log_message "无法备份sshd_config文件"
            wait_for_key_press
            return 1
        }

        sed -i -e 's/^#*PermitRootLogin.*/PermitRootLogin yes/' \
               -e 's/^#*PasswordAuthentication.*/PasswordAuthentication yes/' \
               "$SSH_CONFIG"

        if systemctl restart ssh; then
            echo -e "${GREEN}SSH服务配置已生效${NC}"
            log_message "SSH服务配置已生效"
        else
            echo -e "${RED}服务重启失败，请检查配置${NC}"
            log_message "服务重启失败，请检查配置"
            wait_for_key_press
            return 1
        fi
    else
        echo -e "${YELLOW}已跳过SSH配置${NC}"
        log_message "已跳过SSH配置"
    fi
    wait_for_key_press
    return 0
}

# 软件源配置模块
setup_sources() {
    log_message "开始软件源配置"
    echo -e "${CYAN}\n=== 软件源配置 ===${NC}"

    if ! confirm "是否更换系统软件源？"; then
        echo -e "${YELLOW}已跳过软件源更换${NC}"
        log_message "已跳过软件源更换"
        wait_for_key_press
        return 0
    fi

    # 获取当前Debian版本代码名
    DEBIAN_VERSION=$(grep -oP 'VERSION_CODENAME=\K\w+' /etc/os-release)
    if [ -z "$DEBIAN_VERSION" ]; then
        echo -e "${RED}无法获取Debian版本代码名，使用默认值bookworm${NC}"
        DEBIAN_VERSION="bookworm"
        log_message "无法获取Debian版本代码名，使用默认值bookworm"
    fi
    echo -e "${GREEN}当前Debian版本代码名: $DEBIAN_VERSION${NC}"

    # 镜像源选择菜单
    echo "请选择要使用的镜像源："
    echo " 1) 阿里云"
    echo " 2) 华为云"
    echo " 3) 清华大学开源软件镜像站"
    echo " 4) 中国科学技术大学开源软件镜像站"
    echo " 5) 恢复默认官方源"
    read -p "请输入数字选择 [1-5]: " mirror_choice

    SOURCES_FILE="/etc/apt/sources.list"
    # 检查源文件是否存在且可读写
    if [ ! -f "$SOURCES_FILE" ]; then
        echo -e "${RED}源文件不存在: $SOURCES_FILE${NC}"
        log_message "源文件不存在: $SOURCES_FILE"
        wait_for_key_press
        return 1
    fi
    
    if [ ! -r "$SOURCES_FILE" ]; then
        echo -e "${RED}源文件不可读，请检查权限${NC}"
        log_message "源文件不可读，请检查权限"
        wait_for_key_press
        return 1
    fi
    
    if [ ! -w "$SOURCES_FILE" ]; then
        echo -e "${RED}源文件不可写，请检查权限${NC}"
        log_message "源文件不可写，请检查权限"
        wait_for_key_press
        return 1
    fi
    
    # 备份源文件
    cp "$SOURCES_FILE" "${SOURCES_FILE}.bak" 2>&1 || {
        echo -e "${RED}源文件备份失败: $?${NC}"
        log_message "源文件备份失败: $?"
        wait_for_key_press
        return 1
    }

    case $mirror_choice in
        1)
            cat > "$SOURCES_FILE" << EOL
deb https://mirrors.aliyun.com/debian/ $DEBIAN_VERSION main non-free non-free-firmware
deb-src https://mirrors.aliyun.com/debian/ $DEBIAN_VERSION main non-free non-free-firmware

deb https://mirrors.aliyun.com/debian/ ${DEBIAN_VERSION}-updates main non-free non-free-firmware
deb-src https://mirrors.aliyun.com/debian/ ${DEBIAN_VERSION}-updates main non-free non-free-firmware

deb https://mirrors.aliyun.com/debian/ ${DEBIAN_VERSION}-backports main non-free non-free-firmware
deb-src https://mirrors.aliyun.com/debian/ ${DEBIAN_VERSION}-backports main non-free non-free-firmware

deb https://mirrors.aliyun.com/debian-security/ ${DEBIAN_VERSION}-security main non-free non-free-firmware
deb-src https://mirrors.aliyun.com/debian-security/ ${DEBIAN_VERSION}-security main non-free non-free-firmware
EOL
            ;;
        2)
            cat > "$SOURCES_FILE" << EOL
deb https://repo.huaweicloud.com/debian/ $DEBIAN_VERSION main non-free non-free-firmware
deb-src https://repo.huaweicloud.com/debian/ $DEBIAN_VERSION main non-free non-free-firmware

deb https://repo.huaweicloud.com/debian/ ${DEBIAN_VERSION}-updates main non-free non-free-firmware
deb-src https://repo.huaweicloud.com/debian/ ${DEBIAN_VERSION}-updates main non-free non-free-firmware

deb https://repo.huaweicloud.com/debian/ ${DEBIAN_VERSION}-backports main non-free non-free-firmware
deb-src https://repo.huaweicloud.com/debian/ ${DEBIAN_VERSION}-backports main non-free non-free-firmware

deb https://repo.huaweicloud.com/debian-security/ ${DEBIAN_VERSION}-security main non-free non-free-firmware
deb-src https://repo.huaweicloud.com/debian-security/ ${DEBIAN_VERSION}-security main non-free non-free-firmware
EOL
            ;;
        3)
            cat > "$SOURCES_FILE" << EOL
deb https://mirrors.tuna.tsinghua.edu.cn/debian/ $DEBIAN_VERSION main non-free contrib
deb-src https://mirrors.tuna.tsinghua.edu.cn/debian/ $DEBIAN_VERSION main non-free contrib

deb https://mirrors.tuna.tsinghua.edu.cn/debian/ ${DEBIAN_VERSION}-updates main non-free contrib
deb-src https://mirrors.tuna.tsinghua.edu.cn/debian/ ${DEBIAN_VERSION}-updates main non-free contrib

deb https://mirrors.tuna.tsinghua.edu.cn/debian/ ${DEBIAN_VERSION}-backports main non-free contrib
deb-src https://mirrors.tuna.tsinghua.edu.cn/debian/ ${DEBIAN_VERSION}-backports main non-free contrib
deb https://mirrors.tuna.tsinghua.edu.cn/debian-security ${DEBIAN_VERSION}-security main non-free contrib
deb-src https://mirrors.tuna.tsinghua.edu.cn/debian-security ${DEBIAN_VERSION}-security main non-free contrib
EOL
            ;;
        4)
            cat > "$SOURCES_FILE" << EOL
deb https://mirrors.ustc.edu.cn/debian/ $DEBIAN_VERSION main non-free contrib
deb-src https://mirrors.ustc.edu.cn/debian/ $DEBIAN_VERSION main non-free contrib

deb https://mirrors.ustc.edu.cn/debian/ ${DEBIAN_VERSION}-updates main non-free contrib
deb-src https://mirrors.ustc.edu.cn/debian/ ${DEBIAN_VERSION}-updates main non-free contrib

deb https://mirrors.ustc.edu.cn/debian/ ${DEBIAN_VERSION}-backports main non-free contrib
deb-src https://mirrors.ustc.edu.cn/debian/ ${DEBIAN_VERSION}-backports main non-free contrib
deb https://mirrors.ustc.edu.cn/debian-security ${DEBIAN_VERSION}-security main non-free contrib
deb-src https://mirrors.ustc.edu.cn/debian-security ${DEBIAN_VERSION}-security main non-free contrib
EOL
            ;;
        5)
            cat > "$SOURCES_FILE" << EOL
deb http://deb.debian.org/debian/ $DEBIAN_VERSION main non-free non-free-firmware
deb-src http://deb.debian.org/debian/ $DEBIAN_VERSION main non-free non-free-firmware

deb http://deb.debian.org/debian/ ${DEBIAN_VERSION}-updates main non-free non-free-firmware
deb-src http://deb.debian.org/debian/ ${DEBIAN_VERSION}-updates main non-free non-free-firmware

deb http://deb.debian.org/debian/ ${DEBIAN_VERSION}-backports main non-free non-free-firmware
deb-src http://deb.debian.org/debian/ ${DEBIAN_VERSION}-backports main non-free non-free-firmware
deb http://security.debian.org/debian-security ${DEBIAN_VERSION}-security main non-free non-free-firmware
deb-src http://security.debian.org/debian-security ${DEBIAN_VERSION}-security main non-free non-free-firmware
EOL
            ;;
        *)
            echo -e "${RED}无效选择，已恢复原始配置${NC}"
            mv "${SOURCES_FILE}.bak" "$SOURCES_FILE" 2>/dev/null
            log_message "无效选择，已恢复原始配置"
            wait_for_key_press
            return 1
            ;;
    esac

    echo -e "${PURPLE}正在更新软件列表..."
    if apt update -y > /dev/null 2>&1; then
        echo -e "${GREEN}软件源更新成功${NC}"
        log_message "软件源更新成功"
    else
        echo -e "${RED}软件源更新失败，请检查网络连接${NC}"
        log_message "软件源更新失败，请检查网络连接"
        mv "${SOURCES_FILE}.bak" "$SOURCES_FILE" 2>/dev/null
    fi
    wait_for_key_press
}

# 删除非root用户及其文件夹模块
remove_non_root_users() {
    log_message "开始删除非root用户"
    echo -e "${CYAN}\n=== 删除非root用户及其文件夹 ===${NC}"

    if ! confirm "是否删除除root之外的所有用户及其文件夹？"; then
        echo -e "${YELLOW}已跳过删除用户操作${NC}"
        log_message "已跳过删除用户操作"
        wait_for_key_press
        return 0
    fi

    NON_ROOT_USERS=$(awk -F: '($3 >= 1000) && ($1 != "root") {print $1}' /etc/passwd)
    if [ -z "$NON_ROOT_USERS" ]; then
        echo -e "${YELLOW}未找到需要删除的非root用户${NC}"
        log_message "未找到需要删除的非root用户"
        wait_for_key_press
        return 0
    fi

    echo "以下用户将被删除："
    echo "$NON_ROOT_USERS"

    for user in $NON_ROOT_USERS; do
        if userdel -r "$user" 2>/dev/null; then
            echo -e "${GREEN}用户 $user 及其文件夹已成功删除${NC}"
            log_message "用户 $user 及其文件夹已成功删除"
        else
            echo -e "${RED}用户 $user 删除失败，请检查权限或是否存在活动进程${NC}"
            log_message "用户 $user 删除失败，请检查权限或是否存在活动进程"
        fi
    done
    wait_for_key_press
    return 0
}

# WiFi连接模块
connect_wifi() {
    log_message "开始WiFi连接设置"
    echo -e "${CYAN}\n=== WiFi连接设置 ===${NC}"

    # 检查网络连接
    if ! check_network_connection; then
        return 1
    fi

    # 检查是否安装了 NetworkManager
    if ! command -v nmcli &>/dev/null; then
        echo -e "${YELLOW}未检测到 NetworkManager，正在尝试自动安装..."

        # 提示用户是否继续安装
        if confirm "是否自动安装 NetworkManager？"; then
            apt update -y > /dev/null 2>&1
            if ! apt install -y network-manager > /dev/null 2>&1; then
                echo -e "${RED}安装 NetworkManager 失败，请手动安装：sudo apt install network-manager${NC}"
                log_message "安装 NetworkManager 失败，请手动安装：sudo apt install network-manager"
                wait_for_key_press
                return 1
            fi
            echo -e "${GREEN}NetworkManager 安装完成${NC}"
            log_message "NetworkManager 安装完成"
        else
            echo -e "${RED}请手动安装 NetworkManager：sudo apt install network-manager${NC}"
            log_message "请手动安装 NetworkManager：sudo apt install network-manager"
            wait_for_key_press
            return 1
        fi
    fi

    # 启动 NetworkManager 服务
    if ! systemctl is-active --quiet NetworkManager; then
        systemctl start NetworkManager
        systemctl enable NetworkManager
    fi

    # 列出可用的 WiFi 网络
    echo -e "${PURPLE}正在扫描可用的WiFi网络..."
    WIFI_LIST=$(nmcli dev wifi | awk 'NR>1 {print $2}')
    if [ -z "$WIFI_LIST" ]; then
        echo -e "${RED}未找到可用的WiFi网络，请检查无线网卡状态${NC}"
        log_message "未找到可用的WiFi网络，请检查无线网卡状态"
        wait_for_key_press
        return 1
    fi

    echo "可用的WiFi网络："
    nmcli dev wifi | awk 'NR>1 {print NR-1 ") " $2}'
    read -p "请输入要连接的WiFi编号： " wifi_choice

    # 获取用户选择的WiFi名称
    SELECTED_WIFI=$(nmcli dev wifi | awk -v choice="$wifi_choice" 'NR==choice+1 {print $2}')
    if [ -z "$SELECTED_WIFI" ]; then
        echo -e "${RED}无效选择，请重新输入${NC}"
        log_message "无效选择，请重新输入"
        wait_for_key_press
        return 1
    fi

    # 输入WiFi密码
    read -sp "请输入WiFi密码： " wifi_password
    echo

    # 连接WiFi
    echo -e "${PURPLE}正在连接到WiFi：$SELECTED_WIFI..."
    if nmcli dev wifi connect "$SELECTED_WIFI" password "$wifi_password"; then
        echo -e "${GREEN}WiFi连接成功${NC}"
        log_message "WiFi连接成功"
    else
        echo -e "${RED}WiFi连接失败，请检查密码或网络状态${NC}"
        log_message "WiFi连接失败，请检查密码或网络状态"
    fi
    wait_for_key_press
    return 0
}

# 更新系统模块
update_system() {
    log_message "开始更新系统"
    echo -e "${CYAN}\n=== 更新系统 ===${NC}"

    if confirm "是否更新系统？"; then
        echo -e "${PURPLE}正在更新系统...${NC}"
        if apt update -y && apt upgrade -y && apt dist-upgrade -y; then
            echo -e "${GREEN}系统更新完成${NC}"
            log_message "系统更新完成"
        else
            echo -e "${RED}系统更新失败，请检查网络连接${NC}"
            log_message "系统更新失败，请检查网络连接"
        fi
    else
        echo -e "${YELLOW}已跳过系统更新${NC}"
        log_message "已跳过系统更新"
    fi
    wait_for_key_press
}

# 清理不必要的软件包模块
cleanup_packages() {
    log_message "开始清理不必要的软件包"
    echo -e "${CYAN}\n=== 清理不必要的软件包 ===${NC}"

    if confirm "是否清理不必要的软件包？"; then
        echo -e "${PURPLE}正在清理不必要的软件包..."
        if apt autoremove -y && apt autoclean -y; then
            echo -e "${GREEN}清理完成${NC}"
            log_message "清理完成"
        else
            echo -e "${RED}清理失败${NC}"
            log_message "清理失败"
        fi
    else
        echo -e "${YELLOW}已跳过清理操作${NC}"
        log_message "已跳过清理操作"
    fi
    wait_for_key_press
}

# ==================== 主流程 ====================
main() {
    check_root
    if ! install_dependencies; then
        echo -e "${RED}依赖工具安装失败，但继续执行后续任务...${NC}"
        log_message "依赖工具安装失败，但继续执行后续任务..."
    fi

    while true; do
        display_header
        echo -e "\n${GREEN}请选择要执行的功能模块：${NC}"
        echo -e " ${GREEN}1)${NC} 环境变量配置"
        echo -e " ${GREEN}2)${NC} SSH服务配置"
        echo -e " ${GREEN}3)${NC} 软件源配置"
        echo -e " ${GREEN}4)${NC} 删除其他用户"
        echo -e " ${GREEN}5)${NC} 连接WiFi"
        echo -e " ${GREEN}6)${NC} 更新系统"
        echo -e " ${GREEN}7)${NC} 清理不必要的软件包"
        echo -e " ${GREEN}0)${NC} 完成并退出"
        read -p "请输入数字选择 [0-7]: " choice

        case $choice in
            1)
                setup_environment
                ;;
            2)
                setup_ssh
                ;;
            3)
                setup_sources
                ;;
            4)
                remove_non_root_users
                ;;
            5)
                connect_wifi
                ;;
            6)
                update_system
                ;;
            7)
                cleanup_packages
                ;;
            0)
                echo -e "${GREEN}\n=== 脚本执行完成 ===${NC}"
                log_message "脚本执行完成"
                exit 0
                ;;
            *)
                echo -e "${RED}无效选择，请重新输入${NC}"
                log_message "无效选择，请重新输入"
                ;;
        esac
    done
}

# 启动主程序
main