#!/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.

set -o pipefail

# Print manual if no parameters provided or invalid amount of parameters is provided
if [[ ! -n $1 ]]; then
	cat <<- END
	Usage: arkanium-deploy <action> [target]
	Actions:
	  update		Check for updates, optionally provide a target, if no target provided it defaults to primary
	  deploy		Deploy a new or update an existing deployment
	  init			Initialize Arkanium on a new system
	  teardown		Remove all Arkanium-deploy related files and folders
	  get-available		List available packages in repo
	END
	exit 0
fi

## Set common variables
#
declare -r arkanium_dir='/arkanium/'
declare -r arkanium_deployment_dir="$arkanium_dir/deployments"

## 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

	# Remove temporary btrfs volumes
	# ONLY IF CUSTOM DEFINED?
	#[[ ! -n $arkanium_NO_CLEANUP || ! -n $subvol_created ]] &&
	#	btrfs subvolume delete $workdir 2> /dev/null

	# Quit program if argument provided to function
	[[ -n $1 ]] && exit 1

	# Otherwise just quit, there is no error
	exit 0

}

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

# Check if all dependencies are installed, quit if not
for prog in btrfs wget; 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

## Core functions
#
# Initialize the system for arkanium
init () {

	# Ensure systemd-boot is installed before continuing, for it is the only thing we support
	bootctl is-installed || cleanup_and_quit 'systemd-boot seems to not be installed'

	printf '\e[1;34m-->\e[0m\e[1m Initializing arkanium\e[0m\n'

	[[ -d $arkanium_dir ]] && cleanup_and_quit "$arkanium_dir already exists"

	# Create the /arkanium subvolume
	printf "\e[1;34m-->\e[0m\e[1m Creating $(readlink -m $arkanium_dir) subvolume\e[0m\n"
	btrfs subvolume create $arkanium_dir || cleanup_and_quit "Failed to create btrfs subvolume"

	# Create directory structure
	printf "\e[1;34m-->\e[0m\e[1m Creating directory structure\e[0m\n"
	mkdir -pv $(readlink -m $arkanium_dir/deployments) \
		$(readlink -m $arkanium_dir/deployments) \
		$(readlink -m $arkanium_dir/cache) \
		$(readlink -m $arkanium_dir/templates) \
		$(readlink -m $arkanium_dir/overlay) \
		$(readlink -m $arkanium_dir/shared) ||
		cleanup_and_quit "Failed to create /arkanium and related directories"

	# Add home shared subvolume and make writable
	btrfs subvolume create $(readlink -m $arkanium_dir/shared/home) || cleanup_and_quit "Failed to create home subvolume"
	btrfs property set -f -ts $(readlink -m  $arkanium_dir/shared/home) ro false

	# Write default config file
	printf "\e[1;34m-->\e[0m\e[1m Adding default config file\e[0m\n"
	cat <<- END > $arkanium_dir/config
	# Write /arkanium/overlay overlay to root or etc
	enable_overlay=0

	# Do not install additional packages defined customize
	custom_packages=0

	# URL to image repository, do not add trailing slash
	repo_url='https://repo.arkanelinux.org/arkanium'

	# Default image pulled from repo if nothing defined
	repo_default_image='arkanelinux'
	END

	# Add default bootloader config file
	cat <<- END > $arkanium_dir/templates/systemd-boot
	title Arkane GNU/Linux - arkanium
	linux /vmlinuz-linux
	initrd /amd-ucode.img
	initrd /intel-ucode.img
	initrd /initramfs-linux.img
	options root="LABEL=arkane_root" rootflags=subvol=$arkanium_dir/%TARGET_DEPLOYMENT% rw
	END

	exit 0

}

