#!/bin/bash

# Define the image file name
IMG_NAME="rom.img"
GZ_IMG_NAME="$IMG_NAME.gz"

# Check if the required components exist, otherwise install them using opkg
REQUIRED_COMMANDS=("fdisk" "pv" "gzip")
OPKG_UPDATED_FLAG="/tmp/opkg_updated"

# Update opkg (only if the update flag does not exist)
if [ ! -f "$OPKG_UPDATED_FLAG" ]; then
    echo "Running opkg update..."
    if opkg update; then
        touch "$OPKG_UPDATED_FLAG"
    else
        echo "Error: Failed to update opkg. Please check your network or opkg configuration." >&2
        exit 1
    fi
fi

# Check and install required components
for cmd in "${REQUIRED_COMMANDS[@]}"; do
    if ! command -v "$cmd" &> /dev/null; then
        echo "$cmd is not installed. Installing it using opkg..."
        if ! opkg install "$cmd"; then
            echo "Error: Failed to install $cmd. Please install it manually and re-run the script." >&2
            exit 1
        fi
    fi
done

# Get the end block number and add +1
END_BLOCK=$(fdisk -l /dev/sda | awk '/sda2/ {print $3}')
END_BLOCK_PLUS_ONE=$((END_BLOCK + 1))
echo "End block of /dev/sda2 is $END_BLOCK_PLUS_ONE"

# Prompt the user to input COUNT_SIZE (number of blocks to backup)
read -p "Please enter the number of blocks to backup (default is ${END_BLOCK_PLUS_ONE}, or press Enter to use default): " COUNT_SIZE

# Use the default value if the user does not input anything
COUNT_SIZE=${COUNT_SIZE:-$END_BLOCK_PLUS_ONE}

# Validate that the user input is a valid positive integer
if ! [[ "$COUNT_SIZE" =~ ^[0-9]+$ ]]; then
    echo "Error: Invalid input. Please enter a valid number for COUNT_SIZE." >&2
    exit 1
fi

# Get the directory of the script
SCRIPT_DIR="$(cd "$(dirname "$0")" && pwd)"
cd "$SCRIPT_DIR" || exit 1

# Check if /dev/sda exists and is a block device
if [ ! -b /dev/sda ]; then
    echo "Error: /dev/sda not found or is not a block device." >&2
    exit 1
fi

# Check available space and convert units (MB)
AVAILABLE_SPACE_MB=$(df "$SCRIPT_DIR" | awk 'NR==2 {printf "%.0f", $4 / 1024}')
REQUIRED_SPACE_MB=$(($COUNT_SIZE * 512 / 1024 / 1024))  # Convert to MB
if [ "$AVAILABLE_SPACE_MB" -lt "$REQUIRED_SPACE_MB" ]; then
    echo "Error: Not enough space in $SCRIPT_DIR. Required: ${REQUIRED_SPACE_MB}MB, Available: ${AVAILABLE_SPACE_MB}MB." >&2
    exit 1
fi

# Notify the user to prepare for the backup operation
echo "Prepare to backup your ROM to an image file."

# Ask the user to confirm the backup operation
read -p "Are you sure you want to backup /dev/sda to $IMG_NAME with $COUNT_SIZE blocks? (yes/no): " CONFIRM
if [ "$CONFIRM" != "yes" ]; then
    echo "Backup operation cancelled."
    exit 0
fi

# Check and remove existing files
if [ -f "$IMG_NAME" ]; then
    echo "Found existing $IMG_NAME. Deleting it..."
    rm -f "$IMG_NAME"
fi

if [ -f "$GZ_IMG_NAME" ]; then
    echo "Found existing $GZ_IMG_NAME. Deleting it..."
    rm -f "$GZ_IMG_NAME"
fi

# Perform the backup operation and show progress
if pv -s $((COUNT_SIZE * 512)) /dev/sda | dd bs=512 count=$COUNT_SIZE of="$SCRIPT_DIR/$IMG_NAME" conv=notrunc; then
    echo "Backup successful: $IMG_NAME created."
else
    echo "Error: Backup failed." >&2
    exit 1
fi

# Compress the image file
echo "Compressing $IMG_NAME to $GZ_IMG_NAME..."

# Ask the user to confirm the compression operation
read -p "Are you sure you want to compress $IMG_NAME to $GZ_IMG_NAME? (yes/no): " CONFIRM
if [ "$CONFIRM" != "yes" ]; then
    echo "Compression operation cancelled."
    exit 0
fi

# Perform the compression operation and show progress
if pv "$SCRIPT_DIR/$IMG_NAME" | gzip > "$SCRIPT_DIR/$GZ_IMG_NAME"; then
    echo "Compression successful: $GZ_IMG_NAME created."
else
    echo "Error: Compression failed." >&2
    exit 1
fi