#!/bin/bash

# Exit on error
set -e

# ---FUNCTIONS---

confirm() {
    # call with a prompt string or use a default
    read -r -p "${1:-Are you sure? [y/N]} " response
    case "$response" in
        [yYzZ][eE][sS]|[yYzZ])
            true
            ;;
        *)
            false
            ;;
    esac
}

archchroot(){
  arch-chroot /mnt $1
  # work-around for /mnt/dev sometimes not unmounting properly when leaving the chroot
  while [[ "$(findmnt --target=/dev -o TARGET -n)" == "/mnt/dev" ]]; do
  	umount /mnt/dev || true
  done
}

# Unmount all disks
unmount(){
	if [[ $efi -eq 1 ]]; then
		umount /mnt/tmp/refind_install
		umount /mnt/tmp
		umount /mnt/boot
	fi
	umount /mnt/.snapshots
	umount /mnt

	if [[ "$luks" == "true" ]]; then
		cryptsetup luksClose /dev/mapper/cryptroot
	else
		swapoff "$swapdev"
	fi
}

# Set up pacman & keyring
pacman_setup(){
	pacman -Sy --needed --noconfirm archlinux-keyring
	pacman -S --needed --noconfirm arch-install-scripts
	if [ ! -f /usr/bin/wget ]; then
		pacman -S --needed --noconfirm wget libnewt
	fi
}

# Create disk partitions
partition_disk(){
	parted -s $disk mklabel gpt
	if [[ $efi -eq 1 ]]; then
		sgdisk $disk -n=1:0:+1024M -t=1:ef00
		sgdisk $disk -n=2:0:+$swapsize -t=2:8200
		sgdisk $disk -n=3:0:0
	else
		sgdisk $disk -n=1:0:+1M -t=1:ef02
		sgdisk $disk -n=2:0:+$swapsize -t=2:8200
		sgdisk $disk -n=3:0:0
	fi

	# Special naming for eMMC and NVME devices
	if [[ "${disk::8}" == "/dev/nvm" || "${disk::8}" == "/dev/mmc" ]]; then
	  bootdev=$disk"p1"
	  swapdev=$disk"p2"
	  rootdev=$disk"p3"
	else
	  bootdev=$disk"1"
	  swapdev=$disk"2"
	  rootdev=$disk"3"
	fi

	export bootdev
	export swapdev
	export rootdev

	# Format ESP to fat32
	if [[ $efi -eq 1 ]]; then
		mkfs.fat $bootdev
	fi
}

format_root(){
	if [[ "$luks" == "true" ]]; then
		format_root_luks
		return
	fi

	mkswap $swapdev # format swap
	mkfs.btrfs -f $rootdev # format root to btrfs
	mount $rootdev /mnt
	btrfs subvolume create /mnt/@ # new root subvolume
	btrfs subvolume set-default /mnt/@ # set as mount-default to prevent boot failures
	btrfs subvolume create /mnt/@snapshots # new snapshot subvolume
	mkdir -p /mnt/@/var/cache/pacman # create directory so next cmd works
	btrfs subvolume create /mnt/@/var/cache/pacman/pkg # exclude pacman cache from snapshots
	umount /mnt
	mount $rootdev -o subvol=@ /mnt # remount with new root
	mkdir /mnt/.snapshots
	mount $rootdev -o subvol=@snapshots /mnt/.snapshots # mount snapshots as /.snapshots in new root
}

format_root_luks(){
	echo "$lukspass" | cryptsetup -y -v luksFormat $rootdev # setup LUKS partition
	echo "$lukspass" | cryptsetup open $rootdev cryptroot # unlock it
	mkfs.btrfs -f /dev/mapper/cryptroot # format root to btrfs
	mount /dev/mapper/cryptroot /mnt
	btrfs subvolume create /mnt/@ # new root subvolume
	btrfs subvolume set-default /mnt/@ # set as mount-default to prevent boot failures
	btrfs subvolume create /mnt/@snapshots # new snapshots subvolume
	mkdir -p /mnt/@/var/cache/pacman
	btrfs subvolume create /mnt/@/var/cache/pacman/pkg # exclude pacman cache from snapshots
	umount /mnt
	mount /dev/mapper/cryptroot -o subvol=@ /mnt # remount with new root
	mkdir /mnt/.snapshots
	mount /dev/mapper/cryptroot -o subvol=@snapshots /mnt/.snapshots # mount snapshots as /.snapshots in new root
}

swapsetup(){
	if [[ "$luks" == "true" ]]; then
		# setup encrypted swap
		swapuuid=$(blkid -o value $swapdev)
		echo "swap PARTUUID=$swapuuid /dev/urandom swap,cipher=aes-cbc-essiv:sha256,size=256" >> /mnt/etc/crypttab
	else
		swapon $swapdev
	fi
}

