#!/bin/bash
# 定义颜色变量，让输出更醒目
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[0;33m'
BLUE='\033[0;34m'
MAGENTA='\033[0;35m'
CYAN='\033[0;36m'
NC='\033[0m' # No Color

# 全局变量：是否自动执行
AUTO_MODE=false

# 函数：输出带时间戳的日志
log_message() {
    local log_level=$1
    local message=$2
    local timestamp=$(date "+%Y-%m-%d %H:%M:%S")

    case "$log_level" in
        "INFO")
            echo -e "${BLUE}[INFO]${NC} ${timestamp} - ${message}"
            ;;
        "SUCCESS")
            echo -e "${GREEN}[SUCCESS]${NC} ${timestamp} - ${message}"
            ;;
        "WARNING")
            echo -e "${YELLOW}[WARNING]${NC} ${timestamp} - ${message}"
            ;;
        "ERROR")
            echo -e "${RED}[ERROR]${NC} ${timestamp} - ${message}"
            ;;
        "STEP")
            echo -e "${CYAN}[STEP]${NC} ${timestamp} - ${message}"
            ;;
        *)
            echo -e "${timestamp} - ${message}"
            ;;
    esac
}

# 函数：显示进度条
show_spinner() {
    local pid=$1
    local message=$2
    local spin='-\|/'
    local i=0

    echo -ne "${YELLOW}$message${NC} "

    while kill -0 $pid 2>/dev/null; do
        local temp=${spin:i++%4:1}
        echo -ne "\b$temp"
        sleep 0.1
    done

    echo -e "\b${GREEN}✓${NC}"
}

# 函数：确认是否继续操作
confirm_action() {
    local prompt="$1"
    # 如果是自动模式，直接返回 0 (yes)
    if [ "$AUTO_MODE" = true ]; then
        log_message "INFO" "自动模式：默认选择是"
        return 0
    fi

    while true; do
        read -r -p "$(echo -e "${MAGENTA}$prompt${NC} (y/n): ")" choice
        case "$choice" in
            [yY])
                return 0 # 继续
                ;;
            [nN])
                return 1 # 退出
                ;;
            *)
                log_message "ERROR" "请输入 'y' 或 'n'！"
                ;;
        esac
    done
}

# 函数：检查命令执行结果
check_result() {
    if [ $1 -eq 0 ]; then
        log_message "SUCCESS" "$2"
        return 0
    else
        log_message "ERROR" "$3"
        return 1
    fi
}

# 函数：检测是否在中国
is_china() {
    # 使用 ping 百度来判断，如果 ping 得通，则认为在中国
    ping -c 1 baidu.com > /dev/null 2>&1
    if [ $? -eq 0 ]; then
        return 0 # 在中国
    else
        return 1 # 不在中国
    fi
}

# 函数：写入默认官方源
write_default_repo() {
    log_message "INFO" "正在写入默认官方源..."

    local default_repo="
deb http://deb.debian.org/debian bookworm main contrib non-free non-free-firmware
deb-src http://deb.debian.org/debian bookworm main contrib non-free non-free-firmware

deb http://security.debian.org/debian-security bookworm-security main contrib non-free non-free-firmware
deb-src http://security.debian.org/debian-security bookworm-security main contrib non-free non-free-firmware

deb http://deb.debian.org/debian bookworm-updates main contrib non-free non-free-firmware
deb-src http://deb.debian.org/debian bookworm-updates main contrib non-free non-free-firmware
"
    # 备份 sources.list
    local backup_file="/etc/apt/sources.list.bak.$(date +%Y%m%d%H%M%S)"
    cp /etc/apt/sources.list "$backup_file"
    log_message "INFO" "已备份 /etc/apt/sources.list 到 $backup_file"

    # 替换 sources.list
    echo "$default_repo" > /etc/apt/sources.list
    log_message "INFO" "已替换 /etc/apt/sources.list 为默认官方源"
}

