#!/usr/bin/env bash

#    This program is free software: you can redistribute it and/or modify
#    it under the terms of the GNU General Public License as published by
#    the Free Software Foundation, either version 3 of the License, or
#    (at your option) any later version.

## Common functions
#
# Cleanup and quit if error
cleanup_and_quit () {

	# If any paramters are passed we will assume it to be an error
	[[ -n $1 ]] && printf "\e[1;31m<#>\e[0m $*\e[0m\n" >&2

	if [[ $ARKDEP_NO_CLEANUP -eq 1 ]]; then
		printf 'Cleanup disabled, not running cleanup\n'
		exit 1
	fi

	umount -Rl $build_image_mountpoint

	rm $build_image
	rm -rf $build_image_mountpoint

	# Quit program if argument provided to function
	if [[ -n $1 ]]; then
		exit 1
	fi

	# Otherwise just quit, there is no error
	exit 0

}

## Set common variables
#
declare -r build_image='/var/tmp/arkdep-build.img'
declare -r build_image_mountpoint='/var/tmp/arkdep-build'
declare -r workdir="$build_image_mountpoint/rootfs"
declare -r variant="$1"
declare -r build_image_size='15G'

if [[ -v ARKDEP_CONFIGS ]]; then
	declare -r configsdir="$ARKDEP_CONFIGS"
else
	declare -r configsdir="$(readlink -m ./arkdep-build.d/)"
fi

# Before we continue setting variables, lets first ensure the configsdir actually exist
# We do this now to give better error reporting to the user
if [[ ! -d $configsdir ]]; then
	printf "\e[1;31m<#>\e[0m\e[1m $configsdir does not exist, are you located inside of the configuration directory's parent directory?\n\e[0m"
	exit 1
fi

declare -r variantdir="$configsdir/$variant"
declare -r type="$(cat $(readlink -m $variantdir/type) 2> /dev/null || (printf '\e[1;31m<#>\e[0m\e[1m Failed to get build type\n\e[0m'; exit 1))"

if [[ -v ARKDEP_OUTPUT_TARGET ]]; then
	declare -r output_target="$ARKDEP_OUTPUT_TARGET"
else
	declare -r output_target="$(pwd)/target"
fi

## Set common functions
#
# Generate a 42 character long random string, used for generating psuedo-random image names
# Unless overwritten with $ARKDEP_CUSTOM_NAME
gen_random_string () {

	if [[ -v ARKDEP_CUSTOM_NAME ]]; then
		random=$ARKDEP_CUSTOM_NAME
	else
		random=$(openssl rand -hex 100 | head -c 42)
	fi

	printf "${random}\n"

}

declare -r image_name=$(gen_random_string)

## Error checking
#
# Quit if not root
if [[ ! $EUID -eq 0 ]]; then
	printf '\e[1;31m<#>\e[0m\e[1m This program has to be run as root\n\e[0m'
	exit 1
fi

# Check if all dependencies are installed, quit if not
for prog in btrfs pacstrap; do
	if ! command -v $prog > /dev/null; then
		printf "\e[1;31m<#>\e[0m\e[1m Failed to locate $prog, ensure it is installed\e[0m\n"
		exit 1
	fi
done

# Check if requested variant exists
if [[ ! -d $variantdir ]]; then
	printf '\e[1;31m<#>\e[0m\e[1m The requested variant does not exist\e[0m\n'
	exit 1
fi

## Variants
#
# Build migration type image
if [[ $type == 'migration' ]]; then

	printf '\e[1;34m-->\e[0m\e[1m Started migration image build\e[0m\n'

	# Create output directory structure
	mkdir -p $output_target/$image_name

	# Copy migration contents to compression target dir
	cp -v $variantdir/migration.sh $output_target/$image_name/$image_name-migration.sh ||
		cleanup_and_quit 'No primary migration script found'

	# Copy migration related files if they are provided
	if [[ -d $variantdir/migration ]]; then
		cp -rv $variantdir/migration $output_target/$image_name/$image_name-migration
	fi

	# Compress migration files
	tar -cv -I 'zstd -12 -T0 ' -f $output_target/$image_name.tar.zst -C $output_target/$image_name .

	# We can just exit, no need to do a cleanup
	exit 0

fi

