#!/bin/bash
set -e

qpushd() {
	pushd $@ > /dev/null
}

qpopd() {
	popd $@ > /dev/null
}

quiet() {
	local DUMMY
	set +e
	DUMMY=$($@ 2>&1 > /dev/null)
	if [ $? -ne 0 ]; then
		echo "$DUMMY"
		set -e
		return 1
	fi
	set -e
}

ultra_quiet() {
	local DUMMY
	DUMMY=$($@ 2>&1 > /dev/null)
}


### FORCE ROOT ###
[ $(whoami) != "root" ] && echo "Please run as root" && exit 1

### CD TO MY DIR ###
cd "$(dirname "$0")"

### URLs ###
FZF_DOWNLOAD="$(curl -s https://api.github.com/repos/junegunn/fzf/releases/latest | grep linux_amd64 | sed -nE 's/^\s*"browser_download_url":\s*"(.*)"\s*$/\1/p')"
PARTED_DOWNLOAD="https://archlinux.org/packages/extra/x86_64/parted/download"
POST_INSTALL_SCRIPT="https://raw.githubusercontent.com/augustogunsch/install-arch/master/post-install.sh"
KEYBOARD_MAP="https://raw.githubusercontent.com/augustogunsch/install-arch/master/keyboard-map.csv"

### COLORS ###
RED='\033[0;31m'
LGREEN='\033[1;32m'
YELLOW='\033[1;33m'
NC='\033[0m' # No Color
BOLD=$(tput bold)
NORM=$(tput sgr0)
readonly RED
readonly LGREEN
readonly YELLOW
readonly NC
readonly BOLD
readonly NORM

### VARS ###
CUR_PHASE=1
MAX_PHASE=4

### INFO ###
AVAILABLE_PLATFORMS='Both BIOS and UEFI systems\nOnly x86_64 systems\nDistros:\nArch\nArtix (OpenRC)\n'
readonly AVAILABLE_PLATFORMS

echo "This script can only be run interactively. Make sure you are in a supported platform and have an Internet connection. Available platforms:"
echo -e "$AVAILABLE_PLATFORMS"

### SYSTEM ###
DISTRO=$(cat /etc/os-release | sed -nE 's/^ID=(.*)/\1/p')
INIT_SYS=$(basename $(readlink /bin/init))
set +e
[ -d /sys/firmware/efi ] && UEFI=1 || UEFI=0
set -e
readonly DISTRO
readonly INIT_SYS
readonly UEFI

right_chroot() {
	[ "$DISTRO" = "arch" ] && arch-chroot $@ || artix-chroot $@
}

right_fstabgen() {
	[ "$DISTRO" = "arch" ] && genfstab $@ || fstabgen $@
}

right_basestrap() {
	[ "$DISTRO" = "arch" ] && pacstrap $@ || basestrap $@
}

print_phase() {
	echo -e "${BOLD}${YELLOW}[$CUR_PHASE/$MAX_PHASE] $1 phase${NC}${NORM}"
	CUR_PHASE=$((CUR_PHASE+1))
}

download_fzf() {
	[ -f /usr/bin/fzf ] && return 0
	echo -n "Downloading fzf (for script use only)..."
	curl -sL "$FZF_DOWNLOAD" -o fzf.tar.gz
	tar -xf fzf.tar.gz
	mv ./fzf /usr/bin/fzf
	rm fzf.tar.gz
	echo "done"
}

download_parted() {
	[ -f /usr/bin/parted ] && return 0
	echo -n "Downloading parted (for script use only)..."
	curl -sL "$PARTED_DOWNLOAD" -o parted.tar.zst
	tar -xf parted.tar.zst
	cp -r ./usr /
	rm -r ./usr
	rm parted.tar.zst
	echo "done"
}

prompt() {
	echo -n "$1 [Y/n] "
	[ $NO_CONFIRM ] && echo "y" && return 1
	read ans
	case $ans in
		n|N) return 0 ;;
		*) return 1 ;;
	esac
}

prompt_drive() {
	local DRIVES="$(lsblk -pno NAME,TYPE,MODEL | awk 'BEGIN {count=1} $1 ~ /^\// { 
		if ($2 == "disk") {printf("%i'")"' %s    \"", count, $1); for(i=3;i<NF;i++) printf("%s ", $i); printf("%s\"\n", $NF); count++ } }')" 
	echo "Available drives:"
	printf "   %-12s %s\n" "DISK" "IDENTIFIER"
	echo -e "$DRIVES" 
	echo -n "Choose drive to install $DISTRO into (WARNING: it will be repartitioned and wiped out) (default=1): "
	local drive
	read drive

	[ -z "$drive" ] && drive=1
	DRIVE_TARGET=$(echo -e "$DRIVES" | awk '$1 ~ /^'$drive'\)$/ { print $2 }')
	if [ -z "$DRIVE_TARGET" ]; then
		echo "Invalid target"
		exit 1
	fi

	set +e
	prompt "Installing to $DRIVE_TARGET. Confirm?"
	[ $? -eq 0 ] && exit 0
	set -e
}

