#!/usr/bin/env bash

trap 'pkill -P $$' EXIT INT TERM

# Exit on error
set -e

# Keep sudo session alive
sudo -v
( while true; do sleep 60; sudo -v; done ) &

# Kill it on exit or interrupt

# Check for source argument
if [[ -z "$1" ]]; then
    echo "Usage: $0 /dev/sdX"
    exit 1
fi

sourceDevice="$1"
# Check if the device exists
if [ ! -b "$sourceDevice" ]; then
    echo "Error: $sourceDevice is not a valid block device."
    exit 1
fi

# Ask for architecture
while true; do
    read -p "Is this x86_64 or aarch64? (x/a): " architecture
    architecture="${architecture:0:1}"
    architecture="${architecture^^}"
    case "$architecture" in
        X) arch="x86_64"; break ;;
        A) arch="aarch64"; break ;;
        *) echo "Invalid choice. Please enter 'x' for x86_64 or 'a' for aarch64." ;;
    esac
done

# Generate output file name with current date and architecture
dateStamp=$(date +%Y.%m.%d)
outFile="stormux_gaming_image-${dateStamp}-${arch}.img"

# Get last used sector from partition table
lastSector=$(sudo parted -s "$sourceDevice" unit s print | awk '/^ / {print $3}' | sed 's/s//' | sort -n | tail -n 1)
# Get sector size in bytes
sectorSize=$(cat /sys/block/$(basename "$sourceDevice")/queue/hw_sector_size)
# Calculate total bytes to copy
usedBytes=$(( (lastSector + 1) * sectorSize ))
# Round up to nearest multiple of bs (4 MiB)
blockSizeBytes=$((4 * 1024 * 1024))
roundedBytes=$(( (usedBytes + blockSizeBytes - 1) / blockSizeBytes * blockSizeBytes ))
# Calculate count for dd (bs=4M)
count=$((roundedBytes / blockSizeBytes))
# Build the dd command
ddCommand="sudo dd if=$sourceDevice of=$outFile bs=4M count=$count oflag=sync status=progress"
# Show command and ask for confirmation
echo "About to run the following command:"
echo "$ddCommand"
read -rp "Proceed? [y/N] " confirm
confirm="${confirm:0:1}"
if [[ "${confirm^}" == "Y" ]]; then
    eval "$ddCommand"
else
    echo "Aborted."
    exit 1
fi
if [[ -e "$outFile" ]]; then
    sudo xz -T0 -9 "$outFile"
fi