#!/bin/bash
# Author: Steven Shiau <steven _at_ clonezilla org>
# License: GPL
# A script to expand the GPT partition table by disk size ratio.
# ///NOTE/// This program only works for GPT partition table, not for MBR.
# This file contains code generated by Google Gemini.

# --- Environment Setup ---
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

# --- Defaults ---
chk_tgt_disk_size_bf_mk_pt="yes"
batch_mode="no"
cmd_name="$(basename "$0")"
EXTRA_SFDISK_OPT=${EXTRA_SFDISK_OPT:-"--force --wipe always"}

#
USAGE() {
    echo "$cmd_name: To create a proportional GPT partition table in a disk based on an existing partition table (sfdisk format)"
    echo "Usage:"
    echo "  $cmd_name [OPTION] PARTITION_TABLE_FILE TARGET_DEVICE"
    echo "  TARGET_DEVICE can be with or without /dev/, e.g., /dev/sda or sda."
    echo 
    echo "OPTION:"
    echo "  -b, --batch     Run $cmd_name in batch mode, i.e. without any prompt or wait to press enter. VERY DANGEROUS!"
    echo "   $cmd_name will honor experimental variable EXTRA_SFDISK_OPT and use it as the option for sfdisk."
    echo "  -icds, --ignore-chk-dsk-size-pt  Skip checking destination disk size before creating the partition table on it. By default it will be checked and if the size is smaller than the source disk, quit."
    echo 
    echo "Example:"
    echo "To create a proportional partition table on disk /dev/sda based on /home/partimag/IMAGE/sda-pt.sf, use:"
    echo "$cmd_name /home/partimag/IMAGE/sda-pt.sf /dev/sda"
}         
          
# --- Arg Parsing ---
while [ $# -gt 0 ]; do
    case "$1" in
        -b|--batch) batch_mode="yes"; shift ;;
        -icds|--ignore-chk-dsk-size-pt) chk_tgt_disk_size_bf_mk_pt="no"; shift ;;
        -*) echo "${0}: ${1}: invalid option" >&2
            USAGE >& 2
            exit 2 ;;
        *) break ;;
    esac
done

orig_sf="$1"
target_disk="$(format_dev_name_with_leading_dev "$2")"

# --- Validation ---
[ ! -f "$orig_sf" ] && { echo "Error: SF file missing."; exit 1; }
[ -z "$target_disk" ] && { echo "Error: Target disk missing."; exit 1; }

# --- 0. Sector Size Check & External Conversion ---
# Detect Source Sector Size (from sfdisk file)
src_sector_size=$(grep "sector-size:" "$orig_sf" | awk '{print $2}')
[ -z "$src_sector_size" ] && src_sector_size=512 # Default to 512 if not specified

# Detect Target Sector Size
# Priority: logical_block_size (Standard/Modern) -> hw_sector_size (Legacy) -> 512 (Default)
tgt_dev_name=$(basename "$target_disk")
tgt_sector_size=$(cat "/sys/block/$tgt_dev_name/queue/logical_block_size" 2>/dev/null)
if [ -z "$tgt_sector_size" ]; then
    tgt_sector_size=$(cat "/sys/block/$tgt_dev_name/queue/hw_sector_size" 2>/dev/null)
fi
[ -z "$tgt_sector_size" ] && tgt_sector_size=512

echo "Source Sector: $src_sector_size | Target Sector: $tgt_sector_size"

working_sf="$orig_sf"

if [ "$src_sector_size" != "$tgt_sector_size" ]; then
    echo "Sector size mismatch detected ($src_sector_size -> $tgt_sector_size)."
    
    # Determine mode
    if [ "$src_sector_size" -eq 512 ] && [ "$tgt_sector_size" -eq 4096 ]; then
        mode="512to4k"
    elif [ "$src_sector_size" -eq 4096 ] && [ "$tgt_sector_size" -eq 512 ]; then
        mode="4kto512"
    else
        echo "Error: Unsupported conversion ($src_sector_size -> $tgt_sector_size)"
        exit 1
    fi

    working_sf=$(mktemp /tmp/ocs_converted_sf.XXXXXX)
    
    # Call external converter directly
    # Assumes ocs-pt-512-4k-convert is in PATH or DRBL sbin
    echo "Running conversion: ocs-pt-512-4k-convert $mode $orig_sf $working_sf"
    ocs-pt-512-4k-convert "$mode" "$orig_sf" "$working_sf"
    
    if [ $? -eq 0 ]; then
        echo "Conversion successful. Using temporary SF file."
    else
        echo "Error: Partition table conversion failed."
        rm -f "$working_sf"
        exit 1
    fi
fi

