#!/bin/bash
# License: GPL
# Author: Steven Shiau <steven _at_ clonezilla org>
# Description: Program to restore 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 restore the MDRAID layout from image dir"
    echo "Usage:"
    echo "To run $ocs:"
    echo "$ocs [OPTION] IMAGE_DIR_PATH"
    echo
    echo "Options:"
    echo "-b, --batch-mode   Run image checking in batch mode"
    echo "Ex:"
    echo "To restore the MDRAID device from 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
   -b|--batch) ocs_batch_mode="on"; shift;;
   -*)     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

IMG_DIR="${1%/}"

if [ ! -d "$IMG_DIR" ]; then
    echo "Error: Directory '$IMG_DIR' does not exist."
    exit 1
fi

echo "------------------------------------------"
echo "1: Recreating the MDRAID Arrays"
echo "------------------------------------------"

# Check if there are any config files
if ! ls "$IMG_DIR"/md*-config.env 1> /dev/null 2>&1; then
    echo "Error: Could not find any MDRAID configuration files (e.g., md127-config.env)."
    exit 1
fi

# ====================================================================
# The "udev Auto-Assembly Race Condition" Preventer
# We MUST stop all arrays and zero all superblocks across ALL configs 
# simultaneously before recreating, otherwise udev will lock the disks.
# ====================================================================
echo "  -> Stopping any background arrays auto-assembled by udev..."
ocs-mdraid-stop
udevadm settle
sleep 2

echo "  -> Pre-wiping all MDRAID superblocks to clear stale metadata..."
for ENV_FILE in "$IMG_DIR"/md*-config.env; do
    [ -e "$ENV_FILE" ] || continue
    # Run in subshell to isolate variables
    (
        source "$ENV_FILE"
        for MEMBER in $MD_MEMBERS_ORDERED; do
            if [ "$MEMBER" != "missing" ] && [ -b "$MEMBER" ]; then
                mdadm --zero-superblock "$MEMBER" 2>/dev/null || true
                wipefs -a "$MEMBER" 2>/dev/null || true
            fi
        done
    )
done
udevadm settle
# ====================================================================

