#!/bin/bash

# Debian 12 国内源配置脚本
# 支持多种国内镜像源

set -e

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

# 日志函数
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"
}

# 备份当前源配置
backup_sources() {
    local backup_dir="/root/sources_backup_$(date +%Y%m%d_%H%M%S)"
    mkdir -p "$backup_dir"

    log_info "备份当前源配置到: $backup_dir"

    if [[ -f /etc/apt/sources.list ]]; then
        cp /etc/apt/sources.list "$backup_dir/"
        log_success "已备份 /etc/apt/sources.list"
    fi

    if [[ -d /etc/apt/sources.list.d ]]; then
        cp -r /etc/apt/sources.list.d "$backup_dir/"
        log_success "已备份 /etc/apt/sources.list.d/"
    fi

    echo "$backup_dir" > /tmp/sources_backup_path
    log_success "备份完成，路径已保存到 /tmp/sources_backup_path"
}

# 恢复源配置
restore_sources() {
    local backup_path="$1"

    if [[ -z "$backup_path" ]]; then
        if [[ -f /tmp/sources_backup_path ]]; then
            backup_path=$(cat /tmp/sources_backup_path)
        else
            log_error "未找到备份路径"
            return 1
        fi
    fi

    if [[ ! -d "$backup_path" ]]; then
        log_error "备份目录不存在: $backup_path"
        return 1
    fi

    log_info "从 $backup_path 恢复源配置"

    if [[ -f "$backup_path/sources.list" ]]; then
        cp "$backup_path/sources.list" /etc/apt/sources.list
        log_success "已恢复 /etc/apt/sources.list"
    fi

    if [[ -d "$backup_path/sources.list.d" ]]; then
        cp -r "$backup_path/sources.list.d" /etc/apt/
        log_success "已恢复 /etc/apt/sources.list.d/"
    fi

    log_success "源配置恢复完成"
}

# 配置阿里云源
configure_aliyun() {
    log_info "配置阿里云镜像源..."

    cat > /etc/apt/sources.list << 'EOF'
deb https://mirrors.aliyun.com/debian/ bookworm main non-free non-free-firmware contrib
deb-src https://mirrors.aliyun.com/debian/ bookworm main non-free non-free-firmware contrib
deb https://mirrors.aliyun.com/debian-security/ bookworm-security main
deb-src https://mirrors.aliyun.com/debian-security/ bookworm-security main
deb https://mirrors.aliyun.com/debian/ bookworm-updates main non-free non-free-firmware contrib
deb-src https://mirrors.aliyun.com/debian/ bookworm-updates main non-free non-free-firmware contrib
deb https://mirrors.aliyun.com/debian/ bookworm-backports main non-free non-free-firmware contrib
deb-src https://mirrors.aliyun.com/debian/ bookworm-backports main non-free non-free-firmware contrib
EOF

    log_success "阿里云源配置完成"
}

# 配置清华大学源
configure_tsinghua() {
    log_info "配置清华大学镜像源..."

    cat > /etc/apt/sources.list << 'EOF'
deb https://mirrors.tuna.tsinghua.edu.cn/debian/ bookworm main contrib non-free non-free-firmware
deb-src https://mirrors.tuna.tsinghua.edu.cn/debian/ bookworm main contrib non-free non-free-firmware
deb https://mirrors.tuna.tsinghua.edu.cn/debian-security/ bookworm-security main contrib non-free non-free-firmware
deb-src https://mirrors.tuna.tsinghua.edu.cn/debian-security/ bookworm-security main contrib non-free non-free-firmware
deb https://mirrors.tuna.tsinghua.edu.cn/debian/ bookworm-updates main contrib non-free non-free-firmware
deb-src https://mirrors.tuna.tsinghua.edu.cn/debian/ bookworm-updates main contrib non-free non-free-firmware
deb https://mirrors.tuna.tsinghua.edu.cn/debian/ bookworm-backports main contrib non-free non-free-firmware
deb-src https://mirrors.tuna.tsinghua.edu.cn/debian/ bookworm-backports main contrib non-free non-free-firmware
EOF

    log_success "清华大学源配置完成"
}

# 配置网易源
configure_163() {
    log_info "配置网易镜像源..."

    cat > /etc/apt/sources.list << 'EOF'
deb https://mirrors.163.com/debian/ bookworm main non-free non-free-firmware contrib
deb-src https://mirrors.163.com/debian/ bookworm main non-free non-free-firmware contrib
deb https://mirrors.163.com/debian-security/ bookworm-security main
deb-src https://mirrors.163.com/debian-security/ bookworm-security main
deb https://mirrors.163.com/debian/ bookworm-updates main non-free non-free-firmware contrib
deb-src https://mirrors.163.com/debian/ bookworm-updates main non-free non-free-firmware contrib
deb https://mirrors.163.com/debian/ bookworm-backports main non-free non-free-firmware contrib
deb-src https://mirrors.163.com/debian/ bookworm-backports main non-free non-free-firmware contrib
EOF

    log_success "网易源配置完成"
}

