#!/bin/bash
# ================================================================
#  PVE 7/8 一键换国内源 + 一键还原
#  默认中科大，参数 tsinghua 则换清华
#  支持子命令：install（换源） / restore（还原）
# ================================================================
set -e

MIRROR="${1:-ustc}"
ACTION="${2:-install}"   # 新增：install | restore
BACKUP_DIR="/root/pve_repo_backup/$(date +%F-%T)"
mkdir -p "$BACKUP_DIR"

RED='\033[0;31m'; GREEN='\033[0;32m'; YELLOW='\033[0;33m'; NC='\033[0m'
log()  { echo -e "${GREEN}[INFO]${NC} $1"; }
warn() { echo -e "${YELLOW}[WARN]${NC} $1"; }
error(){ echo -e "${RED}[ERROR]${NC} $1"; }

# 检测系统版本
if grep -q 'VERSION_CODENAME=bookworm' /etc/os-release; then
    CODENAME="bookworm"
elif grep -q 'VERSION_CODENAME=bullseye' /etc/os-release; then
    CODENAME="bullseye"
else
    error "未检测到支持的 PVE7/8 版本"; exit 1
fi
log "系统版本：Debian $CODENAME"

# 镜像地址
case "$MIRROR" in
    tsinghua)
        DEB_MIRROR="https://mirrors.tuna.tsinghua.edu.cn/debian"
        DEB_SEC_MIRROR="https://mirrors.tuna.tsinghua.edu.cn/debian-security"
        PVE_MIRROR="https://mirrors.tuna.tsinghua.edu.cn/proxmox/debian/pve"
        CEPH_MIRROR="https://mirrors.tuna.tsinghua.edu.cn/proxmox/debian/ceph-quincy"
        CT_MIRROR="https://mirrors.tuna.tsinghua.edu.cn/proxmox"
        ;;
    *)
        DEB_MIRROR="https://mirrors.ustc.edu.cn/debian"
        DEB_SEC_MIRROR="https://mirrors.ustc.edu.cn/debian-security"
        PVE_MIRROR="https://mirrors.ustc.edu.cn/proxmox/debian/pve"
        CEPH_MIRROR="https://mirrors.ustc.edu.cn/proxmox/debian/ceph-quincy"
        CT_MIRROR="https://mirrors.ustc.edu.cn/proxmox"
        ;;
esac

# 需要备份/还原的文件列表
FILES_TO_BACKUP=(
    /etc/apt/sources.list
    /etc/apt/sources.list.d/pve-enterprise.list
    /etc/apt/sources.list.d/pve-no-subscription.list
    /etc/apt/sources.list.d/ceph.list
    /usr/share/perl5/PVE/APLInfo.pm
)

# 备份函数
backup() {
    for f in "${FILES_TO_BACKUP[@]}"; do
        [[ -f $f ]] && { cp --parents "$f" "$BACKUP_DIR" && log "已备份 $f"; }
    done
}

# 一键还原函数
restore() {
    local latest
    latest=$(ls -1d /root/pve_repo_backup/*/ 2>/dev/null | tail -n1)
    if [[ -z $latest ]]; then
        error "未找到备份目录，无法还原"; exit 1
    fi
    log "将使用备份目录：$latest"
    for f in "${FILES_TO_BACKUP[@]}"; do
        local bf="$latest${f#/}"
        [[ -f $bf ]] && { cp "$bf" "$f" && log "已还原 $f"; }
    done
    apt-get update
    log "还原完成！"
}

# 换源函数
change_repo() {
    log "开始换源（$MIRROR）..."
    backup

    # sources.list
    cat > /etc/apt/sources.list <<EOF
deb $DEB_MIRROR $CODENAME main contrib non-free non-free-firmware
deb $DEB_MIRROR $CODENAME-updates main contrib non-free non-free-firmware
deb $DEB_SEC_MIRROR $CODENAME-security main contrib non-free non-free-firmware
EOF

    # PVE 源
    rm -f /etc/apt/sources.list.d/pve-enterprise.list
    cat > /etc/apt/sources.list.d/pve-no-subscription.list <<EOF
deb $PVE_MIRROR $CODENAME pve-no-subscription
EOF

    # Ceph 源
    [[ -f /etc/apt/sources.list.d/ceph.list ]] && {
        cat > /etc/apt/sources.list.d/ceph.list <<EOF
deb $CEPH_MIRROR $CODENAME no-subscription
EOF
    }

    # CT 模板源
    sed -i.bak \
        -e "s|http://download.proxmox.com|${CT_MIRROR}|g" \
        -e "s|https://releases.turnkeylinux.org|${CT_MIRROR}/turnkeylinux/metadata|g" \
        /usr/share/perl5/PVE/APLInfo.pm

    apt-get update
    log "换源完成！"
}

# 主流程
case "${ACTION}" in
    install|i)
        change_repo
        ;;
    restore|r)
        restore
        ;;
    *)
        echo "用法："
        echo "  换源  : $0 [ustc|tsinghua] install"
        echo "  还原  : $0 restore"
        ;;
esac