# sets hostname, keymap, locale, timezone, root password
basic_postinstall(){
	echo "$hostname" > /mnt/etc/hostname
	echo "KEYMAP=$keylayout" > /mnt/etc/vconsole.conf

	echo "echo \"EDITOR=nano\" >> /mnt/etc/environment"
	echo "echo \"COLORTERM=truecolor\" >> /mnt/etc/environment"

	echo "EDITOR=nano" >> /mnt/etc/environment
	echo "COLORTERM=truecolor" >> /mnt/etc/environment

	locale=en_US
	echo "echo \"LANG=$locale.UTF-8\" > /mnt/etc/locale.conf"
	echo "LANG=$locale.UTF-8" > /mnt/etc/locale.conf
	echo "sed -i '/"$locale".UTF-8/s/^#//g' /mnt/etc/locale.gen"
	sed -i '/'$locale'.UTF-8/s/^#//g' /mnt/etc/locale.gen
	sed -i 's/<U002E> IGNORE/<U002E> <RES-1>/g' /mnt/usr/share/i18n/locales/iso14651_t1_common
	archchroot locale-gen

	sed -i 's%#Color%Color%g' /mnt/etc/pacman.conf
	sed -i 's%#ParallelDownloads %ParallelDownloads %g' /mnt/etc/pacman.conf

	cat > /mnt/usr/share/libalpm/hooks/locale-fix.hook <<"EOS"
[Trigger]
Operation = Install
Operation = Upgrade
Type = Path
Target = usr/share/i18n/locales/iso14651_t1_common

[Action]
Description = Fixing locale sort...
When = PostTransaction
Exec = /usr/bin/bash -c "sed -i 's/<U002E> IGNORE/<U002E> <RES-1>/g' /usr/share/i18n/locales/iso14651_t1_common && locale-gen"

EOS

	ln -sf /usr/share/zoneinfo/Europe/Berlin /mnt/etc/localtime
	archchroot "hwclock --systohc --utc"
	arch-chroot /mnt /bin/bash -c "echo -e '$password\n$password' | passwd root"
}

# sets up fstab, initramfs as well as encrypted swap if LUKS is enabled
setup_initramfs(){
	genfstab -U -p /mnt > /mnt/etc/fstab
	if [[ "$luks" == "true" ]]; then
		echo "/dev/mapper/swap	none swap defaults 0 0" >> /mnt/etc/fstab
		sed -i "s/HOOKS=(base udev autodetect modconf block filesystems keyboard fsck)/HOOKS=(base udev autodetect keyboard keymap consolefont modconf block encrypt filesystems fsck)/g" /mnt/etc/mkinitcpio.conf
		archchroot "mkinitcpio -p linux"
	fi

	# add nvme module to initramfs if rootdevice is on a nvme drive
	if [ "${disk::8}" == "/dev/nvm" ]; then
		sed -i "s/MODULES=()/MODULES=(nvme)/g" /mnt/etc/mkinitcpio.conf
		archchroot "mkinitcpio -p linux"
	fi
}

# installs either refind or grub
install_bootloader(){
	if [[ "$efi" -eq 1 ]]; then
		install_refind
	else
		install_grub
	fi
}

