#!/bin/bash
# License: GPL
# Author: Steven Shiau <steven _at_ clonezilla org>
# Description: Program to save the MDRAID layout.
# This file contains code generated by Google's Gemini. 

#
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

#
USAGE() {
    echo "$ocs - To save the MDRAID layout to the image dir"
    echo "Usage:"
    echo "To run $ocs:"
    echo "$ocs [OPTION] IMAGE_DIR_PATH"
    echo
    echo "Ex:"
    echo "To save the MDRAID device to the image \"my-image\", which is located in $ocsroot, run"
    echo "   $ocs $ocsroot/my-image"
    echo
} # end of USAGE

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

ocs_file="$0"
ocs=`basename $ocs_file`
#
while [ $# -gt 0 ]; do
 case "$1" in
   -*)     echo "${0}: ${1}: invalid option" >&2
           USAGE >& 2
           exit 2 ;;
   *)      break ;;
 esac
done

#
check_if_root
ask_and_load_lang_set

if [ "$#" -ne 1 ]; then
    USAGE
    exit 1
fi

echo $msg_delimiter_star_line
# Create the backup directory with a timestamp
IMG_DIR="$1"
mkdir -p "$IMG_DIR"

echo "Starting Advanced MDRAID layout backup into: $IMG_DIR"

# Find active MD devices from /proc/mdstat
MD_DEVICES=$(LC_ALL=C awk '/^md/ {print $1}' /proc/mdstat)

if [ -z "$MD_DEVICES" ]; then
    echo "No MDRAID devices found in /proc/mdstat."
    exit 1
fi