# 函数：切换软件源
switch_repo() {
    log_message "INFO" "正在检测是否在中国..."
    if is_china; then
        log_message "INFO" "检测到在中国，切换到中科大镜像源..."

        local ustc_repo="
# 默认注释了源码仓库，如有需要可自行取消注释
deb http://mirrors.ustc.edu.cn/debian bookworm main contrib non-free non-free-firmware
# deb-src http://mirrors.ustc.edu.cn/debian bookworm main contrib non-free non-free-firmware
deb http://mirrors.ustc.edu.cn/debian bookworm-updates main contrib non-free non-free-firmware
# deb-src http://mirrors.ustc.edu.cn/debian bookworm-updates main contrib non-free non-free-firmware

# backports 软件源，请按需启用
# deb http://mirrors.ustc.edu.cn/debian bookworm-backports main contrib non-free non-free-firmware
# deb-src http://mirrors.ustc.edu.cn/debian bookworm-backports main contrib non-free non-free-firmware
"
        # 备份 sources.list
        local backup_file="/etc/apt/sources.list.bak.$(date +%Y%m%d%H%M%S)"
        cp /etc/apt/sources.list "$backup_file"
        log_message "INFO" "已备份 /etc/apt/sources.list 到 $backup_file"

        # 替换 sources.list
        echo "$ustc_repo" > /etc/apt/sources.list
        log_message "INFO" "已替换 /etc/apt/sources.list 为中科大镜像源"
    else
        log_message "INFO" "检测到不在中国，写入默认官方源..."
        write_default_repo
    fi

    # 更新软件源
    apt update
    check_result $? "软件源更新成功！" "软件源更新失败"
}

# 函数：选择最优软件源
select_best_repo() {
    if confirm_action "是否选择最优软件源？"; then
        # 检查 netselect-apt 是否安装
        if ! command -v netselect-apt &> /dev/null; then
            log_message "WARNING" "netselect-apt 未安装，正在安装..."
            apt update  #先更新一下
            apt install -y netselect-apt
            if [ $? -ne 0 ]; then
                log_message "ERROR" "netselect-apt 安装失败，请手动安装后重试"
                return 1
            fi
        fi

        log_message "INFO" "正在测试最优软件源，请稍候..."

        # 获取当前系统代号
        local codename=$(lsb_release -cs)

        # 执行 netselect-apt 命令
        netselect-apt --nonfree --sources "$codename"
        local result=$?

        if [ $result -eq 0 ]; then
            log_message "SUCCESS" "最优软件源选择完成！"

            # 检查当前目录下是否存在 sources.list 文件
            if [ -f ./sources.list ]; then
                # 备份现有的 sources.list 文件
                local backup_file="/etc/apt/sources.list.bak.$(date +%Y%m%d%H%M%S)"
                cp /etc/apt/sources.list "$backup_file"
                log_message "INFO" "已备份 /etc/apt/sources.list 到 $backup_file"

                # 移动 netselect-apt 生成的 sources.list 文件到 /etc/apt 目录下 
                mv ./sources.list /etc/apt/sources.list
                log_message "INFO" "已将 netselect-apt 生成的 sources.list 移动到 /etc/apt/sources.list"

                # 替换 non-free 为 non-free-firmware
                sed -i 's/non-free /non-free-firmware /g' /etc/apt/sources.list
                sed -i 's/non-free$/non-free-firmware/g' /etc/apt/sources.list

                # 更新软件源
                apt update
                check_result $? "软件源更新成功！" "软件源更新失败"
            else
                log_message "ERROR" "在当前目录下未找到 netselect-apt 生成的 sources.list 文件"
            fi
        else
            log_message "ERROR" "最优软件源选择失败，错误代码: $result"
        fi

    else
        log_message "WARNING" "已跳过选择最优软件源！"
    fi
}

# 函数：更新软件源
update_repo() {
    # 尝试更新软件源，如果失败则切换到临时源
    log_message "INFO" "正在尝试更新软件源..."
    apt update
    local update_result=$?

    if [ $update_result -ne 0 ]; then
        log_message "WARNING" "软件源更新失败，尝试切换到临时源..."
        switch_repo
    else
        log_message "SUCCESS" "软件源更新成功！"
    fi
}

