#!/bin/bash
# License: GPL
# Author: Steven Shiau <steven _at_ clonezilla org>
# This file contains code generated by Gemini, created by Google.
# Description: This program runs once to create sorted symlinks for all block devices.
# The format for the created block device: ocs-?d*
# E.g.,
# NVMe device: ocs-nd01, ocs-nd02...
# SATA/SAS/SCSI/USB device: ocs-sd01, ocs-sd02... and USB one will be in low priority
# EMMC device: ocs-md01, ocs-md02...

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

check_if_root

# Make sure only one instance is running.
lock_dir="/var/run/ocs-blkdev-sorter.lock"
# The following check is atomic, it's race-condition-free.
if ! mkdir "$lock_dir" 2>/dev/null; then
  # Another instance is running, exit.
  exit 0
fi
# Make sure the lock dir is removed when this script exits.
trap 'rm -rf "$lock_dir"' EXIT

# Functions
outpu_mapping_devics() {
  echo "The OCS block device mapping:"
  echo "--------------------------"
  echo "Clonezilla alias | Linux kernel name | ID_PATH"

  local ocs_devices
  shopt -s nullglob
  ocs_devices=($ocs_dev_dir/ocs-*)
  shopt -u nullglob

  for dev in "${ocs_devices[@]}"; do
    # Read the target of the symbolic link
    target=$(readlink -e "$dev")

    # Check if the symlink is broken. If so, remove it and skip.
    if [ -z "$target" ] || [ ! -e "$target" ]; then
      rm -f "$dev"
      continue
    fi

    # Get the base names for clean output
    ocs_name=$(basename "$dev")
    kernel_name=$(basename "$target")

    local id_path=""

    # Check if the symlink is a partition
    if [[ "$ocs_name" =~ p[0-9]+$ ]]; then
      # If a partition, get the parent kernel name
      local parent_kernel_name
      if [[ "$kernel_name" =~ ^nvme ]]; then
        # For NVMe, the parent is the part before 'p'
        parent_kernel_name=$(echo "$kernel_name" | sed 's/p[0-9]*$//')
      else
        # For other types, it's the part before the partition number
        parent_kernel_name=$(echo "$kernel_name" | sed 's/[0-9]*$//')
      fi

      # Get the parent's ID_PATH
      local base_id_path
      if [ -n "$parent_kernel_name" ]; then
        base_id_path=$(udevadm info --name="$parent_kernel_name" | grep -Ew "ID_PATH" | cut -d'=' -f2)
      fi

      # Get the partition number from the symlink name
      local part_num
      if [[ "$ocs_name" =~ ^ocs-nd ]]; then
        part_num=$(echo "$ocs_name" | sed 's/^.*p//')
      else
        part_num=$(echo "$ocs_name" | sed 's/^.*[a-zA-Z]//')
      fi

      # Construct the final ID_PATH
      if [ -n "$base_id_path" ] && [ -n "$part_num" ]; then
        id_path="${base_id_path}-part${part_num}"
      fi
    else
      # If a main disk, get the ID_PATH directly
      id_path=$(udevadm info --name="$kernel_name" | grep -Ew "ID_PATH" | cut -d'=' -f2)
    fi

    printf "%-18s %-19s %s\n" "$ocs_name" "$kernel_name" "$id_path"
  done
} # end of outpu_mapping_devics

mkdir -p "$ocs_dev_dir"
rm -f "$ocs_dev_dir/ocs-*"
> $ocs_dev_dir/nvme.tmp
> $ocs_dev_dir/scsi_non_usb.tmp
> $ocs_dev_dir/scsi_usb.tmp
> $ocs_dev_dir/mmc.tmp
> $ocs_dev_dir/dev-mapping.txt

# Collect devices by type and sort key
for dev in /sys/class/block/nvme*n*; do
  [ -e "$dev/partition" ] && continue
  devname=$(basename "$dev")
  id_path=$(udevadm info --name="$devname" | grep -Ew "ID_PATH" | cut -d'=' -f2)
  [ -z "$id_path" ] && continue
  printf "%s %s\n" "$id_path" "$devname" >> $ocs_dev_dir/nvme.tmp
done

for dev in /sys/class/block/sd*; do
  if [ ! -e "$dev" ] || [ -e "$dev/partition" ]; then
    continue
  fi
  devname=$(basename "$dev")
  id_path=$(udevadm info --name="$devname" | grep -Ew "ID_PATH" | cut -d'=' -f2)
  [ -z "$id_path" ] && continue
  if udevadm info --name="$devname" | grep -q "ID_BUS=usb"; then
    printf "%s %s\n" "$id_path" "$devname" >> $ocs_dev_dir/scsi_usb.tmp
  else
    printf "%s %s\n" "$id_path" "$devname" >> $ocs_dev_dir/scsi_non_usb.tmp
  fi
done

