#!/bin/bash

# Ubuntu 24.04 换源和自动升级脚本
# 作者: SOLO Coding
# 版本: 1.0
# 适用于: Ubuntu 24.04 LTS

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

# 日志函数
log_info() {
    echo -e "${BLUE}[INFO]${NC} $1"
}

log_success() {
    echo -e "${GREEN}[SUCCESS]${NC} $1"
}

log_warning() {
    echo -e "${YELLOW}[WARNING]${NC} $1"
}

log_error() {
    echo -e "${RED}[ERROR]${NC} $1"
}

# 检查是否为root用户
check_root() {
    if [[ $EUID -ne 0 ]]; then
        log_error "此脚本需要root权限运行，请使用 sudo 执行"
        exit 1
    fi
}

# 检查Ubuntu版本
check_ubuntu_version() {
    if ! grep -q "Ubuntu 24.04" /etc/os-release; then
        log_warning "检测到非Ubuntu 24.04系统，脚本可能无法正常工作"
        read -p "是否继续执行？(y/N): " -n 1 -r
        echo
        if [[ ! $REPLY =~ ^[Yy]$ ]]; then
            exit 1
        fi
    fi
}

# 备份原始sources.list
backup_sources() {
    local backup_file="/etc/apt/sources.list.backup.$(date +%Y%m%d_%H%M%S)"
    if [[ -f /etc/apt/sources.list ]]; then
        cp /etc/apt/sources.list "$backup_file"
        log_success "已备份原始sources.list到: $backup_file"
    fi
}