# 函数：升级所有软件包，自动选择本地配置
upgrade_packages() {
    if confirm_action "是否升级所有软件包？"; then
        log_message "INFO" "正在升级所有软件包，请稍候..."

        # 先计算需要升级的包数量
        local packages_to_upgrade=$(apt list --upgradable 2>/dev/null | grep -v "Listing..." | wc -l)
        log_message "INFO" "将要升级 $packages_to_upgrade 个软件包"

        # 执行命令,不隐藏输出
        apt upgrade -y --option Dpkg::Options::="--force-confdef" --option Dpkg::Options::="--force-confold"
        local result=$?

        # 输出重要信息
        if [ $result -eq 0 ]; then
            log_message "INFO" "软件包升级成功！"
        else
            log_message "ERROR" "软件包升级失败，错误代码: $result"
        fi

        check_result $result "软件包升级成功！" "软件包升级失败，错误代码: $result"

    else
        log_message "WARNING" "已跳过软件包升级！"
    fi
}

# 函数：安装常用工具
install_common_tools() {
    if confirm_action "是否安装常用工具？"; then
        log_message "INFO" "正在安装常用工具，请稍候..."

        local tools="wget curl sudo vim gcc g++ make git sl lrzsz bash-completion exa duf figlet fd-find htop rsync iotop"
        log_message "INFO" "将安装以下工具: $tools"

        # 执行命令, 不隐藏输出
        apt install -y $tools
        local result=$?

        if [ $result -eq 0 ]; then
            log_message "INFO" "常用工具安装成功！"
        else
            log_message "ERROR" "常用工具安装失败，错误代码: $result"
        fi

        # 处理 sl 命令
        if [ $result -eq 0 ] && [ -f /usr/games/sl ]; then
            log_message "INFO" "正在配置 sl 命令..."
            mv /usr/games/sl /usr/bin/sl 2>/dev/null
            if [ $? -eq 0 ]; then
                log_message "SUCCESS" "sl 命令已移至 /usr/bin/sl"
            else
                log_message "WARNING" "无法移动 sl 命令"
            fi

            # 添加别名
            if ! grep -q "alias sl='sl -e'" /etc/profile; then
                echo "alias sl='sl -e'" >> /etc/profile
                log_message "SUCCESS" "sl 命令别名已添加到 /etc/profile"
            else
                log_message "INFO" "sl 命令别名已存在，跳过添加"
            fi
        fi

        # 安装 dust
        install_dust

        check_result $result "常用工具安装成功！" "常用工具安装失败，错误代码:$result"
    else
        log_message "WARNING" "已跳过常用工具安装！"
    fi
}

# 函数：移除 Python pip 警告
remove_python_pip_warning() {
    log_message "INFO" "正在移除 Python pip 警告..."

    # 检测 Python 版本目录
    PYTHON_LIB_DIR=$(ls /usr/lib/ | grep -E '^python3\.[0-20]+$' | sort -V | tail -n 1)

    if [ -z "$PYTHON_LIB_DIR" ]; then
        log_message "ERROR" "未找到 Python 库目录！"
        return 1
    else
        log_message "INFO" "找到 Python 库目录: /usr/lib/$PYTHON_LIB_DIR"
    fi

    # 备份文件
    if [ -f "/usr/lib/$PYTHON_LIB_DIR/EXTERNALLY-MANAGED" ]; then
        mv /usr/lib/$PYTHON_LIB_DIR/EXTERNALLY-MANAGED /usr/lib/$PYTHON_LIB_DIR/EXTERNALLY-MANAGED.bk 2>/dev/null
        local result=$?
        if [ $result -eq 0 ]; then
            check_result 0 "Python 增强功能已成功备份和移除！" "Python 增强功能移除失败，错误代码: $result"
        else
            log_message "ERROR" "Python 增强功能移除失败，错误代码: $result"
        fi
    else
        log_message "WARNING" "EXTERNALLY-MANAGED 文件不存在，可能已经移除或使用不同的 Python 版本"
    fi
}