for dev in /sys/class/block/mmcblk*; do
  if [ ! -e "$dev" ] || [ -e "$dev/partition" ]; then
      continue
  fi
  devname=$(basename "$dev")
  id_path=$(udevadm info --name="$devname" | grep -Ew "ID_PATH" | cut -d'=' -f2)
  [ -z "$id_path" ] && continue
  printf "%s %s\n" "$id_path" "$devname" >> $ocs_dev_dir/mmc.tmp
done

# Determine the number of digits to use for each type
nvme_count=$(wc -l < $ocs_dev_dir/nvme.tmp)
scsi_count=$(($(wc -l < $ocs_dev_dir/scsi_non_usb.tmp) + $(wc -l < $ocs_dev_dir/scsi_usb.tmp)))
mmc_count=$(wc -l < $ocs_dev_dir/mmc.tmp)

# Set the format string based on the count
nvme_fmt=$( [ "$nvme_count" -ge 100 ] && echo '%03d' || echo '%02d' )
scsi_fmt=$( [ "$scsi_count" -ge 100 ] && echo '%03d' || echo '%02d' )
mmc_fmt=$( [ "$mmc_count" -ge 100 ] && echo '%03d' || echo '%02d' )

# Create numbered symlinks for NVMe disks and their partitions
i=0
sort -V $ocs_dev_dir/nvme.tmp | while read -r line; do
  id_path=$(echo "$line" | cut -d' ' -f1)
  devname=$(echo "$line" | cut -d' ' -f2)
  i=$((i+1))
  ocs_disk_name="$ocs_dev_dir/ocs-nd$(printf "$nvme_fmt" "$i")"
  ln -sf "/dev/disk/by-path/$id_path" "$ocs_disk_name"
  # Find and create symlinks for partitions using lsblk
  lsblk -no KNAME,TYPE "/dev/$devname" | tail -n +2 | while read -r part_kname part_type; do
    # Process only if it's a partition, skip LVM PVs, LVs, etc.
    if [ "$part_type" == "part" ]; then
      if [[ "$part_kname" =~ ^nvme ]]; then
        # For NVMe, the partition number is after 'p'
        part_num=$(echo "$part_kname" | sed 's/^.*p//')
      else
        # For other types, it's at the end
        part_num=$(echo "$part_kname" | sed 's/^.*[a-zA-Z]//')
      fi
      ln -sf "/dev/$part_kname" "${ocs_disk_name}p${part_num}"
    fi
  done
done

# Create numbered symlinks for SATA/SAS/SCSI/USB disks and their partitions
i=0
{ sort -V $ocs_dev_dir/scsi_non_usb.tmp; sort -V $ocs_dev_dir/scsi_usb.tmp; } | while read -r line; do
  id_path=$(echo "$line" | cut -d' ' -f1)
  devname=$(echo "$line" | cut -d' ' -f2)
  i=$((i+1))
  ocs_disk_name="$ocs_dev_dir/ocs-sd$(printf "$scsi_fmt" "$i")"
  ln -sf "/dev/disk/by-path/$id_path" "$ocs_disk_name"
  # Find and create symlinks for partitions using lsblk
  lsblk -no KNAME,TYPE "/dev/$devname" | tail -n +2 | while read -r part_kname part_type; do
    # Process only if it's a partition, skip LVM PVs, LVs, etc.
    if [ "$part_type" == "part" ]; then
      part_num=$(echo "$part_kname" | sed 's/^.*[a-zA-Z]//')
      ln -sf "/dev/$part_kname" "${ocs_disk_name}p${part_num}"
    fi
  done
done

# Create numbered symlinks for eMMC disks and their partitions
i=0
sort -V $ocs_dev_dir/mmc.tmp | while read -r line; do
  id_path=$(echo "$line" | cut -d' ' -f1)
  devname=$(echo "$line" | cut -d' ' -f2)
  i=$((i+1))
  ocs_disk_name="$ocs_dev_dir/ocs-mc$(printf "$mmc_fmt" "$i")"
  ln -sf "/dev/disk/by-path/$id_path" "$ocs_disk_name"
  # Find and create symlinks for partitions using lsblk
  lsblk -no KNAME,TYPE "/dev/$devname" | tail -n +2 | while read -r part_kname part_type; do
    # Process only if it's a partition, skip LVM PVs, LVs, etc.
    if [ "$part_type" == "part" ]; then
      part_num=$(echo "$part_kname" | sed 's/^.*[a-z]//')
      ln -sf "/dev/$part_kname" "${ocs_disk_name}p${part_num}"
    fi
  done
done

rm -f $ocs_dev_dir/nvme.tmp $ocs_dev_dir/scsi_non_usb.tmp $ocs_dev_dir/scsi_usb.tmp $ocs_dev_dir/mmc.tmp

echo "--------------------------"
outpu_mapping_devics | tee "$ocs_dev_dir/dev-mapping.txt"
