#!/bin/bash

BACKUP_SNAPSHOTS_PATH="/backup/snapshots"
EXCLUDE_PATHS="/backup \
        /cdrom \
        /dev \
        /efi \
        /etc/uid_list \
        /data/ghost \
        /ghost \
        */ghost \
        /lost+found \
        /media \
        /mnt   \
        /proc  \
        /run   \
        /swap_file \
        /sys   \
        /tmp   \
        /var/lib/docker/overlay2 \
        /var/libkmre/data \
        /var/lib/kmre/kmre-*-*/data/media/0/0-麒麟* \
        /var/lib/udisk2 \
        /var/log \
        */backup/snapshots \
        /data/home \
        /data/root \
	/etc/.kyhwid"
MSG_PARENT_DIRECOTRY="上一个目录"
MSG_EXIT_BROWSING="Exit directory browsing"

m_filelist=""
m_ghost_img_dir=""
m_ghost_img_dir_created=0
m_is_mount_backup=0
m_backup_uuid=""
m_backup_path_default="/backup"
m_gpg_key_id=""
m_settings_customization=0
M_DEFAULT_GHOST_IMG_DIR="/ghost"
m_ghost_store_partition=$M_DEFAULT_GHOST_IMG_DIR
m_ghost_directory_exclude_path=("/proc" "/var"  "/dev")

show_help() {
    echo "用法: mk-ghost [options]"
    echo "该命令用于快速创建系统镜像。"
    echo
    echo "参数:"
    echo "  -h, --help    显示帮助信息。"
    echo "  -d, --define  自定义创建系统镜像。"
    echo
    echo "举例："
    echo "  mk-ghost      # 使用默认值立即创建系统镜像。"
    echo "  mk-ghost -d   # 可自定义镜像的名称及存放目录。"
    echo ""
}

check_if_root() {
    if [ ! "$UID" = "0" ]; then
        echo
        echo "需要以 root 用户身份运行。"
        echo
        exit 1
    fi
}

check_and_remove_folder() {
    local folder_path="$1"

    if [[ ! -d "$folder_path" ]]; then
        echo "文件夹不存在：$folder_path"
        return
    fi

    if [[ -z "$(ls -A "$folder_path")" ]]; then
        # 文件夹为空，尝试删除
        rmdir "$folder_path"
    fi
}

overwrite() {
    printf "\r%s" "$1"
}

rep_whspc_w_udrsc() {
    local inp_txt="$1"
    local out_txt
    out_txt="$(echo ${inp_txt} | sed -r -e "s/[[:space:]]/_/g")"
    echo "$out_txt"
}

get_directories_list_common() {
    for file in ./*; do
        [ ! -d "$file" ] && continue
    realpath=$(readlink -f "$file")
    if printf "%s\n" "${m_ghost_directory_exclude_path[@]}" | grep -q "^$realpath$"; then
          continue;
        fi
        #fileinfo="$(unalias ls 2>/dev/null; ls -ldh "$file" 2>/dev/null| awk '{ print $6"_"$7; }')"
    fileinfo=$(date -d "$(stat -c "%y" "$realpath")" +"%m_%d")
        #if [ "$(find "${file}" -maxdepth 1 -type d | wc -l)" -eq 1 ]; then
        #    # Mark dir containing no sub dir
        #    fileinfo="${fileinfo}_NO_SUBDIR"
        #fi
        filelist_ND="$filelist_ND $(basename "$file") $fileinfo"
        #numfiles=`expr $numfiles + 1`
    numfiles=$((numfiles + 1))
    done
    if [ $numfiles -eq 0 ]; then
    filelist=".. $(rep_whspc_w_udrsc "$MSG_PARENT_DIRECOTRY")"
    numfiles=1
    else
      filelist="$filelist $filelist_ND"
      # Strip the leading spaces for filelist
      filelist="$(echo $filelist | sed -r -e "s/^[[:space:]]*//g")"
      # Append an abort option in the end, so the menu will show:
      # First is .. (Parent Directory) (if applicable)
      # second is subdirectories
      # third is cz_img directories
      # fourth is Exit/Abort option
      #filelist="$filelist <ABORT> $(rep_whspc_w_udrsc "$MSG_EXIT_BROWSING")"
    fi
    m_filelist="$filelist"
    height="$numfiles"

    filelist=""
    filelist_ND=""
}