# 函数：配置 SSH 允许 root 密码登录
config_ssh_root_login() {
    if confirm_action "是否配置 SSH 允许 root 密码登录？"; then
        log_message "INFO" "正在配置 SSH 允许 root 密码登录..."

        # 备份 sshd_config
        local backup_file="/etc/ssh/sshd_config.bak.$(date +%Y%m%d%H%M%S)"
        cp /etc/ssh/sshd_config "$backup_file"
        log_message "INFO" "已备份 SSH 配置文件到 $backup_file"

        # 修改配置
        sed -i 's/#\?PermitRootLogin .*/PermitRootLogin yes/' /etc/ssh/sshd_config
        local result1=$?

        sed -i 's/#\?PasswordAuthentication .*/PasswordAuthentication yes/' /etc/ssh/sshd_config
        local result2=$?

        # 修改端口号
        sed -i 's/#\?Port .*/Port 20202/' /etc/ssh/sshd_config
        local result3=$?

        if [ $result1 -eq 0 ] && [ $result2 -eq 0 ]; then
            log_message "SUCCESS" "SSH 配置文件修改成功"

            # 检查更改
            log_message "INFO" "当前 SSH 配置状态:"
            grep -E "PermitRootLogin|PasswordAuthentication" /etc/ssh/sshd_config | grep -v "#"

            # 重启 SSH 服务
            log_message "INFO" "正在重启 SSH 服务..."
            systemctl restart sshd
            local restart_result=$?

            check_result $restart_result "SSH 服务重启成功！" "SSH 服务重启失败，错误代码: $restart_result"
        else
            log_message "ERROR" "SSH 配置修改失败，错误代码: $result1, $result2, $result3"
        fi
    else
        log_message "WARNING" "已跳过 SSH 配置！"
    fi
}

# 函数：安装 Python 3
install_python3() {
    log_message "STEP" "正在安装 Python 3，请稍候..."

    # 执行命令，不隐藏输出
    apt install -y python3 python3-pip
    local result=$?

    if [ $result -eq 0 ]; then
        log_message "INFO" "Python 3 安装成功！"
        # 安装成功后，询问是否移除 Python pip 警告
        if confirm_action "是否移除 Python pip 警告？"; then
            remove_python_pip_warning
        fi
    else
        log_message "ERROR" "Python 3 安装失败，错误代码: $result"
    fi

    check_result $result "Python 3 安装成功！" "Python 3 安装失败，错误代码:$result"
}

# 函数：安装 uv
install_uv() {
    log_message "STEP" "正在安装 uv，请稍候..."

    # 检查 curl 是否安装
    if ! command -v curl &> /dev/null; then
        log_message "WARNING" "curl 未安装，正在安装..."
        apt install -y curl > /dev/null 2>&1 #这里还是隐藏curl的安装信息，避免刷屏
    fi

    # 执行命令,不隐藏输出
    log_message "INFO" "正在下载并运行 uv 安装脚本..."
    curl -LsSf https://astral.sh/uv/install.sh | sh
    local result=$?

    if [ $result -eq 0 ]; then
        log_message "INFO" "uv 安装成功！"
    else
        log_message "ERROR" "uv 安装失败，错误代码: $result"
    fi

    check_result $result "uv 安装成功！" "uv 安装失败，错误代码: $result"
}

# 函数：选择安装 Python 3 还是 uv
choose_python_or_uv() {
    log_message "STEP" "请选择安装 Python 3 或 uv"
    while true; do
        read -p "$(echo -e "${MAGENTA}请选择安装 Python 3 (输入 'python') 或 uv (输入 'uv', 默认)${NC}: ")" choice
        case "$choice" in
            "python")
                install_python3
                break
                ;;
            "uv"|"")
                install_uv
                break
                ;;
            *)
                log_message "ERROR" "输入错误，请重新选择！"
                ;;
        esac
    done
}

# 函数：安装 dust
install_dust() {
    if confirm_action "是否安装 dust？"; then
        log_message "INFO" "正在安装 dust，请稍候..."

        # 下载 dust
        log_message "INFO" "正在下载 dust 软件包..."
        wget https://github.com/bootandy/dust/releases/download/v1.1.2/du-dust_1.1.2-1_amd64.deb -O /opt/du-dust.deb
        local download_result=$?

        if [ $download_result -eq 0 ]; then
            log_message "SUCCESS" "dust 软件包下载成功"

            # 安装 dust
            log_message "INFO" "正在安装 dust 软件包..."
            apt install /opt/du-dust.deb -y
            local install_result=$?

            if [ $install_result -eq 0 ]; then
                log_message "INFO" "dust 安装成功！"
            else
                log_message "ERROR" "dust 安装失败，错误代码: $install_result"
            fi
            # 清理下载文件
            log_message "INFO" "正在清理下载的软件包..."
            rm /opt/du-dust.deb

            # 验证安装
            if [ $install_result -eq 0 ] && command -v dust &> /dev/null; then
                log_message "INFO" "dust 版本: $(dust --version 2>&1)"
            fi

            check_result $install_result "dust 安装成功！" "dust 安装失败，错误代码: $install_result"
        else
            log_message "ERROR" "dust 软件包下载失败，错误代码: $download_result"
        fi
    else
        log_message "WARNING" "已跳过 dust 安装！"
    fi
}