partition() {
	print_phase "Disk partitioning"
	set +e
	ultra_quiet swapoff -a
	set -e
	[ -f /bin/parted ] || download_parted

	local rootN
	echo -n "Partitioning drive..."

	if [ $UEFI -eq 0 ]; then
	# Legacy
		rootN=2
		parted --script "$DRIVE_TARGET" \
		mklabel msdos \
		mkpart primary linux-swap 0% 4GiB \
		mkpart primary ext4 4GiB 100%
		echo "done"
	else
	# EFI
		rootN=3
		parted --script "$DRIVE_TARGET" \
		mklabel gpt \
		mkpart swap linux-swap 0% 4GiB \
		mkpart boot fat32 4GiB 5Gib \
		mkpart root ext4 5GiB 100% \
		set 2 esp on
		echo "done"

		echo -n "Configuring BOOT partition..."
		quiet mkfs.fat -F 32 "$DRIVE_TARGET"2
		fatlabel "$DRIVE_TARGET"2 BOOT
		echo "done"
	fi

	echo -n "Configuring SWAP partition..."
	quiet mkswap -L SWAP "$DRIVE_TARGET"1
	quiet swapon "$DRIVE_TARGET"1
	echo "done"

	echo -n "Configuring ROOT partition..."
	quiet mkfs.ext4 -L ROOT "$DRIVE_TARGET"$rootN
	quiet mount "$DRIVE_TARGET"$rootN /mnt
	[ $UEFI -eq 1 ] && mkdir -p /mnt/boot && quiet mount "$DRIVE_TARGET"2 /mnt/boot
	echo "done"
}

install_base() {
	print_phase "System installation"
	echo -n "Installing base system, kernel, bootloader and vi..."
	quiet right_basestrap /mnt base base-devel linux linux-firmware grub vi
	echo "done"

	if [ "$DISTRO" = "artix" ]; then
		if [ "$INIT_SYS" = "openrc-init" ]; then
			echo -n "Installing openrc..."
			quiet right_basestrap /mnt openrc elogind-openrc
			echo "done"
		else	
			echo
			echo "Error: Unsupported init system \"$INIT_SYS\""
			exit 1
		fi
	elif [ "$DISTRO" != "arch" ]; then
		echo
		echo "Error: Unsupported distro \"$DISTRO\""
		exit 1
	fi

	echo -n "Generating fstab..."
	right_fstabgen -U /mnt >> /mnt/etc/fstab
	echo "done"
}

set_timezone() {
	ln -sf "/usr/share/zoneinfo/$TIMEZONE" /mnt/etc/localtime
	quiet right_chroot /mnt hwclock --systohc
}

set_locale() {
	echo -n "Configuring locale..."
	sed "s/^#$LOCALE/$LOCALE/" < /mnt/etc/locale.gen  > /etc/locale.gen
	quiet locale-gen
	cp -f /usr/lib/locale/locale-archive /mnt/usr/lib/locale/locale-archive
	cp -f /etc/locale.gen /mnt/etc/locale.gen

	local short_locale
	local encoding
	read short_locale encoding <<< "$LOCALE"
	echo "export LANG=\"$short_locale\"" > /mnt/etc/locale.conf
	echo "done"

	echo -n "Setting keyboard layout..."
	IFS=, read -r dummy XKBD_LAYOUT XKBD_MODEL XKBD_VARIANT XKBD_OPTIONS <<< "$(grep "^$KBD_LAYOUT," -m1 keyboard-map.csv)"

	# systemd and others may read from here
	echo "KEYMAP=$KBD_LAYOUT" > /mnt/etc/vconsole.conf

	# while openrc and others may read from here
	KBD_LAYOUT_FULL_PATH=$(find /usr/share/kbd/keymaps -name "$KBD_LAYOUT.map.gz")
	mkdir -p /mnt/etc/conf.d
	echo "keymap=\"$KBD_LAYOUT_FULL_PATH\"" > /mnt/etc/conf.d/keymaps

	# and X11 will read from here
	local XKBD_CONF="/mnt/etc/X11/xorg.conf.d/00-keyboard.conf"
	mkdir -p $(dirname $XKBD_CONF)
	echo "Section \"InputClass\"" > $XKBD_CONF
	echo "	Identifier \"system-keyboard\"" >> $XKBD_CONF
	echo "	MatchIsKeyboard \"on\"" >> $XKBD_CONF
	echo "	Option \"XkbLayout\" \"$XKBD_LAYOUT\"" >> $XKBD_CONF
	echo "	Option \"XkbModel\" \"$XKBD_MODEL\"" >> $XKBD_CONF
	[ -n "$XKBD_VARIANT" ] && echo "	Option \"XkbVariant\" \"$XKBD_VARIANT\"" >> $XKBD_CONF
	[ -n "$XKBD_OPTIONS" ] && echo "	Option \"XkbOptions\" \"$XKBD_OPTIONS\"" >> $XKBD_CONF
	echo "EndSection" >> $XKBD_CONF
	
	[ "$INIT_SYS" = "openrc-init" ] && quiet right_chroot /mnt rc-update add keymaps boot
	echo "done"
}