get_directories() {
    # Initialize filelist
    # Keep the parent dir.
    filelist=".. $(rep_whspc_w_udrsc "$MSG_PARENT_DIRECOTRY")"
    numfiles=0
    get_directories_list_common
}

get_sub_directories() {
    filelist=""
    numfiles=0
    get_directories_list_common
}

browser_current_directory() {
    local wk_dir="$1"
    local ANS_TMP="$2"
    local TMP=`mktemp /tmp/kybackup.XXXX`
    local basedir currentdir currentdir_fullpath cmd option_rc
    trap "[ -f "$TMP" ] && rm -f $TMP" HUP INT QUIT TERM EXIT

#    if [ -z "$wk_dir" ]; then
        basedir="/"
#    else
#        basedir="$wk_dir"
#    fi

    cd $m_ghost_store_partition
    currentdir="$(basename `pwd`)"
    while true; do
        if [ "$(pwd)" == "$basedir" ]; then
            get_sub_directories
        else
         get_directories
        fi

        $DIA --title "制作系统镜像" \
             --menu  "请选择镜像的存储位置：`pwd`                 " \
             --ok-button "浏览" --cancel-button "完成" --yes-button "yes" \
             0 90 $height $DIA_ESC $m_filelist   2>$TMP
        rc=$?

        option="$(cat $TMP)"

        if [ "$option" == "<ABORT>" ]; then
            exit;
        fi

        if [ "$rc" -eq 1 ]; then
            break;
        fi

        cmd="cd "$option
        eval $cmd
    done

    echo "$(pwd)" > $ANS_TMP
}

get_backup_file_size() {
    local check_path_args="$1"
    local arg_local_exclude_paths="$2"
    local src_path="$3"
    local check_path="$4"

    rsync $check_path_args $arg_local_exclude_paths $src_path $check_path | while IFS= read -r line; do
        if [[ "$line" == "Number of files:"* ]]; then
            num_of_dirs=$(echo $line|grep -oP '\(.*\)'|cut -d' ' -f4 |tr -d ,)

            # There are also . and .. in each directory, so the total number is *3
            # The directory itself in ext4 format takes up 4K space (4096)
            dir_size=$(expr $num_of_dirs \* 3 \* 4096)

            echo $dir_size>/tmp/kybackup.num.of_dirs
        fi

        if [[ "$line" == "Total transferred file size: "* ]]; then
            _n=$(echo $line | cut -d' ' -f5 |tr -d ,)
            dir_size="$(cat /tmp/kybackup.num.of_dirs)"
            all_size=$(expr $dir_size \+ $_n)
            echo $all_size>/tmp/kybackup.num.of_dirs
        fi
    done
}

validate_image_name() {
    local input=$1
    if [[ ! "$input" =~ ^[a-zA-Z0-9_-]+$ ]]; then
        return 1  # 验证失败
    fi
    return 0  # 验证成功
}

# 解析参数
while [[ $# -gt 0 ]]; do
    case $1 in
        -h|--help)
       show_help
           exit 0
           ;;
        -d|--default)
           m_settings_customization=1
       shift 1
       ;;
    *)
       ;;
    esac
done

check_if_root

#FOLDER_PATH="/"
#local_exclude_paths=$EXCLUDE_PATHS
#arg_local_exclude_paths=$(echo "$local_exclude_paths" | awk '{for (i=1; i<=NF; i++) printf "--exclude=%s ",$i}')
#echo $arg_local_exclude_paths

# 计算大小，排除指定子文件夹，并以 MB 为单位显示
#echo "Calculating size of '$FOLDER_PATH' excluding '$EXCLUDE_FOLDER'..."
#SIZE_KB=$(du -sk "$FOLDER_PATH" --exclude="$EXCLUDE_FOLDER" | cut -f1)
#SIZE_KB=$(du -sk "$FOLDER_PATH" "$arg_local_exclude_path" | cut -f1)
#SIZE_MB=$(echo "scale=2; $SIZE_KB / 1024" | bc)  # 将 KB 转换为 MB，保留两位小数