# --- 1. Size Detection (Using working_sf) ---
# Get Original Size (trust the converted/original SF file header now)
last_lba=$(grep "^last-lba:" "$working_sf" | awk '{print $2}')
[ -n "$last_lba" ] && ori_disk_size=$((last_lba + 1))

# Fallback to .parted file ONLY if no conversion happened
if [ "$working_sf" == "$orig_sf" ] && [ -z "$ori_disk_size" ] && [ -f "${orig_sf%.sf}.parted" ]; then
    ori_disk_size=$(grep -E "^Disk /dev" "${orig_sf%.sf}.parted" | awk -F: '{print $2}' | sed 's/[^0-9]//g')
fi

[ -z "$ori_disk_size" ] && { echo "Error: Could not determine original disk size."; exit 1; }

# Get Target Size
tgt_disk_size=$(parted -s "$target_disk" unit s print | grep -E "^Disk /dev" | awk -F: '{print $2}' | sed 's/[^0-9]//g')
last_usable=$((tgt_disk_size - 34))

echo "Original Size (Converted): $ori_disk_size | Target Size: $tgt_disk_size"

# --- 2. Analyze Partitions ---
new_sf_tmp=$(mktemp /tmp/new_sf_tmp.XXXXXX)
grep -E "^/dev" "$working_sf" > "$new_sf_tmp"

total_static=0
total_expandable=0

while read -r line; do
    # Robust cleanup: remove all syntax sugar to get raw numbers
    clean_line=$(echo "$line" | sed -E 's/ : / /; s/start=//; s/size=//; s/,/ /g')
    read -r dev start size type rest <<< "$clean_line"

    # Identify protected partitions (EFI, MSR, Recovery, Swap)
    if echo "$clean_line" | grep -qiE 'C12A7328|E3C9E316|DE94BBA4|0657FD6D'; then
        total_static=$((total_static + size))
    else
        total_expandable=$((total_expandable + size))
    fi
done < "$new_sf_tmp"

# Calculate Expansion Ratio
# Note: Subtract 34 sectors for GPT Backup Header so we don't truncate the last partition
available_space=$((tgt_disk_size - total_static - 2048 - 34))

if [ "$total_expandable" -gt 0 ]; then
    ratio=$(echo "scale=10; $available_space / $total_expandable" | bc -l)
else
    ratio=0
fi

echo "Static: $total_static | Expandable: $total_expandable | Ratio: $ratio"

# --- 3. Construct New Table ---
new_sf=$(mktemp /tmp/new_sf.XXXXXX)

# Copy SF Headers from working file
grep -E "^label:|^label-id:|^device:|^unit:|^first-lba:" "$working_sf" > "$new_sf"
echo "" >> "$new_sf"

curr_offset=2048 

while read -r line; do
    clean_line=$(echo "$line" | sed -E 's/ : / /; s/start=//; s/size=//; s/,/ /g')
    read -r dev start size type rest <<< "$clean_line"

    # Determine New Size
    if echo "$clean_line" | grep -qiE 'C12A7328|E3C9E316|DE94BBA4|0657FD6D'; then
        new_size=$size
    else
        if [ "$total_expandable" -gt 0 ]; then
            new_size=$(printf "%.0f" "$(echo "$size * $ratio" | bc -l)")
        else
            new_size=$size
        fi
    fi

    # Extract original attributes from the raw line to preserve UUIDs etc.
    suffix=$(echo "$line" | grep -oE "type=.*")

    # Write line
    echo "$dev : start=$curr_offset, size=$new_size, $suffix" >> "$new_sf"

    curr_offset=$((curr_offset + new_size))
done < "$new_sf_tmp"

# --- 4. Final Boundary Fix ---
# Only cap if we somehow exceeded the limit (should be rare now with correct calculation)
last_line=$(tail -n 1 "$new_sf")
last_start=$(echo "$last_line" | sed -E 's/.*start=([0-9]+).*/\1/')
current_end=$((last_start + new_size))

if [ "$current_end" -gt "$last_usable" ]; then
    final_size=$((last_usable - last_start))
    [ "$final_size" -lt 0 ] && { echo "Error: Partitions exceed disk size!"; exit 1; }
    echo "Warning: Adjusting last partition to fit disk boundary."
    sed -i "$ s/size=[0-9]*/size=$final_size/" "$new_sf"
fi

# --- 5. Execution ---
echo "Applying new partition layout to $target_disk..."
eval "LC_ALL=C sfdisk $EXTRA_SFDISK_OPT $target_disk < $new_sf"
rc=$?

# Cleanup
rm -f "$new_sf" "$new_sf_tmp"
[ "$working_sf" != "$orig_sf" ] && rm -f "$working_sf"
exit $rc