install_refind(){
	pacstrap /mnt refind
	archchroot "refind-install --usedefault $bootdev"
	rootuuid=$(blkid -s UUID -o value $rootdev)

	if [[ "$luks" == "true" ]]; then
		echo "\"Arch Linux         \" \"cryptdevice=UUID=$rootuuid:cryptroot root=/dev/mapper/cryptroot rw add_efi_memmap fbcon=nodefer quiet\"" > /mnt/boot/refind_linux.conf
		echo "\"Arch Linux Verbose \" \"cryptdevice=UUID=$rootuuid:cryptroot root=/dev/mapper/cryptroot rw add_efi_memmap fbcon=nodefer\"" >> /mnt/boot/refind_linux.conf
		echo "\"Arch Linux Fallback\" \"cryptdevice=UUID=$rootuuid:cryptroot root=/dev/mapper/cryptroot rw add_efi_memmap initrd=/initramfs-linux-fallback.img fbcon=nodefer quiet\"" >> /mnt/boot/refind_linux.conf
		echo "\"Arch Linux Terminal\" \"cryptdevice=UUID=$rootuuid:cryptroot root=/dev/mapper/cryptroot rw add_efi_memmap systemd.unit=multi-user.target fbcon=nodefer quiet\"" >> /mnt/boot/refind_linux.conf
	else
		echo "\"Arch Linux         \" \"root=UUID=$rootuuid rw add_efi_memmap quiet\"" > /mnt/boot/refind_linux.conf
		echo "\"Arch Linux Verbose \" \"root=UUID=$rootuuid rw add_efi_memmap\"" >> /mnt/boot/refind_linux.conf
		echo "\"Arch Linux Fallback\" \"root=UUID=$rootuuid rw add_efi_memmap initrd=/initramfs-linux-fallback.img quiet\"" >> /mnt/boot/refind_linux.conf
		echo "\"Arch Linux Terminal\" \"root=UUID=$rootuuid rw add_efi_memmap systemd.unit=multi-user.target quiet\"" >> /mnt/boot/refind_linux.conf
	fi

	sed -i 's/timeout 20/timeout 5/g' /mnt/boot/EFI/BOOT/refind.conf
	sed -i 's/install, shell, bootorder, gdisk, memtest, mok_tool, apple_recovery, windows_recovery, about, //g' /mnt/boot/EFI/BOOT/refind.conf
	sed -i 's/#use_graphics_for osx,linux/use_graphics_for osx,linux/g' /mnt/boot/EFI/BOOT/refind.conf

	git clone https://github.com/bobafetthotmail/refind-theme-regular.git /mnt/tmp/refind-theme
	rm -rf /mnt/tmp/refind-theme/{src,.git,install.sh}
	cp -r /mnt/tmp/refind-theme /mnt/boot/efi/EFI/BOOT/refind-theme
	rm -rf /mnt/tmp/refind-theme
	echo "include refind-theme-regular/theme.conf" >> /mnt/boot/EFI/BOOT/refind.conf
}

install_grub(){
	pacstrap /mnt grub
	mkdir /mnt/boot/grub
	archchroot "grub-mkconfig -o /boot/grub/grub.cfg"
    archchroot "grub-install --target=i386-pc --recheck $disk"
}

# runs postinstall of selected profile
profile_postinstall(){
	if [[ "$profile" == "cli" ]]; then
		postinstall_cli
	elif [[ "$profile" == "i3desktop" ]]; then
		postinstall_i3_desktop
	elif [[ "$profile" == "i3mobile" ]]; then
		postinstall_i3_mobile
	elif [[ "$profile" == "i3lite" ]]; then
		postinstall_i3_lite
	elif [[ "$profile" == "kde" ]]; then
		postinstall_kde
	elif [[ "$profile" == "gnome" ]]; then
		postinstall_gnome
	fi
	archchroot "btrfs subvolume snapshot / /.snapshots/postinstall"
}

# install some packages
# adds the user & sets the password
# enables network
# creates backup snapshot
postinstall_cli(){
	# run this entire thing in new root
	archchroot "/bin/bash" <<"EOS"
	pacman -Syu --needed --noconfirm git openssh htop dialog zsh kernel-modules-hook nftables
	systemctl enable haveged sshd

	sed -i 's%#PasswordAuthentication yes%PasswordAuthentication no%g' /etc/ssh/sshd_config

	systemctl mask systemd-resolved

	systemctl enable linux-modules-cleanup
	systemctl enable systemd-timesyncd
	timedatectl set-ntp 1

	groupadd sudo -g 666
	echo "%sudo ALL=(ALL) NOPASSWD: ALL" >> /etc/sudoers

	cat /etc/makepkg.conf | sed "s/xz -c -z -)/xz -c -z - --threads=0)/g" > /tmp/makepkg.conf; mv /tmp/makepkg.conf /etc/makepkg.conf

	useradd "$user"
	echo -e "$password\n$password" | passwd "$user"
	usermod -aG sudo "$user"
	usermod -aG adm "$user"
	mkdir /home/"$user"
	chown "$user":"$user" /home/"$user"

	systemctl enable dhcpcd

	btrfs subvolume snapshot / /.snapshots/basesystem

EOS
}

postinstall_xorg(){
	postinstall_cli
	archchroot "/bin/bash" <<"EOS"
	pacman -S --needed --noconfirm xorg-server xorg-drivers ttf-opensans ttf-dejavu ttf-hack ttf-liberation networkmanager
EOS
	if [[ "$keylayout" == "us" ]]; then
		return
	fi
cat <<EOF > /mnt/etc/X11/xorg.conf.d/00-keyboard.conf
Section "InputClass"
    Identifier "system-keyboard"
    MatchIsKeyboard "on"
    Option "XkbLayout" "de"
    Option "XkbModel" "pc105"
    Option "XkbVariant" "deadacute"
EndSection
EOF
}