# 函数：打印系统信息
print_system_info() {
    log_message "INFO" "系统信息摘要:"

    # OS 信息
    if [ -f /etc/os-release ]; then
        log_message "INFO" "- 操作系统: $(grep PRETTY_NAME /etc/os-release | cut -d= -f2 | tr -d '"')"
    fi

    # 内核版本
    log_message "INFO" "- 内核版本: $(uname -r)"

    # 硬盘使用情况
    log_message "INFO" "- 硬盘使用情况:"
    df -h / | tail -n 1 | awk '{print "  使用: " $5 ", 可用: " $4 ", 总计: "$2}'

    # 内存使用情况
    local mem_total=$(free -m | grep Mem | awk '{print $2}')
    local mem_used=$(free -m | grep Mem | awk '{print $3}')
    local mem_percent=$((mem_used * 100 / mem_total))
    log_message "INFO" "- 内存使用情况: ${mem_percent}% (${mem_used}MB/${mem_total}MB)"
}

# SSH 配置脚本
function configure_ssh() {
  local config_file="/etc/ssh/sshd_config.d/99-custom.conf"

  # 询问并设置端口
  read -p "修改 SSH 端口 (y/n, 留空跳过)? " modify_port
  if [[ "$modify_port" == [Yy] ]]; then
    read -p "新端口 (1024-65535)? " port
    [[ "$port" =~ ^[0-9]+$ && "$port" -ge 1024 && "$port" -le 65535 ]] || { echo "无效端口"; return 1; }
    echo "Port $port" >> "$config_file"
  fi

  # 询问并设置 PermitRootLogin
  read -p "允许 root 登录 (y/n, 强烈建议 n)? " allow_root
  allow_root=${allow_root,,} # 转为小写
  case "$allow_root" in
    y|yes) echo "PermitRootLogin yes" >> "$config_file" ;;
    *)       echo "PermitRootLogin no" >> "$config_file" ;; # 默认 no
  esac

  # 询问并设置 PasswordAuthentication
  read -p "允许密码登录 (y/n, 建议 n)? " password_auth
  password_auth=${password_auth,,} # 转为小写
  case "$password_auth" in
    y|yes) echo "PasswordAuthentication yes" >> "$config_file" ;;
    *)       echo "PasswordAuthentication no" >> "$config_file" ;; # 默认 no
  esac

  systemctl restart sshd

  echo "SSH 配置已更新"
  return 0
}



# 主程序：开始执行
main() {
    # 打印欢迎信息
    figlet -f slant "System Update" 2>/dev/null || echo -e "\n${CYAN}=== System Update ====${NC}\n"
    log_message "INFO" "系统升级脚本启动于 $(date)"

    # 检查是否以 root 权限运行
    if [ "$(id -u)" -ne 0 ]; then
        log_message "ERROR" "此脚本需要 root 权限运行！"
        exit 1
    fi

    # 打印系统信息
    print_system_info

    # 增加 "一键执行" 选项
    read -r -p "$(echo -e "${MAGENTA}是否一键执行所有步骤？ (输入 'A' 启用, 其他任意键继续)${NC}: ")" auto_choice
    case "$auto_choice" in
        [aA])
            AUTO_MODE=true
            log_message "INFO" "已启用一键执行模式，所有步骤将自动进行。"
            ;;
        *)
            AUTO_MODE=false
            log_message "INFO" "未启用一键执行模式，将逐个询问。"
            ;;
    esac

    # 选择最优软件源
    select_best_repo

    # 先尝试更新，如果失败则切换到临时源
    update_repo

    # 执行主要功能
    upgrade_packages
    install_common_tools
    config_ssh_root_login
    choose_python_or_uv

    # 完成提示
    log_message "SUCCESS" "所有操作已完成，系统升级脚本执行结束"

    # 重新打印系统信息
    print_system_info

    log_message "INFO" "建议重启系统以应用所有更改"
}

# 执行主程序
main