# Loop through all saved MDRAID configurations
for ENV_FILE in "$IMG_DIR"/md*-config.env; do
    [ -e "$ENV_FILE" ] || continue
    
    MD_DEV_NAME=$(basename "$ENV_FILE" -config.env)
    MD_DEV_PATH="/dev/$MD_DEV_NAME"
    MD_DUMP="$IMG_DIR/${MD_DEV_NAME}-pt.sf"

    if [ "$ocs_batch_mode" != "on" ]; then
      echo "WARNING: This will overwrite MDRAID layout on target device $MD_DEV_PATH."
      confirm_continue_or_not_default_quit
    fi

    echo "  -> Target MD array: $MD_DEV_PATH"

    MD_UUID_PARAM=""
    MD_NAME_PARAM=""
    MD_META_PARAM=""
    MD_LEVEL_PARAM=""
    MD_DEVICES_PARAM=""
    MD_CHUNK_PARAM=""
    MD_LAYOUT_PARAM=""
    MD_DATA_OFFSET_PARAM=""
    MD_MEMBERS=""

    # Unset variables to prevent leakage from the previous loop iteration
    unset MD_UUID MD_NAME MD_META MD_LEVEL MD_DEVICES_COUNT MD_CHUNK MD_LAYOUT MD_DATA_OFFSET MD_MEMBERS_ORDERED
    
    # Parse all advanced variables from the env config
    source "$ENV_FILE"

    [ -n "${MD_NAME:-}" ] && MD_NAME_PARAM="--name=$MD_NAME"
    [ -n "${MD_UUID:-}" ] && MD_UUID_PARAM="--uuid=$MD_UUID"
    [ -n "${MD_META:-}" ] && MD_META_PARAM="--metadata=$MD_META"
    [ -n "${MD_LEVEL:-}" ] && MD_LEVEL_PARAM="--level=$MD_LEVEL"
    [ -n "${MD_DEVICES_COUNT:-}" ] && MD_DEVICES_PARAM="--raid-devices=$MD_DEVICES_COUNT"
    [ -n "${MD_CHUNK:-}" ] && MD_CHUNK_PARAM="--chunk=$MD_CHUNK"
    [ -n "${MD_LAYOUT:-}" ] && MD_LAYOUT_PARAM="--layout=$MD_LAYOUT"
    [ -n "${MD_DATA_OFFSET:-}" ] && MD_DATA_OFFSET_PARAM="--data-offset=${MD_DATA_OFFSET}s"
    [ -n "${MD_MEMBERS_ORDERED:-}" ] && MD_MEMBERS="$MD_MEMBERS_ORDERED"

    echo "  -> Array Level: ${MD_LEVEL:-Unknown} | Devices: ${MD_DEVICES_COUNT:-Unknown}"
    echo "  -> Exact Device Order: $MD_MEMBERS"

    # ====================================================================
    # Selective MDRAID Assassination
    # Prevent mdadm --create from failing or grabbing the wrong size if 
    # the kernel still holds the physical partitions in a phantom array.
    # ====================================================================
    for MEMBER in $MD_MEMBERS; do
        if [ "$MEMBER" != "missing" ] && [ -b "$MEMBER" ]; then
            MEMBER_KNAME="$(basename "$MEMBER")"
            for md_sys in /sys/class/block/md*; do
                if [ -d "$md_sys" ] && [ -e "$md_sys/slaves/$MEMBER_KNAME" ]; then
                    stale_md="/dev/$(basename "$md_sys")"
                    echo "  -> Target $MEMBER is held by stale array $stale_md. Stopping $stale_md..."
                    umount "$stale_md" 2>/dev/null || true
                    mdadm --stop "$stale_md" 2>/dev/null || true
                fi
            done
            # Ensure the superblock is wiped one last time before creation
            mdadm --zero-superblock "$MEMBER" 2>/dev/null || true
        fi
    done
    udevadm settle
    sleep 2
    # ====================================================================

    # ====================================================================
    # Smart Array Name Retry Logic
    # ====================================================================
    BASE_CMD="mdadm --create \"$MD_DEV_PATH\" --force --run $MD_LEVEL_PARAM $MD_DEVICES_PARAM $MD_META_PARAM $MD_DATA_OFFSET_PARAM $MD_UUID_PARAM $MD_CHUNK_PARAM $MD_LAYOUT_PARAM $MD_MEMBERS --assume-clean"

    echo "  -> Attempting to create array with original name: ${MD_NAME:-None}"
    if ! eval "yes | $BASE_CMD $MD_NAME_PARAM" 2>/dev/null; then
        echo "  -> [Warning] Creation failed! Current mdadm likely rejects non-POSIX characters like ':'."
        
        SAFE_MD_NAME="$(echo "$MD_NAME" | sed -r -e 's/[^a-zA-Z0-9.-]/_/g')"
        echo "  -> Retrying with sanitized name: $SAFE_MD_NAME..."
        
        MD_NAME_PARAM="--name=$SAFE_MD_NAME"
        if ! eval "yes | $BASE_CMD $MD_NAME_PARAM"; then
            echo "  -> [Error] Array creation failed completely!"
            exit 1
        fi
    else
        echo "  -> Array created successfully with original name."
    fi

    udevadm settle
    sleep 3
    # ====================================================================

    if [ -f "$MD_DUMP" ]; then
        echo "------------------------------------------"
        echo "Phase 2: Restoring MD Array Internal Partitions for $MD_DEV_PATH"
        echo "------------------------------------------"
        echo "  -> Applying internal partition table to $MD_DEV_PATH..."
        sfdisk --force "$MD_DEV_PATH" < "$MD_DUMP"
        
        # TARGETED partprobe for MD device
        echo "  -> Reloading partition table for $MD_DEV_PATH..."
        partprobe "$MD_DEV_PATH" 2>/dev/null || true
        udevadm settle
        sleep 3
    fi
done

echo "Restoration Complete (Layout & Metadata Only)!"
echo "Advanced RAID array is now structurally identical to the original."
echo $msg_delimiter_star_line
