#!/bin/sh

CMDLINE="/boot/env_k1-x.txt"
ASK_TO_REBOOT=0
CMDLINE_CONTENT="commonargs=setenv bootargs earlyprintk quiet splash plymouth.ignore-serial-consoles plymouth.prefer-fbcon clk_ignore_unused swiotlb=65536 workqueue.default_affinity_scope=system"

get_bootro_now() {
  findmnt /boot | grep -q " ro,"
  echo $?
}

get_bootro_conf() {
  grep /boot /etc/fstab | grep -q "defaults.*,ro[ ,]"
  echo $?
}

get_overlay_now() {
  grep -q "overlayroot" /proc/cmdline
  echo $?
}

get_overlay_conf() {
  grep -q "overlayroot" $CMDLINE
  echo $?
}

is_emmc() {
  # 获取cmdline中的root UUID
  CMDLINE_ROOT_UUID=$(grep -oP 'root=UUID=\K[^\s]+' /proc/cmdline)
  # 获取p6分区的UUID
  EMMC_ROOT_UUID=$(blkid -s UUID -o value /dev/mmcblk2p6)
  if [ "$CMDLINE_ROOT_UUID" = "$EMMC_ROOT_UUID" ]; then
    return 0
  else
    return 1
  fi
}

is_sd() {
  # 获取cmdline中的root UUID
  CMDLINE_ROOT_UUID=$(grep -oP 'root=UUID=\K[^\s]+' /proc/cmdline)
  # 获取p6分区的UUID
  SD_ROOT_UUID=$(blkid -s UUID -o value /dev/mmcblk0p6)
  if [ "$CMDLINE_ROOT_UUID" = "$SD_ROOT_UUID" ]; then
    return 0
  else
    return 1
  fi
}

is_ssd(){
  # 获取cmdline中的root UUID
  CMDLINE_ROOT_UUID=$(grep -oP 'root=UUID=\K[^\s]+' /proc/cmdline)
  # 获取p6分区的UUID
  SSD_ROOT_UUID=$(blkid -s UUID -o value /dev/nvme0n1p6)
  if [ "$CMDLINE_ROOT_UUID" = "$SSD_ROOT_UUID" ]; then
    return 0
  else
    return 1
  fi
}

do_finish() {
  if [ $ASK_TO_REBOOT -eq 1 ]; then
    while true; do
      read -p "It will take effect after reboot. Do you want to reboot right now? [Y/n]: " answer
      case "$answer" in
        [Yy]|[Yy][Ee][Ss]|"")
          sync
          reboot
          exit 0
          ;;
        [Nn]|[Nn][Oo])
          exit 0
          ;;
        *)
          echo "Error: Please enter Y, N, yes, no, or press Enter for default (Y)"
          ;;
      esac
    done
  fi
  exit 0
}

enable_overlayfs() {
  if [ "$(awk '/MemTotal/{print $2; exit}' /proc/meminfo)" -le 262144 ]; then
    if [ "$INTERACTIVE" = True ]; then
      whiptail --msgbox "At least 512MB of RAM is recommended for overlay filesystem" 20 60 1
    else
      echo "At least 512MB of RAM is recommended for overlay filesystem"
    fi
    return 1
  fi
  # is_installed overlayroot || apt-get install -y overlayroot
  # mount the boot partition as writable if it isn't already
  if [ $(get_bootro_now) -eq 0 ] ; then
    if ! mount -o remount,rw /boot 2>/dev/null ; then
      echo "Unable to mount boot partition as writable - cannot enable"
      return 1
    fi
    BOOTRO=yes
  else
    BOOTRO=no
  fi

  # modify command line
  if ! grep -q "overlayroot" /proc/cmdline ; then
    if is_emmc; then
      NEW_CMDLINE="${CMDLINE_CONTENT} overlayroot=device:dev=/dev/mmcblk2p7,recurse=0"
    elif is_sd; then
      NEW_CMDLINE="${CMDLINE_CONTENT} overlayroot=device:dev=/dev/mmcblk0p7,recurse=0"
    elif is_ssd; then
      NEW_CMDLINE="${CMDLINE_CONTENT} overlayroot=device:dev=/dev/nvme0n1p7,recurse=0"
    fi
    echo "commonargs=setenv bootargs ${NEW_CMDLINE}" >> $CMDLINE
  fi

  if [ "$BOOTRO" = "yes" ] ; then
    if ! mount -o remount,ro /boot 2>/dev/null ; then
      echo "Unable to remount boot partition as read-only"
    fi
  fi
}