if type whiptail >/dev/null 2>&1; then
    DIA="whiptail"
    DIA_ESC="--"
else
    DIA="dialog"
fi

if ! which mksquashfs >/dev/null; then
    $DIA --title "制作系统镜像" \
         --msgbox "当前环境无mksquashufs，请先安装组件包." 10 60
    exit
fi

###############
# 判断gpg是否存在
#if ! command -v gpg >/dev/null 2>&1; then
#    $DIA --title "制作系统镜像" \
#         --msgbox "当前环境无gpg命令，请先安装组件包." 10 60
#    exit
#fi

#for ((i=1; i<=3; i++))
#    do
#        read -p  "将使用GnuPG对备份镜像签名，请输入用户名字：" gpg_key_name
#        m_gpg_key_id=$(gpg --list-keys --with-colons | awk -F: '/^pub:/ { key_id=$5 } /^uid:/ && $10 ~ /'$gpg_key_name/' { print key_id  }')
#        if [ ! -n "$m_gpg_key_id" ]; then
#            continue;
#   fi
#        IF=$'\n' _id_arr=($m_gpg_key_id)
#   _id_arr_len=${#_id_arr[@]}
#   if [ $_id_arr_len -ne 1 ]; then
#       continue;
#        else
#            break;
#        fi
#    done
#gpg_key_name="mock-backup"
#m_gpg_key_id=$(gpg --list-keys --with-colons | awk -F: '/^pub:/ { key_id=$5 } /^uid:/ && $10 ~ /'$gpg_key_name/' { print key_id  }')
#if [ ! -n "$m_gpg_key_id" ]; then
#    $DIA --title "制作系统镜像" \
#         --msgbox "备份环境已被破坏，建议重新安装此安装包." 10 60
#    exit;
#fi
#IF=$'\n' _id_arr=($m_gpg_key_id)
#id_arr_len=${#_id_arr[@]}
#if [ $_id_arr_len -ne 1 ]; then
#    $DIA --title "制作系统镜像" \
#     --msgbox "gpg 用户名字获取失败.请创建gpg后再尝试备份系统." 10 60
#    exit;
#fi

# 1. get removeable Storage uuidgen
m_backup_device_path=""
#for device_path in /dev/*; do
#    lsblk ${device_path} >/dev/null 2>&1
#    if [ $? -ne 0 ]; then
#        continue;
#    fi
#    if [[ $device_path == /dev/dm-* ]];then
#        continue;
#    fi
#    label=`blkid -o value -s LABEL "${device_path}"`
#
#    echo $label | grep -i "backup" >/dev/null 2>&1
#    if [ $? -eq 0 ]; then
#        backup_device_path=${device_path}
#        break;
#    fi
#done
#
#if [ "${backup_device_path}" != "" ]; then
#  m_backup_uuid=$(lsblk -p -o NAME,UUID,TYPE | grep "${backup_device_path}" | awk '{print $2}')
#  if [ ! -e "${m_backup_path_default}" ]; then
#      mkdir -p "${m_backup_path_default}"
#  fi
#  if [ "$m_backup_uuid" = "" ]; then
#    $DIA --title "制作系统镜像" \
#     --msgbox "获取设备\"${backup_device_path}\"的uuid失败，请确认该设备挂载是否正确." 10 60
#    exit
#  fi
#  mount |grep -w $backup_device_path |grep -w "/backup" >/dev/null 2>&1
#  if [ $? -ne 0 ]; then
#      mount -o rw /dev/disk/by-uuid/${m_backup_uuid} /backup
#      if [ $? -eq 0 ]; then
#          m_ghost_store_partition="${m_backup_path_default}"
#          m_is_mount_backup=1
#      else
#          $DIA --title "制作系统镜像" \
#               --msgbox "\"${m_backup_path_default}\"目录挂载失败，请确认该目录是否已被挂载." 10 60
#          exit
#      fi
#  else
#      m_ghost_store_partition="${m_backup_path_default}"
#      m_is_mount_backup=1
#  fi
#else
#  $DIA --title "制作系统镜像" \
#       --msgbox "请接入标签为backup的存储介质后再执行系统备份." 10 60
#   exit
#fi
#sleep 10
# 2. select a storage location for ghost img.
if [ ! -e "${m_ghost_store_partition}" ]; then
    mkdir -p "${m_ghost_store_partition}"