teardown () {

	cat <<- END
	WARNING: Removing arkanium may leave your system in an unbootable state and you
	may have to manually reconfigure your bootloader etc.. Only proceed if you know
	what you are doing!

	The following changes will be made to your system;
	- All subvolumes under $arkanium_dir will be deleted

	END

	read -p 'Type "I KNOW WHAT I AM GOING" in uppercase to confirm that you know what you are doing: ' input_confirm


	if [[ $input_confirm == 'I KNOW WHAT I AM DOING' ]]; then

		printf '\e[1;34m-->\e[0m\e[1m Tearing down arkanium\e[0m\n'

		# Quit with error if $arkanium_dir does not exist
		if [[ ! -d $arkanium_dir ]]; then
			printf "\e[1;31m<#>\e[0m $(readlink -m $arkanium_dir) does not exist, there is nothing to tear down"
			exit 1
		fi

		# Remove all nested subvolumes in $arkanium_dir
		for volume in $(btrfs subvolume list $arkanium_dir | grep -oE '[^ ]+$'); do
			btrfs property set -f -ts $(readlink -m $arkanium_dir/$volume) ro false
			btrfs subvolume delete $(readlink -m $arkanium_dir/$volume)
		done

		# Remove $arkanium_dir itself
		btrfs property set -f -ts $(readlink -m $arkanium_dir) ro false
		btrfs subvolume delete $(readlink -m $arkanium_dir)

	else
		printf '\e[1;34m-->\e[0m\e[1m Teardown canceled, no changes made to system\e[0m\n'
	fi

	exit 0

}

# List all available packages defined in the repo's list file
get_available () {

	printf "\e[1;34m-->\e[0m\e[1m Downloading list file from $repo_url\e[0m\n"
	curl -sf "${repo_url}/list" || cleanup_and_quit 'Failed to download repo file'

}

