#!/bin/bash
# Author: Steven Shiau <steven _at_ clonezilla org>
# License: GPL
# Improved by Google Gemini for performance and reliability.

export LC_ALL=C

# Load DRBL/OCS settings and functions
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

#
prog="$(basename $0)"
USAGE() {
   echo "Usage: $prog [OPTION] DEVICE VALUE1 VALUE2..."
   echo "DEVICE can be partition or disk with or without /dev/. e.g. /dev/sda1 or sda1"
   echo "VALUE is one of these values: start, end, size, type, filesystem, UUID, PARTUUID, PTUUID, PARTLABEL, SERIALNO, LABEL or flags"
   echo
   echo "OPTION:"
   echo "-l, --long-serial-no  Output the long serial number of a hard drive. By default the short one from udevadm will be given."   
   echo "-u, --unit UNIT       When the VALUE1, VALUE2... is about size, show the size with unit UNIT (s, kB, MB, GB... which parted accepts)."
   echo
   echo "Ex:" 
   echo "To list the /dev/sda1 size with unit sector"
   echo "  $prog -u s /dev/sda1 size"
}

#
# --- Internal Helper Functions ---
#

# A lightweight, dependency-free function to check if a device is a whole disk.
_is_whole_disk() {
  local dev_name
  dev_name=$(basename "$1")
  # Whole disks exist as top-level directories in /sys/block. Partitions do not.
  if [[ -d "/sys/block/$dev_name" ]]; then
    return 0 # It's a whole disk
  else
    return 1 # It's a partition or does not exist
  fi
}

get_fs_from_partclone_fstype() {
  local part_="$1"
  if type partclone.fstype &>/dev/null; then
    LC_ALL=C partclone.fstype "$part_" 2>/dev/null | grep -o -E 'TYPE="[^[:space:]]*"' | cut -d'"' -f2
  fi
}

get_fs_from_blkid() {
  local part_="$1"
  local TYPE
  TYPE=$(LC_ALL=C blkid -o export -- "$part_" 2>/dev/null | grep -E '^TYPE=' | cut -d'=' -f2)

  if [[ -n "$TYPE" ]] || _is_whole_disk "$part_"; then
    echo "$TYPE"
    return
  fi

  # Partition-Specific Checks
  if [[ "$TYPE" == "vfat" ]] && [[ "$(file -Ls "$part_" 2>/dev/null)" == *"FAT (12 bit)"* ]]; then
    TYPE="fat12"
  elif [[ "$TYPE" =~ (isw_raid_member|linux_raid_member) ]]; then
    # Exclude non-imageable RAID member types
    TYPE=""
  elif [[ "$TYPE" == "ufs" ]]; then
    local hdtmp="/dev/$(get_diskname "${part_}")"; local ptno="$(get_part_number "${part_}")"
    local p_type="$(get_partition_table_type_from_disk "$hdtmp")"
    case "$p_type" in
      mbr)
        local sfdisk_opt="--part-type"; [[ "$(sfdisk --help 2>&1)" != *--part-type* ]] && sfdisk_opt="--print-id"
        local pt_id; pt_id="$(sfdisk "$sfdisk_opt" "$hdtmp" "$ptno" 2>/dev/null)"
        if [[ "$pt_id" =~ (a5|a6|a9) ]]; then TYPE="BSD_slice"; fi
        ;;
      gpt)
        local pt_info; pt_info="$(sgdisk -i "$ptno" "$hdtmp" 2>/dev/null)"
        if [[ "$pt_info" == *"FreeBSD boot"* ]]; then TYPE="BSD_boot";
        elif [[ "$pt_info" == *"FreeBSD swap"* ]]; then TYPE="BSD_part"; fi
        ;;
    esac
  elif [[ "$TYPE" =~ ext ]] && [[ "$(blkid -c /dev/null "$part_")" == *'LABEL="ROOT-A"'* ]]; then
    if e2fsck -n "$part_" 2>/dev/null | grep -q "has unsupported feature(s)"; then TYPE="chromium_ext4"; fi
  elif [[ "$TYPE" == "apfs" ]]; then
    if fsapfsinfo "$part_" 2>/dev/null | grep -q "unable to initialize encryption context"; then TYPE="apfs_encrypted"; fi
  fi
  echo "$TYPE"
}

