#!/bin/bash
# Author: Steven Shiau <steven _at_ clonezilla org>
# License: GPL
#
# To solve the small partition image restored to larger partition problem.
# Improved by Google's Gemini for safety and reliability.

DRBL_SCRIPT_PATH="${DRBL_SCRIPT_PATH:-/usr/share/drbl}"

. $DRBL_SCRIPT_PATH/sbin/drbl-conf-functions
. /etc/drbl/drbl-ocs.conf
. $DRBL_SCRIPT_PATH/sbin/ocs-functions

# Settings
# By default we won't run ntfsfix before resizing it.
run_ntfsfix="no"

#
export LC_ALL=C

USAGE() {
   echo "Usage: [OPTION] DEVICE"
   echo "DEVICE name can be with or without /dev/, e.g., /dev/sda1 or sda1."
   echo
   echo "OPTION:"
   echo "-b, --batch    Run program in batch mode, i.e. without any prompt or wait to press enter."
   echo "-n, --ntfsfix  Run ntfsfix before resizing a NTFS partition."
   echo
   echo "Example: $0 /dev/sdc1"
}

cleanup() {
  # If the partition variable is set and currently mounted, unmount it.
  if [ -n "$partition" ] && findmnt --source "$partition" >/dev/null 2>&1; then
     # Only unmount if it looks like we mounted it in a temp dir
     if findmnt --source "$partition" -n -o TARGET | grep -q "tmp/.*mnt\."; then
         echo "Interrupted! Unmounting $partition..."
         umount "$partition" 2>/dev/null
     fi
  fi
  # Remove any temporary mount directories created by this script
  rm -rf /tmp/*mnt.?????? 2>/dev/null
}
trap cleanup EXIT INT TERM

# Handles the logic for FS types that must be mounted to be resized (XFS, JFS, BTRFS, NILFS2)
resize_via_mount() {
  local fs_type="$1"
  local resize_cmd_template="$2" # Use __MNT__ as placeholder for mountpoint
  local tool_name="${resize_cmd_template%% *}" # Get the first word (command name)

  if ! type "$tool_name" &>/dev/null; then
      echo "$tool_name: command not found! Skip growing this partition ${partition}"
      return 1
  fi

  local TMP_MNT
  TMP_MNT="$(mktemp -d /tmp/${fs_type}mnt.XXXXXX)"
  
  # Try to mount
  mount -t "$fs_type" "$partition" "$TMP_MNT"
  local rc=$?
  
  if [ "$rc" -eq 0 ]; then
    echo "Running resize operation for $fs_type on $partition..."
    
    # Replace __MNT__ with the actual temp path and execute
    local final_cmd="${resize_cmd_template//__MNT__/$TMP_MNT}"
    echo "Executing: $final_cmd"
    eval "$final_cmd"
    
    sleep 1
    umount "$partition"
    [ -d "$TMP_MNT" ] && rmdir "$TMP_MNT"
  else
    echo "Failed to mount $fs_type partition $partition!"
    echo "Without mounting, we cannot increase the filesystem size to fit partition size!"
    echo "Program terminated!"
    exit 1
  fi
}

# Parse command-line options
while [ $# -gt 0 ]; do
  case "$1" in
    -b|--batch)
            batch_mode="yes"
            shift;;
    -n|--ntfsfix)
            run_ntfsfix="yes"
            shift;;
    -*)     echo "${0}: ${1}: invalid option" >&2
            USAGE >& 2
            exit 2 ;;
    *)      break ;;
  esac
done

[ $# -ne 1 ] && USAGE && exit 1
# No matter the input is like /dev/sda or sda, format it as /dev/sda
partition="$(format_dev_name_with_leading_dev $1)"

###################
### Main script ###
###################
check_if_root
ask_and_load_lang_set
#
if [ ! -b "$partition" ]; then
  [ "$BOOTUP" = "color" ] && $SETCOLOR_FAILURE
  echo "Block device \"$partition\" not found!"
  [ "$BOOTUP" = "color" ] && $SETCOLOR_NORMAL
  echo "$msg_program_stop!"
  my_ocs_exit 1
fi

if findmnt --source "$partition" >/dev/null; then
  [ "$BOOTUP" = "color" ] && $SETCOLOR_FAILURE
  echo "Partition $partition is already mounted! You must unmount it first! Program terminated!"
  [ "$BOOTUP" = "color" ] && $SETCOLOR_NORMAL
  [ -f "$part_info" ] && rm -f $part_info
  exit 1
fi

# Check partition, if no such partition exist, exit, otherwise "parted $partition p" will hang
part_fs="$(ocs-get-dev-info $partition filesystem)"
if [ -z "$part_fs" ]; then
  [ "$BOOTUP" = "color" ] && $SETCOLOR_WARNING
  echo "Unknown or unsupported partition ($partition) found! Skip this partition ${partition}."
  [ "$BOOTUP" = "color" ] && $SETCOLOR_NORMAL
  exit 1
fi

case "$part_fs" in
   reiserfs)
              if ! command -v resize_reiserfs &> /dev/null; then
                 echo "resize_reiserfs: command not found! Skip resizing this partition ${partition}"
                 exit 1
              fi
              [ "$batch_mode" = "yes" ] && resize_reiserfs_opt="-f"
              echo "resize_reiserfs $resize_reiserfs_opt ${partition}"
              resize_reiserfs $resize_reiserfs_opt ${partition}
              ;;
   vfat|fat16|fat32)
              if ! command -v fatresize&> /dev/null; then
                 [ "$BOOTUP" = "color" ] && $SETCOLOR_WARNING
                 echo "Program fatresize not found! Skip this partition ${partition}"
                 [ "$BOOTUP" = "color" ] && $SETCOLOR_NORMAL
                 exit 1
              fi
              fat_fs_size="$(LC_ALL=C lsblk -b -n -d -o size ${partition})"
              # A workaround to avoid fatresize 1.1.0 bug. We do not extend it to the boundary.
              # Just try to resize it smaller.
              # If we run fatresize with "-s max", it will complain:
              # .Error: Can't have overlapping partitions.
              # Ref: https://github.com/ya-mouse/fatresize/issues/18
              #      https://gitlab.tails.boum.org/tails/tails/-/issues/18699
              [ "$batch_mode" = "yes" ] && fatresize_opt="-f"
              fatresize $fatresize_opt -p --verbose -s $((fat_fs_size - 1)) ${partition}
              ;;
   ext2|ext3|ext4)
              if ! command -v e2fsck &> /dev/null || ! command -v resize2fs &> /dev/null; then
                 echo "e2fsck: command not found! Skip resizing this partition ${partition}"
                 exit 1
              fi
              [ "$batch_mode" = "yes" ] && resize2fs_opt="-f"
              echo "e2fsck -f -y ${partition}; resize2fs -p $resize2fs_opt ${partition}"
              e2fsck -f -y ${partition}
              resize2fs -p $resize2fs_opt ${partition}
              ;;
   ntfs)
              if ! command -v ntfsresize &> /dev/null; then
                 echo "ntfsresize: command not found! Skip resizing this partition ${partition}"
                 exit 1
              fi
              [ "$batch_mode" = "yes" ] && ntfsresize_opt="-f -f"
              if [ "$run_ntfsfix" = "yes" ]; then
                ntfsfix ${partition}
              fi
              echo "ntfsresize $ntfsresize_opt ${partition}"
              ntfsresize $ntfsresize_opt ${partition}
              ;;
   xfs)
              # Mount, then run xfs_growfs on the mountpoint
              resize_via_mount "xfs" "xfs_growfs __MNT__"
              ;;
   jfs)
              # Ref: http://jfs.sourceforge.net/project/pub/faq.txt
              # Mount, then remount with resize option
              resize_via_mount "jfs" "mount -o remount,resize __MNT__"
              ;;
   btrfs)
              # Mount, then run btrfs filesystem resize on the mountpoint
              resize_via_mount "btrfs" "btrfs filesystem resize max __MNT__"
              ;;
   nilfs2)
              # Mount, then run nilfs-resize on the DEVICE (but it requires the device to be mounted)
              resize_via_mount "nilfs2" "nilfs-resize -y $partition"
              ;;
   *)
              echo "\"$part_fs\" is an unknown or unsupported filesystem... Skip resizing that."
              exit 1
esac

[ -f "$part_info" ] && rm -f $part_info
