#! /bin/bash -e

#
# rpi-lvm-imager.sh - initialize NVME with Raspberry Pi OS (Lite) on LVM
#
# Version 1.0, latest version, documentation and bugtracker available at:
#               https://gitea.lindenaar.net/scripts/raspberrypi
#
# Copyright (c) 2025 Frederik Lindenaar
#
# This script is free software: you can redistribute and/or modify it under the
# terms of version 3 of the GNU General Public License as published by the Free
# Software Foundation, or (at your option) any later version of the license.
#
# This script is distributed in the hope that it will be useful but WITHOUT ANY
# WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS FOR
# A PARTICULAR PURPOSE.  See the GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License along with
# this program.  If not, visit <http://www.gnu.org/licenses/> to download it.
#

# Inspired by: https://raspberrypi.stackexchange.com/questions/7159/can-the-raspberry-boot-to-an-lvm-root-partition

# Fetch the list of Raspberry Pi OS images, if required, and flatten it
# next return the url to download the specified image or list of images
PIOS_LIST_DOWNLOAD_URL=https://downloads.raspberrypi.com/os_list_imagingutility_v4.json
get_raspios() {
	# Download the Raspberry Pi Imager catalog (if a newer version exists)
	wget --quiet --timestamping --directory-prefix "${WORKDIR}" "${PIOS_LIST_DOWNLOAD_URL}" 
	RPI_IMAGER_CONFIG="${WORKDIR}/$(basename "${PIOS_LIST_DOWNLOAD_URL}" )"

	# Extract and flatten the list of entries, if required
	PIOS_LIST_FILE="${WORKDIR}/os_list.json"
	if [ "${RPI_IMAGER_CONFIG}" -nt "${PIOS_LIST_FILE}" ]; then
		jq '.os_list 
        | map(select( .url )) 
        + map(select( .subitems ).subitems 
                  | map(select( .url )) 
                  + map(select( .subitems ).subitems 
                           | map(select( .url )) 
                           + map(select( .subitems ).subitems
                                    | map(select( .url )) 
                                    + map(select( .subitems ).subitems) 
                                    | flatten)
                           | flatten)
                  | flatten)
        | flatten' "${RPI_IMAGER_CONFIG}" > "${PIOS_LIST_FILE}"
        if fgrep -q '"subitems"' "${PIOS_LIST_FILE}"; then
        	echo >&2 Warning: structure of the Raspberry Pi Imager catalog has changed
        fi
	fi

	if [ -z "$1" -o "$1" = list ]; then
	    jq --raw-output "map(.name ) | .[]" "${PIOS_LIST_FILE}" | sort -u
	else
		echo $(jq --raw-output "map(select( .name == \"${1}\"))[0] | .url,.extract_size,.extract_sha256" "${PIOS_LIST_FILE}")
	fi
}

# defaults for command line options
opt_raspios_lite=true
opt_raspios_bits=64
opt_raspios_image=
opt_raspios_checksum=
opt_check_image_checksum=true
opt_raspios_checksum=
opt_lvm_group_name=
opt_lvm_size=+
opt_boot_size=
opt_root_size=16G
opt_lvm_partitions=
opt_nvme_pcie_gen3=false

# parse command line parameters
while [ $# -gt 0 ]
do
	case "$1" in
        -h|--help)           cat << EOT
Raspberry Pi LVM Imager - write Raspberry Pi OS image to an LVM-enabled device

Usage: $(baesename $0) [options] device

