#!/bin/bash

# Enable strict error handling
set -e

# Function to log messages with timestamps
log() {
	echo "$(date '+%F %T') - $@"
}

# Function to display usage information
usage() {
	echo "Usage: $0 [-d DEVICE] [-f FARM] [-e ENVIRONMENT]"
	echo "  -d, --device DEVICE      Specify the block device to use (e.g., sda, sdb)."
	echo "  -f, --farm FARM          Specify the FARM ID."
	echo "  -e, --environment ENV    Specify the environment (dev, qa, test, prod)."
	echo "  -h, --help               Display this help message."
	exit 1
}

# Function to find suitable block devices
find_blockdev() {
	# Minimum size in bytes, e.g., 10 GiB
	MIN_SIZE=$((10 * 1024 * 1024 * 1024))

	BLOCKDEVICES=()

	# Exclude the boot disk
	BOOT_DISK=$(findmnt -n -o SOURCE / | sed 's/[0-9]*$//;s|^/dev/||')
	log "Boot disk detected: /dev/${BOOT_DISK}"

	# Search for NVMe devices first
	log "Searching for NVMe devices..."
	for DEV in /dev/nvme*n1; do
		DEV_BASE=$(basename "$DEV")
		if [ -b "$DEV" ] && [ "$DEV_BASE" != "$BOOT_DISK" ]; then
			SIZE_BYTES=$(blockdev --getsize64 "$DEV")
			if [ "$SIZE_BYTES" -ge "$MIN_SIZE" ]; then
				log "Found NVMe device: $DEV ($(numfmt --to=iec $SIZE_BYTES))"
				BLOCKDEVICES+=("$DEV_BASE")
			else
				log "$DEV is smaller than $(numfmt --to=iec $MIN_SIZE). Skipping."
			fi
		else
			[ "$DEV_BASE" == "$BOOT_DISK" ] && log "Skipping boot disk: $DEV"
		fi
	done

	# Search for SSD devices
	log "Searching for SSD devices..."
	for DEV_PATH in /sys/block/*; do
		DEV_NAME=$(basename "$DEV_PATH")
		# Exclude unwanted devices and the boot disk
		case "$DEV_NAME" in
		loop* | ram* | sr* | fd* | md* | dm-*)
			continue
			;;
		"$BOOT_DISK")
			log "Skipping boot disk: /dev/$DEV_NAME"
			continue
			;;
		esac
		DEVICE="/dev/$DEV_NAME"
		if [ -b "$DEVICE" ] && [ -e "/sys/block/$DEV_NAME/queue/rotational" ] && [ "$(cat /sys/block/$DEV_NAME/queue/rotational)" == "0" ]; then
			SIZE_BYTES=$(blockdev --getsize64 "$DEVICE")
			if [ "$SIZE_BYTES" -ge "$MIN_SIZE" ]; then
				log "Found SSD device: $DEVICE ($(numfmt --to=iec $SIZE_BYTES))"
				BLOCKDEVICES+=("$DEV_NAME")
			else
				log "$DEVICE is smaller than $(numfmt --to=iec $MIN_SIZE). Skipping."
			fi
		fi
	done

	if [ ${#BLOCKDEVICES[@]} -eq 0 ]; then
		log "No suitable block devices found."
		return 1
	fi

	log "List of block devices to be wiped:"
	for DEV in "${BLOCKDEVICES[@]}"; do
		log "/dev/$DEV"
	done

	# Select the first device as the block device to partition
	BLOCKDEV="${BLOCKDEVICES[0]}"
}

# Main script execution
main() {
	# Default values
	BLOCKDEV_SPECIFIED=""
	FARM="1234"
	ENVIRONMENT="prod"

	# Parse command-line arguments
	while [[ $# -gt 0 ]]; do
		case "$1" in
		-d | --device)
			shift
			if [ -z "$1" ]; then
				log "Error: '--device' requires a non-empty option argument."
				usage
			fi
			BLOCKDEV_SPECIFIED="$1"
			shift
			;;
		-f | --farm)
			shift
			if [ -z "$1" ]; then
				log "Error: '--farm' requires a non-empty option argument."
				usage
			fi
			FARM="$1"
			shift
			;;
		-e | --environment)
			shift
			if [ -z "$1" ]; then
				log "Error: '--environment' requires a non-empty option argument."
				usage
			fi
			ENVIRONMENT="$1"
			shift
			;;
		-h | --help)
			usage
			;;
		*)
			log "Error: Unknown option: $1"
			usage
			;;
		esac
	done

	# Validate environment
	case "$ENVIRONMENT" in
	dev | qa | test | prod)
		log "Using environment: $ENVIRONMENT"
		;;
	*)
		log "Error: Invalid environment '$ENVIRONMENT'. Valid options are 'dev', 'qa', 'test', 'prod'."
		exit 1
		;;
	esac

	if [ -n "$BLOCKDEV_SPECIFIED" ]; then
		# Use the specified block device
		if [[ "$BLOCKDEV_SPECIFIED" != /dev/* ]]; then
			BLOCKDEV_SPECIFIED="/dev/$BLOCKDEV_SPECIFIED"
		fi
		if [ ! -b "$BLOCKDEV_SPECIFIED" ]; then
			log "Error: Specified device $BLOCKDEV_SPECIFIED does not exist or is not a block device."
			exit 1
		fi
		BLOCKDEV="${BLOCKDEV_SPECIFIED##*/}"
		BLOCKDEVICES=("$BLOCKDEV")
		log "Using specified block device: $BLOCKDEV_SPECIFIED"
	else
		# Find suitable block devices
		if ! find_blockdev; then
			log "Error: No suitable block devices found."
			exit 1
		fi
	fi

	# Confirm with the user before proceeding
	log "The following disks will be wiped:"
	for DEV in "${BLOCKDEVICES[@]}"; do
		log "/dev/$DEV"
	done
	read -p "This will erase all data on the above disks. Are you sure you want to proceed? Type 'WIPE ALL' to confirm: " CONFIRM
	if [ "$CONFIRM" != "WIPE ALL" ]; then
		log "Operation cancelled by user."
		exit 0
	fi

	# Wipe filesystem signatures on all disks
	log "Wiping filesystem signatures on all selected disks..."
	for DEV in "${BLOCKDEVICES[@]}"; do
		log "Wiping /dev/$DEV..."
		wipefs -a -f /dev/$DEV || {
			log "Failed to wipe /dev/$DEV"
			exit 1
		}
	done

	# Proceed with partitioning and formatting the first disk
	# Determine partition naming convention
	if [[ $BLOCKDEV =~ [0-9]$ ]]; then
		PARTPREFIX="${BLOCKDEV}p"
	else
		PARTPREFIX="${BLOCKDEV}"
	fi

	# Create GPT partition table on the selected block device
	log "Creating GPT partition table on /dev/${BLOCKDEV}..."
	parted -s /dev/${BLOCKDEV} mklabel gpt || {
		log "Failed to create GPT partition table."
		exit 1
	}

	# Create BIOS Boot Partition
	log "Creating BIOS Boot partition on /dev/${BLOCKDEV}..."
	parted -s /dev/${BLOCKDEV} mkpart primary 1MiB 2MiB || {
		log "Failed to create BIOS Boot partition."
		exit 1
	}
	parted -s /dev/${BLOCKDEV} set 1 bios_grub on

	# Create additional partitions as needed
	log "Creating data partitions on /dev/${BLOCKDEV}..."
	parted -s /dev/${BLOCKDEV} mkpart primary 2MiB 100MiB || {
		log "Failed to create partition."
		exit 1
	}
	parted -s /dev/${BLOCKDEV} mkpart primary 100MiB 100% || {
		log "Failed to create partition."
		exit 1
	}

	# Format partitions
	log "Formatting /dev/${PARTPREFIX}2 as FAT16..."
	if ! mkfs.vfat -F16 /dev/${PARTPREFIX}2; then
		log "Failed to format /dev/${PARTPREFIX}2 as FAT16."
		exit 1
	fi

	log "Formatting /dev/${PARTPREFIX}3 as BTRFS..."
	if ! mkfs.btrfs -f /dev/${PARTPREFIX}3; then
		log "Failed to format /dev/${PARTPREFIX}3 as BTRFS."
		exit 1
	fi

	# Install iPXE bootloader to MBR
	log "Installing iPXE bootloader to MBR of /dev/${BLOCKDEV}..."

	# Download the iPXE bootloader suitable for BIOS
	log "Downloading iPXE bootloader..."
	if ! wget -q -O ipxe.lkrn http://boot.ipxe.org/ipxe.lkrn; then
		log "Failed to download iPXE bootloader."
		exit 1
	fi

	# Write iPXE bootloader to the MBR
	log "Writing iPXE bootloader to MBR of /dev/${BLOCKDEV}..."
	dd if=ipxe.lkrn of=/dev/${BLOCKDEV} bs=440 count=1 conv=notrunc || {
		log "Failed to write iPXE to MBR."
		exit 1
	}

	# Clean up the downloaded file
	rm -f ipxe.lkrn

	# Prepare the iPXE script
	log "Preparing iPXE script..."

	# Mount the FAT16 partition to store the iPXE script
	MOUNT_POINT="zospxe"
	log "Creating and mounting directory ./${MOUNT_POINT}..."
	mkdir -p "${MOUNT_POINT}"
	if ! mount /dev/${PARTPREFIX}2 -t vfat "${MOUNT_POINT}"; then
		log "Failed to mount /dev/${PARTPREFIX}2 to ${MOUNT_POINT}."
		exit 1
	fi

	# Download the BOOTX64.EFI file (optional, depending on your setup)
	# Since this is a BIOS system, we might not need this step
	# However, if required, adjust the script accordingly

	# Alternatively, create an embedded iPXE script
	log "Creating iPXE script..."
	cat >"${MOUNT_POINT}/boot.ipxe" <<EOF
#!ipxe
dhcp
set farm ${FARM}
chain https://bootstrap.grid.tf/uefi/${ENVIRONMENT}/\${FARM}/debug/zero-os-development-zos-v4-debug-7d2de62033.efi?version=v4"
# chain https://bootstrap.grid.tf/ipxe/${ENVIRONMENT}/\${farm}
EOF

	log "iPXE script created at ${MOUNT_POINT}/boot.ipxe"

	# Unmount the partition
	log "Unmounting ${MOUNT_POINT}..."
	umount "${MOUNT_POINT}" || {
		log "Failed to unmount ${MOUNT_POINT}."
		exit 1
	}

	log "Operation completed successfully."
}

# Execute the main function
main "$@"