#!/bin/sh
# MT-3000 软件包更换分区助手
# by @wukongdaily

# 定义颜色输出函数
red() { echo -e "\033[31m\033[01m$1\033[0m"; }
green() { echo -e "\033[32m\033[01m$1\033[0m"; }
yellow() { echo -e "\033[33m\033[01m$1\033[0m"; }
cyan() { echo -e "\033[38;2;0;255;255m$1\033[0m"; }

install_depends_apps() {
    cyan "正在安装必备工具..."
    opkg update >/dev/null 2>&1
    for pkg in lsblk fdisk; do
        if opkg list-installed | grep -qw "$pkg"; then
            cyan "$pkg 已安装。"
        else
            if opkg install "$pkg" >/dev/null 2>&1; then
                green "$pkg 安装成功。"
            else
                red "$pkg 安装失败。"
                exit 1
            fi
        fi
    done
}

# 卸载USB设备
unmount_usb_device() {
    for mount in $(mount | grep "$1" | awk '{print $3}'); do
        yellow "正在尝试卸载U盘挂载点：$mount"
        umount $mount || {
            red "警告：无法卸载挂载点 $mount。可能有文件正在被访问或权限不足。"
            exit 1
        }
        green "U盘挂载点 $mount 卸载成功。"
    done
}

create_and_format_partitions() {
    local device=$1
    local total_bytes=$(fdisk -l $device | grep "Disk $device:" | awk '{print $5}')
    local total_gb=$(echo "$total_bytes" | awk '{print int($1/(1024*1024*1024))}')

    if [ -n "$CUSTOM_OPKG_SIZE" ]; then
        part1_gb=$CUSTOM_OPKG_SIZE
        yellow "U盘总容量约为 $total_gb GB,您设置的自定义软件包大小为 ${part1_gb}GB。"
    else
        part1_gb=$((total_gb / 10))
        yellow "U盘总容量约为 $total_gb GB,第一分区大小设置为U盘容量的10% 即 ${part1_gb}GB。"
    fi
    green "计划将第一分区分配给软件包 其大小为 ${part1_gb}GB"
    cyan "没错～你没有看错,让我们任性的告别容量焦虑！"

    {
        echo g
        echo n
        echo p
        echo 1
        echo
        echo +${part1_gb}G
        echo n
        echo p
        echo 2
        echo
        echo
        echo w
    } | fdisk $device >/dev/null 2>&1

    sleep 5

    local new_partition1="${device}1"
    cyan "正在将 $new_partition1 格式化为EXT4文件系统..."
    mkfs.ext4 -F $new_partition1 >/dev/null 2>&1
    green "$new_partition1 已成功格式化为EXT4文件系统。"
    yellow "第2分区 ${device}2 暂不格式化,未来您可自行分配使用"
}

# 换区到U盘
change_overlay_usb() {
    install_depends_apps
    cyan "现在开始查找USB设备分区 请稍后..."
    local USB_PARTITION=$(lsblk -dn -o NAME,TRAN | awk '$2=="usb" {print "/dev/"$1; exit}')
    if [ -z "$USB_PARTITION" ]; then
        red "未找到USB磁盘。"
        exit 1
    fi
    yellow "找到USB磁盘 $USB_PARTITION"
    dd if=/dev/zero of=$USB_PARTITION bs=1M count=10
    sync
    unmount_usb_device "$USB_PARTITION"
    red "正在将U盘${USB_PARTITION}分为2个区..."
    create_and_format_partitions "$USB_PARTITION"

    MOUNT_POINT="/mnt/usb_overlay"
    TMP_DIR="/tmp/overlay_backup"
    mkdir -p $MOUNT_POINT
    cyan "重新挂载第一分区 ${USB_PARTITION}1 到 $MOUNT_POINT"
    mount ${USB_PARTITION}1 $MOUNT_POINT >/dev/null 2>&1
    mkdir -p $TMP_DIR
    cp -a /overlay/. $TMP_DIR
    green "正在拷贝当前系统文件到U盘"
    cp -a $TMP_DIR/. $MOUNT_POINT
    cyan "正在更新启动时的配置文件"
    uci set fstab.overlay=mount
    uci set fstab.overlay.uuid="$(blkid -o value -s UUID ${USB_PARTITION}1)"
    uci set fstab.overlay.target="/overlay"
    uci commit fstab
    rm -rf $TMP_DIR
    cp /etc/config/fstab $MOUNT_POINT/fstab.bak
    green "overlay更换分区完成 重启验证是否成功。"
    red "是否立即重启？(y/n)"
    read -r answer
    if [ "$answer" = "y" ] || [ -z "$answer" ]; then
        red "正在重启..."
        reboot
    else
        yellow "您选择了不重启"
    fi
}

# 重新绑定
rebind_usb_overlay() {
    cyan "正在重新绑定U盘设备..."
    if ! opkg list-installed | grep -qw "lsblk"; then
        opkg update >/dev/null 2>&1
        opkg install "lsblk" >/dev/null 2>&1 || {
            red "lsblk 安装失败。"
            exit 1
        }
    fi
    local USB_DEVICE=$(lsblk -dn -o NAME,RM,TYPE | awk '$2=="1" && $3=="disk" {print "/dev/"$1; exit}')
    if [ -z "$USB_DEVICE" ]; then
        red "未找到USB磁盘。"
        exit 1
    fi
    uci set fstab.overlay=mount
    uci set fstab.overlay.uuid="$(blkid -o value -s UUID ${USB_DEVICE}1)"
    uci set fstab.overlay.target="/overlay"
    uci commit fstab
    green "重新绑定成功！重启后生效"
    red "正在重启..."
    reboot
}

# 自定义软件包的大小，默认为U盘容量的10%
custom_package_size() {
    while :; do
        echo "请输入想分配的软件包的大小(数字,单位:GB):"
        read size
        if [[ $size =~ ^[0-9]+$ ]]; then
            CUSTOM_OPKG_SIZE=$size
            yellow "已设置软件包大小为:$CUSTOM_OPKG_SIZE GB"
            green "接下来,您可以执行第一项啦"
            break
        else
            red "错误: 请输入一个有效的数字。"
        fi
    done
}

while true; do
    clear
    echo "***********************************************************************"
    green "      MT-3000 软件包更换分区助手"
    echo "**********************************************************************"
    echo
    cyan " 1. MT-3000 一键更换overlay分区到U盘"
    cyan " 2. 自定义设置软件包大小(GB)"
    yellow " 3. 重新绑定U盘"
    echo
    echo " Q. 退出本程序"
    echo
    read -p "请选择一个选项: " choice
    echo

    case $choice in
    1)
        change_overlay_usb
        ;;
    2)
        custom_package_size
        ;;
    3)
        rebind_usb_overlay
        ;;
    q | Q)
        echo "退出"
        exit 0
        ;;
    *)
        echo "无效选项，请重新选择。"
        ;;
    esac
    read -p "按 Enter 键继续..."
done