setup_grub() {
	echo -n "Configuring boot loader..."
	if [ $UEFI -eq 1 ]; then
		quiet right_basestrap /mnt efibootmgr
		quiet right_chroot /mnt grub-install --target=x86_64-efi --efi-directory=/boot --bootloader-id=GRUB
	else
		quiet right_chroot /mnt grub-install --target=i386-pc "$DRIVE_TARGET"
	fi
	quiet right_chroot /mnt grub-mkconfig -o /boot/grub/grub.cfg
	echo "done"
}

setup_users() {
	echo -n "Configuring users..."

	set +e
	# There might be a group with the user's name
	awk -F: '$1 ~ /^'$PERSONAL_USER'$/ { exit 1 }' /etc/group
	[ $? -eq 0 ] && right_chroot /mnt useradd --badnames -m "$PERSONAL_USER" || \
		right_chroot /mnt useradd --badnames -m -g "$PERSONAL_USER" "$PERSONAL_USER"
	set -e

	echo -e "root:$ROOT_PASSWORD\n$PERSONAL_USER:$PERSONAL_PASSWORD" | chpasswd -R /mnt
	echo "done"
}

setup_network() {
	echo -n "Configuring hostname and network..."
	echo "$MACHINE_HOSTNAME" > /mnt/etc/hostname
	echo "127.0.0.1	localhost" > /mnt/etc/hosts
	echo "::1	localhost" >> /mnt/etc/hosts
	echo "127.0.1.1	$MACHINE_HOSTNAME.localdomain	$MACHINE_HOSTNAME" >> /mnt/etc/hosts

	quiet right_basestrap /mnt dhcpcd wpa_supplicant
	if [ "$DISTRO" = "artix" ]; then
		if [ "$INIT_SYS" = "openrc-init" ]; then
			echo "hostname=\"$MACHINE_HOSTNAME\"" > /mnt/etc/conf.d/hostname
			quiet right_basestrap /mnt connman-openrc
			quiet right_chroot /mnt rc-update add connmand
		fi
	else
		quiet right_chroot /mnt systemctl enable dhcpcd
	fi
	echo "done"
}

ask_password() {
	echo -n "Type password for $1: "
	stty -echo
	read USER_PASSWORD
	stty echo
	echo
	echo -n "Confirm password: "
	stty -echo
	local PASSWORD_CONFIRM
	read PASSWORD_CONFIRM
	stty echo
	echo
	if [ "$USER_PASSWORD" != "$PASSWORD_CONFIRM" ]; then
		echo "Wrong passwords. Please try again."
		ask_password $1
	fi
}

prompt_all() {
	download_fzf

	prompt_drive

	echo -n "Choose timezone:"
	qpushd /usr/share/zoneinfo
	TIMEZONE="$(fzf --layout=reverse --height=20)"
	qpopd
	echo "Choose timezone: $TIMEZONE"

	echo -n "Choose locale:"
	LOCALE=$(sed '/^#\s/D' < /etc/locale.gen | sed '/^#$/D' | sed 's/^#//' | fzf --layout=reverse --height=20)
	echo "Choose locale: $LOCALE"

	[ -f keyboard-map.csv ] || curl -sL "$KEYBOARD_MAP" -o keyboard-map.csv
	echo -n "Choose keyboard layout:"
	KBD_LAYOUT="$(awk -F, '{print $1}' keyboard-map.csv | fzf --layout=reverse --height=20)"
	echo "Choose keyboard layout: $KBD_LAYOUT"

	ask_password root
	ROOT_PASSWORD="$USER_PASSWORD"

	echo -n "Type your personal username: "
	read PERSONAL_USER
	ask_password "$PERSONAL_USER"
	PERSONAL_PASSWORD="$USER_PASSWORD"

	echo -n "Type the machine hostname: "
	read MACHINE_HOSTNAME
}

post_install() {
	curl -sL "$POST_INSTALL_SCRIPT" -o post-install.sh
	mv post-install.sh /mnt/root
	chmod +x /mnt/root/post-install.sh
	print_phase "Post installation"
	echo "From now on, the post installation script will be running"
	right_chroot /mnt /root/post-install.sh -nu "$PERSONAL_USER"
}

configure() {
	print_phase "System configuration"

	set_timezone
	set_locale
	setup_grub
	setup_users
	setup_network
}

main() {
	prompt_all
	partition
	install_base
	configure
	post_install

	umount -R /mnt
	echo "ALSA must be configured manually. You can use 'alsamixer' to configure and then 'alsactl store' to save the changes."
	echo "Most likely you just have to unmute the Main channel."
	echo "Graphics drivers must also be installed manually. See Arch Wiki: Xorg#Driver_installation."
	echo -n "Ready to reboot. Press any key to continue..."
	read dummy
	reboot
}

main