#!/usr/bin/sh

# Copyright (C) 2014-2024 Red Hat Inc.
# SPDX-License-Identifier:  GPL-2.0+

# Automate Media Creation for Fedora ARM
# Current version
VERSION=5.3

# Cleanup function for loopback devices
cleanup_loopback() {
	if [ -n "$LOOP_DEVICE" ] && [ -b "$LOOP_DEVICE" ]; then
		echo "= Cleaning up loopback device $LOOP_DEVICE ..."
		losetup -d "$LOOP_DEVICE" 2>/dev/null
	fi
}

# Set trap to cleanup on exit
trap cleanup_loopback EXIT INT TERM

# usage message
usage() {
    echo "
Usage: $(basename ${0}) <options>

	--image=IMAGE   - xz compressed image file name
	--media=DEVICE  - media device file (/dev/[sdX|mmcblkX]) or path to disk image file

Optional
	--add-copr      - Install u-boot from COPR (pbrobinson/u-boot) for rk35xx boards
	--addconsole    - Add system console kernel parameter for the target
	--addkey        - /path/to/ssh-public-key
	--args          - Optional kernel parameters listed in quotes
	--ign-url	- url for ignition configuration file (IoT)
	--ignition	- local path for ignition configuration file (IoT)
	--norootpass    - Set the root password to the empty string
	--relabel       - SELinux relabel root filesystem on first boot
	--resizefs      - Resize root filesystem to fill media device
	--showboot	- Show boot messages, removes 'rhgb quiet' from kargs
	--sysrq         - Enable System Request debugging of the kernel
	--target=TARGET - target board for uboot
	--wifi-ssid=SSID        - Wi-Fi SSID to configure
	--wifi-pass=PASS        - Wi-Fi password to configure
	--wifi-security=TYPE    - Wi-Fi security type (wpa-psk, sae)
	-y              - Assumes yes, will not wait for confirmation

Help
	--supported     - List of hardware we support writing out boot firmware
	--version       - Display version and exit

Example: $(basename ${0}) --image=Fedora-Rawhide.xz --target=pine64_plus --media=/dev/mmcblk0 --wifi-ssid=MY_SSID --wifi-pass=MY_PASSWORD

"
}

# Set some global variables for the command directory, target board directory,
# and valid targets.
DIR=$(dirname $0)
# Prefer local checkout over system installation
if [ -d "${DIR}/boards.d" ]; then
	BOARDDIR="${DIR}/boards.d"
	DOC_DIR="${DIR}/"
elif [ -d "/usr/share/arm-image-installer/boards.d" ]; then
	BOARDDIR="/usr/share/arm-image-installer/boards.d"
	DOC_DIR="/usr/share/doc/arm-image-installer/"
else
	BOARDDIR="${DIR}/boards.d"
	DOC_DIR="${DIR}/"
fi

