#!/bin/bash
#
# Script to flash the Auracaster System image to an SD card
# This script must be run with sudo privileges

set -e

# Configuration
IMAGE_FILE="./work/iot-system/export-image/2025-04-23-iot-system-lite.img"
AUTO_YES=false

# Parse command line arguments
while getopts "y" opt; do
  case $opt in
    y)
      AUTO_YES=true
      ;;
    \?)
      echo "Invalid option: -$OPTARG" >&2
      exit 1
      ;;
  esac
done

shift $((OPTIND -1))

# Require image file argument
if [ $# -lt 1 ]; then
    echo "Usage: sudo $0 [-y] <path-to-image-file.img>"
    exit 1
fi

IMAGE_FILE="$1"

# Check if script is run as root
if [ "$(id -u)" -ne 0 ]; then
    echo "Error: This script must be run as root (sudo)."
    exit 1
fi

# Check if image exists
if [ ! -f "$IMAGE_FILE" ]; then
    echo "Error: Image file not found at $IMAGE_FILE"
    echo "You need to build the project first."
    exit 1
fi

# Find SD cards (both removable devices and mmcblk devices)
echo "Looking for SD cards..."
echo

# Get list of potential SD cards (both removable devices and mmcblk devices)
SD_CARDS=$(lsblk -d -o NAME,SIZE,MODEL,VENDOR,TRAN,RM | grep -v "loop" | grep -E '(1$|mmcblk)' | awk '{print $1}')
SD_CARD_COUNT=$(echo "$SD_CARDS" | grep -v "^$" | wc -l)

if [ "$SD_CARD_COUNT" -eq 0 ]; then
    echo "No SD cards found. Please insert an SD card and try again."
    exit 1
fi

# Display available SD cards
echo "Available SD cards:"
echo
lsblk -d -o NAME,SIZE,MODEL,VENDOR,TRAN,RM | head -n 1
lsblk -d -o NAME,SIZE,MODEL,VENDOR,TRAN,RM | grep -v "loop" | grep -E '(1$|mmcblk)'
echo

# Safety check - get list of internal disks to avoid (exclude mmcblk devices from this check)
INTERNAL_DISKS=$(lsblk -d -o NAME,RM | grep -E '0$' | grep -v "mmcblk" | awk '{print $1}')

# Select SD card
DEVICE=""
if [ "$SD_CARD_COUNT" -eq 1 ]; then
    DEVICE=$(echo "$SD_CARDS" | tr -d '[:space:]')
    echo "Found single SD card: /dev/$DEVICE"
else
    echo "Multiple SD cards found. Please select which one to use:"
    select DEVICE in $SD_CARDS; do
        if [ -n "$DEVICE" ]; then
            break
        else
            echo "Invalid selection. Please try again."
        fi
    done
fi

# Safety check - ensure we're not flashing to an internal disk
if echo "$INTERNAL_DISKS" | grep -q "$DEVICE"; then
    echo "ERROR: Selected device /dev/$DEVICE appears to be an internal disk!"
    echo "This script will only flash to removable devices or mmcblk devices for safety."
    exit 1
fi

# Validate device exists
if [[ "$DEVICE" == mmcblk* ]]; then
    DEV_PATH="/dev/$DEVICE"
else
    DEV_PATH="/dev/$DEVICE"
fi

if [ ! -b "$DEV_PATH" ]; then
    echo "Error: Device $DEV_PATH does not exist or is not a block device."
    exit 1
fi

# Double-check with clear warning
echo
echo "You are about to erase ALL DATA on $DEV_PATH and write the pi-gen built image ($IMAGE_FILE) to it."
echo "This operation cannot be undone."
echo

# Check if pv is available for better progress display
if command -v pv >/dev/null 2>&1; then
    # Get image size for pv
    IMAGE_SIZE=$(stat -c %s "$IMAGE_FILE")
    
    # Show the command that will be executed with pv
    FLASH_COMMAND="pv -s $IMAGE_SIZE \"$IMAGE_FILE\" | dd of=\"$DEV_PATH\" bs=4M conv=fsync oflag=direct"
    echo "The following command will be executed (with progress bar):"
    echo "$FLASH_COMMAND"
else
    # Fallback to dd with status=progress
    FLASH_COMMAND="dd if=\"$IMAGE_FILE\" of=\"$DEV_PATH\" bs=4M status=progress conv=fsync oflag=direct"
    echo "The following command will be executed:"
    echo "$FLASH_COMMAND"
    echo "For better progress visualization, consider installing 'pv' (Pipe Viewer)."
fi
echo

if [ "$AUTO_YES" = true ]; then
    CONFIRM="y"
    echo "Auto-confirming with -y flag"
else
    read -p "Are you ABSOLUTELY SURE you want to continue? (y/n): " CONFIRM
    
    if [ "$CONFIRM" != "y" ]; then
        echo "Operation cancelled."
        exit 0
    fi
fi

# Unmount any partitions on the device, reporting failures
echo "Unmounting any mounted partitions on $DEV_PATH..."
UNMOUNT_FAILED=0
if [[ "$DEVICE" == mmcblk* ]]; then
    for partition in $(lsblk -n -o NAME | grep "^$DEVICE"); do
        if [ "$partition" != "$DEVICE" ]; then
            if mount | grep -q "/dev/$partition"; then
                umount "/dev/$partition" 2>/dev/null || {
                    echo "Warning: Failed to unmount /dev/$partition" >&2
                    UNMOUNT_FAILED=1
                }
            fi
        fi
    done
else
    for partition in $(lsblk -n -o NAME /dev/$DEVICE | grep -v "^$DEVICE$"); do
        if mount | grep -q "/dev/$partition"; then
            umount "/dev/$partition" 2>/dev/null || {
                echo "Warning: Failed to unmount /dev/$partition" >&2
                UNMOUNT_FAILED=1
            }
        fi
    done
fi
if [ $UNMOUNT_FAILED -eq 1 ]; then
    echo "Some partitions could not be unmounted. Please close any open files or applications using the SD card and try again." >&2
    exit 1
fi
# Check for open files on the device
if lsof | grep -q "/dev/$DEVICE"; then
    echo "Warning: Open files detected on /dev/$DEVICE. Please close them before flashing to avoid corruption." >&2
    lsof | grep "/dev/$DEVICE"
    exit 1
fi

# Flash the image
echo "Flashing image to $DEV_PATH..."
echo "This may take several minutes. Please wait..."
echo

# Execute the flashing command
eval "$FLASH_COMMAND"

# Sync to ensure all writes are complete
sync

# Remind user to eject the card safely
echo "Flash complete! SD card is ready with the pi-gen built image."
echo
# Re-read the partition table to ensure the system recognizes the new layout
echo "Running partprobe to re-read the partition table on $DEV_PATH..."
if partprobe "$DEV_PATH"; then
    echo "Partition table re-read successfully. The SD card is ready for use."
else
    echo "Warning: partprobe failed. If partitions do not show up, try reinserting the SD card."
fi
echo
# Instructions for next steps
echo "1. Insert the SD card into your Raspberry Pi 3"
echo "2. Power on the device"
echo

exit 0