Where <<options>> can be one or more of:

  -h, --help           show this help and exit

  -l, --list-images    output the list of available Raspberry Pi images and exit
  -L, --lite           install Raspberry Pi OS Lite (default)
  -R, --regular        install the regular Raspberry Pi OS (with destktop)
  -6, --64bit          install a 64-bit version of Raspberry Pi OS
  -3, --32bit          install a 32-bit version of Raspberry Pi OS
  -i, --image-name     name of the image to install 
  -U, --url            provide the download URL for Raspberry Pi OS
                       see https://www.raspberrypi.com/software/operating-systems
  -C, --checksum       provide sha256 checksum of image
  -n, --no-checksum    disable checking the sha256 checksum of the downloaded image

  -g, --lvm-group      name of the LVM volume group to create (required and must not exist!)
  -s, --lvm-size       size of the LVM volume group to create (defaults to the rest of the disk)
  -b, --boot-size      size of the boot filesystem (defaults to the size in the image)
  -r, --root-size      size of the root filesystem (defaults to 16G)
  -p, --lvm-partition  comma separated /path:size list to move to separate LVMs

  -P, --nvme-pcie-gen3 enable PCIe Gen 3 mode for the NVMe

  -w, --workdir        (persistent) work directory to store files, not cleaned up
EOT
                             exit
					         ;;
        -l|--list-images)    opt_raspios_image=list
                             ;;
        -i|--image-name)     opt_raspios_image="$2"
                             shift
                             ;;
        -L|--lite)           opt_raspios_lite=true
                             ;;
        -R|--regular)        opt_raspios_lite=false
                             ;;
        -3|--32-bit)         opt_raspios_bits=32
                             ;;
        -6|--64-bit)         opt_raspios_bits=64
                             ;;
        -U|--url)            IMAGE_DOWNLOAD_URL="$2"
                             shift
                             ;;
        -C|--checksum)       opt_raspios_checksum="$2"
                             shift
                             ;;
        -n|--no-checksum)    opt_check_image_checksum=false
                             ;;
        -g|--lvm-group)      opt_lvm_group_name="$2"
                             shift
                             ;;
        -s|--lvm-size)       opt_lvm_size="$2"
                             shift
                             ;;
        -b|--boot-size)      opt_boot_size="$2"
                             shift
                             ;;
        -r|--root-size)      opt_root_size="$2"
                             shift
                             ;;
        -p|--lvm-partition)  opt_lvm_partitions="$2"
                             shift
                             ;;
        -P|--nvme-pcie-gen3) opt_nvme_pcie_gen3=true
                             ;;
        -w|--workdir)        WORKDIR="$2"
                             shift
                             ;;
        -*)                  echo Error: unknown option $1, run $0 -h for supported options
                             exit 2
                             ;;
		*)                   if [ -z "$opt_target_device" ]; then
	                             opt_target_device=$1
	                         else
	                             echo Error: only one device is supported, run $0 -h for usage
	                             exit 2
                             fi
                             ;;
	esac
	shift
done

# Ensure the script is run as root
if [ $(id -u) -gt 0 ]; then
	echo "Error: $(basename $0) must be run as root"
	exit 1
fi

# Ensure cleanup is done when we exist, irrespective how
cleanup() {
    echo cleaning up
    # Unmount the partitions under the rootfs in reverse order
    if [ -n "${TARGET_ROOT_MOUNTPOINT}" ]; then
        findmnt --raw --noheadings --output target --submounts "${TARGET_ROOT_MOUNTPOINT}" | sort -r | while read mountpoint
        do 
            umount "${mountpoint}"
        done
    fi

    # Unmount the image from the loopback when it was mounted
    if [ -n "$RASPIOS_DEV" ]; then
        losetup --detach ${RASPIOS_DEV}
    fi

    # Remove the temporary working directory (if created)
    if [ -n "$TEMPDIR" ]; then
        rm -rf $TEMPDIR
    fi
}
trap cleanup EXIT

# Create a temporary working directory, if required
if [ -z "$WORKDIR" ]; then
	TEMPDIR=$(mktemp --directory --tmpdir $(basename $0 .sh).XXXXXXXX)
	WORKDIR=${TEMPDIR}
elif [ ! -d "${WORKDIR}" ]; then
    mkdir "${WORKDIR}"
fi

# If we need to print the list of available images, print it and exit
if [ "${opt_raspios_image}" = "list" ]; then
	get_raspios list
	exit
fi

# Ensure the target is a valid disk
if [ -z "${opt_target_device}" ]; then
    echo Error: target device is required, run $0 -h for usage
    exit 2