# Return help for no args
if [ $# -eq 0 ]; then
	usage
	exit 0
fi

# check the args
while [ $# -gt 0 ]; do
	case $1 in
		--debug)
			set -x
			;;
		-h|--help)
			usage
			exit 0
			;;
		--target*)
			if echo $1 | grep '=' >/dev/null ; then
				TARGET=$(echo $1 | sed 's/^--target=//')
			elif [ -n "$2" ]; then
				TARGET=$2
				shift
			else
				echo "$(basename ${0}): Error - '--target' expects an argument"
				usage
				exit 1
			fi
			;;
		--image*)
			if echo $1 | grep '=' >/dev/null ; then
				IMAGE=$(echo $1 | sed 's/^--image=//')
			elif [ -n "$2" ]; then
				IMAGE=$2
				shift
			else
				echo "$(basename ${0}): Error - '--image' expects an argument"
				usage
				exit 1
			fi
                        # Explicitly expand the tilde if it's the first character
                        if [ "$(echo "$IMAGE" | cut -c1)" = "~" ]; then
                                IMAGE="${HOME}${IMAGE#~}"
                        fi
			;;
		--media*)
			if echo $1 | grep '=' >/dev/null ; then
				MEDIA=$(echo $1 | sed 's/^--media=//')
			elif [ -n "$2" ]; then
				MEDIA=$2
				shift
			else
				echo "$(basename ${0}): Error - '--media' expects an argument"
				usage
				exit 1
			fi
			;;
		--add-copr)
			ADD_COPR=1
			;;
		--addkey*)
			if echo $1 | grep '=' >/dev/null ; then
				SSH_KEY=$(echo $1 | sed 's/^--addkey=//')
			elif [ -n "$2" ]; then
				SSH_KEY=$2
				shift
			else
				echo "$(basename ${0}): Error - '--addkey' expects an argument"
				usage
				exit 1
			fi
			# Explicitly expand the tilde if it's the first character
			if [ "$(echo "$SSH_KEY" | cut -c1)" = "~" ]; then
				SSH_KEY="${HOME}${SSH_KEY#~}"
			fi
			;;
		--selinux*)
			if echo $1 | grep '=' >/dev/null ; then
				SELINUX=$(echo $1 | sed 's/^--selinux=//')
			elif [ -n "$2" ]; then
				SELINUX=$2
				shift
			else
				echo "$(basename ${0}): Error - '--selinux' expects an argument"
				usage
				exit 1
			fi
			;;
		--args*)
			if echo $1 | grep '=' >/dev/null ; then
				OPT_ARGS=$(echo $1 | sed 's/^--args=//')
			elif [ -n "$2" ]; then
				OPT_ARGS=$2
				shift
			else
				echo "$(basename ${0}): Error - '--args' expects an argument"
				usage
				exit 1
			fi
			;;
		--wifi-ssid*)
			if echo $1 | grep '=' >/dev/null ; then
				WIFI_SSID=$(echo $1 | sed 's/^--wifi-ssid=//')
			elif [ -n "$2" ]; then
				WIFI_SSID=$2
				shift
			else
				echo "$(basename ${0}): Error - '--wifi-ssid' expects an argument"
				usage
				exit 1
			fi
			;;
		--wifi-pass*)
			if echo $1 | grep '=' >/dev/null ; then
				WIFI_PASS=$(echo $1 | sed 's/^--wifi-pass=//')
			elif [ -n "$2" ]; then
				WIFI_PASS=$2
				shift
			else
				echo "$(basename ${0}): Error - '--wifi-pass' expects an argument"
				usage
				exit 1
			fi
			;;
		--wifi-security*)
			if echo $1 | grep '=' >/dev/null ; then
				WIFI_SECURITY=$(echo $1 | sed 's/^--wifi-security=//')
			elif [ -n "$2" ]; then
				WIFI_SECURITY=$2
				shift
			else
				echo "$(basename ${0}): Error - '--wifi-security' expects an argument"
				usage
				exit 1
			fi
			# Validate the security type
			if [[ "$WIFI_SECURITY" != "wpa-psk" && "$WIFI_SECURITY" != "sae" ]]; then
				echo "$(basename ${0}): Error - Invalid Wi-Fi security type: $WIFI_SECURITY"
				usage
				exit 1
			fi
			;;
		--ign-url*)
                        if echo $1 | grep '=' >/dev/null ; then
                                IGN_URL=$(echo $1 | sed 's/^--ign-url=//')
                        elif [ -n "$2" ]; then
                                IGN_URL=$2
                                shift
                        else
                                echo "$(basename ${0}): Error - '--ign-url' expects an argument"
                                usage
                                exit 1
                        fi
                        ;;
		--ignition*)
                        if echo $1 | grep '=' >/dev/null ; then
                                IGNITION=$(echo $1 | sed 's/^--ignition=//')
                        elif [ -n "$2" ]; then
                                IGNITION=$2
                                shift
                        else
                                echo "$(basename ${0}): Error - '--ignition' expects an argument"
                                usage
                                exit 1
                        fi
			# Explicitly expand the tilde if it's the first character
                        if [ "$(echo "$IGNITION" | cut -c1)" = "~" ]; then
                                IGNITION="${HOME}${IGNITION#~}"
                        fi
                        ;;
		--norootpass)
			NOROOTPASS=1
			;;
		--sysrq)
			SYSRQ=1
			;;
		--resizefs)
			RESIZEFS=1
			;;
		--addconsole)
			CONSOLE=1
			;;
		--blacklistvc4)
			FIX_RPI=1
			;;
		--supported)
			cat $DOC_DIR/SUPPORTED-BOARDS
			exit 0
			;;
		--relabel)
			RELABEL=1
			;;
		--showboot)
			SHOWBOOT=1
			;;
		--version)
			echo "$(basename ${0})-$VERSION"
			exit 0
			;;
		-y)
			NOASK=1
			;;
		*)
			echo "$(basename ${0}): Error - ${1}"
			usage
			exit 1
			;;
	esac
	shift
done

contains() {
	string="$1"
	substring="$2"
	if test "${string#*$substring}" != "$string"; then
		return 0    # $substring is in $string
	else
		return 1    # $substring is not in $string
	fi
}

# ensure sudo user
if [ "$(whoami)" != "root" ]; then
	echo "Error: This script requires 'sudo' privileges in order to write to disk & mount media."
	exit 1
fi

# check to make sure populated
if [ "$MEDIA" = "" ]; then
	usage
	exit 1
fi

ROOTDISK="$(mount | grep "on / " | awk '{printf $1"\n"}')"
case "$ROOTDISK" in  
  *nvme*)
    ROOTDISK="$(echo $ROOTDISK | head -c 10)"
    ;;
  *sd*)
    ROOTDISK="$(echo $ROOTDISK | head -c 8)"
    ;;
