#!/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

# 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 the target device $MD_DEV_PATH."
      confirm_continue_or_not_default_quit
    fi

    if [ -e "$MD_DEV_PATH" ]; then
      mdadm --stop $MD_DEV_PATH 2>/dev/null || true
    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_UUID:-}" ] && MD_UUID_PARAM="--uuid=$MD_UUID"
    [ -n "${MD_NAME:-}" ] && MD_NAME_PARAM="--name=$MD_NAME"
    [ -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"

    # Silently clear old superblocks from target disks before creating to prevent interactive prompts
    for MEMBER in $MD_MEMBERS; do
        if [ "$MEMBER" != "missing" ] && [ -b "$MEMBER" ]; then
            mdadm --zero-superblock "$MEMBER" 2>/dev/null || true
        fi
    done

    # Construct the create command dynamically. MD_MEMBERS handles the physical disk paths and any 'missing' keywords.
    CREATE_CMD="yes | mdadm --create \"$MD_DEV_PATH\" --force --run $MD_LEVEL_PARAM $MD_DEVICES_PARAM $MD_META_PARAM $MD_DATA_OFFSET_PARAM $MD_UUID_PARAM $MD_NAME_PARAM $MD_CHUNK_PARAM $MD_LAYOUT_PARAM $MD_MEMBERS --assume-clean"

    echo "  -> Running: $CREATE_CMD"
    eval "$CREATE_CMD"

    udevadm settle
    sleep 3

    # ====================================================================
    # CRITICAL COMPATIBILITY FIX: Strip modern feature flags (like BBL)
    # ====================================================================
    echo "  -> Stripping modern mdadm feature flags..."
    mdadm --stop "$MD_DEV_PATH" 2>/dev/null || true
    if ! mdadm --assemble "$MD_DEV_PATH" --update=no-bbl $MD_MEMBERS 2>/dev/null; then
        mdadm --assemble "$MD_DEV_PATH" $MD_MEMBERS 2>/dev/null || true
    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
