#!/bin/bash

# Call:
# $0 --tmpfs base tools-base xfce4 browser-debian
# to build with all configs and tools. Will build rootfs in memory, so
# run on machine with enough ram (32gb, 16 might work too).
#
# Otherwise use --tmpdir <path>
#
# <path> will be /tmp/mltk-work for --tmpfs
#
# Final kernel + initrd will be in <path>/out-*
# Final qcow2 is /tmp/compressed.qcow2
#
# !!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
# !!                                        !!
# !!  Designed to be run on non-persistent  !!
# !!  worker nodes only. Will mess with     !!
# !!  the running OS!                       !!
# !!  Must be run as root.                  !!
# !!                                        !!
# !!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!

# Config
# TODO: Make configurable via external include, check required vars are set

# Supports deb/apt based distros for now
#distro="ubuntu"
#release="focal"
#mirror="ftp.halifax.rwth-aachen.de"
#pkg_sources="main restricted universe"
distro="debian"
release="bookworm"
mirror="ftp.halifax.rwth-aachen.de"
pkg_sources="main contrib"

# if you don't put a kernel.config in $ROOT_DIR, this will be used
# fallback would be the running kernel's config
kernel_base_config="https://github.com/archlinux/svntogit-packages/raw/packages/linux/trunk/config"

# https://git.openslx.org/bwlp/ansible-bwlp.git/tree/desktop-common/tasks/main.yml

MLTK_CONFIG='
export http_proxy="http://132.230.4.234:8123/"
sourceforge_mirror="netcologne"

CONFIG_NFS_CACHE="10.4.180.32:/escience-bwlp01"

NVIDIA_VERSIONS="550.135 390.157"
CONFIG_KERNEL_VERSION="6.6.62"
CONFIG_VMWARE_VERSION="17.6.1"
CONFIG_VBOX_VERSION="7.1.4"
CONFIG_QEMU_VERSION="v9.1.1"
CONFIG_VIRTMANAGER_VERSION="4.1.0"
CONFIG_LIBTPMS_VERSION="v0.9.6"
CONFIG_SWTPM_VERSION="v0.9.0"
'

##################
#                #
#   End config   #
#                #
##################

disabled_services=
pkgs_dummy= # "libwayland-dev libwayland-client0"
pkgs_full=
pkgs_lean=
tmp_dir=

perror () {
	echo "[ERROR] $*" >&2
	kill "$ppid"
	exit 1
}

load_config() {
	local i
	local dir="configs/$1"
	[ -d "$dir" ] || perror "Could not find directory '$dir'"
	for i in disabled_services pkgs_dummy pkgs_full pkgs_lean; do
		[ -s "$dir/$i" ] || continue
		declare -g "$i=${!i} $( cat "$dir/$i" )"
	done
}

# Parse options, load configs
while [ $# -gt 0 ]; do
	case "$1" in
	--tmpfs)
		tmp_dir=tmpfs
		;;
	--tmpdir)
		tmp_dir="$2"
		shift
		;;
	--*)
		perror "Unknown option '$1'"
		;;
	*)
		load_config "$1"
		;;
	esac
	shift
done

[ -z "$tmp_dir" ] && perror "No temp dir set. use --tmpdir <dir> or --tmpfs"

ppid="$$"

export DEBIAN_FRONTEND="noninteractive"

apt update
# Breaks on current MaxiLinux because of missing kernel
apt remove -y initramfs-tools
# Essential tools
apt install -y systemd-container debootstrap equivs gdisk \
	|| perror "Cannot install nspawn or debootstrap"

run () {
	systemd-nspawn -E DEBIAN_FRONTEND="noninteractive" -D "${root}/" "$@"
}

fix_resolv () {
	# resolv.conf used during build process is just copied from host
	unlink "${root}/etc/resolv.conf"
	cp -L "/etc/resolv.conf" "${root}/etc/resolv.conf" || perror "No resolv.conf"
	[ -L "${root}/etc/resolv.conf" ] && perror "resolv.conf is still a link"
}