esac

if [ "$MEDIA" = "$ROOTDISK" ]; then
	echo ""
	echo " ***********************************************************"
	echo " ** WARNING: You have requested the image be written to sda."
	echo " ** $ROOTDISK is mounted as the root filesystem of the host."
	echo " ***********************************************************"
	echo " ** Do you wish to continue? (type 'yes' to continue)"
	echo " ***********************************************************"
	# wait for agreement
	printf " = Continue? "
	read AGREE
	if [ "$(echo ${AGREE} | tr '[:lower:]' '[:upper:]')" != "YES" ]; then
		echo "User exit, no image written."
		exit 0
	fi
fi

# check for boards
if [ "$TARGET" != "" ] && ! [ -e "${BOARDDIR}/${TARGET}" ]; then
	# If --add-copr is set, check if it's an rk35xx board
	if [ "$ADD_COPR" = "1" ]; then
		TARGET_SOC=$(echo "$TARGET" | grep -oE 'rk35[0-9]{2}')
		if [ -n "$TARGET_SOC" ]; then
			# rk35xx board with --add-copr, will validate SoC later
			:
		else
			echo "Error: You must choose a supported board or none at all."
			echo "With --add-copr, only rk35xx boards are supported."
			usage
			exit 1
		fi
	else
		echo "Error: You must choose a supported board or none at all."
		usage
		exit 1
	fi
fi

# image exists
if [ ! -f "$IMAGE" ] && [ "$IMAGE" != "" ]; then
	echo "Error: $IMAGE not found! Please choose an existing image."
	exit 1
fi

# device or file exists
if [ ! -e "$MEDIA" ]; then
	echo "Error: $MEDIA not found! Please choose an existing device or file."
	exit 1
fi

# If MEDIA is a regular file, set up a loopback device
LOOP_DEVICE=""
if [ -f "$MEDIA" ] && [ ! -b "$MEDIA" ]; then
	# Check if file exists and has content, or if we need to create it
	if [ ! -s "$MEDIA" ] && [ "$IMAGE" != "" ]; then
		# File is empty or doesn't exist - we'll write an image to it
		# No need to set up loop device yet, we'll do it after writing
		MEDIA_FILE="$MEDIA"
		echo "= Will create disk image file: $MEDIA"
	else
		# File already has content - set up loop device now
		echo "====================================================="
		echo "====================================================="
		echo "= Setting up loopback device for $MEDIA ..."
		LOOP_DEVICE=$(losetup --show --partscan -f "$MEDIA")
		if [ $? -ne 0 ] || [ -z "$LOOP_DEVICE" ]; then
			echo "Error: Failed to set up loopback device for $MEDIA"
			exit 1
		fi
		echo "= Loopback device created: $LOOP_DEVICE"
		# Save original MEDIA path and use loop device for all operations
		MEDIA_FILE="$MEDIA"
		MEDIA="$LOOP_DEVICE"
	fi
fi

# check only one ignition source was added
if [ -n "$IGNITION" ] && [ -n "$IGN_URL" ]; then
        echo "Error: Please pick one ignition source, local or url."
        exit 1
fi

# Validate COPR board support if requested
if [ "$ADD_COPR" = "1" ]; then
	# List of supported boards for COPR (rk35xx series: rk3566, rk3568, rk3576, rk3588)
	COPR_SUPPORTED_BOARDS="rk3566 rk3568 rk3576 rk3588"

	# Extract the SoC from the target name (e.g., rock5b-rk3588 -> rk3588)
	TARGET_SOC=$(echo "$TARGET" | grep -oE 'rk35[0-9]{2}')

	if [ -z "$TARGET_SOC" ]; then
		echo "Error: --add-copr requires a target board with rk35xx SoC."
		echo "Supported: rk3566, rk3568, rk3576, rk3588 based boards"
		echo "Your target: $TARGET"
		exit 1
	fi

	# Check if the SoC is in the supported list
	if ! echo "$COPR_SUPPORTED_BOARDS" | grep -qw "$TARGET_SOC"; then
		echo "Error: Target SoC $TARGET_SOC is not supported by COPR."
		echo "Supported SoCs: $COPR_SUPPORTED_BOARDS"
		exit 1
	fi
fi

# Last chance to back out
echo "====================================================="

# check to see if host system uses LVM and VG named fedora for root
# if writing the aarch64 server image rename to fedora-server
if [ -b /dev/fedora/root ] && [ "$(echo $IMAGE|awk '/aarch64/ && /Server/')" != "" ]; then
	echo "**************************************************"
	echo "= NOTE: This host system uses the same VG name as "
	echo "= the AArch64 disk image. To avoid issues, the VG "
	echo "= on the image will be renamed to 'fedora-server'."
	echo "**************************************************"
	RENAME_LVM="1"
fi