# Build archlinux type image
if [[ $type == 'archlinux' ]]; then

	# Ensure bootstrap.list exists, if not error and quit
	# This is the mimimum requirement to perform a build and this checked first
	if [[ ! -e $variantdir/bootstrap.list ]]; then
		printf "\e[1;31m<#>\e[0m\e[1m The required file 'bootstrap.list' is not preset in $variantdir\e[0m\n"
		exit 1
	fi

	printf '\e[1;34m-->\e[0m\e[1m Started Arch linux image build\e[0m\n'

	printf '\e[1;34m-->\e[0m\e[1m Creating disk image\e[0m\n'
	fallocate -l $build_image_size $build_image || cleanup_and_quit "Failed to create disk image at $build_image"
	mkfs.btrfs -f $build_image || cleanup_and_quit "Failed to partition $build_image"

	printf "\e[1;34m-->\e[0m\e[1m Mounting $build_image at $workdir\e[0m\n"
	# The data is compressed to minimize writes to the disk, the actual export does not maintain this compression
	mount -m -t btrfs -o loop,compress=zstd $build_image $build_image_mountpoint || cleanup_and_quit "Failed to mount disk image to $workdir"

	# Create temporary Btrfs subvolume
	printf "\e[1;34m-->\e[0m\e[1m Creating temporary Btrfs subvolumes at $workdir\e[0m\n"
	btrfs subvolume create $workdir/ || cleanup_and_quit "Failed to create btrfs subvolume $workdir)"
	mount --bind $workdir $workdir || cleanup_and_exit "Failed to bind mount disk $workdir"
	btrfs subvolume create $workdir/etc || cleanup_and_quit "Failed to create btrfs subvolume $workdir/etc)"
	btrfs subvolume create $workdir/var || cleanup_and_quit "Failed to create btrfs subvolume $workdir/var)"

	# Read base package list and install base system
	readarray bootstrap_packages < $variantdir/bootstrap.list
	printf '\e[1;34m-->\e[0m\e[1m Installing base packages\e[0m\n'
	# If pacman.conf is available in overlay, use it
	if [[ -f $variantdir/pacman.conf ]]; then
		pacstrap -c -C $variantdir/pacman.conf $workdir ${bootstrap_packages[*]} || cleanup_and_quit 'Failed to install secondary package list'
		cp -v $variantdir/pacman.conf $workdir/etc/pacman.conf
	else
		pacstrap -c $workdir ${bootstrap_packages[*]} || cleanup_and_quit 'Failed to bootstrap system'
	fi

	# If overlay directory exists in variant copy it's contents to the temporary subvolume
	if [[ -d $variantdir/overlay/post_bootstrap ]]; then
		printf '\e[1;34m-->\e[0m\e[1m Copying overlay/post_bootstrap to root\e[0m\n'
		cp -rv $variantdir/overlay/post_bootstrap/* $workdir/
	fi

	# Run post_bootstrap script if exists
	if [[ -f $variantdir/extensions/post_bootstrap.sh ]]; then
		printf '\e[1;34m-->\e[0m\e[1m Running post_bootstrap extension\e[0m\n'
		(source $variantdir/extensions/post_bootstrap.sh)
	fi

	# Read package list and install secondary system components, skip if not used
	if [[ -e $variantdir/package.list ]]; then
		printf '\e[1;34m-->\e[0m\e[1m Installing secondary packages\e[0m\n'

		# Mount the pacman cache
		mount --bind /var/cache/pacman/pkg $workdir/var/cache/pacman/pkg

		# Read package list and install
		readarray packages < $variantdir/package.list
		arch-chroot $workdir pacman -S --noconfirm ${packages[*]} || cleanup_and_quit 'Failed to install packages'

		# Unmount pacman cache
		umount -l $workdir/var/cache/pacman/pkg
	fi

	# If postinstall overlay directory exists in variant copy it's contents to the temporary subvolume
	if [[ -d $variantdir/overlay/post_install ]]; then
		printf '\e[1;34m-->\e[0m\e[1m Copying overlay/post_install to root\e[0m\n'
		cp -rv $variantdir/overlay/post_install/* $workdir/
	fi

	# Run post_install script if exists
	if [[ -f $variantdir/extensions/post_install.sh ]]; then
		printf '\e[1;34m-->\e[0m\e[1m Running post_install extension\e[0m\n'
		(source $variantdir/extensions/post_install.sh)
	fi

	# Clear pacman cache
	printf '\e[1;34m-->\e[0m\e[1m Clearing pacman cache\e[0m\n'
	arch-chroot $workdir pacman -Scc <<< Y <<< Y
	# Insert a new line to clean, the previous command fails to do so
	printf '\n'

	# Remove subvolumes created by systemd
	[[ -d $workdir/var/lib/portables ]] &&
		printf '\e[1;34m-->\e[0m\e[1m Removing systemd subvolume var/lib/portables\e[0m\n'
		btrfs subvolume delete $workdir/var/lib/portables
	[[ -d $workdir/var/lib/machines ]] &&
		printf '\e[1;34m-->\e[0m\e[1m Removing systemd subvolume var/lib/machines\e[0m\n'
		btrfs subvolume delete $workdir/var/lib/machines

	# Make /usr/local symlink in var
	printf '\e[1;34m-->\e[0m\e[1m Moving dirs to var and creating symlinks\e[0m\n'
	mv $workdir/usr/local $workdir/var/usrlocal || cleanup_and_quit 'Failed to move usr/local to var/usrlocal'
	ln -sv ../var/usrlocal $workdir/usr/local || cleanup_and_quit 'Failed to create usrlocal symlink'

	# locale symlink
	mv $workdir/usr/lib/locale $workdir/var/usrliblocale || cleanup_and_quit 'Failed to move usr/local to var/usrlocal'
	ln -sv ../../../var/usrliblocale $workdir/usr/lib/locale || cleanup_and_quit 'Failed to create usrlocal symlink'

	# Opt symlink
	mv $workdir/opt $workdir/var/ || cleanup_and_quit 'Failed to move opt to var/opt'
	ln -sv var/opt $workdir/opt || cleanup_and_quit 'Failed to create opt symlink'

	# srv symlink
	mv $workdir/srv $workdir/var/srv || cleanup_and_quit 'Failed to move srv to var/srv'
	ln -sv var/srv $workdir/srv || cleanup_and_quit 'Failed to create srv symlink'

	# mnt symlink
	mv $workdir/mnt $workdir/var/mnt || cleanup_and_quit 'Failed to move mnt to var/mnt'
	ln -sv var/mnt $workdir/mnt || cleanup_and_quit 'Failed to create mnt symlink'

	# NetworkManager system-connections symlink if installed
	if [[ -d $workdir/etc/NetworkManager ]]; then
		mv $workdir/etc/NetworkManager/system-connections $workdir/var/nm-system-connections || cleanup_and_quit 'Failed to move etc/NetworkManager/system-connections to var/nm-system-connections'
		ln -sv ../../var/nm-system-connections $workdir/etc/NetworkManager/system-connections || cleanup_and_quit 'Failed to create nm-system-connections symlink'
	fi

	printf '\e[1;34m-->\e[0m\e[1m Creating mountpoints for shared subvolumes\e[0m\n'

	# Remove the folders to ensure they are empty
	rm -rf $workdir/root
	rm -rf $workdir/var/lib/flatpak

	# Ensure these folder exist
	mkdir -pv $workdir/root
	mkdir -pv $workdir/arkdep
	mkdir -pv $workdir/var/lib/flatpak

	printf '\e[1;34m-->\e[0m\e[1m Moving passwd, shadow and group files to usr/lib\e[0m\n'

	# Create second passwd, group and shadow file in usr/lib and configure
	for file in passwd group shadow; do
		grep -v "^root:" $workdir/etc/$file > $workdir/usr/lib/$file
	done

	# Remove all users except for root, is typically overwritten by user overlay but
	# may be used during os installation as a template
	for file in passwd group shadow; do
		grep "^root:" $workdir/etc/$file > $workdir/etc/$file-tmp
		mv $workdir/etc/$file-tmp $workdir/etc/$file
	done

	# Ensure passwd/group/shadow permissions are set properly
	chmod 600 $workdir/etc/shadow $workdir/usr/lib/shadow
	chmod 644 $workdir/etc/{passwd,group} $workdir/usr/lib/{passwd,group}

	#
	# nss-switch.conf is added using the overlay
	#

	# Remove passwd/group/shadow backup files
	rm $workdir/etc/{passwd-,shadow-,group-}

	printf '\e[1;34m-->\e[0m\e[1m Moving CPU microcode to usr/lib\e[0m\n'
	# Move CPU firmware to /usr/lib if present
	mv $workdir/boot/*-ucode.img $workdir/usr/lib/

	# Make subvolume read-only
	printf '\e[1;34m-->\e[0m\e[1m Adding read-only property to subvolumes\e[0m\n'
	btrfs property set -ts $workdir ro true || cleanup_and_quit 'Failed to set root to read-only'
	btrfs property set -ts $workdir/etc ro true || cleanup_and_quit 'Failed to set etc to read-only'
	btrfs property set -ts $workdir/var ro true || cleanup_and_quit 'Failed to set var to read-only'

	# Create dir for storing the images
	mkdir -p $output_target/$image_name

	# Generate package list
	pacman -Q --root=$workdir > $output_target/$image_name.pkgs

	# Write subvolume to image
	printf '\e[1;34m-->\e[0m\e[1m Creating images\e[0m\n'
	btrfs send -f $output_target/$image_name/$image_name-rootfs.img $workdir
	btrfs send -f $output_target/$image_name/$image_name-etc.img $workdir/etc
	btrfs send -f $output_target/$image_name/$image_name-var.img $workdir/var

	if [[ ! -v ARKDEP_NO_TAR ]]; then
		printf '\e[1;34m-->\e[0m\e[1m Compressing images\e[0m\n'
		tar -cv -I 'zstd -12 -T0 ' -f $output_target/$image_name.tar.zst -C $output_target/$image_name .
	fi

	if [[ -f $variantdir/extensions/post_build.sh ]]; then
		printf '\e[1;34m-->\e[0m\e[1m Running post_build extension\e[0m\n'
		(source $variantdir/extensions/post_build.sh)
	fi

	cleanup_and_quit

fi