# 彻底清理所有第三方源文件
clean_all_sources() {
    log_info "正在彻底清理所有第三方源文件..."
    
    # 清理sources.list.d目录中的所有源文件
    if [[ -d /etc/apt/sources.list.d ]]; then
        # 备份sources.list.d目录
        local sources_d_backup="/etc/apt/sources.list.d.backup.$(date +%Y%m%d_%H%M%S)"
        cp -r /etc/apt/sources.list.d "$sources_d_backup" 2>/dev/null || true
        log_info "已备份sources.list.d目录到: $sources_d_backup"
        
        # 禁用所有第三方源文件
        find /etc/apt/sources.list.d -name "*.list" -exec mv {} {}.disabled \; 2>/dev/null || true
        find /etc/apt/sources.list.d -name "*.sources" -exec mv {} {}.disabled \; 2>/dev/null || true
        
        log_success "已清理sources.list.d目录中的所有第三方源文件"
    fi
    
    # 清理可能的缓存
    rm -rf /var/lib/apt/lists/* 2>/dev/null || true
    log_info "已清理APT缓存"
}

# 验证当前源配置
verify_single_source() {
    local expected_mirror="$1"
    local source_name="$2"
    
    log_info "正在验证 $source_name 源配置..."
    
    # 检查sources.list文件
    if [[ ! -f /etc/apt/sources.list ]]; then
        log_error "sources.list文件不存在"
        return 1
    fi
    
    # 统计不同镜像源的数量
    local active_sources=$(grep -v '^#' /etc/apt/sources.list | grep -v '^$' | grep -c '^deb' || echo "0")
    local expected_sources=$(grep -v '^#' /etc/apt/sources.list | grep -v '^$' | grep -c "$expected_mirror" || echo "0")
    
    if [[ $active_sources -eq $expected_sources && $expected_sources -gt 0 ]]; then
        log_success "源验证通过：当前使用的是纯净的 $source_name 镜像源（共 $expected_sources 个源）"
        return 0
    else
        log_warning "源验证失败：检测到混合源配置（总源数：$active_sources，$source_name 源数：$expected_sources）"
        return 1
    fi
}

# 换源到阿里云镜像
change_to_aliyun() {
    log_info "正在切换到阿里云镜像源..."
    backup_sources
    
    # 彻底清理所有第三方源
    clean_all_sources
    
    # 清空并重写sources.list文件
    cat > /etc/apt/sources.list << 'EOF'
# ============================================================
# 阿里云Ubuntu 24.04 LTS镜像源 - 纯净单一源配置
# 此文件由ubuntu_manager.sh自动生成，请勿手动编辑
# 镜像提供商: 阿里云 (Alibaba Cloud)
# 镜像地址: http://mirrors.aliyun.com/ubuntu/
# 配置时间: $(date '+%Y-%m-%d %H:%M:%S')
# ============================================================

# 主要软件仓库
deb http://mirrors.aliyun.com/ubuntu/ noble main restricted universe multiverse

# 安全更新仓库
deb http://mirrors.aliyun.com/ubuntu/ noble-security main restricted universe multiverse

# 推荐更新仓库
deb http://mirrors.aliyun.com/ubuntu/ noble-updates main restricted universe multiverse

# 后移植软件仓库
deb http://mirrors.aliyun.com/ubuntu/ noble-backports main restricted universe multiverse
EOF
    
    log_success "已切换到阿里云镜像源"
    
    # 验证源配置
    if verify_single_source "mirrors.aliyun.com" "阿里云"; then
        update_package_list
    else
        log_error "源配置验证失败，请检查配置"
        return 1
    fi
}

# 换源到清华大学镜像
change_to_tsinghua() {
    log_info "正在切换到清华大学镜像源..."
    backup_sources
    
    # 彻底清理所有第三方源
    clean_all_sources
    
    # 清空并重写sources.list文件
    cat > /etc/apt/sources.list << 'EOF'
# ============================================================
# 清华大学Ubuntu 24.04 LTS镜像源 - 纯净单一源配置
# 此文件由ubuntu_manager.sh自动生成，请勿手动编辑
# 镜像提供商: 清华大学 (Tsinghua University)
# 镜像地址: https://mirrors.tuna.tsinghua.edu.cn/ubuntu/
# 配置时间: $(date '+%Y-%m-%d %H:%M:%S')
# ============================================================

# 主要软件仓库
deb https://mirrors.tuna.tsinghua.edu.cn/ubuntu/ noble main restricted universe multiverse

# 安全更新仓库
deb https://mirrors.tuna.tsinghua.edu.cn/ubuntu/ noble-security main restricted universe multiverse

# 推荐更新仓库
deb https://mirrors.tuna.tsinghua.edu.cn/ubuntu/ noble-updates main restricted universe multiverse

# 后移植软件仓库
deb https://mirrors.tuna.tsinghua.edu.cn/ubuntu/ noble-backports main restricted universe multiverse
EOF
    
    log_success "已切换到清华大学镜像源"
    
    # 验证源配置
    if verify_single_source "mirrors.tuna.tsinghua.edu.cn" "清华大学"; then
        update_package_list
    else
        log_error "源配置验证失败，请检查配置"
        return 1
    fi
}

# 换源到中科大镜像
change_to_ustc() {
    log_info "正在切换到中科大镜像源..."
    backup_sources
    
    # 彻底清理所有第三方源
    clean_all_sources
    
    # 清空并重写sources.list文件
    cat > /etc/apt/sources.list << 'EOF'
# ============================================================
# 中科大Ubuntu 24.04 LTS镜像源 - 纯净单一源配置
# 此文件由ubuntu_manager.sh自动生成，请勿手动编辑
# 镜像提供商: 中国科学技术大学 (USTC)
# 镜像地址: https://mirrors.ustc.edu.cn/ubuntu/
# 配置时间: $(date '+%Y-%m-%d %H:%M:%S')
# ============================================================

# 主要软件仓库
deb https://mirrors.ustc.edu.cn/ubuntu/ noble main restricted universe multiverse

# 安全更新仓库
deb https://mirrors.ustc.edu.cn/ubuntu/ noble-security main restricted universe multiverse

# 推荐更新仓库
deb https://mirrors.ustc.edu.cn/ubuntu/ noble-updates main restricted universe multiverse

# 后移植软件仓库
deb https://mirrors.ustc.edu.cn/ubuntu/ noble-backports main restricted universe multiverse
EOF
    
    log_success "已切换到中科大镜像源"
    
    # 验证源配置
    if verify_single_source "mirrors.ustc.edu.cn" "中科大"; then
        update_package_list
    else
        log_error "源配置验证失败，请检查配置"
        return 1
    fi
}

# 换源到华为云镜像
change_to_huawei() {
    log_info "正在切换到华为云镜像源..."
    backup_sources
    
    # 彻底清理所有第三方源
    clean_all_sources
    
    # 清空并重写sources.list文件
    cat > /etc/apt/sources.list << 'EOF'
# ============================================================
# 华为云Ubuntu 24.04 LTS镜像源 - 纯净单一源配置
# 此文件由ubuntu_manager.sh自动生成，请勿手动编辑
# 镜像提供商: 华为云 (Huawei Cloud)
# 镜像地址: https://mirrors.huaweicloud.com/ubuntu/
# 配置时间: $(date '+%Y-%m-%d %H:%M:%S')
# ============================================================

# 主要软件仓库
deb https://mirrors.huaweicloud.com/ubuntu/ noble main restricted universe multiverse

# 安全更新仓库
deb https://mirrors.huaweicloud.com/ubuntu/ noble-security main restricted universe multiverse

# 推荐更新仓库
deb https://mirrors.huaweicloud.com/ubuntu/ noble-updates main restricted universe multiverse

# 后移植软件仓库
deb https://mirrors.huaweicloud.com/ubuntu/ noble-backports main restricted universe multiverse
EOF
    
    log_success "已切换到华为云镜像源"
    
    # 验证源配置
    if verify_single_source "mirrors.huaweicloud.com" "华为云"; then
        update_package_list
    else
        log_error "源配置验证失败，请检查配置"
        return 1
    fi
}

# 更新软件包列表
update_package_list() {
    log_info "正在更新软件包列表..."
    if apt update; then
        log_success "软件包列表更新成功"
    else
        log_error "软件包列表更新失败"
        return 1
    fi
}

# 系统升级
system_upgrade() {
    log_info "开始系统升级..."
    
    # 更新软件包列表
    if ! update_package_list; then
        return 1
    fi
    
    # 升级已安装的软件包
    log_info "正在升级已安装的软件包..."
    if apt upgrade -y; then
        log_success "软件包升级完成"
    else
        log_error "软件包升级失败"
        return 1
    fi
    
    # 完整系统升级（包括内核等）
    log_info "正在进行完整系统升级..."
    if apt full-upgrade -y; then
        log_success "完整系统升级完成"
    else
        log_error "完整系统升级失败"
        return 1
    fi
    
    # 清理不需要的软件包
    log_info "正在清理不需要的软件包..."
    apt autoremove -y
    apt autoclean
    
    log_success "系统升级完成！"
    
    # 检查是否需要重启
    if [[ -f /var/run/reboot-required ]]; then
        log_warning "系统升级完成，建议重启系统以应用所有更改"
        read -p "是否现在重启？(y/N): " -n 1 -r
        echo
        if [[ $REPLY =~ ^[Yy]$ ]]; then
            log_info "系统将在5秒后重启..."
            sleep 5
            reboot
        fi
    fi
}

# 显示当前源信息
show_current_sources() {
    log_info "当前软件源配置:"
    echo "----------------------------------------"
    if [[ -f /etc/apt/sources.list ]]; then
        grep -v '^#' /etc/apt/sources.list | grep -v '^$'
    else
        log_warning "未找到sources.list文件"
    fi
    echo "----------------------------------------"
}

# 恢复原始源
restore_original_sources() {
    log_info "查找备份文件..."
    local backup_files=($(ls /etc/apt/sources.list.backup.* 2>/dev/null | sort -r))
    
    if [[ ${#backup_files[@]} -eq 0 ]]; then
        log_error "未找到备份文件"
        return 1
    fi
    
    echo "找到以下备份文件:"
    for i in "${!backup_files[@]}"; do
        echo "$((i+1)). ${backup_files[i]}"
    done
    
    read -p "请选择要恢复的备份文件编号 (1-${#backup_files[@]}): " choice
    
    if [[ $choice -ge 1 && $choice -le ${#backup_files[@]} ]]; then
        local selected_backup="${backup_files[$((choice-1))]}"
        cp "$selected_backup" /etc/apt/sources.list
        log_success "已恢复备份: $selected_backup"
        update_package_list
    else
        log_error "无效的选择"
        return 1
    fi
}

# 显示主菜单
show_menu() {
    clear
    echo -e "${BLUE}================================================${NC}"
    echo -e "${BLUE}        Ubuntu 24.04 换源和升级工具${NC}"
    echo -e "${BLUE}================================================${NC}"
    echo
    echo -e "${GREEN}1.${NC} 换源到阿里云镜像"
    echo -e "${GREEN}2.${NC} 换源到清华大学镜像"
    echo -e "${GREEN}3.${NC} 换源到中科大镜像"
    echo -e "${GREEN}4.${NC} 换源到华为云镜像"
    echo -e "${GREEN}5.${NC} 系统升级"
    echo -e "${GREEN}6.${NC} 显示当前软件源"
    echo -e "${GREEN}7.${NC} 恢复原始软件源"
    echo -e "${GREEN}8.${NC} 一键换源并升级（推荐阿里云）"
    echo -e "${GREEN}0.${NC} 退出"
    echo
    echo -e "${BLUE}================================================${NC}"
}

# 一键换源并升级
quick_setup() {
    log_info "开始一键换源并升级..."
    change_to_aliyun
    if [[ $? -eq 0 ]]; then
        system_upgrade
    else
        log_error "换源失败，跳过升级步骤"
    fi
}

# 主程序
main() {
    # 检查权限和系统版本
    check_root
    check_ubuntu_version
    
    while true; do
        show_menu
        read -p "请选择操作 (0-8): " choice
        
        case $choice in
            1)
                change_to_aliyun
                ;;
            2)
                change_to_tsinghua
                ;;
            3)
                change_to_ustc
                ;;
            4)
                change_to_huawei
                ;;
            5)
                system_upgrade
                ;;
            6)
                show_current_sources
                ;;
            7)
                restore_original_sources
                ;;
            8)
                quick_setup
                ;;
            0)
                log_info "感谢使用，再见！"
                exit 0
                ;;
            *)
                log_error "无效的选择，请重新输入"
                ;;
        esac
        
        echo
        read -p "按回车键继续..." -r
    done
}

# 脚本入口点
if [[ "${BASH_SOURCE[0]}" == "${0}" ]]; then
    main "$@"
fi