# Image if included
if [ "$IMAGE" != "" ]; then
	echo "= Selected Image:                                 "
	echo "= $IMAGE"
fi
echo "= Selected Media : $MEDIA"
# target hardware platform
if [ "$TARGET" != "" ]; then
	echo "= U-Boot Target : $TARGET"
fi
# SE Linux On/Off
if [ "$SELINUX" != "" ]; then
	echo "= SELINUX = $SELINUX"
fi
if [ "$RELABEL" != "" ]; then
	echo "= SELinux relabel will be completed on first boot."
fi
# Empty the root password
if [ "$NOROOTPASS" != "" ]; then
	echo "= Root Password will be set to the empty string."
fi
# Enable System Request debugging of the kernel
if [ "$SYSRQ" != "" ]; then
	echo "= System Request debugging of the kernel will be enabled."
fi
# Resize root filesystem to fill media device
if [ "$RESIZEFS" != "" ]; then
	echo "= Root partition will be resized"
fi
# Console to be added
if [ "$CONSOLE" != "" ]; then
	# don't add a console param for target none
	if echo "$TARGET" | grep -q 'none'; then
		CONSOLE=0
	else
		echo "= Console for $TARGET will be added."
	fi
fi
# wifi ssid to be added
if [ "$WIFI_SSID" != "" ]; then
	echo "= Wifi ssid: $WIFI_SSID will autoconnect on boot."
fi
# User ssh key
if [ "$SSH_KEY" != "" ]; then
	echo "= SSH Public Key $SSH_KEY will be added."
fi
# show boot messages
if [ "$SHOWBOOT" != "" ]; then
        echo "= Boot messages will be shown onscreen."
fi
if [ "$IGN_URL" != "" ]; then
        echo "= Ignition URL: $IGN_URL "
fi
if [ "$IGNITION" != "" ]; then
	echo "= Ignition: $IGNITION "
fi
# COPR repository warning
if [ "$ADD_COPR" != "" ]; then
	echo "= COPR repository (pbrobinson/u-boot) will be enabled on this system"
fi
echo "= Version: $VERSION"
echo "====================================================="
echo " "
echo "*****************************************************"
echo "*****************************************************"
echo "******** WARNING! ALL DATA WILL BE DESTROYED ********"
echo "*****************************************************"
echo "*****************************************************"
if [ "$NOASK" != 1 ]; then
	echo " "
	echo " Type 'YES' to proceed, anything else to exit now "
	echo " "
	# wait for agreement
	printf "= Proceed? "
	read PROCEED
	if [ "$(echo ${PROCEED} | tr '[:lower:]' '[:upper:]')" != "YES" ]; then
		echo "User exit, no image written."
		exit 0
	fi
fi

# Install COPR repository and arm-images-copr package if requested
if [ "$ADD_COPR" = "1" ]; then
	# Check if COPR repo is already enabled
	if dnf repolist --enabled | grep -q 'copr.*pbrobinson.*u-boot'; then
		echo "= COPR repository pbrobinson/u-boot already enabled"
	else
		echo "= Enabling COPR repository: pbrobinson/u-boot"
		# Enable COPR repo on local system (use fedora-44-aarch64 explicitly)
		if ! dnf copr enable -y pbrobinson/u-boot fedora-44-aarch64; then
			echo "Error: Failed to enable COPR repository"
			echo "Please ensure dnf-plugins-core is installed: dnf install dnf-plugins-core"
			exit 1
		fi
	fi

	# Check if uboot-images-copr is already installed
	if rpm -q uboot-images-copr >/dev/null 2>&1; then
		echo "= Package uboot-images-copr already installed"
	else
		echo "= Installing uboot-images-copr package..."
		if ! dnf install -y uboot-images-copr; then
			echo "Error: Failed to install uboot-images-copr package"
			exit 1
		fi
	fi

	echo "= Using u-boot from COPR (installed on host system)"
fi

# umount before starting
umount $MEDIA* > /dev/null 2>&1
# Note: we can't deactivate LVM here because ROOTPART isn't set yet
# The partition table hasn't been read at this point

# Write the disk image to media
if [ "$IMAGE" != "" ]; then
	echo "= Writing: "
	echo "= $IMAGE "
	echo "= To: $MEDIA ...."
	if [ -n "$MEDIA_FILE" ]; then
		# Writing to a file or loop device backed by a file - don't use oflag=direct
		# (direct I/O doesn't work well with sparse files and loop devices)
		xzcat $IMAGE | dd of=$MEDIA bs=4M status=progress iflag=fullblock; sync; sleep 3
		echo "= Writing image complete!"
		# Set up loop device if we haven't already
		if [ -z "$LOOP_DEVICE" ]; then
			echo "= Setting up loopback device for $MEDIA_FILE ..."
			LOOP_DEVICE=$(losetup --show --partscan -f "$MEDIA_FILE")
			if [ $? -ne 0 ] || [ -z "$LOOP_DEVICE" ]; then
				echo "Error: Failed to set up loopback device for $MEDIA_FILE"
				exit 1
			fi
			echo "= Loopback device created: $LOOP_DEVICE"
			MEDIA="$LOOP_DEVICE"
		fi
	else
		# Writing to a real block device - use oflag=direct for better performance
		xzcat $IMAGE | dd of=$MEDIA oflag=direct bs=4M status=progress iflag=fullblock; sync; sleep 3
		echo "= Writing image complete!"
	fi
