#!/bin/bash
# License: GPL
# Author: Steven Shiau <steven _at_ clonezilla org>
# Description: Program to get partition size from the sfdisk's dump file.
# This file contains code generated by Google's Gemini.

# Default values
UNIT="B"

# Helper function for usage
usage() {
    echo "Usage: $0 [-u|--unit B|MB|GB|TB] <sfdisk_dump_file> <partition_device>"
    echo "Example: $0 -u GB nvme0n1-pt.sf nvme0n1p6"
    exit 1
}

# 1. Parse optional unit parameter if it appears first
if [[ "$1" == "-u" || "$1" == "--unit" ]]; then
    UNIT="$2"
    shift 2
fi

# 2. Check if the mandatory positional arguments exist
if [ "$#" -lt 2 ]; then
    usage
fi

SF_FILE=$1
PART_INPUT=$2

# Normalize the partition name to match sfdisk format (/dev/...)
if [[ "$PART_INPUT" == /dev/* ]]; then
    PART_PATH="$PART_INPUT"
else
    PART_PATH="/dev/$PART_INPUT"
fi

# 3. Extract the sector-size from the file header
SECTOR_SIZE="$(grep "sector-size:" "$SF_FILE" | awk -F': ' '{print $2}' | tr -d ' ')"

# 4. Extract the sector count for the specific partition
PART_SECTORS="$(grep "^$PART_PATH" "$SF_FILE" | sed -n 's/.*size=[[:space:]]*\([0-9]*\).*/\1/p')"

# 5. Validate extracted data
if [ -z "$SECTOR_SIZE" ] || [ -z "$PART_SECTORS" ]; then
    echo "Error: Could not find partition or sector size info." >&2
    exit 1
fi

# 6. Calculate total size in Bytes
TOTAL_BYTES="$(($SECTOR_SIZE * $PART_SECTORS))"

# 7. Convert unit if requested (using 1024-based binary units)
case ${UNIT^^} in
    B)
        echo "$TOTAL_BYTES"
        ;;
    MB)
        # Using 'bc' for floating point if you want decimals, 
        # otherwise bash $(()) for integers.
        echo "scale=2; $TOTAL_BYTES / 1024 / 1024" | bc -l | sed 's/\.00$//'
        ;;
    GB)
        echo "scale=2; $TOTAL_BYTES / 1024 / 1024 / 1024" | bc -l | sed 's/\.00$//'
        ;;
    TB)
        echo "scale=2; $TOTAL_BYTES / 1024 / 1024 / 1024 / 1024" | bc -l | sed 's/\.00$//'
        ;;
    *)
        echo "Error: Unsupported unit '$UNIT'. Use B, MB, GB, or TB." >&2
        exit 1
        ;;
esac
