#!/bin/bash
# License: GPL 
# Author: Steven Shiau <steven _at_ clonezilla org>
# Description: Program to update initramfs for the deployed OS in the hard drive.

#
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

# Load the config in ocs-live.conf. This is specially for Clonezilla live. It will overwrite some settings of /etc/drbl/drbl-ocs.conf, such as $DIA...
[ -e "/etc/ocs/ocs-live.conf" ] && . /etc/ocs/ocs-live.conf

# Provide get_restored_os_extra_root_path if not already loaded from ocs-functions
if ! type get_restored_os_extra_root_path &>/dev/null; then
  get_restored_os_extra_root_path() {
    # Function to find the subvolume root path under the mount point for btrfs or similar layouts
    # Returns extra path like "/@/.snapshots/1/snapshot", "/@", "/root", or ""
    local target_mnt="$1"
    local def_subvol_out def_path isnap cand sub

    # 1. Standard flat root: if /boot and (/etc/os-release or /etc/fstab) exist directly at root
    if [ -d "$target_mnt/boot" ] && [ -f "$target_mnt/etc/os-release" -o -f "$target_mnt/etc/fstab" ]; then
      echo ""
      return 0
    fi

    # 2. btrfs subvolume get-default
    if type btrfs &>/dev/null; then
      def_subvol_out="$(btrfs subvolume get-default "$target_mnt" 2>/dev/null)"
      if [ -n "$def_subvol_out" ] && [ -z "$(echo "$def_subvol_out" | grep "(FS_TREE)")" ]; then
        def_path="$(echo "$def_subvol_out" | sed -r -e 's/.*path[[:space:]]+//')"
        if [ -n "$def_path" ] && [ -d "$target_mnt/$def_path/boot" -o -f "$target_mnt/$def_path/etc/os-release" -o -f "$target_mnt/$def_path/etc/fstab" ]; then
          echo "/$def_path"
          return 0
        fi
      fi
    fi

    # 3. openSUSE Snapper layout: @/.snapshots/*/snapshot
    if [ -d "$target_mnt/@/.snapshots" ]; then
      for isnap in $(ls -vd "$target_mnt"/@/.snapshots/*/snapshot 2>/dev/null | tac); do
        if [ -d "$isnap/boot" -o -f "$isnap/etc/os-release" -o -f "$isnap/etc/fstab" ]; then
          echo "${isnap#$target_mnt}"
          return 0
        fi
      done
    fi

    # 4. Ubuntu / Debian / openSUSE flat / Arch @ layout
    if [ -d "$target_mnt/@/boot" -o -f "$target_mnt/@/etc/os-release" -o -f "$target_mnt/@/etc/fstab" ]; then
      echo "/@"
      return 0
    fi

    # 5. Fedora >= 38 root layout
    if [ -d "$target_mnt/root/boot" -o -f "$target_mnt/root/etc/os-release" -o -f "$target_mnt/root/etc/fstab" ]; then
      echo "/root"
      return 0
    fi

    # 6. Fallback: check any snapshot or subvolume with boot
    cand="$(find "$target_mnt" -maxdepth 4 -type d -name "boot" 2>/dev/null | head -n 1)"
    if [ -n "$cand" ]; then
      sub="${cand%/boot}"
      echo "${sub#$target_mnt}"
      return 0
    fi

    echo ""
    return 0
  }
fi

# Override do_run_update_initrd_from_restored_os to support openSUSE Tumbleweed and Btrfs subvolumes
do_run_update_initrd_from_restored_os() {
  local mnt_pnt mkinitrd_run_sh mkinitrd_cmd ival rc_ extra_root_path
  local grub2_prt="$1"  # e.g. /dev/sda1, the partition where /boot/grub exists
  local root_prt="$2"   # e.g. /dev/sda2, the partition where / exists
  local diff_grub_pt_flag kver_initrd initrd_tool initrd_2_be_updated tmpdir_opt
  for i in grub2_prt root_prt; do
    eval ivar=\$${i}
    if [ -z "$ivar" ]; then
      echo "No \"$i\" assigned in function do_run_update_initrd_from_restored_os." | tee --append ${OCS_LOGFILE}
      echo "Program terminated." | tee --append ${OCS_LOGFILE}
      [ "$save_restore_error_log" = "yes" ] && copy_error_log
      exit 1
    fi
  done
  echo "Re-creating initrd on boot dir in partition $grub2_prt and root partition $root_prt... "
  mnt_pnt="$(mktemp -d /tmp/root_reinst.XXXXXX)"
  mount $root_prt $mnt_pnt
  rc_=$?
  if [ "$rc_" -gt 0 ]; then
    echo "Failed to mount $root_prt! Maybe it does not exist or the file system is not supported in the kernel. Skip running grub2 reinstallation!" 
    return 1
  else
    extra_root_path="$(get_restored_os_extra_root_path "$mnt_pnt")"
    echo "Test if we can chroot the restored OS partition $root_prt (subpath: $extra_root_path)..."
    # Actually we'd better to test /bin/bash instead of ls since we will run a shell script to for grub-install. However, if we test /bin/bash, we will enter the shell if it's successful.
    LC_ALL=C chroot $mnt_pnt/$extra_root_path ls &>/dev/null
    rc_=$?
    # Fedora >= 38 uses an extra dir "root" for root file system, i.e., /root/ under the root partition. 
    if [ "$rc_" -ne 0 ] && [ -z "$extra_root_path" ]; then
      LC_ALL=C chroot $mnt_pnt/root/ ls &>/dev/null
      rc_=$?
      [ "$rc_" -eq 0 ] && extra_root_path="/root"
    fi
    if [ "$rc_" -eq 0 ]; then
      [ "$BOOTUP" = "color" ] && $SETCOLOR_WARNING
      echo "Yes, we are able to chroot the restored OS partition $root_prt."
      [ "$BOOTUP" = "color" ] && $SETCOLOR_NORMAL
    else
      [ "$BOOTUP" = "color" ] && $SETCOLOR_WARNING
      echo "No, we are not able to chroot the restored OS partition $root_prt."
      [ "$BOOTUP" = "color" ] && $SETCOLOR_NORMAL
      echo "Skip updating initrd from restored OS."
      unmount_wait_and_try $mnt_pnt
      return 199
    fi
  fi
  # Extra root path detection for btrfs subvolumes (openSUSE snapper, openSUSE/Ubuntu @, Fedora root)
  if [ -z "$extra_root_path" ]; then
    extra_root_path="$(get_restored_os_extra_root_path "$mnt_pnt")"
    if [ -z "$extra_root_path" ] && [ -d "$mnt_pnt/root/boot" ]; then
      extra_root_path="/root"
    fi
  fi
  local target_root="${mnt_pnt%/}/${extra_root_path#/}"
  target_root="${target_root%/}"

  # /var/tmp and /var/lib/initramfs-tools are required by update-initramfs.
  # This is especially when /var/ is in different partition with kernel.
  mkdir -p $target_root/dev $target_root/var/tmp $target_root/var/lib/initramfs-tools
  mount --bind /dev $target_root/dev
  mount --bind /proc $target_root/proc
  mount --bind /sys $target_root/sys
  if [ -d "$target_root/run" ]; then
    mount --bind /run $target_root/run
  fi

  # If boot partition is not in the same partition as root partition, we have to mount grub partition in the root partition.
  diff_grub_pt_flag="no"
  if [ "$grub2_prt" != "$root_prt" ]; then
     chroot $target_root/ mount $grub2_prt /boot/
     diff_grub_pt_flag="yes"
  fi

  # TODO: for each kernel version, update
  # Now only for the latest kernel version
  # Initramfs file name examples:
  # Debian: initrd.img-4.18.0-0.bpo.3-amd64
  # CentOS: initramfs-3.10.0-957.12.2.el7.x86_64.img
  # openSUSE: initrd-6.11.0-default
  kver_initrd="$(LC_ALL=C find $target_root/boot/ -name "vmlinuz*" -print 2>/dev/null | sort -r -V | head -n 1 | xargs basename | sed -r -e "s/^vmlinuz-//g")"
  # Get the file name for initramfs
  initrd_2_be_updated="$(LC_ALL=C find $target_root/boot/ -name "init*${kver_initrd}*" -print 2>/dev/null | grep -v kdum)"
  # Decide it's update-initramfs or dracut
  mkinitrd_run_sh="$(mktemp $target_root/initrd-sh-XXXXXX)"
  tmpdir_opt=""
  # CentOS 6: /sbin/dracut, CentOS >=7: /usr/sbin/dracut, openSUSE: /usr/bin/dracut
  if [ -e "$target_root/usr/bin/dracut" ]; then
    # CentOS >=7: /usr/sbin/dracut
    # Newer dracut supports --tmpdir
    initrd_tool="dracut"
    if [ -n "$(LC_ALL=C strings $target_root/usr/bin/dracut | \
	    grep -Eo -- --tmpdir)" ]; then
      tmpdir_opt="--tmpdir /tmp"
    fi
  elif [ -e "$target_root/usr/sbin/dracut" ]; then
    initrd_tool="dracut"
    if [ -n "$(LC_ALL=C strings $target_root/usr/sbin/dracut | \
	    grep -Eo -- --tmpdir)" ]; then
      tmpdir_opt="--tmpdir /tmp"
    fi
  elif [ -e "$target_root/sbin/dracut" ]; then
    # CentOS 6: /sbin/dracut
    initrd_tool="dracut"
    if [ -n "$(LC_ALL=C strings $target_root/sbin/dracut | \
	    grep -Eo -- --tmpdir)" ]; then
      tmpdir_opt="--tmpdir /tmp"
    fi
  elif [ -e "$target_root/usr/sbin/update-initramfs" ]; then
    initrd_tool="update-initramfs"
  elif [ -e "$target_root/usr/sbin/mkinitcpio" ]; then
    initrd_tool="mkinitcpio"
  fi
  case "$initrd_tool" in
    "update-initramfs")
      mkinitrd_cmd="update-initramfs -u -k all" ;;
    "dracut")
      local target_initrd="${initrd_2_be_updated#$target_root}"
      if [ -z "$target_initrd" ]; then
        if [ -f "$target_root/boot/initrd-$kver_initrd" ]; then
          target_initrd="/boot/initrd-$kver_initrd"
        elif [ -f "$target_root/boot/initramfs-$kver_initrd.img" ]; then
          target_initrd="/boot/initramfs-$kver_initrd.img"
        else
          target_initrd="/boot/initrd-$kver_initrd"
        fi
      fi
      # Here we assign tmpdir as /tmp, since by default it's /var/tmp, but in chroot the /var might be in another partition.
      mkinitrd_cmd="dracut $tmpdir_opt -f $target_initrd $kver_initrd" ;;
    "mkinitcpio")
      mkinitrd_cmd=" mkinitcpio -P" ;;
    *)
      [ "$BOOTUP" = "color" ] && $SETCOLOR_FAILURE
      echo "The tool to update initramfs was NOT found in function do_run_update_initrd_from_restored_os!"
      [ "$BOOTUP" = "color" ] && $SETCOLOR_NORMAL
      unmount_wait_and_try $target_root/boot
      if [ -d "$target_root/sys/firmware/efi/efivars" ]; then
        unmount_wait_and_try $target_root/sys/firmware/efi/efivars
      fi
      unmount_wait_and_try $target_root/sys
      unmount_wait_and_try $target_root/proc
      unmount_wait_and_try $target_root/dev
      if [ -d "$target_root/run" ]; then
        unmount_wait_and_try $target_root/run
      fi
      unmount_wait_and_try $mnt_pnt
      return 1
      ;;
  esac
  cat <<-SH_END > $mkinitrd_run_sh
#!/bin/sh
export LC_ALL=C
export PATH=/sbin/:/usr/sbin:/bin/:/usr/bin
$mkinitrd_cmd
SH_END
  chmod 755 $mkinitrd_run_sh
  # Run mkinitrd
  echo "Running the command in the restored GNU/Linux chroot env:"
  echo "$mkinitrd_cmd"
  echo "This might take a few minutes..."
  LC_ALL=C chroot $target_root/ /"$(basename $mkinitrd_run_sh)"
  rc_=$?
  [ -e "$mkinitrd_run_sh" ] && rm -f $mkinitrd_run_sh

  #
  if [ "$diff_grub_pt_flag" = "yes" ]; then
    chroot $target_root/ umount /boot/
  fi
  if [ -d "$target_root/sys/firmware/efi/efivars" ]; then
    unmount_wait_and_try $target_root/sys/firmware/efi/efivars
  fi
  unmount_wait_and_try $target_root/sys
  unmount_wait_and_try $target_root/proc
  unmount_wait_and_try $target_root/dev
  if [ -d "$target_root/run" ]; then
    unmount_wait_and_try $target_root/run
  fi
  unmount_wait_and_try $mnt_pnt
  [ -d "$target_root/dev" ] && rmdir $target_root/dev
  [ -d "$mnt_pnt" -a -n "$mnt_pnt" ] && rmdir $mnt_pnt

  return $rc_
} # end of do_run_update_initrd_from_restored_os

# Settings
verbose="no"
grub_partition="" # e.g. /dev/sda1
selected_hd=""    # e.g. sda
selected_parts="" # e.g. sda1, sda2
grub_ver_in_restored_os=""
grub2_boot_dir=""

#
USAGE() {
    echo "$ocs - To update initramfs on a hard drive."
    echo "Usage:"
    echo "To run $ocs:"
    echo "$ocs [OPTION] BOOT_PARTITION"
    echo "BOOT_PARTITION is the partition has /boot dir where Linux kernel, initramfs (and grub files) are located."
    echo "It can be with or without /dev/, e.g., /dev/sda1 or sda1. If \"auto\" is assigned, $ocs will try to find the boot partition."
    echo 
    echo "OPTION:"
    echo "-s, --selected-hd  DISK Assign the disk to search for _only_, i.e. not to search those not in the lists. This is used for not to touch those disk which are not restored. It can be with or without /dev/, e.g, /dev/sda or sda." 
    echo "-p, --selected-parts PARTITION  Assign the partitions to search for _only_, i.e. not to search those not in the lists. This is used for not to touch those partitions which are not restored. It can be with or without /dev/, e.g., /dev/sda1 or sda1. For multiple partitions, use something like \"sda1 sda2\""
    echo 
    echo "Ex:"
    echo "To update initramfs on /dev/sda with possible boot partitions sda1 and sda2, run"
    echo "   $ocs -p \"sda1 sda2\" -s sda auto"
    echo "To update initramfs on /dev/sda1, run"
    echo "   $ocs /dev/sda1"
    echo "To update initramfs and let $ocs to find the boot partition, run"
    echo "   $ocs auto"
    echo
} # end of USAGE


####################
### Main program ###
####################

ocs_file="$0"
ocs=`basename $ocs_file`
#
# -s or -p option is used to pass the selected device, which means only these device will be searched for grub partition. If no assigning, all the available partitions on the system will be searched.
while [ $# -gt 0 ]; do
  case "$1" in
    -s|--selected-hd)
       shift
       if [ -z "$(echo $1 |grep ^-.)" ]; then
         # skip the -xx option, in case 
	 selected_hd="$(strip_leading_dev $1)"
         shift
       fi
       if [ -z "$selected_hd" ]; then
         echo "-s is used, but no selected_hd assigned."
         echo "$msg_program_stop"
         [ "$save_restore_error_log" = "yes" ] && copy_error_log
         exit 1
       fi
       ;;
    -p|--selected-parts)
       shift
       if [ -z "$(echo $1 |grep ^-.)" ]; then
         # skip the -xx option, in case 
	 selected_parts="$(strip_leading_dev $1)"
         shift
       fi
       if [ -z "$selected_parts" ]; then
         echo "-p is used, but no selected_parts assigned."
         echo "$msg_program_stop"
         [ "$save_restore_error_log" = "yes" ] && copy_error_log
         exit 1
       fi
       ;;
    -*)     echo "${0}: ${1}: invalid option"
            USAGE >& 2
            echo "$msg_program_stop"
            [ "$save_restore_error_log" = "yes" ] && copy_error_log
            exit 2 ;;
    *)      break ;;
  esac
done
grub_partition="$(format_dev_name_with_leading_dev $1)"

if [ -z "$grub_partition" ]; then
  USAGE
  exit 1
fi

#
check_if_root

# If grub_partition is not set, set default. Here grub partition is the same as
# boot partition where grub-related files, Linux kernel and initramfs files are put.
# Hence grub partition is the same with boot partition. We keep the similar codes with that in ocs-install-grub.
if [ "$grub_partition" = "auto" ]; then
   grub_partition=""   # reset it as none
   # //NOTE// "selected_hd" could be sda, sdb, or "sda sdb" (more than 1)
   # If it's not assigned, all the partition(s) on all the disk(s) will be searched.
   check_grub_partition "$selected_hd"  
   # We get $found_grub_partition from check_grub_partition, it might be 1 or 2 or more partitions, e.g. /dev/sda1 /dev/hda1
   # 3 cases:
   # (1) selected_parts is assigned
   # From the selected_parts, we can decide which one got from check_grub_partition is really for the target disk
   # (2) selected_hd is assigned
   # found_grub_partition e.g.: /dev/sda1 /dev/hda1
   # selected_hd e.g.: sda sdb
   # (3) no selected_parts or selected_hd is assigned
   if [ -n "$selected_parts" ]; then
     # found_grub_partition e.g.: /dev/sda1 /dev/hda1 /dev/cciss/c0d0p1
     # selected_parts e.g.: sda1 sda2 sda5 cciss/c0d0p1 cciss/c0d0p2
     for i in $found_grub_partition; do
       for j in $selected_parts; do
         if [  "${i#/dev/*}" = "$j" ]; then
          grub_partition="$i"  # e.g. /dev/sda1
          break
         fi
       done
       [ -n "$grub_partition" ] && break
     done
   elif [ -n "$selected_hd" ]; then
     # found_grub_partition e.g.: /dev/sda1 /dev/hda1
     # selected_hd e.g.: sda sdb
     for i in $found_grub_partition; do
       for j in $selected_hd; do
         if [ "$(get_diskname $i)" = "$j" ]; then
          grub_partition="$i"  # e.g. /dev/sda1
          break
         fi
       done
       [ -n "$grub_partition" ] && break
     done
   else
     # No assignment, use the one found. However, maybe there are more than 1. We only take the first one.
     if [ "$(LC_ALL=C echo $found_grub_partition | wc -w)" -gt 1 ]; then
       grub_partition="$(echo $found_grub_partition | awk -F" " '{print $1}')"
       echo "More than one boot and grub config partitions exist: $found_grub_partition. Choose the 1st one: $grub_partition."
     else
       grub_partition="$found_grub_partition"
     fi
   fi

   # ========================================================
   # Fallback for Btrfs subvolumes (openSUSE Snapper, @, Fedora root)
   # ========================================================
   if [ -z "$grub_partition" ]; then
     echo "Standard boot/grub discovery found nothing. Attempting Btrfs subvolume boot discovery..."
     btrfs_cand_parts="$selected_parts"
     if [ -z "$btrfs_cand_parts" -a -n "$selected_hd" ]; then
       get_partition_list false "$selected_hd"
       btrfs_cand_parts="$partition_list"
     fi
     if [ -z "$btrfs_cand_parts" ]; then
       get_partition_list false ""
       btrfs_cand_parts="$partition_list"
     fi
     for ipart in $btrfs_cand_parts; do
       ipart="${ipart#/dev/}"
       p_type="$(LC_ALL=C ocs-get-dev-info "/dev/$ipart" type filesystem parttype 2>/dev/null)"
       [ -n "$(echo "$p_type" | grep -Eiw "(swap|extended|0x5|0xf)")" ] && continue
       hd_test="$(mktemp -d /tmp/hd_btrfs.XXXXXX)"
       mount -o ro "/dev/$ipart" "$hd_test" >/dev/null 2>&1
       mrc=$?
       if [ "$mrc" -eq 0 ]; then
         if [ -n "$(ls -d "$hd_test"/@/.snapshots/*/snapshot/boot 2>/dev/null)" -o \
              -d "$hd_test/@/boot" -o \
              -d "$hd_test/root/boot" -o \
              -n "$(LC_ALL=C find "$hd_test/boot" -name "vmlinuz*" 2>/dev/null)" -o \
              -n "$(LC_ALL=C find "$hd_test" -maxdepth 5 -name "vmlinuz*" 2>/dev/null)" ]; then
           grub_partition="/dev/$ipart"
           unmount_wait_and_try "$hd_test"
           [ -d "$hd_test" -a -n "$hd_test" ] && rmdir "$hd_test" 2>/dev/null
           echo "Found boot partition on Btrfs/subvolume: $grub_partition"
           break
         fi
         unmount_wait_and_try "$hd_test"
         [ -d "$hd_test" -a -n "$hd_test" ] && rmdir "$hd_test" 2>/dev/null
       fi
     done
   fi
   # ========================================================

   if [ -n "$grub_partition" ]; then
     echo "Found boot and grub partition: $grub_partition"
   else
     echo "The boot and grub directory was NOT found. Maybe it does not exist (so other boot manager exists) or the file system is not supported in the kernel. Skip running $grub_install_exec."
     exit 1
   fi
   # Check if grub_partition is on the list of target partitions
   if [ -n "$selected_parts" ]; then
     # Only if we assign the selected partitions, we have to check...
     # E.g.
     # grub_partition=/dev/sda1
     # echo "sda1 sda3" | grep -Ewo "${grub_partition#/dev/*}"
     if [ -n "$(echo $selected_parts | grep -Ewo "${grub_partition#/dev/*}")" ]; then
       echo "Found grub partition \"$grub_partition\", which is on the restored partitions list ($selected_parts)."
     else
       echo "The found grub partition \"$grub_partition\" is _NOT_ on the restored partitions list ($selected_parts)."
       exit 3
     fi
   fi
fi
# grub_partition is like "/dev/hda1" or "/dev/sda1"
# So grub_hd is like "hda" or "sda"
# Remove the leading space, just in case.
grub_partition="$(LC_ALL=C echo "$grub_partition" | sed -r -e "s/^[[:space:]]*//g")"
grub_hd="$(get_diskname $grub_partition)"
# active this device
fdisk -l /dev/$grub_hd &>/dev/null

# process the boot loader
hd_img="$(mktemp -d /tmp/hd_img.XXXXXX)"
mount $grub_partition $hd_img

# Check if /boot is in its own partition, if so, different mount point.
# grub root-directory must have "/boot/". If no "/boot/", it must be the grub partition itself (i.e. a separate partition will be mounted as /boot, and it conains /grub/ in the partition.)
# For btrfs (e.g. openSUSE, Ubuntu, Fedora), /boot might reside inside a subvolume (@/.snapshots/*/snapshot/boot, @/boot, root/boot).
boot_in_grub_pt="no"
if [ -d "$hd_img/boot/" ]; then
   boot_in_grub_pt="yes"
elif [ -n "$(ls -d "$hd_img"/@/.snapshots/*/snapshot/boot 2>/dev/null)" -o \
       -d "$hd_img/@/boot/" -o \
       -d "$hd_img/root/boot/" -o \
       -n "$(LC_ALL=C find "$hd_img" -maxdepth 5 -name "vmlinuz*" 2>/dev/null)" ]; then
   boot_in_grub_pt="yes"
fi

if [ "$boot_in_grub_pt" = "no" ]; then
   # In this case, /boot must exist in its own partition, 
   # not same partition with / (root partition).

   # remount it as /tmp/hd_img.XXXXXX/boot
   unmount_wait_and_try $hd_img
   [ "$BOOTUP" = "color" ] && $SETCOLOR_WARNING
   echo "There is a separate boot partition in target device. Trying to mount root partition for $grub_install_exec to work with that..."
   [ "$BOOTUP" = "color" ] && $SETCOLOR_NORMAL
   # Important! We have to mount the root dir first. Otherwise in FC6, CentOS 5 or later, grub-install will refuse to do that without root device. It will show the error message like: "Could not find device for /mnt"... For Debian or Ubuntu, no such problem. 
   # Actually we can modify /sbin/grub-install (grub-0.97-13) in FC6 or Centos5:
   # -----------
   # ...
   # # Get the root drive.
   # root_device=`find_device ${rootdir}` || exit 1
   # ...
   # -----------
   # as 
   # root_device=`find_device ${rootdir}`
   # i.e. remove "|| exit 1". However, it's better to follow the way distribution does. Therefore we modify ocs-function instead of modifying grub-install.
   if [ -n "$selected_parts" ]; then
     mount_root_partition_for_separate_boot_part -p "$selected_parts" $hd_img  # return variable $found_root_partition
   else
     mount_root_partition_for_separate_boot_part $hd_img  # return variable $found_root_partition
   fi

   # ========================================================
   # Fallback for MDRAID devices / Bare file systems
   # ========================================================
   if [ -z "$found_root_partition" ]; then
     echo "Standard root discovery failed. Attempting aggressive mount discovery..."
     fallback_parts="$selected_parts"
     [ -z "$fallback_parts" ] && fallback_parts="$(awk '/^md/ {print $1}' /proc/mdstat 2>/dev/null)"
     for ipart in $fallback_parts; do
       ipart="${ipart#/dev/}"
       mrc=32
       # Standard mount fallback if explicit types fail
       mount /dev/$ipart $hd_img >/tmp/mnt_err_$ipart 2>&1
       mrc=$?
       if [ "$mrc" -eq 0 ]; then
         # Look for deterministic OS files, including Btrfs subvolumes
         if [ -f "$hd_img/etc/fstab" -o -f "$hd_img/etc/os-release" -o \
              -f "$hd_img/@/etc/fstab" -o -f "$hd_img/@/etc/os-release" -o \
              -n "$(ls "$hd_img"/@/.snapshots/*/snapshot/etc/os-release 2>/dev/null)" -o \
              -n "$(ls "$hd_img"/@/.snapshots/*/snapshot/etc/fstab 2>/dev/null)" -o \
              -f "$hd_img/root/etc/fstab" -o -f "$hd_img/root/etc/os-release" ] && [ ! -e "$hd_img/Clonezilla-Live-Version" ]; then
           found_root_partition="/dev/$ipart"
           unmount_wait_and_try $hd_img
           echo "Aggressive discovery found root on $found_root_partition!"
           break
         fi
         unmount_wait_and_try $hd_img
       else
         echo "Failed to mount /dev/$ipart in aggressive discovery. Error logs:"
         cat /tmp/mnt_err_${ipart}_xfs 2>/dev/null
         cat /tmp/mnt_err_$ipart 2>/dev/null
       fi
     done
   fi
   # ========================================================
   unmount_wait_and_try $grub_partition 2>/dev/null
else
   # In this case, /boot must co-exist with root system in same partition, 
   # Then found_root_partition is the grub partition
   found_root_partition="$grub_partition"
fi
unmount_wait_and_try $hd_img 2>/dev/null

# Let do_run_update_initrd_from_restored_os do the real job
do_run_update_initrd_from_restored_os "$grub_partition" "$found_root_partition"
rc="$?"

if [ "$rc" -eq 199 ]; then
    # Just show warning
    [ "$BOOTUP" = "color" ] && $SETCOLOR_WARNING
    echo "Failed to create initrd in the restored OS due to different system architecture (e.g., i686 vs x86-64)."
    [ "$BOOTUP" = "color" ] && $SETCOLOR_NORMAL
    [ "$debug_mode" = "on" ] && sulogin
elif [ "$rc" -ne 0 ]; then
    # Show error
    [ "$BOOTUP" = "color" ] && $SETCOLOR_FAILURE
    echo "Failed to create initrd in the restored OS."
    [ "$BOOTUP" = "color" ] && $SETCOLOR_NORMAL
    [ "$debug_mode" = "on" ] && sulogin
else
    echo "done!"
fi

exit $rc
