#!/usr/bin/env bash
set -euo pipefail

############################################
#  Helper Functions
############################################

confirm() {
    echo "WARNING: This will ERASE ALL DATA on $1"
    read -rp "Type YES to continue: " ans
    [[ $ans == "YES" ]]
}

ask() {
    local prompt=$1
    local var
    read -rp "$prompt: " var
    echo "$var"
}

pause() {
    read -rp "Press ENTER to continue..."
}

############################################
#  Begin
############################################

echo "== Arch Linux FIDO2-Ready Installer =="

lsblk
DRIVE=$(ask "Enter install drive (e.g., /dev/sda)")
confirm "$DRIVE" || exit 1

# Required packages
pacman -Syd --noconfirm parted cryptsetup libfido2 pam-u2f systemd-ukify

############################################
#  Partitioning
############################################

RAM_GB=$(free --giga | awk '/Mem/ {print $2}')
DISK_GB=$(lsblk -dn -o SIZE -b "$DRIVE" | awk '{print int($1/1024/1024/1024)}')

EFI_SIZE=5
SWAP_SIZE=$RAM_GB
ROOT_SIZE=$((DISK_GB - SWAP_SIZE - EFI_SIZE - 1))

if (( ROOT_SIZE < 8 )); then
    echo "ERROR: Disk too small for layout."
    exit 1
fi

echo "EFI=${EFI_SIZE}G, Root=${ROOT_SIZE}G, Swap=${SWAP_SIZE}G"

parted -s "$DRIVE" mklabel gpt \
    mkpart EFI fat32 1MiB "${EFI_SIZE}GiB" \
    set 1 esp on \
    mkpart ROOT "${EFI_SIZE}GiB" "$((EFI_SIZE + ROOT_SIZE))GiB" \
    mkpart SWAP "$((EFI_SIZE + ROOT_SIZE))GiB" 100%

EFI_PART="${DRIVE}1"
ROOT_PART="${DRIVE}2"
SWAP_PART="${DRIVE}3"

mkfs.fat -F32 "$EFI_PART"
mkswap "$SWAP_PART"
swapon "$SWAP_PART"

############################################
#  Encryption
############################################

echo
read -rp "Enable FIDO2 for unlocking root? (YES/NO): " ENABLE_FIDO_ROOT

echo "Formatting LUKS2 root..."
cryptsetup luksFormat "$ROOT_PART" --type luks2
cryptsetup open "$ROOT_PART" cryptroot

if [[ $ENABLE_FIDO_ROOT == "YES" ]]; then
    echo "Enroll FIDO2 key for LUKS2"
    pause
    systemd-cryptenroll "$ROOT_PART" --fido2-device=auto --fido2-with-client-pin=no
fi

# Add fallback password
echo "Adding fallback LUKS password (recommended)"
cryptsetup luksAddKey "$ROOT_PART"

############################################
#  Filesystem
############################################

mkfs.btrfs /dev/mapper/cryptroot

mount /dev/mapper/cryptroot /mnt
btrfs subvolume create /mnt/@
btrfs subvolume create /mnt/@home
umount /mnt

mount -o subvol=@ /dev/mapper/cryptroot /mnt
mkdir -p /mnt/home
mount -o subvol=@home /dev/mapper/cryptroot /mnt/home

mkdir -p /mnt/boot
mount "$EFI_PART" /mnt/boot

############################################
#  Base System Install
############################################

GPU_INFO=$(lspci | grep -E "VGA|3D")
GPU_PKGS=""

if echo "$GPU_INFO" | grep -qi nvidia; then
    #GPU_PKGS="nvidia nvidia-utils"
    GPU_PKGS="nvidia-open"
elif echo "$GPU_INFO" | grep -qi amd; then
    GPU_PKGS="xf86-video-amdgpu"
elif echo "$GPU_INFO" | grep -qi intel; then
    GPU_PKGS="xf86-video-intel"
fi

KERNEL=$(ask "Kernel (linux, linux-lts, linux-zen)")
HOSTNAME=$(ask "Hostname")
USERNAME=$(ask "Username")
read -rsp "Password for $USERNAME: " USERPASS; echo

read -rp "Enable FIDO2 for user login? (YES/NO): " ENABLE_FIDO_USER

pacstrap /mnt \
    base "$KERNEL" linux-firmware vim zsh git networkmanager grub efibootmgr \
    btrfs-progs cryptsetup libfido2 pam-u2f sudo "$GPU_PKGS"

genfstab -U /mnt >> /mnt/etc/fstab

############################################
#  CHROOT Configuration
############################################

ROOT_UUID=$(blkid -s UUID -o value "$ROOT_PART")


arch-chroot /mnt /usr/bin/env \
  HOSTNAME="$HOSTNAME" \
  USERNAME="$USERNAME" \
  USERPASS="$USERPASS" \
  ENABLE_FIDO_ROOT="$ENABLE_FIDO_ROOT" \
  ENABLE_FIDO_USER="$ENABLE_FIDO_USER" \
  ROOT_UUID="$ROOT_UUID" \
  /bin/bash <<'EOF'

set -euo pipefail

echo "en_US.UTF-8 UTF-8" > /etc/locale.gen
locale-gen
echo "LANG=en_US.UTF-8" > /etc/locale.conf

ln -sf /usr/share/zoneinfo/Europe/Vienna /etc/localtime
hwclock --systohc

echo "$HOSTNAME" > /etc/hostname
systemctl enable NetworkManager

useradd -m -G wheel -s /bin/zsh "$USERNAME"
echo "$USERNAME:$USERPASS" | chpasswd

echo "%wheel ALL=(ALL) ALL" >> /etc/sudoers

# Initramfs + GRUB for FIDO2
if [[ "$ENABLE_FIDO_ROOT" == "YES" ]]; then
    sed -i 's/^HOOKS=.*/HOOKS=(base systemd autodetect modconf block sd-encrypt filesystems keyboard fsck)/' /etc/mkinitcpio.conf
else
    sed -i 's/^HOOKS=.*/HOOKS=(base udev autodetect modconf block encrypt filesystems keyboard fsck)/' /etc/mkinitcpio.conf
fi

mkinitcpio -P

if [[ "$ENABLE_FIDO_ROOT" == "YES" ]]; then
    GRUB_CMDLINE="rd.luks.name=$ROOT_UUID=cryptroot rd.luks.options=fido2-device=auto root=/dev/mapper/cryptroot"
else
    GRUB_CMDLINE="cryptdevice=UUID=$ROOT_UUID:cryptroot root=/dev/mapper/cryptroot"
fi

sed -i "s|^GRUB_CMDLINE_LINUX=.*|GRUB_CMDLINE_LINUX=\"$GRUB_CMDLINE\"|" /etc/default/grub

grub-install --target=x86_64-efi --efi-directory=/boot --bootloader-id=GRUB
grub-mkconfig -o /boot/grub/grub.cfg

# User login FIDO2
if [[ "$ENABLE_FIDO_USER" == "YES" ]]; then
    echo "Enrolling FIDO2 for user login"
    mkdir -p /home/$USERNAME/.config/Yubico
    chown $USERNAME:$USERNAME /home/$USERNAME/.config/Yubico
    sudo -u "$USERNAME" bash -c "pamu2fcfg >> /home/$USERNAME/.config/Yubico/u2f_keys"
    echo "auth required pam_u2f.so" >> /etc/pam.d/system-auth
fi
EOF

echo "Installation complete!"
echo "Run: umount -R /mnt && reboot"