#!/bin/bash

# 
# ARTIX LINUX INSTALLATION
# 

# See the README for details

# FIXME: DHCP client from iwd does not work
# TODO: add BIOS support
# TODO: add pacman repository options (https://wiki.artixlinux.org/Main/Repositories)

source config

# Ensure nothing mounted
swapoff -a &> /dev/null
cryptsetup close cryptroot &> /dev/null
umount -R /mnt &> /dev/null

# Init shell environment
set -e

# Checks
if [[ ! -d /sys/firmware/efi/efivars ]]; then
    echo "Only UEFI systems are currently supported"
    exit
fi
if [[ "${drive}" == "/dev/DRIVE" ]]; then
    echo "You forgot to set the DRIVE option!"
    exit
fi

echo "Checking for internet connection..."
ping -c 3 artixlinux.org &> /dev/null \
    || { echo "No internet connection found"; exit; }

# Sync clock
dinitctl start ntpd

# Read password
echo "Enter a password for ${user}"
while true; do
    read -sr -p "Password: " password
    printf "\n"
    read -sr -p "Confirm password: " password2
    printf "\n"
    [[ "${password}" == "${password2}" ]] && break
    echo "Incorrect password!";
    read -rp "Press ENTER to try again...";
done

# Get system RAM size
pacman --needed --noconfirm -Sy bc
ram_kB=$(awk 'FNR==1 {print $2}' /proc/meminfo)
ram_gb=$(bc <<< "${ram_kB} / 1000^2")

# Check there is at least 1GB RAM for swap
if [[ "${ram_gb}" -lt 1 ]]; then
    echo "Not enough ram for SWAP"
    exit
fi

# Calculate SWAP size
if [[ "${swap_size}" == auto ]]; then
    swap_size="$(bc <<< "sqrt(${ram_gb}) * 4")G"
fi

# Request confirmation
drive_bytes=$(blockdev --getsize64 "${drive}")
drive_size="$(bc <<< "${drive_bytes} / 1000000000")G"

# Get a list of options
[[ $encrypt == true ]] && config_options+="encrypt "
[[ $autologin == true ]] && config_options+="autologin "

echo "
================ CONFIRM INSTALLATION ================
Drive: ${drive} (size: ${drive_size})
BOOT Partition: ${boot}, Size: ${boot_size}
ROOT Partition: ${root}, Size: MAX
SWAP Size: ${swap_size}
------------------------------------------------------"
if [[ $config_options != "" ]]; then
    echo "Options: ${config_options}
------------------------------------------------------"
fi
echo "CAUTION: ALL data from ${drive} will be erased !!!
------------------------------------------------------"

echo "Are you sure you want continue?"
unset input
read -rp "Type YES (in uppercase letters) to begin installation: " input
[[ "${input}" != "YES" ]] && exit

# Wipe file-system
wipefs -a "${drive}"

# Create UEFI boot & root partition
printf ',%s,U,*\n,+,L\n' "${boot_size}" | sfdisk -qf -X gpt ${drive}

# Encryption setup
if [[ $encrypt == true ]]; then

    # Read crypt password
    echo "Enter a password for encryption"
    while true; do
        read -sr -p "Password (crypt): " cryptpassword
        printf "\n"
        read -sr -p "Confirm encryption password: " cryptpassword2
        printf "\n"
        [[ "${cryptpassword}" == "${cryptpassword2}" ]] && break
        echo "Incorrect password!";
        read -rp "Press ENTER to try again...";
    done

    # Create encrypted drive
    echo "${cryptpassword}" | cryptsetup --hash sha512 \
                                         --pbkdf pbkdf2 \
                                         --label LUKS \
					 luksFormat "${root}"

    # Open encrypted drive
    echo "${cryptpassword}" | cryptsetup luksOpen ${root} cryptroot

    # Change root to mapper
    root="/dev/mapper/cryptroot"
fi

# Make filesystem
mkfs.fat -n BOOT -F 32 "${boot}"
mkfs.btrfs -qfL ROOT "${root}"

# Mount btrfs ROOT drive
mount "${root}" /mnt

# Create BTRFS subvolumes
btrfs -q subvolume create /mnt/@
btrfs -q subvolume create /mnt/@home
btrfs -q subvolume create /mnt/@tmp
btrfs -q subvolume create /mnt/@var
btrfs -q subvolume create /mnt/@snapshots
btrfs -q subvolume create /mnt/@swap