disable_overlayfs() {
  # mount the boot partition as writable if it isn't already
  if [ $(get_bootro_now) -eq 0 ] ; then
    if ! mount -o remount,rw /boot 2>/dev/null ; then
      echo "Unable to mount boot partition as writable - cannot disable"
      return 1
    fi
    BOOTRO=yes
  else
    BOOTRO=no
  fi

  # modify command line
  sed -i '/^commonargs=setenv bootargs/d' $CMDLINE

  if [ "$BOOTRO" = "yes" ] ; then
    if ! mount -o remount,ro /boot 2>/dev/null ; then
      echo "Unable to remount boot partition as read-only"
    fi
  fi
}

enable_bootro() {
  if [ $(get_overlay_now) -eq 0 ] ; then
    echo "Overlay in use; cannot update fstab"
    return 1
  fi
  sed -i /etc/fstab -e "s#\(.*/boot.*\)defaults\(.*\)#\1defaults,ro\2#"
}

disable_bootro() {
  if [ $(get_overlay_now) -eq 0 ] ; then
    echo "Overlay in use; cannot update fstab"
    return 1
  fi
  sed -i /etc/fstab -e "s#\(.*/boot.*\)defaults,ro\(.*\)#\1defaults\2#"
}

update_resize_partition_script() {
  local TARGET_GB="$1"
  local FS_GB=$((TARGET_GB+2))
  local PART_GB=$((TARGET_GB+3))
  local SCRIPT_PATH="/etc/initramfs-tools/scripts/init-premount/resize_partition"
  if is_emmc; then
    sed -i "s/resize2fs -f \/dev\/mmcblk2p6 [0-9]\+G/resize2fs -f \/dev\/mmcblk2p6 ${FS_GB}G/g" "$SCRIPT_PATH"
    sed -i "s/parted \/dev\/mmcblk2 resizepart 6 [0-9]\+G/parted \/dev\/mmcblk2 resizepart 6 ${PART_GB}G/g" "$SCRIPT_PATH"
    sed -i "s/parted --script \/dev\/mmcblk2 mkpart userdata ext4 [0-9]\+GB 100%/parted --script \/dev\/mmcblk2 mkpart userdata ext4 ${PART_GB}GB 100%/g" "$SCRIPT_PATH"
  elif is_sd; then
    sed -i "s/parted \/dev\/mmcblk0 resizepart 6 [0-9]\+G/parted \/dev\/mmcblk0 resizepart 6 ${PART_GB}G/g" "$SCRIPT_PATH"
    sed -i "s/parted --script \/dev\/mmcblk0 mkpart userdata ext4 [0-9]\+GB 100%/parted --script \/dev\/mmcblk0 mkpart userdata ext4 ${PART_GB}GB 100%/g" "$SCRIPT_PATH"
  elif is_ssd; then
    sed -i "s/resize2fs -f \/dev\/nvme0n1p6 [0-9]\+G/resize2fs -f \/dev\/nvme0n1p6 ${FS_GB}G/g" "$SCRIPT_PATH"
    sed -i "s/parted \/dev\/nvme0n1 resizepart 6 [0-9]\+G/parted \/dev\/nvme0n1 resizepart 6 ${PART_GB}G/g" "$SCRIPT_PATH"
    sed -i "s/parted --script \/dev\/nvme0n1 mkpart userdata ext4 [0-9]\+GB 100%/parted --script \/dev\/nvme0n1 mkpart userdata ext4 ${PART_GB}GB 100%/g" "$SCRIPT_PATH"
  fi
}

backup_boot() {
  echo "Creating backup directory and backing up /boot..."
  
  # 创建/backup目录
  mkdir -p /backup
  
  # 使用rsync备份/boot目录到/backup
  rsync -av --delete /boot/ /backup/boot/
  
  if [ $? -eq 0 ]; then
    echo "Backup completed successfully."
  else
    echo "Backup failed!"
    return 1
  fi
}