# https://git.openslx.org/bwlp/ansible-bwlp.git/tree/dummy-package/scripts/dummy-package.sh
# <tmpdir>
dummy_package () {
	[ "$#" -eq 2 ] || return 1
	[ -d "$1" ] || return 2
	rm -f -- "$2" || return 3

	cd "$1" || return 4

	equivs-control "$2" || return 5

	sed -r -i \
		-e "s/^(#\s)?(Maintainer).*/\\2: support@bwlehrpool.de/" \
		-e "s/^(#\s)?(Package).*/\\2: ${2}/" \
		-e "s/^(#\s)?(Version).*/\\2: 99.9.9/" \
		-e "s/^(#\s)?(Description).*/\\2: Dummy package to provide $2/" \
		-e "/^Description.*/q" \
		"$2"

	cat >> "$2" <<-EOF
	 Long description
	 .
	 with
	 some
	 more
	 lines
	EOF

	equivs-build "$2"
}

declare -rg ARG0="$0"
declare -rg SELF="$(readlink -f "$ARG0")"
declare -rg ROOT_DIR="$(dirname "${SELF}")"

modprobe -a overlay nbd nfs nfsv4 || perror "Could not load overlay and nbd and nfs"

if [ "$tmp_dir" = "tmpfs" ]; then
	base="/tmp/mltk-work"
	mkdir -p "${base}"
	if mountpoint "${base}"; then
		umount "${base}" || perror "Could not unmount old workdir"
	fi
	# Generous 100G tmpfs, should be enough...
	mount -t tmpfs -o size=100G mltk-build "${base}" || perror "Tmpfs fail"
elif [ -d "${tmp_dir}" ]; then
	perror "${tmp_dir} must not exist!"
else
	base="${tmp_dir}"
	mkdir -p "${base}"
fi

cd "${base}" || perror "Cannot cd to '${base}'"
root="${base}/fstree"
mkdir -p "${root}" "${base}/mnt" || perror "mkdir root"

# TODO: Hard-coded apt-cacher-ng
debootstrap --variant=minbase --arch=amd64 \
	--include="build-essential,dbus,binutils,lsb-release,wget,rsync,gpg" \
	"${release}" "${root}" \
	"http://10.4.9.64:3142/${mirror}/${distro}/" || perror "debootstrap failed"

fix_resolv

# Static
rsync -avHAX --chown=0:0 "${ROOT_DIR}/data/" "${root}/" || perror "Could not sync data dir"

# Create and install fake packages
mkdir -p /tmp/dummypkg || perror "Could not create tmp dir for dummy packages"
for pkg in $pkgs_dummy; do
	dummy_package /tmp/dummypkg "$pkg" || perror "Could not create dummy package $pkg"