elif [ "$(lsblk --nodeps --noheadings  ${opt_target_device} -o type)" != "disk" ]; then
    echo Error: target device ${opt_target_device} is not a disk
    exit 2  
fi

# Ensure no partitions of the target device are still mounted
if [ -n "$(findmnt -D| fgrep "${opt_target_device}")" ]; then
    echo Error: partitions of target device ${opt_target_device} are still mounted:
    findmnt -D | egrep "^SOURCE|${opt_target_device}"
    exit 1
fi

# Ensure the LVM group (if one is to be setup) is provided and does not exist yet
if [ -z "${opt_lvm_group_name}" ]; then
    echo Error: target LVM volume group name is required, run $0 -h for usage
    exit 2
elif [ -n "${opt_lvm_group_name}" -a -n "$(pvdisplay --columns --options pv_name,vg_name --noheading | egrep "${opt_lvm_group_name}|${opt_target_device}")" ]; then
    echo Existing LVM configuration found, cleaning up conflicting setup
    TTY=$(tty)
    pvdisplay --columns --options pv_name,vg_name --noheading | while read dev vg
    do
        if [ "${vg}" = "${opt_lvm_group_name}" ]; then
            wipefs --quiet --all /dev/${vg}/*
            vgremove --quiet ${vg} < ${TTY}
        fi
        if [[ "${dev}" == ${opt_target_device}* ]]; then
            if ! [ "${vg}" != "${opt_lvm_group_name}" ] || vgreduce --quiet ${vg} ${dev} < ${TTY}; then
                pvremove --quiet ${dev} < ${TTY}
            fi
        fi
    done
fi
lvm_dev=

# Ensure the partition table of the taget device is empty
if [ "$(sudo sfdisk --dump "${opt_target_device}" 2>/dev/null )" ]; then
    echo Warning: target device ${opt_target_device} may contain data:
    sfdisk --quiet --list "${opt_target_device}"

    echo
    read -p "erase everything and image ${opt_target_device}? (yes/no) [no]"
    if [ "${REPLY}" != "yes" ]; then
        echo Aborted.
        exit 2
    fi

    echo
    echo Erasing data on "${opt_target_device}"
    wipefs --quiet --all ${opt_target_device}*
    sfdisk --quiet --wipe always --wipe-partitions always --delete "${opt_target_device}" && echo OK || echo NOK
    # dd if=/dev/zero "of=${opt_target_device}" bs=512 count=1 2> /dev/null
fi

echo partitioning target disk
echo 'label: dos' | sfdisk  --quiet "${opt_target_device}"

# Ensure a Raspberry Pi OS download URL and local image name are set
if [ -z "${IMAGE_DOWNLOAD_URL}" ]; then
    if [ -z "${opt_raspios_image}" ]; then
    	if ${opt_raspios_lite}; then
    		opt_raspios_image="Raspberry Pi OS Lite (${opt_raspios_bits}-bit)"
    	else
    		opt_raspios_image="Raspberry Pi OS (${opt_raspios_bits}-bit)"
    	fi
    fi

    read IMAGE_DOWNLOAD_URL IMAGE_EXTRACT_SIZE IMAGE_EXTRACT_SHA256 <<< $(get_raspios "${opt_raspios_image}")
	if [ -z "${IMAGE_DOWNLOAD_URL}" -o "${IMAGE_DOWNLOAD_URL}" = "null" ]; then
		echo "Error: no URL found for image name '${opt_raspios_image}'"
		exit 1
	fi
fi
RASPIOS_IMG="${WORKDIR}/$(basename "${IMAGE_DOWNLOAD_URL}" .xz)"

# Download the image, if required, and/or verify its checksum (if possible)
if [ ! -f "${RASPIOS_IMG}" -o -n "${IMAGE_EXTRACT_SIZE}" -a "$(du -b "${RASPIOS_IMG}" 2>/dev/null | sed "s/[^0-9].*//")" != "${IMAGE_EXTRACT_SIZE}" ]; then
    if $opt_check_image_checksum && [ -n "${opt_raspios_checksum:-$IMAGE_EXTRACT_SHA256}" ]; then
        echo "${opt_raspios_checksum:-$IMAGE_EXTRACT_SHA256} -" > "${RASPIOS_IMG}.$$.sha256sum"
        if ! wget --quiet --show-progress --output-document - "${IMAGE_DOWNLOAD_URL}" | xz --decompress | tee ${RASPIOS_IMG} | sha256sum --check --strict --quiet "${RASPIOS_IMG}.$$.sha256sum"; then
            echo Error: downloading, extracting and verifying the image failed 
            exit 2    
        fi
        rm ${RASPIOS_IMG}.*.sha256sum
    else
        if ! wget --quiet --show-progress --output-document - "${IMAGE_DOWNLOAD_URL}" | xz --decompress > ${RASPIOS_IMG}; then
            echo Error: downloading the image failed 
            exit 2    
        fi
    fi