restore_boot() {
  echo "Restoring /boot from backup..."
  
  # 检查备份是否存在
  if [ ! -d "/backup/boot" ]; then
    echo "Backup directory /backup/boot not found!"
    return 1
  fi

  # 将备份的/boot目录复制到根目录
  rsync -av --delete /backup/boot/ /boot/
  
  if [ $? -eq 0 ]; then
    echo "Boot directory restored successfully."
  else
    echo "Restore failed!"
    return 1
  fi
}

do_overlayfs() {

  echo "WARNING: !!!Enabling the read-only root filesystem will resize the root filesystem partition, Please operate with caution!!!"

  while true; do
    read -p "Do you want to set the root directory to read-only? [Y/n]: " answer
    case "$answer" in
      [Yy]|[Yy][Ee][Ss])
        if grep -q "overlayroot" /proc/cmdline; then
          echo "The root filesystem is already read-only."
          exit 0
        fi
        # 执行备份操作
        if ! backup_boot; then
          echo "Backup failed, exiting."
          return 1
        fi
        sync
        if is_emmc; then
          ROOT_PART=/dev/mmcblk2p6
        elif is_sd; then
          ROOT_PART=/dev/mmcblk0p6
        elif is_ssd; then
          ROOT_PART=/dev/nvme0n1p6
        fi
        ROOT_USED=$(df -T -B1 | awk -v part="$ROOT_PART" '$1==part{print $4}')
        # echo "ROOT_USED=$ROOT_USED"
        # 向上取整到GB (1GB = 1073741824 bytes)
        TOTAL_GB=$(( (ROOT_USED + 1073741823) / 1073741824 ))

        if enable_overlayfs; then
          update_resize_partition_script $TOTAL_GB
          chmod a+x /etc/initramfs-tools/scripts/init-premount/resize_partition
          update-initramfs -u
          echo "Read-only filesystem has been enabled"
          ASK_TO_REBOOT=1
        else
          echo "Failed to enable read-only filesystem"
          return 1
        fi
        break
        ;;
      [Nn]|[Nn][Oo])
        if ! grep -q "overlayroot" /proc/cmdline; then
          echo "The root filesystem is already writable."
          exit 0
        fi
        # 执行恢复操作
        if ! restore_boot; then
          echo "Restore failed, exiting."
          return 1
        fi
        if disable_overlayfs; then
          update-initramfs -u
          echo "Read-only filesystem has been disabled"
          ASK_TO_REBOOT=1
        else
          echo "Failed to disable read-only filesystem"
          return 1
        fi
        break
        ;;
      *)
        echo "Error: Please enter Y, N, yes or no."
        ;;
    esac
  done

  do_finish
}

# Everything else needs to be run as root
if [ $(id -u) -ne 0 ]; then
  printf "Script must be run as root. Try 'sudo read-only-rootfs-config'\n"
  exit 1
fi

# 在调用do_overlayfs之前，判断root分区的已用空间加3GB并向上取整后不超过磁盘总大小
if is_emmc; then
  DISK=/dev/mmcblk2
  ROOT_USED=$(df -T -B1 | awk '$1=="/dev/mmcblk2p6"{print $4}')
elif is_sd; then
  DISK=/dev/mmcblk0
  ROOT_USED=$(df -T -B1 | awk '$1=="/dev/mmcblk0p6"{print $4}')
elif is_ssd; then
  DISK=/dev/nvme0n1
  ROOT_USED=$(df -T -B1 | awk '$1=="/dev/nvme0n1p6"{print $4}')
else
  echo "Unknown disk type"
  exit 1
fi
ADD_SIZE=3221225472
TOTAL_SIZE=$((ROOT_USED + ADD_SIZE))
# 向上取整到GB
TOTAL_GB=$(( (TOTAL_SIZE + 1073741823) / 1073741824 ))
TOTAL_SIZE_ALIGN=$(( TOTAL_GB * 1073741824 ))
DISK_SIZE=$(blockdev --getsize64 $DISK)
if [ $TOTAL_SIZE_ALIGN -gt $DISK_SIZE ]; then
  echo "Not enough total disk size to expand root filesystem. Exiting."
  exit 1
else
  do_overlayfs
fi