fi

if [ $m_settings_customization -eq 1 ] ; then
    sel_dir="$(mktemp /tmp/sel_dir.XXXXXX)"
    browser_current_directory $m_ghost_store_partition $sel_dir
    m_ghost_img_dir=`eval "cat $sel_dir"`
else
    m_ghost_img_dir=$M_DEFAULT_GHOST_IMG_DIR
fi

part_type=`df -T $m_ghost_img_dir | awk 'NR==2 {print $2}'`
if [ $part_type == "vfat" ]; then
    $DIA --title "制作系统镜像" \
         --msgbox "当前存储设备格式为vFat,建议使用ext4格式存储. " 10 60
    #exit
fi

rw_test_file="$(mktemp $m_ghost_img_dir/rw_test.XXXXXX)"
rc=$?
if [ "$rc" -ne 0 ]; then
    $DIA --title "制作系统镜像" \
         --msgbox "当前选择的存储目录不可写，请确认可写权限. " 10 60
    exit
fi
[ -f "$rw_test_file" ] && rm -f $rw_test_file

if [[ "$m_ghost_img_dir" != *"/ghost" ]]; then
    m_ghost_img_dir="${m_ghost_img_dir}""/ghost"
fi

# 3. cacl backup's size.
 check_path_args="-aAHXrn --stats --ignore-missing-args"
 #1.2 exclude some path
 local_exclude_paths=$EXCLUDE_PATHS
 arg_local_exclude_paths=$(echo "$local_exclude_paths" | awk '{for (i=1; i<=NF; i++) printf "--exclude=%s ",$i}')

# check_path="${m_ghost_img_dir}""${BACKUP_SNAPSHOTS_PATH}""/check/data/"
 check_path="${m_ghost_img_dir}"
 if [ ! -d "$check_path" ]; then
    m_ghost_img_dir_created=1
    mkdir -p "$check_path"
 fi
 overwrite "检测环境..."

# backup the / directory of the current system and package it into ghost img.
 src_path="/"

 get_backup_file_size "${check_path_args}" "${arg_local_exclude_paths}" "${src_path}" "${check_path}"
 backup_files_size="$(cat /tmp/kybackup.num.of_dirs)"
# Some temporary files will be generated during the backup process and will occupy some space, so we reserve 500M of space
 temporary_files_size=524288000   # 500*1024*1024=52428800
 m_ghost_size=$(expr $backup_files_size \+ $temporary_files_size)
 m_ghost_size=$(awk "BEGIN {print int($m_ghost_size * 0.67)}")
 available_space=$(df -B 1 $check_path | awk 'NR==2 {print $4}')
 if [ $available_space -lt $m_ghost_size ]; then
     if [ $m_ghost_img_dir_created -eq 1 ]; then
       check_and_remove_folder $m_ghost_img_dir
     fi
     _size=$(((m_ghost_size-available_space)/1024/1024))
     $DIA --title "制作系统镜像" \
          --msgbox "当前选择的存储空间不足，还需要${_size}MB空间. " 10 60
     overwrite "当前选择的存储空间不足，还需${_size}MB空间."
     echo
     exit
 fi
overwrite "检测环境... 完成"
echo ""
m_ghost_img_name_default="ghost-"`date +%Y%m%d%H%M%S`"-"`arch`
m_ghost_img_name=$m_ghost_img_name_default