done
mkdir -p "${root}/dummypkg"
mv -f /tmp/dummypkg/*.deb "${root}/dummypkg/" || perror "Could not move dummy packages to ${root}"
cd "${base}" || perror "Doof dir"

run /bin/sh -c 'dpkg -i /dummypkg/*.deb && rm -rf -- /dummypkg' || perror "Could not install dummy packages"
run apt-mark hold $pkgs_dummy || perror "Could not apt-mark hold"

# TODO: WTF?
run apt-key adv --keyserver keyserver.ubuntu.com --recv-keys 112695A0E562B32A
run apt-key adv --keyserver keyserver.ubuntu.com --recv-keys 54404762BBB6E853

mkdir -p "${root}/etc/apt/apt.conf.d"
# Use our apt cache while building, but remove afterwards (TODO: Configurable as above)
cat > "${root}/etc/apt/apt.conf.d/01proxy" <<END
Acquire::http { Proxy "http://10.4.9.64:3142"; };
Acquire::https { Proxy "https://"; };
END

if [ "$distro" = "ubuntu" ]; then
	cat > "${root}/etc/apt/sources.list" <<-END
	deb http://${mirror}/${distro} ${release} ${pkg_sources}
	deb http://${mirror}/${distro} ${release}-updates ${pkg_sources}
	deb http://${mirror}/${distro} ${release}-security ${pkg_sources}
	END
else
	cat > "${root}/etc/apt/sources.list" <<-END
	deb http://${mirror}/${distro} ${release} ${pkg_sources}
	deb http://${mirror}/${distro} ${release}-updates ${pkg_sources}
	deb http://security.${distro}.org/${distro}-security ${release}-security ${pkg_sources}
	END
fi

if [ "$distro" = "ubuntu" ]; then
	# Untested since ~2022, when firefox became a snap too and we had enough
	# TODO: Use ungoogled-chromium
	# Non-snap chromium
	cat > "${root}/etc/apt/sources.list.d/xalt7x-${distro}-chromium-deb-vaapi-${release}.list" <<-EOF
	deb http://ppa.launchpad.net/xalt7x/chromium-deb-vaapi/${distro} ${release} main
	# deb-src http://ppa.launchpad.net/xalt7x/chromium-deb-vaapi/${distro} ${release} main
	EOF

	# pin chromium
	cat > "${root}/etc/apt/preferences.d/pin-xalt7x-chromium-deb-vaapi" <<-EOF
	Package: *
	Pin: release o=LP-PPA-xalt7x-chromium-deb-vaapi
	Pin-Priority: 1337
	EOF

	# MESA
	cat > "${root}/etc/apt/sources.list.d/mesa-new.list" <<-EOF
	deb http://ppa.launchpad.net/ernstp/mesarc/ubuntu ${release} main
	EOF
fi

# As some scripts might do commits
cat > "${root}/root/.gitconfig" <<END
[user]
	email = you@example.com
	name = Your Name
END

# Now that we fiddled with package sources, refresh db
run apt update || perror "apt update after resetting package sources failed"
# And run an upgrade just to be sure
run apt -o Dpkg::Options::='--force-confold' -o Dpkg::Options::='--force-confdef' -y full-upgrade \
	|| perror "Could not full-upgrade"
# Then install all packages read from the configs passed on cmdline
run apt -o Dpkg::Options::='--force-confold' -o Dpkg::Options::='--force-confdef' install -y $pkgs_full \
	|| perror "Could not install apt list"
run apt -o Dpkg::Options::='--force-confold' -o Dpkg::Options::='--force-confdef' install -y --no-install-recommends $pkgs_lean \
	|| perror "Could not install no recommends list"

fix_resolv

# Prepare mltk
run git clone --depth 1 git://git.openslx.org/openslx-ng/mltk.git /mltk \
	|| perror "Could not clone mltk"

# Insert our overrides for mltk
echo "$MLTK_CONFIG" > "${root}/mltk/config"

mkdir -p "${root}/boot"

# Figure out which kernel config to use as a base
if [ -n "${kernel_base_config}" ]; then
	if [ -s "${ROOT_DIR}/kernel.config" ]; then
		echo "Already have an override kernel config, not downloading ${kernel_base_config}"
		sleep 2
	else
		wget -O "${ROOT_DIR}/kernel.config" "${kernel_base_config}" \
			|| perror "Could not get kernel base config from ${kernel_base_config}"
	fi
fi

cp "${ROOT_DIR}/kernel.config" "${root}/boot/config-mltk" \
	|| cp "/boot/config-$(uname -r)" "${root}/boot/config-$(uname -r)" \
	|| echo "Did not copy any kernel config over..."

# Finally, build mltk stuff
run /mltk/mltk stage4 -i -b kernel || perror "Could not mltk kernel" # without -d for noninteractive
run /mltk/mltk stage4 -b -d -i || perror "Could not mltk stage 4"
run /mltk/mltk vmware-addon -b -d -i || perror "Could not mltk vmware"
run /mltk/mltk vmware-legacy-addon -b -d -i || perror "Could not mltk vmware-legacy"
run /mltk/mltk nvidia-libs@NVIDIA_VERSIONS -b -d -i || perror "Could not mltk nvidia-libs"
run /mltk/mltk qemu -b -d -i || perror "Could not mltk qemu"

# Locale is messed up by this point
# TODO: Configurable?
cat > "${root}/etc/locale.gen" <<EOF
de_DE.UTF-8 UTF-8
en_US.UTF-8 UTF-8
C.UTF-8 UTF-8
EOF
run update-locale LANG=C.UTF-8
run locale-gen
run fc-cache -v -f -r -s

# Another stupid hack: Force our wallpaper for xfce by replacing every wallpaper with ours
for f in "${root}/usr/share/backgrounds/xfce/"*; do
	[ -f "$f" ] || continue
	ln -nfs "/usr/share/wallpapers/bwLehrpool/contents/images/1920x1080.png" "$f"
done

## Not needed currently/anymore?
# Use old iptables for now, not the nft wrapper, as physdev matching is broken there
#run update-alternatives --set iptables /usr/sbin/iptables-legacy

# Disable annoying motd stuff, is unfortunately part of base pam package on ubuntu
# and the firstlogin disclaimer cannot be disabled?
run sed -i '/pam_motd/d' /etc/pam.d/sshd /etc/pam.d/login || perror "Could not disable pam_motd"

# Copy static stuff we ship
rsync -avHAX --chown=0:0 "${ROOT_DIR}/data/" "${root}/" || perror "Could not sync data dir"

# Disable services we don't want or need
# TODO: udisks2 and numad are currently hard-coded here, because we don't want to mask them.
# They should still be triggerable when they're actually needed, i.e. a native session
run systemctl disable $disabled_services udisks2.service numad.service
run systemctl mask $disabled_services
for s in $disabled_services udisks2.service; do
	if [[ "$s" == *.* ]]; then
		rm -- "${root}/"etc/systemd/system/*.wants/"$s"
	else
		rm -- "${root}/"etc/systemd/system/*.wants/"$s".service
	fi
done

# Now create disk image

# Get size of FS
shopt -s extglob
fs_size=$( du -s -BM "${root}"/!(mltk|systemd-init|boot|proc|sys|dev|run|tmp) | awk '{a+=$1}END{print a}' )

if (( fs_size < 1000 )) || (( fs_size > 50000 )); then
	perror "implausible rootfs size: $fs_size MB"
fi

# Leave 500MB buffer space
un="${base}/uncompressed.qcow2"
cmp="/tmp/compressed.qcow2"
qemu-img create -f qcow2 "$un" "$(( fs_size + 500 ))M" \
	|| perror "Could not create uncompressed qcow2"

qemu-nbd -c /dev/nbd3 --discard=unmap --detect-zeroes=unmap "$un" || perror "qemu-nbd fail"
echo -e "n\n\n\n\n\nc\nSLX_SYS\nw\ny\n" | gdisk /dev/nbd3 || perror "gdisk failed"
partprobe /dev/nbd3
sleep 1
[ -b "/dev/nbd3p1" ] || perror "NBD partition not found"

mkfs.ext4 /dev/nbd3p1 || perror "mkfs.ext4 failed"
mount "/dev/nbd3p1" "${base}/mnt" || perror "Mount failed"

# TODO: Configurable blacklist
rsync -avHAX \
	--exclude="/dev/*" \
	--exclude="/sys/*" \
	--exclude="/proc/*" \
	--exclude="/run/*" \
	--exclude="/boot" \
	--exclude="/snap/*" \
	--exclude="/mltk" \
	--exclude="/systemd-init" \
	--exclude="*~" \
	--exclude="*.tmp" \
	--exclude=".*.swp" \
	--include="/var/log/**/" \
	--exclude="/var/log/**" \
	--include="/var/cache/**/" \
	--include="/var/cache/fontconfig/**" \
	--include="/var/cache/ldconfig/**" \
	--exclude="/var/cache/**" \
	--include="/var/spool/**/" \
	--exclude="/var/spool/**" \
	--exclude="/addon-init" \
	--exclude="/etc/apt/apt.conf.d/01proxy" \
	--exclude="/etc/resolv.conf" \
	--include="/root/.bashrc" \
	--exclude="/root/**" \
	--exclude="/etc/init.d/kexec" \
	--exclude="/etc/init.d/kexec-load" \
	--exclude="/usr/share/xsessions/i3-with-shmlog.desktop" \
	--exclude="/usr/share/xsessions/lightdm-xsession.desktop" \
	--exclude="/usr/lib/udev/rules.d/*-hwclock.rules" \
	--exclude="/usr/lib/udev/rules.d/*-alsa-restore.rules" \
	--exclude="/tmp/**" \
	--exclude="/etc/krb5.conf" \
	"${root}/" "${base}/mnt/" \
		|| perror "rsync failed"

echo "Unmounting container"
umount "${base}/mnt" || perror "Unmount Failed"
sync
echo "Shutting down qemu-nbd"
qemu-nbd -d /dev/nbd3 || perror "closing qemu-nbd failed"
sync

# Convert in background since it's slow
echo "Compressing qcow2 in background job"
qcow_progress="$base/qcow-log"
# -o compression_type=zstd
qemu-img convert -W -m 16 -p -O qcow2 -c "$un" "$cmp" &> "$qcow_progress" &
qcow_pid="$!"

# Dracut
# Build initramfs

# Do we have zstd?
compress=
if grep -qF 'CONFIG_RD_ZSTD=y' "${root}/mltk/tmp/work/kernel/ksrc/.config" \
		&& command -v zstd; then
	compress="zstd -19 -q -T0"
elif grep -qF 'CONFIG_RD_LZ4=y' "${root}/mltk/tmp/work/kernel/ksrc/.config" \
		&& command -v lz4; then
	compress="lz4"
elif grep -qF 'CONFIG_RD_GZIP=y' "${root}/mltk/tmp/work/kernel/ksrc/.config" \
		&& command -v gzip; then
	compress="gzip"
fi

if [ -n "$compress" ]; then
	echo "Using compression for initrd: $compress"
	compress="--compress $compress"
else
	echo "Will not compress initrd. Either not supported by kernel, or compression tool missing."
	compress="--no-compress"
fi

mkdir -p "${root}/systemd-init"
wget -O "${root}/systemd-init/build-initramfs.sh" \
	"https://git.openslx.org/openslx-ng/systemd-init.git/plain/build-initramfs.sh" \
	|| perror "Could not download systemd-init script"
chmod +x "${root}/systemd-init/build-initramfs.sh"

long_debug= # "--debug"
for ver in "${root}/lib/"modules/*-openslx*; do
	# Workaround for xloop build *sometimes* failing. It sometimes claims it can't copy
	# a file that's definitely there, but that's cmake for you I guess.
	for tries in 1 2 FAIL; do
		if [ "$tries" = "FAIL" ]; then
			perror "dracut stuff failed"
		fi
		run /systemd-init/build-initramfs.sh $long_debug --update --file-path "/systemd-init/initramfs-${ver##*/}" \
			--kernel-version "${ver##*/}" --kernel-headers /mltk/tmp/work/kernel/ksrc/ \
			--qcow-handler xloop --all-microcode \
			- \
			--add 'slx-clock slx-addons slx-runmode slx-uuid slx-splash slx-drm slx-ssl' \
			--install '/usr/sbin/mii-tool /usr/sbin/ethtool' \
			--omit crypt --omit-drivers nvidiafb "$compress" \
			&& break
	done
	mkdir -p "${base}/out-${ver##*/}"
	# Move initrd to destination
	mv -f -- "${root}/systemd-init/initramfs-${ver##*/}" \
		"${base}/out-${ver##*/}/initramfs-stage31" \
		|| perror "Error copying initrd to destination"
	# Move kernel to destination
	mv -f -- "${root}/mltk/var/builds/kernel/kernel" \
		"${base}/out-${ver##*/}/kernel" \
		|| perror "Error moving kernel to destination"
done

# Wait for qemu-img to finish compression
tail -f "$qcow_progress" &
cat_pid="$!"
wait "$qcow_pid" || pwarning "Compressing final qcow2 failed"
kill "$cat_pid"
rm -f -- "$qcow_progress"
wait

exit 0