#!/bin/bash
#
# Copyright 2023 The qm Authors
#
# This program is free software; you can redistribute it and/or
# modify it under the terms of the GNU General Public License
# as published by the Free Software Foundation; either version 2
# 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 General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program; If not, see <http://www.gnu.org/licenses/>.

select_disk_to_partition() {
  ###########################################################################
  # Description:                                                            #
  # Select a free partition with size in "G" and update the system           #
  #                                                                         #
  # Arguments: Should be exported by calling script                         #
  #   DEVICE:  Set free /dev/ device                                        #
  #   PART_ID: Set device partition                                         #
  ###########################################################################

  # Fetch the disk table, listing devices and partitions
  local disk_table
  disk_table=$(lsblk --noheadings --raw)

  # Extract all disk names from the disk table
  local disks_arr
  disks_arr=$(echo "$disk_table" | awk '$1~// && $6=="disk" {print $1}')

  # Log found disks for reference
  info_message "select_disk_to_partition, found ${disks_arr}"
  info_message "=============================="

  # Iterate over the array of disks to find a suitable partition
  for disk in $disks_arr; do
    # Check if the disk is "vda" or a valid disk (not zram0), and whether it exists in disk_table
    if [[ ${disk} == "vda" || \
      $(echo "${disk_table}" | grep -c "${disk}") -eq 1 && ${disk} != "zram0" ]]; then

      # Check if any partitions of the disk are currently mounted
      #mountpoints=$(grep "/dev/${disk}" /proc/mounts | awk '{print $2}')

      # Debugging statement to print mountpoints status
      #echo "Mountpoints for /dev/${disk}: ${mountpoints}"

      # If any partitions are mounted, unmount them to avoid disk busy errors
      #echo "Unmounting partitions on /dev/${disk}..."
      #umount "/dev/${disk}"* || {
      #  echo "Failed to unmount /dev/${disk}, skipping."
      #    continue  # Skip to the next disk if unmount fails
      #}

      # Sync the system to flush any cached I/O operations to the disk
      # This ensures that all file system operations are complete before modifying partitions
      echo "Calling sync before partitioning /dev/${disk}..."
      sync

      # Proceed with partitioning using fdisk
      # fdisk is used here in non-interactive mode to create a new primary partition
      echo "Creating new partition on /dev/${disk}..."
      new_part="$( (echo n; echo p; echo ; echo ; echo ; echo w) | fdisk "/dev/${disk}")" || {
        echo "Failed to create partition on /dev/${disk}, skipping."
        continue
      }

      # Extract the partition number from the fdisk output
      part_id="$(echo "${new_part}" | grep -Po "new partition \K([0-9])")"
      DISK="${disk}"

      # Check if the disk is an NVMe device (which has partition naming like /dev/nvme0n1p1)
      # Append the partition ID correctly for NVMe disks
      if [[ $(echo "${DISK}" | grep -c nvme) -eq 1 ]]; then
        PART_ID="p${part_id}"
        info_message "select_disk_to_partition, found /dev/${DISK}${PART_ID}"
        info_message "=============================="
        break  # Stop further processing as we have found and partitioned the desired disk
      else
        PART_ID="${part_id}"
        info_message "select_disk_to_partition, found /dev/${DISK}${PART_ID}"
        info_message "=============================="
        break  # Stop further processing after partitioning
      fi
    fi

    # Additional logic for SOC-specific disk selection (if necessary)
    if [ -e "${SOC_DISTRO_FILE}" ]; then
      if grep -qi "${QC_SOC}" "${SOC_DISTRO_FILE}"; then
        if [[ "${disk}" == "${QC_SOC_DISK}" ]]; then
          DISK=${disk}

          # Get the partition ID for the SOC-specific disk (partition in GB)
          part=$(echo "${disk_table}" | grep "G" | grep "${QC_SOC_DISK}[0-9]." | grep -v "/" | cut -d" " -f1)
          PART_ID="${part:3}"  # Extract partition ID from the result
          break  # Stop further processing after selecting the SOC-specific disk
        fi
      fi
    fi
  done

  # Once a disk is found and partitioned, we need to notify the kernel about the changes
  if [[ -n "${DISK}" ]]; then
    echo "Rescanning partition table on /dev/${DISK}..."
    if ! rpm -qf qm &>/dev/null; then
       dnf install -y parted
    fi
    # Run partprobe to inform the kernel about partition changes
    # This avoids the need for a reboot, ensuring that partition changes are recognized immediately
    partprobe || {
      echo "Warning: Failed to inform the kernel about partition changes on /dev/${DISK}"
    }

    echo "Partitioning and kernel update completed for /dev/${DISK}${PART_ID}"
  else
    echo "No suitable disk found for partitioning."
  fi
}