elif $opt_check_image_checksum && [ -n "${opt_raspios_checksum:-$IMAGE_EXTRACT_SHA256}" ]; then
    if ! echo "${opt_raspios_checksum:-$IMAGE_EXTRACT_SHA256} ${RASPIOS_IMG}" | sha256sum --check --strict --quiet; then
        echo Error: image checksum does not match with the Raspberry Pi catalog
        exit 2    
    fi 
fi

# Mount the image as a loopback device and get it's device name
RASPIOS_DEV=$(losetup --find --show --partscan --read-only ${RASPIOS_IMG} )
sleep 1

# Get the partition sizes of the image
TTY=$(tty)
sfdisk --dump ${RASPIOS_DEV} | egrep "^/" | sed -E "s/[:,][^=]+=/ /g; s/  +/ /g" | while read part start size type
do
    label=$(lsblk "${part}" --noheadings --output label)
	echo Found $label partition $part with $start $size and type $type

    if [ "${type}" = 'c' ]; then
        echo ",${size},${type}" | sfdisk --quiet --append --wipe always --wipe-partition always "${opt_target_device}"
        target_part=$(sfdisk --dump "${opt_target_device}"  | tail -1 | cut -f1 -d\ )
    elif [ "${type}" = '83' ]; then
        if [ -z "${lvm_dev}" ]; then
            echo ",${opt_lvm_size},8e" | sfdisk --quiet --append --wipe always "${opt_target_device}" 2>/dev/null
            lvm_dev="$(sfdisk --dump "${opt_target_device}"  | tail -1 | cut -f1 -d\ )"
            pvcreate --quiet "${lvm_dev}"
            vgcreate --quiet "${opt_lvm_group_name}" "${lvm_dev}"
        fi

        lvcreate  --name "${label}" --size "${opt_root_size}" --wipesignatures y "${opt_lvm_group_name}" <$TTY
        target_part="/dev/${opt_lvm_group_name}/${label}"

    else
        echo Error: image contains unsupported partition ${part} with type 0x$type
        exit 1
    fi

    echo Copying $label to ${target_part}
    dd if=${part} of=$target_part 2>/dev/null
done

echo Setting up the root filesystem for its initial boot

TARGET_BOOTFS_DEV=$(lsblk --raw --noheadings --output path,label "${opt_target_device}" | fgrep bootfs | cut -d\  -f1)
TARGET_ROOTFS_DEV=$(lsblk --raw --noheadings --output path,label "${opt_target_device}" | fgrep rootfs | cut -d\  -f1)

# Mount the new root filesystem for post-processing
TARGET_ROOT_MOUNTPOINT="${WORKDIR}/rootfs"
echo mount "${filesystems[rootfs]}" "${TARGET_ROOT_MOUNTPOINT}"
mount -o noatime "${TARGET_ROOTFS_DEV}" "${TARGET_ROOT_MOUNTPOINT}"
mount "${TARGET_BOOTFS_DEV}" "${TARGET_ROOT_MOUNTPOINT}/boot/firmware"
mount -t proc /proc "${TARGET_ROOT_MOUNTPOINT}/proc"
mount -o bind /dev "${TARGET_ROOT_MOUNTPOINT}/dev"
mount -o bind /dev/pts "${TARGET_ROOT_MOUNTPOINT}/dev/pts"
mount -t sysfs /sys "${TARGET_ROOT_MOUNTPOINT}/sys"