# Deploy a new or update an existing deployment
deploy () {

	# Load defaults from config file
	source $arkanium_dir/config

	# target and version are optional, if not defined default to primary as defined in
	# /arkanium/config and latest
	if [[ -n $2 ]]; then
		declare -r deploy_target=$2
	else
		declare -r deploy_target="$repo_default_image"
	fi

	if [[ -n $3 ]]; then
		declare -r deploy_version=$3
	else
		declare -r deploy_version='latest'
	fi

	printf "\e[1;34m-->\e[0m\e[1m Deploying $deploy_target $deploy_version\e[0m\n"

	# If latest is requested grab database and get first line
	if [[ $deploy_version == 'latest' ]]; then
		declare curl_data=$(curl -sf "${repo_url}/${deploy_target}/database" | head -n 1)
	else
		declare curl_data=$(curl -sf "${repo_url}/${deploy_target}/database" | grep $3)
	fi

	# Split latest_version at the delimiter, creating an array with data.1=package ver, data.2=compression method, data.3=sha1 hash
	readarray -d : -t data <<< "$curl_data"
	# A carriage feed is inserted for some reason, lets remove it
	data[2]=${data[2]//[$'\t\r\n']}


	# Check if requested version is already downloaded
	if [[ -e $(readlink -m $arkanium_dir/cache/${data[0]}.tar.${data[1]}) ]]; then

		printf "\e[1;34m-->\e[0m\e[1m ${data[0]} already in cache, skipping download\e[0m\n"

	else

		# Download the tarball if not yet downloaded
		if [[ ! -e $arkanium_dir/cache/${data[0]}.tar.${data[1]} ]]; then
			wget -P $(readlink -m $arkanium_dir/cache/) "$repo_url/$repo_default_image/${data[0]}.tar.${data[1]}" ||
				cleanup_and_quit 'Failed to download tarball'
		fi

	fi

	printf "\e[1;34m-->\e[0m\e[1m Validating integrity\e[0m\n"
	sha1sum "$(readlink -m $arkanium_dir/cache/${data[0]}.tar.${data[1]})" |
		grep "${data[3]}" ||
		cleanup_and_quit "Checksum does not match repo file, got $chksum\e[0m\n"

	# Create directory using unique deployment name
	mkdir -pv $arkanium_dir/deployments/${data[0]} || cleanup_and_quit 'Failed to create deployment directory'

	# Extract the root image if not yet extracted
	printf "\e[1;34m-->\e[0m\e[1m Writing root\e[0m\n"
	if [[ ! -e $(readlink -m $arkanium_dir/cache/${data[0]}-rootfs.img) ]]; then
		tar -xf $(readlink -m $arkanium_dir/cache/${data[0]}.tar.${data[1]}) -C $(readlink -m $arkanium_dir/cache/) "./${data[0]}-rootfs.img" ||
			cleanup_and_quit 'Failed to extract root'
	fi

	# Write the root image
	btrfs receive -f $(readlink -m $arkanium_dir/cache/${data[0]}-rootfs.img) $(readlink -m $arkanium_dir/deployments/${data[0]}) ||
		cleanup_and_quit 'Failed to receive root'

	# Cleanup root image
	rm $(readlink -m $arkanium_dir/cache/${data[0]}-rootfs.img)

	# Extract the etc image if not yet extracted
	printf "\e[1;34m-->\e[0m\e[1m Writing etc\e[0m\n"
	if [[ ! -e $(readlink -m $arkanium_dir/cache/${data[0]}-etc.img) ]]; then
		tar -xf $(readlink -m $arkanium_dir/cache/${data[0]}.tar.${data[1]}) -C $(readlink -m $arkanium_dir/cache/) "./${data[0]}-etc.img" ||
			cleanup_and_quit 'failed to extract etc'
	fi

	# Write the etc image and create var directory, we have to unlock rootfs temporarily to do this
	btrfs property set -f -ts $(readlink -m $arkanium_dir/deployments/${data[0]}/rootfs) ro false ||
		cleanup_and_quit 'Failed to unlock root to write etc'

	btrfs receive -f $(readlink -m $arkanium_dir/cache/${data[0]}-etc.img) $(readlink -m $arkanium_dir/deployments/${data[0]}/rootfs/) ||
		cleanup_and_quit 'Failed to receive etc'

	mkdir -pv $(readlink -m $arkanium_dir/deployments/${data[0]}/rootfs/var)

	btrfs property set -f -ts $(readlink -m $arkanium_dir/deployments/${data[0]}/rootfs) ro true ||
		cleanup_and_quit 'Failed to lock root'

	# Unlock the etc deployment
	btrfs property set -f -ts $(readlink -m $arkanium_dir/deployments/${data[0]}/rootfs/etc) ro false ||
		cleanup_and_quit 'Failed to unlock root to write etc'

	# Cleanup etc image
	rm $(readlink -m $arkanium_dir/cache/${data[0]}-etc.img)

	# Overlay customizations to etc
	# TODO: overlay here if enabled in config file

	# Extract the var image if not yet extracted
	printf "\e[1;34m-->\e[0m\e[1m Writing var\e[0m\n"
	if [[ ! -e $(readlink -m $arkanium_dir/cache/${data[0]}-var.img) ]]; then
		tar -xf $(readlink -m $arkanium_dir/cache/${data[0]}.tar.${data[1]}) -C $(readlink -m $arkanium_dir/cache/) "./${data[0]}-var.img" ||
			cleanup_and_quit 'failed to extract var'
	fi

	# Write the var image
	btrfs receive -f $(readlink -m $arkanium_dir/cache/${data[0]}-var.img) $(readlink -m $arkanium_dir/shared/) ||
		cleanup_and_quit 'Failed to receive var'

	# Make var writable
	btrfs property set -f -ts $(readlink -m $arkanium_dir/shared/var) ro false ||
		cleanup_and_quit 'Failed to unlock var'

	# Cleanup var image
	rm $(readlink -m $arkanium_dir/cache/${data[0]}-var.img)

}

[[ $1 == 'init' ]] && init
[[ $1 == 'teardown' ]] && teardown
[[ $1 == 'update' ]] && check_for_updates
[[ $1 == 'get-available' ]] && get_available
[[ $1 == 'deploy' ]] && deploy $1 $2 $3