fi
# check to see how many partitions on the image
partprobe "$MEDIA"
sleep 2

get_lvm_name () {
	# Try to get VG name directly from the root partition using pvs
	# This works because ROOTPART is set to the LVM physical volume partition
	# Note: --devices must include both the base device AND the partition
	LVM_NAME=$(pvs --devicesfile "" --devices "$MEDIA" --devices "$ROOTPART" -o vg_name --noheadings "$ROOTPART" 2>/dev/null | tr -d ' ')
}

get_lv_name () {
	# Get the LV name from the VG - assumes there's only one LV in the VG
	# Note: --devices must include both the base device AND the partition
	LV_NAME=$(lvs --devicesfile "" --devices "$MEDIA" --devices "$ROOTPART" -o lv_name --noheadings "$LVM_NAME" 2>/dev/null | tr -d ' ')
}

add_bls_parameter()
{
	for bls in /tmp/boot/loader/entries/*.conf; do
		sed -i "s|options|& $1|" ${bls}
	done
}

add_kernel_parameter () {
	if [ -f /tmp/boot/extlinux/extlinux.conf ]; then
		sed -i "s|append|& $1 |" /tmp/boot/extlinux/extlinux.conf
	elif [ -f /tmp/fw/EFI/fedora/grub.cfg ]; then
		sed -i "s|GRUB_CMDLINE_LINUX=\"|& $1 |" ${PREFIX}/etc/default/grub
		if [ -f /tmp/fw/EFI/fedora/grubenv ]; then
			sed -i "s|kernelopts=|& $1 |" /tmp/fw/EFI/fedora/grubenv
		else
			add_bls_parameter "$1"
		fi
	fi
}

# Detect partitions - check with 'p' separator first (nvme, loop), then without (sd, mmcblk)
if [ -e "$MEDIA"p5 ]; then
	export FIRMPART="${MEDIA}p1"
	BOOTPART="${MEDIA}p2"
	ROOTPART="${MEDIA}p5"
	PARTNUM=5
elif [ -e "$MEDIA"p4 ]; then
	export FIRMPART="${MEDIA}p1"
	BOOTPART="${MEDIA}p2"
	ROOTPART="${MEDIA}p4"
	PARTNUM=4
elif [ -e "$MEDIA"p3 ]; then
	export FIRMPART="${MEDIA}p1"
	BOOTPART="${MEDIA}p2"
	ROOTPART="${MEDIA}p3"
	PARTNUM=3
elif [ -e "$MEDIA"5 ]; then
	export FIRMPART="${MEDIA}1"
	BOOTPART="${MEDIA}2"
	ROOTPART="${MEDIA}5"
	PARTNUM=5
elif [ -e "$MEDIA"4 ]; then
	export FIRMPART="${MEDIA}1"
	BOOTPART="${MEDIA}2"
	ROOTPART="${MEDIA}4"
	PARTNUM=4
else
	export FIRMPART="${MEDIA}1"
	BOOTPART="${MEDIA}2"
	ROOTPART="${MEDIA}3"
	PARTNUM=3
fi

FS_TYPE=$(blkid "$ROOTPART" -s TYPE -o value)

if [ "$FS_TYPE" = "LVM2_member" ]; then FS_TYPE="xfs"; fi

if [ "$FS_TYPE" = "btrfs" ]; then
	BTRFS=1
fi

# resize root filesystem before mounting
if [ "$RESIZEFS" != "" ]; then
	echo "= Resizing $MEDIA ...."
	sync
	count=0
	if [ "$PARTNUM" = "4" ] || [ "$PARTNUM" = "5" ]; then
		echo ", +" | sfdisk -N "4" "$MEDIA"
		while [ $? != '0' ]; do
			sleep 5
			echo ", +" | sfdisk -N "4" "$MEDIA"
			count=$((count + 1))
			if [ $count -gt 5 ]; then
				echo "= Partition Resize Failed."
				continue
			fi
		done
		partprobe "$MEDIA"
		if [ "$PARTNUM" = "5" ]; then
			echo ", +" | sfdisk -N "$PARTNUM" "$MEDIA"
			while [ $? != '0' ]; do
				sleep 5
				echo ", +" | sfdisk -N "$PARTNUM" "$MEDIA"
				count=$((count + 1))
				if [ $count -gt 5 ]; then
					echo "= Partition Resize Failed."
					continue
				fi
			done
		fi
		partprobe "$MEDIA"
	fi

	if [ "$PARTNUM" = "3" ]; then
		if [ "$RENAME_LVM" = "1" ]; then
		# rename VG if host system uses conflicting names
			LVM_UUID=$(pvs --devicesfile "" --devices "$MEDIA" --devices "$ROOTPART" -o +vguuid | grep $MEDIA | awk '{print $7}')
			vgrename --devicesfile "" --devices "$MEDIA" --devices "$ROOTPART" $LVM_UUID fedora-server
			LVM_NAME=fedora-server
		fi

		get_lvm_name
		vgchange --devicesfile "" --devices "$MEDIA" --devices "$ROOTPART" -a n $LVM_NAME > /dev/null 2>&1

		echo ", +" | sfdisk -N "$PARTNUM" "$MEDIA"
		while [ $? != '0' ]; do
			sleep 5
			echo ", +" | sfdisk -N "$PARTNUM" "$MEDIA"
			count=$((count + 1))
			if [ $count -gt 5 ]; then
				echo "= Partition Resize Failed."
				continue
			fi
		done
		sleep 5
		partprobe "$MEDIA"
		if [ "$LVM_NAME" != "" ]; then
			vgchange --devicesfile "" --devices "$MEDIA" --devices "$ROOTPART" -a y $LVM_NAME > /dev/null 2>&1
			pvresize --devicesfile "" --devices "$MEDIA" --devices "$ROOTPART" "$ROOTPART"
			get_lv_name
			lvextend --devicesfile "" --devices "$MEDIA" --devices "$ROOTPART" -l +100%FREE /dev/$LVM_NAME/$LV_NAME
			ROOTLV="/dev/$LVM_NAME/$LV_NAME"
		fi
	fi

	if [ "$FS_TYPE" = "xfs" ] && [ "$LVM_NAME" != "" ]; then
		mkdir /tmp/root > /dev/null 2>&1
		mount "$ROOTLV" /tmp/root > /dev/null 2>&1
		xfs_growfs /tmp/root
	elif [ "$FS_TYPE" = "btrfs" ]; then
		mkdir /tmp/root > /dev/null 2>&1
		mount "$ROOTPART" /tmp/root > /dev/null 2>&1
		btrfs filesystem resize max /tmp/root
	elif [ "$FS_TYPE" = "ext4" ]; then
		fsck.ext4 -fy "$ROOTPART"
		resize2fs "$ROOTPART"
	fi
fi

sleep 5

get_lvm_name

if [ "$RENAME_LVM" = "1" ] && [ "$LVM_NAME" != "fedora-server" ]; then
# rename VG if host system uses conflicting names
  LVM_UUID=$(pvs --devicesfile "" --devices "$MEDIA" --devices "$ROOTPART" -o +vguuid | grep $MEDIA | awk '{print $7}')
  vgrename --devicesfile "" --devices "$MEDIA" --devices "$ROOTPART" $LVM_UUID fedora-server
  LVM_NAME=fedora-server
fi


# make temp mount points
mkdir /tmp/boot /tmp/root /tmp/fw > /dev/null 2>&1
mount "$BOOTPART" /tmp/boot > /dev/null 2>&1
if [ $? -ne 0 ]; then
	echo "Error: mount $BOOTPART /tmp/boot failed"
	exit 1
fi
mount "$FIRMPART" /tmp/fw
if [ $? -ne 0 ]; then
	echo "Error: mount $FIRMPART /tmp/fw failed"
	exit 1
fi

if [ "$RESIZEFS" = "" ]; then
	get_lvm_name
	vgchange --devicesfile "" --devices "$MEDIA" --devices "$ROOTPART" -a y $LVM_NAME > /dev/null 2>&1
fi

if [ "$(grep /tmp/root /proc/mounts)" = "" ]; then
	if [ "$LVM_NAME" != "" ]; then
		get_lv_name
		mount "/dev/$LVM_NAME/$LV_NAME" /tmp/root > /dev/null 2>&1
	else
        	mount "$ROOTPART" /tmp/root
	fi
fi

if [ "$(echo $IMAGE | grep IoT)" != "" ]; then
	IOT_IMAGE="1"
	OSTREE_ROOT_HOME="/tmp/root/ostree/deploy/fedora-iot/var/roothome"
	OSTREE_PREFIX="/tmp/root/ostree/deploy/fedora-iot/deploy/*/"
fi

# fix up grub.cfg to reflect the new vg name
if [ "$RENAME_LVM" != "" ]; then
	if [ -d /tmp/boot/loader/entries ] && [ -f /tmp/boot/grub2/grub.cfg ]; then
		sed -i 's|/dev/mapper/fedora-root|/dev/mapper/fedora--server-root|g; s|rd.lvm.lv=fedora/root|rd.lvm.lv=fedora-server/root|g' /tmp/boot/loader/entries/*.conf
		sed -i 's|/dev/mapper/fedora-root|/dev/mapper/fedora--server-root|g; s|rd.lvm.lv=fedora/root|rd.lvm.lv=fedora-server/root|g' /tmp/boot/grub2/grub.cfg
	fi
fi

if [ "$IOT_IMAGE" = "1" ]; then
	# dd doesnt support wildcards, echo to expand
	PREFIX=$(echo $OSTREE_PREFIX)
elif [ "$BTRFS" = "1" ]; then
	PREFIX="/tmp/root/root/"
else
	PREFIX=/tmp/root
fi

# determine uboot and write to disk
if [ "$TARGET" != "" ]; then
	# IMAGE_ROOT always points to mounted image, PREFIX is temporarily set for u-boot sourcing
	IMAGE_ROOT="$PREFIX"

	# Use COPR files from host system if requested, otherwise use image files
	if [ "$ADD_COPR" = "1" ]; then
		UBOOT_PREFIX=""
	else
		UBOOT_PREFIX="$IMAGE_ROOT"
	fi

	if echo "$TARGET" | grep -q 'rpi[34]' || [ "$TARGET" = "beagleplay" ]; then
		PREFIX="$UBOOT_PREFIX" . "${BOARDDIR}/${TARGET}"
	elif [ -d "${UBOOT_PREFIX}/usr/share/uboot/${TARGET}" ]; then
		PREFIX="$UBOOT_PREFIX" . "${BOARDDIR}/${TARGET}"
	else
		echo "= No U-Boot files found for $TARGET."
	fi

	# Restore PREFIX to image mount point for remaining operations
	PREFIX="$IMAGE_ROOT"
else
	echo "= No U-boot will be written."
	TARGET="board"
fi

# turn off selinux
if [ "$SELINUX" != "" ]; then
	if [ "$(echo ${SELINUX} | tr '[:lower:]' '[:upper:]')" = "OFF" ]; then
		echo "= Turning SELinux off ..."
		sed -i 's/SELINUX=enforcing/SELINUX=permissive/' ${PREFIX}/etc/selinux/config
		# turn on selinux
	elif [ "$(echo ${SELINUX} | tr '[:lower:]' '[:upper:]')" = "ON" ]; then
		echo "= Turning SELinux on ..."
		sed -i 's/SELINUX=permissive/SELINUX=enforcing/' ${PREFIX}/etc/selinux/config
	fi
fi
# Empty the root password
if [ "$NOROOTPASS" = "1" ]; then
	echo "= Setting the root password to the empty string."
	sed -i 's/root:x:/root::/' ${PREFIX}/etc/passwd
fi
# Enable System Request debugging of the kernel
if [ "$SYSRQ" != "" ]; then
	echo "= Enabling System Request debugging of the kernel."
	cat >> ${PREFIX}/etc/sysctl.d/arm-image-installer-sysrq.conf <<-EOH
		# Controls the System Request debugging functionality of the kernel
		kernel.sysrq = 1
	EOH
fi
# Add ssh key to the image
if [ "$SSH_KEY" != "" ]; then
	if [ -f $SSH_KEY ]; then
		echo "= Adding SSH key to authorized keys."
		if [ "$IOT_IMAGE" = "1" ]; then
			mkdir $OSTREE_ROOT_HOME/.ssh/ > /dev/null 2>&1
			echo "# ssh key added by arm-image-installer" >> $OSTREE_ROOT_HOME/.ssh/authorized_keys
			cat $SSH_KEY >> $OSTREE_ROOT_HOME/.ssh/authorized_keys
			echo "# end arm-image-installer key" >> $OSTREE_ROOT_HOME/.ssh/authorized_keys
			chmod -R u=rwX,o=,g= $OSTREE_ROOT_HOME/.ssh/
		else
			mkdir ${PREFIX}/root/.ssh/ > /dev/null 2>&1
			echo "# ssh key added by arm-image-installer" >> ${PREFIX}/root/.ssh/authorized_keys
			cat $SSH_KEY >> ${PREFIX}/root/.ssh/authorized_keys
			echo "# end arm-image-installer key" >> ${PREFIX}/root/.ssh/authorized_keys
			chmod -R u=rwX,o=,g= ${PREFIX}/root/.ssh/
		fi
	else
		echo "= SSH key $SSH_KEY : Not Found!"
		echo "= WARNING: No SSH Key Added."
	fi
fi

# Configure Wi-Fi if SSID is provided
if [ "$WIFI_SSID" != "" ]; then
    echo "= Configuring Wi-Fi: $WIFI_SSID"
    # Use the SSID as the filename (replace spaces and special characters with underscores)
    FILENAME=$(echo "$WIFI_SSID" | tr ' ' '_' | tr -cd '[:alnum:]._-').nmconnection
    cat > ${PREFIX}/etc/NetworkManager/system-connections/$FILENAME <<-EOF
[connection]
id=${WIFI_SSID}
type=wifi

[wifi]
mode=infrastructure
ssid=${WIFI_SSID}
EOF

    # Add security configuration if a password is provided
    if [ -n "$WIFI_PASS" ]; then
        cat >> ${PREFIX}/etc/NetworkManager/system-connections/$FILENAME <<-EOF
[wifi-security]
key-mgmt=${WIFI_SECURITY:-wpa-psk}
psk=${WIFI_PASS}
EOF
    fi

    # Add common IPv4 and IPv6 settings
    cat >> ${PREFIX}/etc/NetworkManager/system-connections/$FILENAME <<-EOF
[ipv4]
method=auto

[ipv6]
method=auto
addr-gen-mode=default
EOF

    # Set file permissions
    chmod 600 ${PREFIX}/etc/NetworkManager/system-connections/$FILENAME
fi

# Add console
if [ "$CONSOLE" = "1" ]; then
	if [ "$SYSCON" = "" ]; then
		SYSCON="ttyS0,115200"
		echo "= No console listed for $TARGET, adding default $SYSCON ."
	fi

	echo "= Adding console $SYSCON to kernel parameters ..."
	add_kernel_parameter "console=$SYSCON console=tty0"

	if echo "$TARGET" | grep -q 'rpi[34]'; then
		sed -i "s|# enable_uart=1|enable_uart=1|" /tmp/fw/config.txt
	fi
fi

# check if host system has selinux disabled, if it does autorelabel is required
if [ "$(getenforce)" = "Disabled" ]; then
	echo "= NOTE: System Relabel required on first boot."
	RELABEL="1"
fi

# touch /.autorelabel
if [ "$RELABEL" != "" ]; then
	if [ "$IOT_IMAGE" = "1" ]; then
		echo "= SELinux relabel not supported on IoT images."
	else
		echo "= Touch /.autorelabel on rootfs."
		touch ${PREFIX}/.autorelabel
	fi
fi
# add option kernel parameters
if [ "$OPT_ARGS" != "" ] ; then
	echo "= Adding optional kernel parameters for $TARGET : "
	echo "= Parameter: $OPT_ARGS"
	add_kernel_parameter "$OPT_ARGS"
fi
# remove quiet from kargs
if [ "$SHOWBOOT" != "" ]; then
		 sed -i 's|rhgb quiet ||g' /tmp/boot/loader/entries/*.conf
fi

if [ -n "$IGN_URL" ]; then
  if echo "$IGN_URL" | grep -q "://"; then
    echo "= Adding ignition url for configuraion on firstboot."
    sed -i "s|true|true ignition.firstboot=1 ignition.config.url=${IGN_URL}|g" /tmp/boot/ignition.firstboot
  else
    echo "Error: \"$IGN_URL\" is not a valid url."
    exit 1
  fi
elif [ -n "$IGNITION" ]; then
  if [ -f "$IGNITION" ]; then
    echo "= Adding ignition file for configuraion on firstboot."
    mkdir -p /tmp/boot/ignition
    cp "$IGNITION" /tmp/boot/ignition/config.ign
    sed -i "s|true|true ignition.firstboot=1 ignition.config.file=/ignition/config.ign|g" /tmp/boot/ignition.firstboot
  else
    echo "Error: Local Ignition file '$IGNITION' not found!"
    exit 1
  fi
fi

sync

umount /tmp/root $BOOTPART $FIRMPART > /dev/null 2>&1
if [ "$LVM_NAME" != "" ]; then
	vgchange --devicesfile "" --devices "$MEDIA" --devices "$ROOTPART" -a n $LVM_NAME > /dev/null 2>&1
fi
rmdir  /tmp/root /tmp/boot /tmp/fw > /dev/null 2>&1

# Clean up loopback device if we created one
if [ -n "$LOOP_DEVICE" ]; then
	echo "= Detaching loopback device $LOOP_DEVICE ..."
	losetup -d "$LOOP_DEVICE"
	if [ $? -eq 0 ]; then
		echo "= Loopback device detached successfully"
		echo "= Custom image created: $MEDIA_FILE"
	else
		echo "= Warning: Failed to detach loopback device $LOOP_DEVICE"
		echo "= You may need to manually detach it with: sudo losetup -d $LOOP_DEVICE"
	fi
fi

if [ "$URL" != "" ]; then
	echo
	echo "= NOTE"
	echo "= Additional instructions for $TARGET can be found at:"
	echo "= $URL "
	echo
fi

echo ""
if [ -n "$LOOP_DEVICE" ]; then
	echo "= Installation Complete! Custom image ready: $MEDIA_FILE"
else
	echo "= Installation Complete! Insert into the $TARGET and boot."
fi
exit 0