for MD in $MD_DEVICES; do
    MD_PATH="/dev/$MD"
    echo "------------------------------------------"
    echo "Processing array: $MD_PATH"

    # 3. Backup raw RAID details for reference
    LC_ALL=C mdadm --detail "$MD_PATH" > "$IMG_DIR/${MD}-detail.info" || true
    
    # 4. Extract Advanced Metadata (Level, Chunk, Layout, Devices)
    MD_UUID="$(LC_ALL=C mdadm --detail "$MD_PATH" | awk '/UUID/ {print $3}')"
    MD_NAME="$(LC_ALL=C mdadm --detail "$MD_PATH" | awk '/Name/ {print $3}')"
    MD_META="$(LC_ALL=C mdadm --detail "$MD_PATH" | awk '/Version/ {print $3}')"
    MD_LEVEL="$(LC_ALL=C mdadm --detail "$MD_PATH" | awk -F: '/Raid Level/ {print $2}' | tr -d ' ')"
    MD_DEVICES_COUNT="$(LC_ALL=C mdadm --detail "$MD_PATH" | awk -F: '/Raid Devices/ {print $2}' | tr -d ' ')"
    
    # Chunk size usually only applies to RAID 0, 4, 5, 6, 10. Extract numbers only (defaults to KB)
    MD_CHUNK="$(LC_ALL=C mdadm --detail "$MD_PATH" | awk -F: '/Chunk Size/ {print $2}' | tr -d 'a-zA-Z ')"
    MD_LAYOUT="$(LC_ALL=C mdadm --detail "$MD_PATH" | awk -F: '/[ \t]Layout[ \t]*:/ {print $2}' | tr -d ' ')"

    # 5. Extract strict device order (Handles 5+ disks, missing drives, AND Hot Spares)
    # The regex matches the first 3 columns (Number, Major, Minor). 
    MD_MEMBERS_ORDERED="$(LC_ALL=C mdadm --detail "$MD_PATH" | awk '/^[ \t]*[0-9]+[ \t]+[0-9]+[ \t]+[0-9]+/ {
        slot = $4;
        dev = $NF;
        if (dev ~ /removed|faulty|missing/) {
            dev = "missing";
        }
        if (slot == "-") {
            # It is a spare drive, assign it a high artificial slot number so it sorts to the very end
            slot = 9999 + $1; 
        }
        print slot, dev
    }' | sort -n | awk '{print $2}')"
    
    # ====================================================================
    # Extract Data Offset
    # ====================================================================
    FIRST_MEMBER="$(echo "$MD_MEMBERS_ORDERED" | grep -v "missing" | head -n 1)"
    MD_DATA_OFFSET=""
    if [ -n "$FIRST_MEMBER" ] && [ -b "$FIRST_MEMBER" ]; then
        MD_DATA_OFFSET="$(LC_ALL=C mdadm --examine "$FIRST_MEMBER" 2>/dev/null | grep -i 'Data Offset' | grep -oE '[0-9]+' | head -n 1)"
    fi

    # Write configuration to env file for the restore script
    ENV_FILE="$IMG_DIR/${MD}-config.env"
    echo "MD_UUID=$MD_UUID" > "$ENV_FILE"
    echo "MD_NAME=$MD_NAME" >> "$ENV_FILE"
    echo "MD_META=$MD_META" >> "$ENV_FILE"
    echo "MD_LEVEL=$MD_LEVEL" >> "$ENV_FILE"
    echo "MD_DEVICES_COUNT=$MD_DEVICES_COUNT" >> "$ENV_FILE"
    echo "MD_CHUNK=$MD_CHUNK" >> "$ENV_FILE"
    echo "MD_LAYOUT=$MD_LAYOUT" >> "$ENV_FILE"
    echo "MD_DATA_OFFSET=$MD_DATA_OFFSET" >> "$ENV_FILE"
    echo "MD_MEMBERS_ORDERED=\"$(echo $MD_MEMBERS_ORDERED)\"" >> "$ENV_FILE"
    
    echo "  -> Saved Level: $MD_LEVEL | Devices: $MD_DEVICES_COUNT"
    [ -n "$MD_CHUNK" ] && echo "  -> Saved Chunk Size: ${MD_CHUNK}K"
    [ -n "$MD_LAYOUT" ] && echo "  -> Saved Layout: $MD_LAYOUT"
    [ -n "$MD_DATA_OFFSET" ] && echo "  -> Saved Data Offset: ${MD_DATA_OFFSET} sectors"
    echo "  -> Saved Order: $(echo $MD_MEMBERS_ORDERED | tr '\n' ' ')"

    # 6. Dump the internal partition table of the RAID device using Clonezilla naming scheme
    echo "Dumping partition table for $MD_PATH..."
    if ! is_block_device_with_fs $MD_PATH; then
      sfdisk -d "$MD_PATH" > "$IMG_DIR/${MD}-pt.sf" 2>/dev/null || true
    fi

    # 7. Backup physical disk layouts
    for MEMBER in $MD_MEMBERS_ORDERED; do
        if [[ "$MEMBER" == "missing" ]] || [ ! -b "$MEMBER" ]; then
            continue
        fi
        
        MEMBER_NAME="$(LC_ALL=C basename "$MEMBER")"
        LC_ALL=C mdadm --examine "$MEMBER" > "$IMG_DIR/${MEMBER_NAME}-examine.info" || true

        # Trace physical hard drive using sysfs
        if [ -L "/sys/class/block/$MEMBER_NAME" ]; then
            PARENT_SYSFS="$(LC_ALL=C readlink -f "/sys/class/block/$MEMBER_NAME/..")"
            PARENT_KNAME="$(LC_ALL=C basename "$PARENT_SYSFS")"
            
            if [ -f "/sys/class/block/$PARENT_KNAME/dev" ]; then
                PARENT_PATH="/dev/$PARENT_KNAME"
            else
                PARENT_PATH="$MEMBER"
                PARENT_KNAME="$MEMBER_NAME"
            fi
        else
            PARENT_PATH="$MEMBER"
            PARENT_KNAME="$MEMBER_NAME"
        fi
        
        # Backup the GPT partition table
        if [ -b "$PARENT_PATH" ]; then
            if ! is_block_device_with_fs $PARENT_PATH; then
              if [ ! -f "$IMG_DIR/${PARENT_KNAME}-pt.sf" ]; then
                  echo "  -> Dumping physical partition table for $PARENT_PATH..."
                  sfdisk -d "$PARENT_PATH" > "$IMG_DIR/${PARENT_KNAME}-pt.sf" 2>/dev/null || true
              fi
            fi
        fi
    done
done

echo "MDRAID info saving completes in folder: $IMG_DIR"
echo $msg_delimiter_star_line
