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

# Safety warning
echo "WARNING: This will ERASE all data on the selected drive!"
read -rp "Type 'YES' to continue: " confirm
[[ "$confirm" == "YES" ]] || { echo "Aborted."; exit 1; }

# Ensure required packages
pacman -Sy --noconfirm parted cryptsetup libfido2 pam-u2f

# Ask for drive
lsblk
read -rp "Enter target drive (e.g., /dev/sda): " DRIVE

# Get RAM size in GB
RAM_GB=$(free --giga | awk '/^Mem:/ {print $2}')

# Calculate partition sizes
BOOT_SIZE=15G
SWAP_SIZE="${RAM_GB}G"
DISK_SIZE=$(lsblk -b -dn -o SIZE "$DRIVE" | awk '{print int($1/1024/1024/1024)}')
ROOT_SIZE=$((DISK_SIZE - RAM_GB - 15))

echo "Partitioning $DRIVE: Boot=$BOOT_SIZE, Root=${ROOT_SIZE}G, Swap=$SWAP_SIZE"

# Partition the disk
parted "$DRIVE" --script mklabel gpt \
    mkpart ESP fat32 1MiB 15GiB \
    set 1 boot on \
    mkpart ROOT 15GiB "$((15+ROOT_SIZE))GiB" \
    mkpart SWAP "$((15+ROOT_SIZE))GiB" 100%

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

# Format boot and swap
mkfs.fat -F32 "$BOOT_PART"
mkswap "$SWAP_PART"
swapon "$SWAP_PART"

# Ask if FIDO2 unlocking should be enabled for LUKS
read -rp "Enable FIDO2 unlocking for root partition? (YES/NO): " FIDO_ROOT

# Encrypt root partition
echo "Encrypting root partition with LUKS..."
cryptsetup luksFormat "$ROOT_PART" --type luks2
cryptsetup open "$ROOT_PART" cryptroot


# Add multiple FIDO2 keys for LUKS
if [[ "$FIDO_ROOT" == "YES" ]]; then
    echo "ENROLLING FIDOKEY"
    read
    systemd-cryptenroll "$ROOT_PART" --fido2-device=auto --fido2-with-client-pin=no
#    read -rp "How many FIDO2 keys for LUKS unlocking? " FIDO_LUKS_COUNT
#    for ((i=1; i<=FIDO_LUKS_COUNT; i++)); do
#        echo "Insert FIDO2 key #$i and touch it when prompted..."
#        syste "$ROOT_PART" --fido2-device=auto
#    done
fi

# Format encrypted root as Btrfs and create subvolumes
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 /mnt/home
mount -o subvol=@home /dev/mapper/cryptroot /mnt/home
mkdir /mnt/boot
mount "$BOOT_PART" /mnt/boot

# Detect GPU and set driver packages
GPU_PKGS=""
GPU_INFO=$(lspci | grep -E "VGA|3D")
echo "Detected GPU: $GPU_INFO"
if echo "$GPU_INFO" | grep -qi "NVIDIA"; then
    GPU_PKGS="nvidia nvidia-utils"
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

# Ask for kernel
read -rp "Enter kernel package (e.g., linux, linux-lts): " KERNEL

# Prompt for user setup
read -rp "Enter hostname: " HOSTNAME
read -rp "Enter username: " USERNAME
read -rsp "Enter password for $USERNAME: " USERPASS
echo

# Ask if FIDO2 should be enabled for user login
read -rp "Enable FIDO2 authentication for user login? (YES/NO): " FIDO_USER

# Install base system
pacstrap /mnt base "$KERNEL" linux-firmware vim bash zsh git less btop fastfetch networkmanager grub cryptsetup libfido2 pam-u2f efibootmgr sudo  $GPU_PKGS

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

# Chroot configuration
arch-chroot /mnt /bin/bash <<EOF
# Locale setup
echo "en_US.UTF-8 UTF-8" > /etc/locale.gen
locale-gen
echo "LANG=en_US.UTF-8" > /etc/locale.conf

# Timezone and hostname
ln -sf /usr/share/zoneinfo/Europe/Vienna /etc/localtime
hwclock --systohc
echo "$HOSTNAME" > /etc/hostname

# Network
systemctl enable NetworkManager

# Create user (with sudo privileges and zsh as default shell)
useradd -m -G wheel -s /bin/zsh "$USERNAME"
echo "$USERNAME:$USERPASS" | chpasswd
echo "%wheel ALL=(ALL) ALL" >> /etc/sudoers

if [[ "$FIDO_ROOT" == "YES" ]]; then

    # Initramfs hooks for LUKS + FIDO2
    sed -i 's/^HOOKS=.*/HOOKS=(base udev autodetect modconf block encrypt filesystems keyboard fsck)/' /etc/mkinitcpio.conf
    mkinitcpio -P

    # GRUB config for encrypted root
    UUID=\$(blkid -s UUID -o value $ROOT_PART)
    sed -i "s|GRUB_CMDLINE_LINUX=.*|GRUB_CMDLINE_LINUX=\"cryptdevice=UUID=\$UUID:cryptroot root=/dev/mapper/cryptroot\"|" /etc/default/grub
fi
grub-install --target=x86_64-efi --efi-directory=/boot --bootloader-id=GRUB
grub-mkconfig -o /boot/grub/grub.cfg

# FIDO2 for user login
if [[ "$FIDO_USER" == "YES" ]]; then
    echo "ENROLLING FIDOKEY"
    read
    mkdir -p /home/$USERNAME/.config/Yubico
    read -rp "How many FIDO2 keys for user login? " FIDO_USER_COUNT
    for ((i=1; i<=FIDO_USER_COUNT; i++)); do
        echo "Insert FIDO2 key #\$i and touch it when prompted..."
        pamu2fcfg >> /home/$USERNAME/.config/Yubico/u2f_keys
    done
    chown $USERNAME:$USERNAME /home/$USERNAME/.config/Yubico/u2f_keys
    echo "auth required pam_u2f.so" >> /etc/pam.d/system-auth
fi
EOF
echo "==============================================="
echo "Input Summary"
echo "==============================================="
echo $FIDO_USER
echo $FIDO_ROOT
echo "==============================================="

echo "Installation complete! Unmount and reboot."