# patch /boot/firmware/config.txt so tat it contains the correct root filesystem
sed -Ei "s|root=[^ \t]+|root=${TARGET_ROOTFS_DEV}|" "${TARGET_ROOT_MOUNTPOINT}/boot/firmware/cmdline.txt"

# patch /etc/fstab so that it contains the right devices
sed -Ei "s|^[^# \t]+([ \t]+/boot/firmware[ \t])|${TARGET_BOOTFS_DEV}\1|" "${TARGET_ROOT_MOUNTPOINT}/etc/fstab"
sed -Ei "s|^[^# \t]+([ \t]+/[ \t])|${TARGET_ROOTFS_DEV}\1|" "${TARGET_ROOT_MOUNTPOINT}/etc/fstab"

# move filesystems to separate LVMs and ensure these are automaticaly mounted
for p in  ${opt_lvm_partitions//,/ }
do
    IFS=: read path size <<< "${p}"
    lvname=${path//*\//}fs
    echo moving ${path} to a LVM ${lvmname} of ${size}

    lvcreate  --name "${lvname}" --size "${size}" --wipesignatures y "${opt_lvm_group_name}" <$TTY
    lvpath=$(lsblk --raw --noheadings --output path,label "${opt_target_device}" | fgrep "${lvname}" | cut -d\  -f1)
    mke2fs -L "${lvname}" "${lvpath}"

    if ! [ -d "${TARGET_ROOT_MOUNTPOINT}${path}" ]; then
        echo Error: direcory ${vpath} does not exist in the target environment
        exit 1
        # mkdir "${TARGET_ROOT_MOUNTPOINT}/${path}"
    else
        mount "${lvpath}" "${TARGET_ROOT_MOUNTPOINT}/mnt"
        chroot "${TARGET_ROOT_MOUNTPOINT}" rsync -aAXH "${path}/" "/mnt/"
        umount "${TARGET_ROOT_MOUNTPOINT}/mnt"
        chroot "${TARGET_ROOT_MOUNTPOINT}" bash -c "rm -rf ${path}/* ${path}/.*"
    fi

    if egrep -Eq "^[^# \t]+[ \t]+${path}[ \t]" "${TARGET_ROOT_MOUNTPOINT}/etc/fstab" ; then
        sed -Ei "s|^[^# \t]+([ \t]+${path}[ \t])|${lvpath}\1|" "${TARGET_ROOT_MOUNTPOINT}/etc/fstab"
    else
        echo "${lvpath}    ${path}    ext4    defaults,noatime    0    2" >> "${TARGET_ROOT_MOUNTPOINT}/etc/fstab"
    fi

    mount "${lvpath}" "${TARGET_ROOT_MOUNTPOINT}${path}"
done

cat "${TARGET_ROOT_MOUNTPOINT}/etc/fstab"

# install required packages for things to work
chroot "${TARGET_ROOT_MOUNTPOINT}" apt install lvm2

# generate password
# openssl passwd -6

# setup initial (defaultt) user
echo 'jfl:$6$EyEFwTRqBv16ws7j$/5qABGo7Bm3p44mml1/LmFp9b/VdkCURRHNsGxl9YOC3IdzZ7N5LdYFOF5nZjGFQONWGgwjBBIn.JGhGNdnP70' > "${TARGET_ROOT_MOUNTPOINT}/boot/firmware/userconf"

# enable SSH 
touch "${TARGET_ROOT_MOUNTPOINT}/boot/firmware/ssh"

# Enable PCIe Gen 3  for NVMe, if requested
if ${opt_nvme_pcie_gen3}; then
    if ! egrep -q "^\s*dtparam\s*=\s*pciex1_gen\s*=" "${TARGET_ROOT_MOUNTPOINT}/boot/firmware/config.txt"; then
        echo dtparam=pciex1_gen=3 >> "${TARGET_ROOT_MOUNTPOINT}/boot/firmware/config.txt"
    fi
fi

echo Done!