#!/bin/bash
export PATH=/bin:/sbin:/usr/bin:/usr/sbin:/usr/local/bin:/usr/local/sbin:$PATH

: '
Copyright (C) 2022 - now  HaoZi Technology Co., Ltd.

This program is free software: you can redistribute it and/or modify
it under the terms of the GNU Affero General Public License as published
by the Free Software Foundation, either version 3 of the License, or
(at your option) any later version.

This program is distributed in the hope that it will be useful,
but WITHOUT ANY WARRANTY; without even the implied warranty of
MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
GNU Affero General Public License for more details.

You should have received a copy of the GNU Affero General Public License
along with this program.  If not, see <https://www.gnu.org/licenses/>.
'

source <(curl -f -s --connect-timeout 10 --retry 3 https://dl.cdn.haozi.net/panel/public.sh)
if [ $? -ne 0 ]; then
    echo "下载 public.sh 失败，请检查网络或稍后重试。"
    echo "Download public.sh failed, please check the network or try again later."
    exit 1
fi

LOGO="+----------------------------------------------------\n| 耗子面板硬盘自动分区挂载脚本\n| Rat Panel disk automatic partition mounting script\n+----------------------------------------------------\n| 本脚本具有超级牛力\n| This script has Super Cow Powers\n+----------------------------------------------------\n| Copyright © 2022-"$(date +%Y)" 耗子科技 All rights reserved.\n+----------------------------------------------------"

echo -e $LOGO

echo "正在安装分区软件，请稍后..."
echo "Installing partition software, please wait..."
if [ ${OS} == "rhel" ]; then
    dnf makecache
    dnf install -y xfsprogs e2fsprogs util-linux
elif [ ${OS} == "debian" ] || [ ${OS} == "ubuntu" ]; then
    apt-get update
    apt-get install -y xfsprogs e2fsprogs fdisk util-linux bsdmainutils
fi
if [ $? -ne 0 ]; then
    error "安装分区软件失败"
fi

echo "分区软件安装完成"
echo "Partition software installation completed"
echo "当前系统中的硬盘设备列表："
echo "Disk list in the current system:"
lsblk -dno NAME,SIZE,TYPE | grep 'disk' | column -t

while true; do
    read -p "请选择硬盘设备(Please select a disk)[e.g. sdb,vdb,nvme1n1]: " disk
    if [ -b /dev/$disk ] && [[ $disk != *a ]]; then
        echo "您选择的设备是 /dev/$disk："
        echo "The device you selected is /dev/$disk:"
        lsblk /dev/$disk
        break
    elif [[ $disk == *a ]]; then
        echo "不允许选择系统盘，请重新输入"
        echo "System disk is not allowed, please re-enter"
    else
        echo "设备 /dev/$disk 不存在，请重新输入"
        echo "Device /dev/$disk does not exist, please re-enter"
    fi
done

while true; do
    read -p "请输入挂载点(Please enter the mount point)[e.g. /www]: " mount_point
    mount_point=${mount_point:-/www}
    if [ "${mount_point:0:1}" != "/" ]; then
        echo "请输入绝对路径(Please enter an absolute path)"
    else
        break
    fi
done

if [ ! -d $mount_point ]; then
    mkdir -p $mount_point
    echo "已创建挂载点 $mount_point"
    echo "Mount point $mount_point created"
fi

# 判断挂载点是否存在数据
if [ "$(ls -A $mount_point)" ]; then
    echo "挂载点 $mount_point 中存在数据，请先手动清空"
    echo "Mount point $mount_point contains data, please empty it manually first"
    exit 1
fi

# 检查是否已分区
partition='y'
if [ -b /dev/${disk}1 ]; then
    echo "设备 /dev/$disk 已分区，是否重新分区？"
    echo "Device /dev/$disk is partitioned, do you want to repartition?"
    echo "注意：重新分区将删除所有数据，输入 n 则仅挂载不分区"
    echo "Note: Repartitioning will delete all data, enter n to mount without partitioning"
    while true; do
        read -p "请输入您的选择(Please enter your choice)[y/n]：" partition
        case $partition in
        [Yy]*)
            break
            ;;
        [Nn]*)
            break
            ;;
        *)
            echo "无效选择，请重新输入"
            echo "Invalid choice, please re-enter"
            ;;
        esac
    done
fi

if [ $partition == 'y' ]; then
    # 卸载分区
    umount /dev/${disk}1 >/dev/null 2>&1
    # 选择文件系统
    echo "请选择文件系统:"
    echo "Please select the file system:"
    echo "1) ext4（推荐）(Recommended)"
    echo "2) xfs"

    while true; do
        read -p "请输入您的选择(Please enter your choice)[1-3]：" fs_choice
        case $fs_choice in
        1)
            fs_type="ext4"
            break
            ;;
        2)
            fs_type="xfs"
            break
            ;;
        *)
            echo "无效选择，请重新输入"
            echo "Invalid choice, please re-enter"
            ;;
        esac
    done

    echo "正在删除 /dev/$disk 上的所有分区..."
    echo "Deleting all partitions on /dev/$disk..."
    (
        partitions=$(lsblk -no NAME /dev/$disk | grep -v "^${disk}$" | wc -l)
        for ((i = 1; i <= $partitions; i++)); do
            echo d
        done
        echo wq
    ) | fdisk /dev/$disk

    echo "正在对 /dev/$disk 进行分区..."
    echo "Partitioning /dev/$disk..."
    (
        echo g  # 新分区表
        echo n  # 新建分区
        echo 1  # 分区号 1
        echo    # 默认开始扇区
        echo    # 默认结束扇区
        echo y  # 确认删除
        echo wq # 写入分区表
    ) | fdisk /dev/$disk

    echo "分区完成"
    echo "Partitioning completed"

    echo "正在格式化 /dev/${disk}1 为 $fs_type 文件系统..."
    echo "Formatting /dev/${disk}1 to $fs_type file system..."
    case $fs_type in
    ext4)
        mkfs.ext4 -F /dev/${disk}1
        ;;
    xfs)
        mkfs.xfs -f /dev/${disk}1
        ;;
    esac
fi

if [ $partition == 'n' ]; then
    fs_type=$(lsblk -no FSTYPE /dev/${disk}1)
    if [ $? -ne 0 ]; then
        error "无法获取 /dev/${disk}1 的文件系统类型，可能需要重新分区"
    fi
fi

# 挂载分区
echo "正在挂载 /dev/${disk}1 到 $mount_point..."
echo "Mounting /dev/${disk}1 to $mount_point..."
systemctl daemon-reload

mount /dev/${disk}1 $mount_point
if [ $? -ne 0 ]; then
    error "挂载失败"
fi

echo "正在更新 /etc/fstab 以自动挂载..."
echo "Updating /etc/fstab for auto-mounting..."
uuid=$(blkid -s UUID -o value /dev/${disk}1)
if [ -z "$uuid" ]; then
    error "无法获取 /dev/${disk}1 的 UUID"
fi

sed -i "\|/dev/${disk}1|d" /etc/fstab
sed -i "\|$mount_point|d" /etc/fstab
echo "UUID=$uuid    $mount_point    $fs_type    defaults    0    0" >>/etc/fstab

systemctl daemon-reload
mount -a
if [ $? -ne 0 ]; then
    error "挂载失败"
fi

# 清理挂载点下的文件
rm -rf $mount_point/*
echo "硬盘分区挂载完成"
echo "Disk partition mounting completed"