# Mount BTRFS subvolumes
umount /mnt
options="noatime,compress=zstd"
mount -o "${options},subvol=@" "${root}" /mnt
mkdir /mnt/{boot,home,tmp,var,.snapshots,.swap}
mount -o "${options},subvol=@home" "${root}" /mnt/home
mount -o "${options},subvol=@tmp" "${root}" /mnt/tmp
mount -o "${options},subvol=@var" "${root}" /mnt/var
mount -o "${options},subvol=@snapshots" "${root}" /mnt/.snapshots \
    && chmod 750 /mnt/.snapshots
mount -o "nodatacow,subvol=@swap" "${root}" /mnt/.swap

# Create swap file
btrfs filesystem mkswapfile \
      --size "$swap_size" \
      --uuid clear \
      /mnt/.swap/swapfile
btrfs property set /mnt/.swap compression none
swapon /mnt/.swap/swapfile

# Mount boot partition.
mount "${boot}" /mnt/boot

# Sync packages
pacman -Syy

# Base packages
basestrap /mnt base base-devel dinit elogind-dinit

# Linux & utilities
basestrap /mnt --needed \
          linux linux-firmware \
          grub efibootmgr os-prober \
          mkinitcpio btrfs-progs \
          alsa-utils curl less \
          vim git man-{db,pages} \
          bc acpi btop fastfetch tmux \
          openresolv tree unzip

if [[ "${extrapkgs}" == true ]]; then
# Extra packages
basestrap /mnt --needed \
          aspell-en jq \
          artools-base pacman-contrib namcap \
          calcurse rclone rsync snapper stow \
          ttf-{hack,dejavu,liberation} \
          ttf-{hack-nerd,font-awesome} \
          brightnessctl
fi

# Get CPU type for microcode
ucode=amd-ucode
if [[ $(grep "vendor_id" /proc/cpuinfo) == *Intel* ]]; then
    ucode=intel-ucode
fi
basestrap /mnt "${ucode}" --overwrite "/mnt/boot/${ucode}.img"

# Install crypt service
if [[ "${encrypt}" == true ]]; then
    basestrap /mnt cryptsetup-dinit
fi

# Install services
basestrap /mnt --needed \
          {iwd,dhcpcd,openntpd,cronie,openssh,ufw,dbus}-dinit

# Enable services
services="dbus ufw iwd dhcpcd ntpd"

# Generate file-system table
fstabgen -L /mnt >> /mnt/etc/fstab

# Set swappiness levels
[ -d /mnt/etc/sysctl.d/ ] || mkdir -p /mnt/etc/sysctl.d/
echo vm.swappiness=10 > /mnt/etc/sysctl.d/99-swappiness.conf

# SETUP SYSTEM
# Set locale
echo "${locale}.UTF-8 UTF-8
${locale} ISO-8859-1" >> /mnt/etc/locale.gen
echo "LANG=${locale}.UTF-8
export LC_COLLATE=C" > /mnt/etc/locale.conf
artix-chroot /mnt bash -c "locale-gen"

# Set timezone
artix-chroot /mnt bash -c \
             "ln -sf /usr/share/zoneinfo/${timezone} /etc/localtime"
artix-chroot /mnt bash -c "hwclock -w"

# Set default text editor
echo "export EDITOR=vim" >> /mnt/etc/profile

# Set hostname
echo "${hostname}" > /mnt/etc/hostname

# Setup localhost
tee -a /mnt/etc/hosts <<EOF
127.0.0.1 localhost
::1       localhost
127.0.1.1 ${hostname}.localdomain ${hostname}
EOF

# Set root password
artix-chroot /mnt bash -c "echo root:artix | chpasswd"

# Add new user
artix-chroot /mnt bash -c "useradd -mG ${user_groups} ${user}"
artix-chroot /mnt bash -c "echo \"${user}:${password}\" | chpasswd"

# Set user privileges
echo "
Cmnd_Alias STAT = /usr/bin/ufw status
Cmnd_Alias PACMAN = /usr/bin/checkupdates
Cmnd_Alias POWER = /usr/bin/halt,/usr/bin/reboot
Defaults pwfeedback
%wheel ALL=(ALL) ALL
${user} ALL=(ALL) NOPASSWD: PACMAN,POWER,STAT
" >> /mnt/etc/sudoers