if [ $m_settings_customization -eq 1 ] ; then
    attempts=3
    while [ $attempts -gt 0 ]; do
        input=$($DIA --title "制作系统镜像" --inputbox "请设置系统镜像的名称:（只能包含字母、数字、下划线_和中划线-）:" 10 60 $m_ghost_img_name 3>&1 1>&2 2>&3)
        # 检查用户是否取消
        exitstatus=$?
        if [ $exitstatus -ne 0 ]; then
            echo "用户取消输入镜像名称"
	    if [ $m_ghost_img_dir_created -eq 1 ]; then
	      check_and_remove_folder $m_ghost_img_dir
	    fi
            exit 1
        fi

	_ghost_img_name=$(echo $input | tr -d '[:space:]')
	if [ "$_ghost_img_name" == "" ]; then
	    break;
        fi
        # 验证输入
        if validate_image_name "$_ghost_img_name"; then
            m_ghost_img_name=$_ghost_img_name
            break;
        else
            attempts=$((attempts-1))
            if [ $attempts -gt 0 ]; then
                $DIA --title "制作系统镜像" --msgbox "输入包含非法字符！还剩 $attempts 次尝试机会。\n只允许字母、数字、下划线 _ 和中划线 -" 12 60
            else
                $DIA --title "错误" --msgbox "输入错误次数过多！程序退出。" 10 60
		if [ $m_ghost_img_dir_created -eq 1 ]; then
		  check_and_remove_folder $m_ghost_img_dir
		fi
                exit 1
            fi
        fi
    done
fi

if [[ "$m_ghost_img_name" != *".kyimg" ]]; then
    m_ghost_img_name="${m_ghost_img_name}"".kyimg"
fi
# 4. make squash
#UUID="uid"
#if [ "${m_backup_uuid}" != "" ]; then
#  UUID="${m_backup_uuid}"
#else
#  if command -v uuidgen >/dev/null 2>&1;then
#    UUID=$(uuidgen)
#  fi
#fi
#m_ghost_img_dir_0="${m_ghost_img_dir}""${BACKUP_SNAPSHOTS_PATH}""/"$UUID"/data"
m_ghost_img_dir_0="${m_ghost_img_dir}"
if [ ! -d "$m_ghost_img_dir_0" ]; then
    mkdir -p "$m_ghost_img_dir_0"
fi
m_ghost_img_fullname="${m_ghost_img_dir_0}""/""${m_ghost_img_name}"
mk_squashfs_args=$(echo "$local_exclude_paths" | awk '{for (i=1; i<=NF; i++) printf "-e %s ",$i}')

mk_squashfs_args=$mk_squashfs_args" ""-e "$m_ghost_img_dir

kyimg_list=$(find / -name "*.kyimg" -a  ! -path "/run/*" 2>/dev/null| sed 's/^/-e /' | tr '\n' ' ')
if [[ -n "$kyimg_list" ]]; then
    mk_squashfs_args=$mk_squashfs_args" "$kyimg_list 
fi

dst_img_list=$(find / -name "dst.img" 2>/dev/null| sed 's/^/-e /' | tr '\n' ' ')
if [[ -n "$dst_img_list" ]]; then
    mk_squashfs_args=$mk_squashfs_args" "$dst_img_list 
fi

backup_snapshots_list=$(find / -type d -path "*/backup/snapshots" -a  ! -path "/run/*" -print 2>/dev/null | sed 's/^/-e /' | tr '\n' ' ')
if [[ -n "$backup_snapshots_list" ]]; then
    mk_squashfs_args=$mk_squashfs_args" "$backup_snapshots_list
fi

touch /etc/.ghost

echo "开始创建镜像..."
/usr/sbin/mksquashfs $src_path $m_ghost_img_fullname $mk_squashfs_args -comp lz4 2>/dev/null
result=$?
if [ $result -eq 0 ]; then
    echo "成功创建系统镜像."
else
    if [ $m_ghost_img_dir_created -eq 1 ]; then
      check_and_remove_folder $m_ghost_img_dir
    fi
    echo "创建系统镜像失败."
fi
rm -f /etc/.ghost

echo 1 > /proc/sys/vm/drop_caches

if [ $result -eq 0 ]; then
    echo "镜像目录："$m_ghost_img_fullname
fi

exit