#
# --- Core Data Gathering Function ---
#
_gather_all_info() {
  local device="$1"
  local assignments

  # 1. Use jq to get info from lsblk.
  assignments=$(lsblk -p --json -o NAME,SIZE,TYPE,FSTYPE,UUID,PARTUUID,PTUUID,PARTLABEL,LABEL,MODEL "$device" 2>/dev/null |
    jq -r '
      .blockdevices[] | select(.name == "'"$device"'") |
      "info_size=\((.size // "") | @sh); " +
      "info_type=\((.type // "") | @sh); " +
      "info_fstype=\((.fstype // "") | @sh); " +
      "info_uuid=\((.uuid // "") | @sh); " +
      "info_partuuid=\((.partuuid // "") | @sh); " +
      "info_ptuuid=\((.ptuuid // "") | @sh); " +
      "info_partlabel=\((.partlabel // "") | @sh); " +
      "info_label=\((.label // "") | @sh); " +
      "info_model=\((.model // "") | @sh);"
    ')

  if [[ -n "$assignments" ]]; then
    local info_size info_type info_fstype info_uuid info_partuuid info_ptuuid info_partlabel info_label info_model
    eval "$assignments"
  fi

  # 2. Get serial numbers specifically from udevadm for the parent disk.
  local parent_disk="/dev/$(get_diskname "$device")"
  local udev_output
  udev_output=$(udevadm info -q env -n "$parent_disk" 2>/dev/null)
  local info_serial_short
  info_serial_short=$(echo "$udev_output" | grep '^ID_SERIAL_SHORT=' | cut -d= -f2-)
  local info_serial_long
  info_serial_long=$(echo "$udev_output" | grep '^ID_SERIAL=' | cut -d= -f2-)

  # 3. Refine filesystem type.
  local final_fs
  final_fs=$(get_fs_from_partclone_fstype "$device")
  if [[ -z "$final_fs" ]]; then
    final_fs=$(get_fs_from_blkid "$device")
  fi
  local _dev_filesystem="${final_fs:-$info_fstype}"

  # If after all checks the filesystem is still empty, check the PARTLABEL from lsblk
  # to see if it's a Microsoft Reserved Partition.
  if [[ -z "$_dev_filesystem" ]] && [[ "$info_partlabel" == "Microsoft reserved partition" ]]; then
    _dev_filesystem="MS_Reserved_Partition"
  fi

  # 4. Output all variables for the main script to eval.
  echo "_dev_filesystem='${_dev_filesystem//\'/\'\\\'\'}'"
  echo "_dev_size='${info_size//\'/\'\\\'\'}'"
  echo "_dev_type='${info_type//\'/\'\\\'\'}'"
  echo "_dev_uuid='${info_uuid//\'/\'\\\'\'}'"
  echo "_dev_partuuid='${info_partuuid//\'/\'\\\'\'}'"
  echo "_dev_ptuuid='${info_ptuuid//\'/\'\\\'\'}'"
  echo "_dev_partlabel='${info_partlabel//\'/\'\\\'\'}'"
  echo "_dev_label='${info_label//\'/\'\\\'\'}'"
  echo "_dev_model='${info_model//\'/\'\\\'\'}'"
  echo "_dev_serial_short='${info_serial_short//\'/\'\\\'\'}'"
  echo "_dev_serial_long='${info_serial_long//\'/\'\\\'\'}'"
}

#
# --- Main Program ---
#
serial_type="short"
unit=""
while [ $# -gt 0 ]; do
  case "$1" in
    -u|--unit) shift; unit="$1"; shift ;;
    -l|--long-serial-no) shift; serial_type="long" ;;
    -*)
      echo "${0}: ${1}: invalid option" >&2
      USAGE >&2
      exit 2
      ;;
    *) break ;;
  esac
done

if [ "$#" -lt 2 ]; then
  USAGE >&2
  exit 1
fi

part=$(format_dev_name_with_leading_dev "$1"); shift
request="$*"

info_gathered=false
for wanted in $request; do
  if ! $info_gathered; then
    eval "$(_gather_all_info "$part")"; info_gathered=true
  fi

  case "$wanted" in
    filesystem|fs) echo "$_dev_filesystem" ;;
    size)
      if [[ -n "$unit" ]]; then
        parted -s "$part" unit "$unit" print | grep -m 1 -E "^Disk " | awk '{print $3}'
      else echo "$_dev_size"; fi ;;
    type) if [[ "$_dev_filesystem" == "swap" ]]; then echo "swap"; else echo "$_dev_type"; fi ;;
    UUID|uuid) echo "$_dev_uuid" ;;
    PTUUID|ptuuid) echo "$_dev_ptuuid" ;;
    PARTUUID|partuuid) echo "$_dev_partuuid" ;;
    PARTLABEL|partlabel) echo "$_dev_partlabel" ;;
    LABEL|label) echo "$_dev_label" ;;
    SERIALNO|serialno)
      if [[ "$serial_type" == "long" ]]; then
        echo "$_dev_serial_long"
      else
        echo "$_dev_serial_short"
      fi ;;
    *) echo "Value '$wanted' is not supported." >&2 ;;
  esac
done