# 配置中科大源
configure_ustc() {
    log_info "配置中科大镜像源..."

    cat > /etc/apt/sources.list << 'EOF'
deb https://mirrors.ustc.edu.cn/debian/ bookworm main contrib non-free non-free-firmware
deb-src https://mirrors.ustc.edu.cn/debian/ bookworm main contrib non-free non-free-firmware
deb https://mirrors.ustc.edu.cn/debian-security/ bookworm-security main contrib non-free non-free-firmware
deb-src https://mirrors.ustc.edu.cn/debian-security/ bookworm-security main contrib non-free non-free-firmware
deb https://mirrors.ustc.edu.cn/debian/ bookworm-updates main contrib non-free non-free-firmware
deb-src https://mirrors.ustc.edu.cn/debian/ bookworm-updates main contrib non-free non-free-firmware
deb https://mirrors.ustc.edu.cn/debian/ bookworm-backports main contrib non-free non-free-firmware
deb-src https://mirrors.ustc.edu.cn/debian/ bookworm-backports main contrib non-free non-free-firmware
EOF

    log_success "中科大源配置完成"
}

# 配置华为云源
configure_huawei() {
    log_info "配置华为云镜像源..."

    cat > /etc/apt/sources.list << 'EOF'
deb https://mirrors.huaweicloud.com/debian/ bookworm main non-free non-free-firmware contrib
deb-src https://mirrors.huaweicloud.com/debian/ bookworm main non-free non-free-firmware contrib
deb https://mirrors.huaweicloud.com/debian-security/ bookworm-security main
deb-src https://mirrors.huaweicloud.com/debian-security/ bookworm-security main
deb https://mirrors.huaweicloud.com/debian/ bookworm-updates main non-free non-free-firmware contrib
deb-src https://mirrors.huaweicloud.com/debian/ bookworm-updates main non-free non-free-firmware contrib
deb https://mirrors.huaweicloud.com/debian/ bookworm-backports main non-free non-free-firmware contrib
deb-src https://mirrors.huaweicloud.com/debian/ bookworm-backports main non-free non-free-firmware contrib
EOF

    log_success "华为云源配置完成"
}

# 测试源连接
test_sources() {
    log_info "测试源连接速度..."

    local sources=("https://mirrors.aliyun.com"
                   "https://mirrors.tuna.tsinghua.edu.cn"
                   "https://mirrors.163.com"
                   "https://mirrors.ustc.edu.cn"
                   "https://mirrors.huaweicloud.com")

    local fastest_source=""
    local fastest_time=999

    for source in "${sources[@]}"; do
        log_info "测试 $source ..."
        local start_time=$(date +%s%N)
        if curl -s --connect-timeout 5 --max-time 10 "${source}/debian/dists/bookworm/Release" > /dev/null 2>&1; then
            local end_time=$(date +%s%N)
            local time_diff=$(( (end_time - start_time) / 1000000 ))
            log_success "$source 响应时间: ${time_diff}ms"

            if (( time_diff < fastest_time )); then
                fastest_time=$time_diff
                fastest_source=$source
            fi
        else
            log_warning "$source 连接失败"
        fi
    done

    if [[ -n "$fastest_source" ]]; then
        log_success "最快源: $fastest_source (${fastest_time}ms)"
        echo "$fastest_source"
    else
        log_error "所有源都无法连接"
        return 1
    fi
}

# 更新包索引
update_packages() {
    log_info "更新包索引..."
    apt update
    log_success "包索引更新完成"
}

# 显示帮助信息
show_help() {
    cat << EOF
Debian 12 国内源配置工具

用法: $0 [选项] [源名称]

源名称:
    aliyun      阿里云镜像源
    tsinghua    清华大学镜像源
    163         网易镜像源
    ustc        中科大镜像源
    huawei      华为云镜像源
    test        测试所有源速度并选择最快
    backup      备份当前源配置
    restore     恢复源配置

选项:
    -h, --help  显示此帮助信息
    -n, --no-update  配置完成后不更新包索引

示例:
    $0 aliyun           # 配置阿里云源
    $0 test             # 测试并选择最快源
    $0 backup           # 备份当前配置
    $0 restore /path/to/backup  # 从指定路径恢复

EOF
}

# 主函数
main() {
    local no_update=false
    local action=""

    # 参数解析
    while [[ $# -gt 0 ]]; do
        case $1 in
            -h|--help)
                show_help
                exit 0
                ;;
            -n|--no-update)
                no_update=true
                shift
                ;;
            -*)
                log_error "未知选项: $1"
                show_help
                exit 1
                ;;
            *)
                action="$1"
                shift
                break
                ;;
        esac
    done

    case $action in
        aliyun)
            backup_sources
            configure_aliyun
            ;;
        tsinghua)
            backup_sources
            configure_tsinghua
            ;;
        163)
            backup_sources
            configure_163
            ;;
        ustc)
            backup_sources
            configure_ustc
            ;;
        huawei)
            backup_sources
            configure_huawei
            ;;
        test)
            local fastest=$(test_sources)
            if [[ $? -eq 0 ]]; then
                case $fastest in
                    *aliyun*)
                        configure_aliyun
                        ;;
                    *tuna.tsinghua*)
                        configure_tsinghua
                        ;;
                    *163*)
                        configure_163
                        ;;
                    *ustc*)
                        configure_ustc
                        ;;
                    *huaweicloud*)
                        configure_huawei
                        ;;
                esac
            fi
            ;;
        backup)
            backup_sources
            exit 0
            ;;
        restore)
            restore_sources "$1"
            exit 0
            ;;
        "")
            log_error "未指定操作"
            show_help
            exit 1
            ;;
        *)
            log_error "未知操作: $action"
            show_help
            exit 1
            ;;
    esac

    # 更新包索引
    if [[ "$no_update" != "true" ]]; then
        update_packages
    fi

    log_success "源配置完成！"
}

# 执行主函数
main "$@"