#!/bin/bash

#==========================
# Set up the environment
#==========================
set -e                  # exit on error
set -o pipefail         # exit on pipeline error
set -u                  # treat unset variable as error

#==========================
# Basic Information
#==========================
export LC_ALL=C
export LANG=en_US.UTF-8
export DEBIAN_FRONTEND=noninteractive
export SCRIPT_DIR="$(dirname "$(readlink -f "$0")")"

#==========================
# Color
#==========================
Green="\033[32m"
Red="\033[31m"
Yellow="\033[33m"
Blue="\033[36m"
Font="\033[0m"
GreenBG="\033[42;37m"
RedBG="\033[41;37m"
OK="${Green}[  OK  ]${Font}"
ERROR="${Red}[FAILED]${Font}"
WARNING="${Yellow}[ WARN ]${Font}"

#==========================
# Print Colorful Text
#==========================
function print_ok() {
  echo -e "${OK} ${Blue} $1 ${Font}"
}

function print_error() {
  echo -e "${ERROR} ${Red} $1 ${Font}"
}

function print_warn() {
  echo -e "${WARNING} ${Yellow} $1 ${Font}"
}

#==========================
# Judge function
#==========================
function judge() {
  if [[ 0 -eq $? ]]; then
    print_ok "$1 succeeded"
    sleep 0.2
  else
    print_error "$1 failed"
    exit 1
  fi
}

function burn()
{
    # If the user want to burn /dev/sda, device should be 'sda'
    device=$1
    iso=$2

    if [ -z "$device" ]; then
        print_error "Usage: $0 <device> <iso>"
        exit 1
    fi

    # Ensure /dev/$device exists
    if [ ! -e "/dev/$device" ]; then
        print_error "Device /dev/$device does not exist."
        exit 1
    fi

    # Unmount all partitions on /dev/$device
    for mountpoint in $(lsblk -o MOUNTPOINT "/dev/$device" | tail -n +2); do
        print_ok "Unmounting $mountpoint"
        sudo umount "$mountpoint"
        judge "Unmount $mountpoint"
    done

    # Ensure the iso exists
    if [ ! -e "$iso" ]; then
        print_error "ISO $iso does not exist."
        exit 1
    fi

    # Ensure the iso is a file
    if [ ! -f "$iso" ]; then
        print_error "ISO $iso is not a file."
        exit 1
    fi

    # Ensure the iso is readable
    if [ ! -r "$iso" ]; then
        print_error "ISO $iso is not readable."
        exit 1
    fi

    # Ensure the disk is larger than the iso
    iso_size=$(stat -c %s "$iso")
    device_size=$(sudo blockdev --getsize64 "/dev/$device")
    if [ "$iso_size" -gt "$device_size" ]; then
        print_error "ISO $iso is larger than /dev/$device."
        exit 1
    else
        print_ok "ISO $iso is $iso_size bytes, /dev/$device is $device_size bytes. Device is large enough."
    fi

    print_ok "Burning $iso to /dev/$device. This will take a while..."
    sudo dd if="$iso" of="/dev/$device" bs=4M status=progress oflag=sync
    sync && sync && sync

    # Verify the burn
    print_ok "Verifying the burn. This will take a while."

    # Foreach the mount point under $device, unmount it
    for mountpoint in $(lsblk -o MOUNTPOINT "/dev/$device" | tail -n +2); do
        print_ok "Unmounting $mountpoint"
        sudo umount "$mountpoint"
        judge "Unmount $mountpoint"
    done

    # Mount the device to /tmp/burn/$device
    print_ok "Mounting /dev/$device to /tmp/burn/$device"
    mkdir -p /tmp/burn/$device
    sudo mount "/dev/$device" /tmp/burn/$device
    judge "Mount /dev/$device to /tmp/burn/$device"

    # Calculate the md5sum of the iso
    print_ok "Calculating the md5sum of the files in the iso"
    (
        cd /tmp/burn/$device && \
        sudo md5sum -c md5sum.txt && \
        print_ok "First pass verification succeeded." || \
        print_error "Burn failed."
    )
    (
        cd /tmp/burn/$device && \
        sudo md5sum -c md5sum.txt && \
        print_ok "Second pass verification succeeded." || \
        print_error "Burn failed."
    )

    # Unmount the device
    print_ok "Unmounting /tmp/burn/$device"
    sudo umount /tmp/burn/$device
    judge "Unmount /tmp/burn/$device"

    # Clean up
    print_ok "Cleaning up"
    rm -rf /tmp/burn/$device
    judge "Clean up"

    print_ok "Done. You can now remove the USB stick."
}

burn "$@"