# Login
# when using greetd
if [[ "${greeter}" == greetd ]]; then
    basestrap /mnt --needed greetd-{dinit,tuigreet}
    services+=" greetd"
    # disable tty1 console for greetd login
    sed 's/\[.*\]/\[2-6\]/' -i /mnt/etc/dinit.d/config/console.conf
    # set tuigreet as greetd command
    tuigreet_args="\
--cmd '/bin/tmux new -s TMUX' \
--power-shutdown 'sudo poweroff' \
--power-reboot 'sudo reboot' \
--power-no-setsid \
--asterisks \
--remember \
--time"
    sed "s,^command =.*,command = \"tuigreet ${tuigreet_args}\"," \
        -i /mnt/etc/greetd/config.toml
    # Allow greeter to run power commands
    printf "greeter ALL=(ALL) NOPASSWD: POWER\n" >> /mnt/etc/sudoers

    if [[ "${autologin}" == true ]]; then
        echo "
[initial_session]
command = \"bash\"
user = \"${user}\"
" >> /mnt/etc/greetd/config.toml
    fi
else
    # when using agetty
    if [[ "${autologin}" == true ]]; then
        cp /mnt/etc/dinit.d/config/agetty-default.conf \
           /mnt/etc/dinit.d/config/agetty-tty1.conf
        sed "s/GETTY_ARGS=.*/GETTY_ARGS=\"--noclear --autologin ${user}\"/" \
            -i /mnt/etc/dinit.d/config/agetty-tty1.conf
    fi
fi

# Setup dinit user log directory
mkdir /mnt/var/log/dinit/user
chgrp log /mnt/var/log/dinit/user
chmod g+rw /mnt/var/log/dinit/user

# Setup user log directory (for manual logs)
mkdir /mnt/var/log/user
chgrp log /mnt/var/log/user

# Enable services
# NOTE: do not quote 'services' variable or space is ignored
for service in ${services}; do
    artix-chroot /mnt bash -c \
                 "ln -s /etc/dinit.d/$service /etc/dinit.d/boot.d/"
done

# Enable firewall
sed "s/ENABLED=no/ENABLED=yes/" -i /mnt/etc/ufw/ufw.conf

# Enable iwd's DHCP client (instead of dhcpcd)
[[ ! -d /mnt/etc/iwd ]] && mkdir /mnt/etc/iwd
echo "
[General]
EnableNetworkConfiguration=true

[Network]
NameResolvingService=resolvconf
" > /mnt/etc/iwd/main.conf

unlink /mnt/etc/dinit.d/boot.d/dhcpcd

# Set MAKEFLAGS to match CPU threads for faster compiling
cp /etc/makepkg.conf /etc/makepkg.conf.bak
sed "s/#MAKEFLAGS=\".*\"/MAKEFLAGS=\"-j$(nproc)\"/" -i /mnt/etc/makepkg.conf

# Configure mkinitcpio.conf
sed "s/^MODULES=(.*)/MODULES=(btrfs)/" -i /mnt/etc/mkinitcpio.conf

if [[ "${encrypt}" == true ]]; then
    hooks="base udev autodetect modconf kms keyboard keymap block encrypt resume filesystems fsck"
    sed "s/^HOOKS=(.*)/HOOKS=(${hooks})/" -i /mnt/etc/mkinitcpio.conf
fi

# Rebuild ram-disk environment for Linux kernel
artix-chroot /mnt bash -c "mkinitcpio -p linux"

# SETUP BOOT LOADER
# --------------------------------------------------------------------
# Add swap device
devices="resume=LABEL=ROOT"

# BTRFS swap files requires an offset for hibernation to work
# https://man.archlinux.org/man/btrfs.5#HIBERNATION
offset=$(btrfs inspect-internal map-swapfile -r /mnt/.swap/swapfile)
devices+=" resume_offset=$offset"

# Add crypt device if enabled
if [[ "${encrypt}" == true ]]; then
    devices+=" cryptdevice=LABEL=LUKS:$(basename $root)"
fi

# Set command options
grub_cmds="quiet loglevel=3 net.iframes=0 splash"

# Replace default grub commands
sed "s/^GRUB_CMDLINE_LINUX_DEFAULT=.*/GRUB_CMDLINE_LINUX_DEFAULT=\"$grub_cmds $devices\"/" \
    -i /mnt/etc/default/grub

# Install grub bootloader
grub_options="--target=x86_64-efi --efi-directory=/boot --bootloader-id=GRUB"
artix-chroot /mnt bash -c "grub-install ${grub_options}"
artix-chroot /mnt bash -c "grub-mkconfig -o /boot/grub/grub.cfg"

# FINISH
swapoff -a
umount -R /mnt
[[ "${encrypt}" == true ]] && cryptsetup close cryptroot
set +x

echo "
======================================================================
                        Installation Finished
======================================================================
"
echo "You can now reboot and log into system"