remove_unmounted_subdirs() {
  ###########################################################################
  # Description:                                                            #
  #  Removes only the subdirectories of the provided directory that are not #
  #  mounted.                                                               #
  #                                                                         #
  # Arguments:                                                              #
  #  $1: The directory whose subdirectories should be checked and removed   #
  #                                                                         #
  ###########################################################################

  local parent_dir="$1"

  # Ensure that the directory exists
  if [[ ! -d "${parent_dir}" ]]; then
    echo "Directory ${parent_dir} does not exist."
    return 1
  fi

  # Loop over all subdirectories in the given directory
  for dir in "${parent_dir}"/*; do
    # Check if the subdirectory is mounted
    if findmnt -r "$dir" > /dev/null; then
      echo "Skipping $dir because it is mounted."
    else
      echo "Removing $dir because it is not mounted."
      rm -rf "$dir" || echo "Failed to remove $dir"
    fi
  done
}

create_qm_var_part() {
  ###########################################################################
  # Description:                                                            #
  #  Create relevant filesystem                                             #
  #    support fstype, non ostree, xfs (c9s vm), ext4(soc)                  #
  # ARGUMENTS: should be be exported by calling script                      #
  #   DEVICE,  set free /dev/ device                                        #
  #   PART_ID, set device partition                                         #
  ###########################################################################

  select_disk_to_partition

  local slash_var
  slash_var="/var/qm"

  # Check is SoC aboot exist, else regular image
  if [ -e "${SOC_DISTRO_FILE}" ]; then
        if grep -qi "${QC_SOC}" "${SOC_DISTRO_FILE}"; then
          mkfs.ext4 "/dev/${DISK}${PART_ID}"
          if_error_exit "Error: mkfs.ext4 /dev/${DISK}${PART_ID} failed on ${QC_SOC}"
          info_message "Create_qm_disks, prepare SoC ext4 fs"
          info_message "=============================="
        fi
  else
  # VM use case
       mkfs.xfs "/dev/${DISK}${PART_ID}"
       if_error_exit "Error: mkfs.xfs /dev/${DISK}${PART_ID} failed on VM"
       info_message "Create_qm_disks, prepare regular xfs fs"
       info_message "=============================="

  fi

  mkdir -p /new_var
  mount "/dev/${DISK}${PART_ID}" /new_var
  # Test if ${slash_var} exist
  if test -d  ${slash_var}; then
    rsync -aqxXP ${slash_var}/* /new_var
    #if_error_exit "Error: rsync failed"
    systemctl stop var-lib-nfs-rpc_pipefs.mount
    systemctl stop qm
    systemctl stop podman
    umount /new_var
    remove_unmounted_subdirs "${slash_var}"
  else
    mkdir -p "${slash_var}"
  fi
  info_message "Create_qm_disks, prepare and mount ${slash_var}"
  info_message "=============================="
  mount "/dev/${DISK}${PART_ID}" "${slash_var}"
  if [ -n "$(ls -A ${slash_var})" ]; then
    #if_error_exit "Error: mount /dev/${DISK}${PART_ID} on ${slash_var} failed"
    systemctl start var-lib-nfs-rpc_pipefs.mount
    systemctl status var-lib-nfs-rpc_pipefs.mount -l
    systemctl start qm
    systemctl status qm
    journalctl -xeu qm.service --no-pager
  fi
}



