#!/bin/bash

# This was I input to o1-preview with this prompt:
#
# BLOCKDEV=nvme0n1
# wipefs -a -f /dev/${BLOCKDEV}
# parted -s /dev/${BLOCKDEV} "mklabel gpt mkpart zosboot fat32 1 10M set 1 esp on mkpart zoscache btrfs 10M 100%"
# mkfs.vfat -F32 /dev/${BLOCKDEV}p1
# mkfs.btrfs -f /dev/${BLOCKDEV}p2
# mkdir zospxe && mount /dev/${BLOCKDEV}p1 -t vfat zospxe && mkdir -p zospxe/efi/boot
# wget -q https://bootstrap.grid.tf/uefi/prod/1234 -O zospxe/efi/boot/BOOTX64.EFI
#
# find a BLOCKDEV that is at least an nvme, ssd (in that order)
# and make this script a bit more log-expressive, and failing gracefully.
# Also, can you add the if -d $device is specified, it will use that and
# not search for a disk, and also make farmerid a parameter
# for https://bootstrap.grid.tf/uefi/prod/$farmerid}

# Exit immediately if a command exits with a non-zero status.
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 FARMERID]"
	echo "  -d, --device DEVICE      Specify the block device to use (e.g., sda, nvme0n1)."
	echo "  -f, --farmerid FARMERID  Specify the farmer ID for the download URL."
	echo "  -h, --help               Display this help message."
	exit 1
}

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

	# First try NVMe devices
	log "Searching for NVMe devices..."
	for DEV in /dev/nvme*n1; do
		if [ -b "$DEV" ]; then
			SIZE_BYTES=$(blockdev --getsize64 "$DEV")
			if [ "$SIZE_BYTES" -ge "$MIN_SIZE" ]; then
				log "Found NVMe device: $DEV ($(numfmt --to=iec $SIZE_BYTES))"
				BLOCKDEV="${DEV##*/}"
				return 0
			else
				log "$DEV is smaller than $(numfmt --to=iec $MIN_SIZE). Skipping."
			fi
		fi
	done

	log "No suitable NVMe devices found. Searching for SSD devices..."

	for DEV_PATH in /sys/block/*; do
		DEV_NAME=$(basename "$DEV_PATH")
		# Exclude unwanted devices
		case "$DEV_NAME" in
		loop* | ram* | sr* | fd* | md* | dm-*)
			continue
			;;
		esac
		DEVICE="/dev/$DEV_NAME"
		if [ -b "$DEVICE" ] && [ "$(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))"
				BLOCKDEV="$DEV_NAME"
				return 0
			else
				log "$DEVICE is smaller than $(numfmt --to=iec $MIN_SIZE). Skipping."
			fi
		fi
	done

	log "No suitable block devices found."
	return 1
}

# Main script execution
main() {
	# Default values
	BLOCKDEV_SPECIFIED=""
	FARMERID="1234" # Default farmer ID

	# 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 | --farmerid)
			shift
			if [ -z "$1" ]; then
				log "Error: '--farmerid' requires a non-empty option argument."
				usage
			fi
			FARMERID="$1"
			shift
			;;
		-h | --help)
			usage
			;;
		*)
			log "Error: Unknown option: $1"
			usage
			;;
		esac
	done

	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##*/}"
		log "Using specified block device: $BLOCKDEV_SPECIFIED"
	else
		# Find a suitable block device
		if ! find_blockdev; then
			log "Error: No suitable block device found."
			exit 1
		fi
		log "Selected block device: /dev/${BLOCKDEV}"
	fi

	# Confirm with the user before proceeding
	read -p "This will erase all data on /dev/${BLOCKDEV}. Are you sure you want to proceed? (yes/[no]): " CONFIRM
	if [ "$CONFIRM" != "yes" ]; then
		log "Operation cancelled by user."
		exit 0
	fi

	# Determine partition naming convention
	if [[ $BLOCKDEV =~ [0-9]$ ]]; then
		PARTPREFIX="${BLOCKDEV}p"
	else
		PARTPREFIX="${BLOCKDEV}"
	fi

	# Wipe filesystem signatures
	log "Wiping filesystem signatures on /dev/${BLOCKDEV}..."
	wipefs -a -f /dev/${BLOCKDEV}

	# Create GPT partition table
	log "Creating GPT partition table on /dev/${BLOCKDEV}..."
	parted -s /dev/${BLOCKDEV} mklabel gpt

	# Create partitions
	log "Creating partition 1 (zosboot) on /dev/${BLOCKDEV}..."
	parted -s /dev/${BLOCKDEV} mkpart zosboot fat32 1MiB 10MiB

	log "Setting partition 1 as ESP..."
	parted -s /dev/${BLOCKDEV} set 1 esp on

	log "Creating partition 2 (zos-cache) on /dev/${BLOCKDEV}..."
	parted -s /dev/${BLOCKDEV} mkpart zos-cache btrfs 10MiB 100%

	# Format partitions
	log "Formatting /dev/${PARTPREFIX}1 as FAT32..."
	mkfs.vfat -F32 /dev/${PARTPREFIX}1

	log "Formatting /dev/${PARTPREFIX}2 as BTRFS..."
	mkfs.btrfs -f /dev/${PARTPREFIX}2

	# Mount and prepare directories
	MOUNT_POINT="zospxe"
	log "Creating mount point at ./${MOUNT_POINT}..."
	mkdir -p "${MOUNT_POINT}"
	log "Mounting /dev/${PARTPREFIX}1 to ./${MOUNT_POINT}..."
	mount /dev/${PARTPREFIX}1 "${MOUNT_POINT}"
	log "Creating directory structure in ${MOUNT_POINT}..."
	mkdir -p "${MOUNT_POINT}/efi/boot"

	# Download the required BOOTX64.EFI file
	EFI_FILE_URL="https://bootstrap.grid.tf/uefi/prod/${FARMERID}"
	log "Downloading BOOTX64.EFI from ${EFI_FILE_URL}..."
	wget -q "${EFI_FILE_URL}" -O "${MOUNT_POINT}/efi/boot/BOOTX64.EFI"

	# Create zos-cache subvolume
	log "Creating mount point at ./${MOUNT_POINT}root..."
	mkdir -p "${MOUNT_POINT}root"
	log "Mounting /dev/${PARTPREFIX}2 to ./${MOUNT_POINT}root..."
	mount /dev/${PARTPREFIX}2 "${MOUNT_POINT}root"
	log "Creating subvolume \"zos-cache\" on  ${MOUNT_POINT}root"
	btrfs subvolume create ${MOUNT_POINT}root/zos-cache

	# Umount mountpoints
	log "Umounting ${MOUNT_POINT}/*"
	umount -l ${MOUNT_POINT}
	umount -l ${MOUNT_POINT}root

	log "Operation completed successfully."
}

# Execute the main function
main "$@"