postinstall_kde(){
	postinstall_xorg
	archchroot "/bin/bash" <<"EOS"
	pacman -S --needed --noconfirm plasma plasma-meta konsole dolphin kdialog sddm
	systemctl enable sddm
	echo -e "[Autologin]\nRelogin=false\nSession=\nUser=\n\n[General]\nHaltCommand=/usr/bin/systemctl poweroff\nRebootCommand=/usr/bin/systemctl reboot\n\n[Theme]\nCurrent=breeze\nCursorTheme=Adwaita\n\n[Users]\nMaximumUid=60000\nMinimumUid=1000" > /etc/sddm.conf
	systemctl disable dhcpcd
	systemctl enable NetworkManager
EOS
}

postinstall_gnome(){
	postinstall_xorg
	archchroot "/bin/bash" <<"EOS"
	pacman -S --needed --noconfirm gnome gnome-extra
	systemctl enable gdm
	systemctl disable dhcpcd
	systemctl enable NetworkManager
EOS
}

# Main install function
install(){
	clear
	echo "---ArchLinux install start---"
	pacman_setup
	partition_disk
	format_root
	mkdir /mnt/{boot,home} 2>/dev/null
	if [[ $efi -eq 1 ]]; then
		mount $bootdev /mnt/boot
	fi
	pacstrap /mnt base base-devel wget btrfs-progs cryptsetup device-mapper dhcpcd e2fsprogs inetutils linux-firmware linux logrotate lvm2 man-db man-pages mdadm nano netctl perl sysfsutils texinfo usbutils vim xfsprogs # start the actual install
	swapsetup
	basic_postinstall
	setup_initramfs
	install_bootloader
	profile_postinstall
	unmount
	echo "Installation complete! Please reboot now."
	exit 0
}

# ---MAIN SCRIPT---

clear

# cli arguments
if [[ "$#" -ne 1 ]]; then
	echo "Invalid number of arguments."
	echo "Syntax: bash $0 </dev/sdX>"
	exit 1
fi

export disk="$1"

if [[ -d /sys/firmware/efi ]]; then
	export efi=1
else
	export efi=0
fi

# Keyboard layout selection
export keylayout="de-latin1"
loadkeys de-latin1
confirm "Do you want to change the keymap from de-latin1 to us? [y/N]" && export keylayout="us" && loadkeys us
echo "Keymap set to $keylayout."
clear

# profile selection
echo "Profiles:"
echo "cli:        barebones system, no Xorg, no DE, no WM. Primarily for servers."
echo "kde:        KDE desktop environment + big software selection"
echo "gnome:      GNOME desktop environment + software selection [WIP]"
echo
read -p "Profile Selection >" profile

if [[ "$profile" != "cli" && "$profile" != "kde" && "$profile" != "gnome" ]]; then
	echo "Invalid profile."
	exit 1
fi

# LUKS selection
clear
if [[ "$efi" -eq 0 ]]; then
	echo "archx does not support LUKS on legacy BIOS systems, skipping..."
	export luks="false"
else
	echo "Encryption:"
	export luks="false"
	confirm "Do you want to encrypt your system and swap partitions? [y/N]" && export luks="true"
fi

# Hostname selection
clear
read -p "Hostname >" hostname
export hostname

# Swap size selection
clear
swapsize=$(cat /proc/meminfo | grep MemTotal | awk '{ print $2 }')
swapsize=$(($swapsize/1000))"M"
echo "Swap size: $swapsize"
confirm "Per default swap is sized equal to RAM. Customize? [y/N]" && read -p "New swap size >" swapsize
export swapsize

# User/Pass
clear
echo "NOTICE: Keyboard layout is $keylayout!"

if [ "$luks" == "true" ]; then
	while [[ "$lukspass" != "$lukspass_verify" ]] || [[ "$lukspass" == "" ]]; do
		read -s -p "LUKS password: " lukspass
		echo
		read -s -p "LUKS password (confirm): " lukspass_verify
		echo
	done
	export lukspass
fi

read -p "Username: " user
export user

while [[ "$password" != "$password_verify" ]] || [[ "$password" == "" ]]; do
	read -s -p "User password: " password
	echo
	read -s -p "User password (confirm): " password_verify
	echo
done
export password

# Summary
clear
echo "---Summary---"
echo "Target disk: $disk"

if [[ $efi -eq 1 ]]; then
	echo "Mode: UEFI"
else
	echo "Mode: Legacy/CSM/BIOS"
fi
echo "LUKS: $luks"
echo "Swap size: $swapsize"
echo "Hostname: $hostname"
echo
echo "Keymap: $keylayout"
echo "Profile: $profile"
echo "Username: $user"
echo

confirm "Everything correct? [FINAL PROMPT] [y/N]" && install
echo "Canceling setup."

