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

############################################
# LOGGING
############################################
LOGFILE="$HOME/arch-autoinstall.log"

{
    echo
    echo "############################################"
    echo " Arch Auto-Install Log - Started $(date)"
    echo "############################################"
    echo
} >> "$LOGFILE"

exec > >(tee -a "$LOGFILE") 2>&1


############################################
# 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; }

############################################
# REQUIRED PACKAGES FOR INSTALL ENVIRONMENT
############################################
pacman -Sy --noconfirm parted cryptsetup libfido2 pam-u2f

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

############################################
# RAM / PARTITION SIZING
############################################
RAM_GB=$(free --giga | awk '/^Mem:/ {print $2}')
BOOT_SIZE=15GiB
SWAP_SIZE="${RAM_GB}GiB"

DISK_SIZE=$(lsblk -b -dn -o SIZE "$DRIVE")
DISK_GIB=$((DISK_SIZE / 1024 / 1024 / 1024))

ROOT_GIB=$((DISK_GIB - RAM_GB - 15))

echo "Partition plan:"
echo "  Boot: ${BOOT_SIZE}"
echo "  Root: ${ROOT_GIB}GiB"
echo "  Swap: ${SWAP_SIZE}"

############################################
# PARTITION DISK
############################################
parted "$DRIVE" --script mklabel gpt \
    mkpart ESP fat32 1MiB 15GiB \
    set 1 boot on \
    mkpart ROOT 15GiB "$((15 + ROOT_GIB))"GiB \
    mkpart SWAP "$((15 + ROOT_GIB))"GiB 100%

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

############################################
# FORMAT BOOT + SWAP
############################################
mkfs.fat -F32 "$BOOT_PART"
mkswap "$SWAP_PART"
swapon "$SWAP_PART"

############################################
# ASK ABOUT FIDO2 LUKS ENROLLMENT
############################################
read -rp "Enable FIDO2 unlocking for root partition? (YES/NO): " FIDO_ROOT

############################################
# LUKS ENCRYPT ROOT
############################################
echo "Encrypting root partition..."
cryptsetup -v luksFormat "$ROOT_PART"
cryptsetup open "$ROOT_PART" root

############################################
# OPTIONAL FIDO2 ENROLLMENT
############################################
if [[ "$FIDO_ROOT" == "YES" ]]; then
    echo "Insert FIDO2 key for LUKS and touch when prompted..."
    systemd-cryptenroll "$ROOT_PART" --fido2-device=auto --fido2-with-client-pin=no
fi

############################################
# BTRFS 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

############################################
# GPU DETECTION
############################################
GPU_INFO=$(lspci | grep -E "VGA|3D")
GPU_PKGS=""
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

echo "Detected GPU: $GPU_INFO"

############################################
# USER INPUT
############################################
read -rp "Enter kernel package (e.g., linux, linux-lts): " KERNEL
read -rp "Enter hostname: " HOSTNAME
read -rp "Enter username: " USERNAME
read -rsp "Enter password for $USERNAME: " USERPASS
echo
read -rp "Enable FIDO2 authentication for user login? (YES/NO): " FIDO_USER

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

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

############################################
# PASS VARIABLES INTO CHROOT
############################################
export HOSTNAME USERNAME USERPASS ROOT_PART FIDO_ROOT FIDO_USER

############################################
# CHROOT CONFIGURATION
############################################
arch-chroot /mnt /bin/bash <<'EOF'
set -euo pipefail

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

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

# NetworkManager
systemctl enable NetworkManager

# User
useradd -m -G wheel -s /bin/zsh "$USERNAME"
echo "$USERNAME:$USERPASS" | chpasswd
echo "%wheel ALL=(ALL:ALL) ALL" >> /etc/sudoers

###################################################
# INITRAMFS CONFIG
###################################################
if [[ "$FIDO_ROOT" == "YES" ]]; then
    sed -i 's/^HOOKS=.*/HOOKS=(base udev systemd microcode kms autodetect modconf consolefont block sd-encrypt encrypt lvm2 filesystems keyboard keymap fsck)/' /etc/mkinitcpio.conf
else
    sed -i 's/^HOOKS=.*/HOOKS=(base udev systemd microcode kms autodetect modconf consolefont block encrypt encrypt lvm2 filesystems keyboard keymap fsck)/' /etc/mkinitcpio.conf
fi

mkinitcpio -p "$KERNEL"

###################################################
# GRUB CONFIG
###################################################
UUID=$(blkid -s UUID -o value "$ROOT_PART")

if [[ "$FIDO_ROOT" == "YES" ]]; then
    KERNEL_CMD="rd.luks.name=${UUID}=cryptroot root=/dev/mapper/cryptroot"
else
    KERNEL_CMD="cryptdevice=UUID=${UUID}:cryptroot root=/dev/mapper/cryptroot"
fi

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

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

###################################################
# USER FIDO2 LOGIN
###################################################
if [[ "$FIDO_USER" == "YES" ]]; then
    mkdir -p /home/$USERNAME/.config/Yubico
    echo "Insert FIDO2 key for user login and touch when prompted..."
    sudo -u "$USERNAME" pamu2fcfg -u "$USERNAME" > /home/$USERNAME/.config/Yubico/u2f_keys
    chown "$USERNAME":"$USERNAME" /home/$USERNAME/.config/Yubico/u2f_keys

    echo "auth required pam_u2f.so" >> /etc/pam.d/system-local-login
fi

EOF

############################################
# SUMMARY OUTPUT (NEW)
############################################
echo
echo "############################################"
echo " INSTALL SUMMARY"
echo "############################################"
echo "Drive:            $DRIVE"
echo "Boot partition:   $BOOT_PART"
echo "Root partition:   $ROOT_PART"
echo "Swap partition:   $SWAP_PART"
echo
echo "Hostname:         $HOSTNAME"
echo "Username:         $USERNAME"
echo "Kernel:           $KERNEL"
echo "GPU detected:     $GPU_INFO"
echo
echo "FIDO2 root unlock: $FIDO_ROOT"
echo "FIDO2 user login:  $FIDO_USER"
echo
echo "Boot size:         $BOOT_SIZE"
echo "Root size:         ${ROOT_GIB}GiB"
echo "Swap size:         $SWAP_SIZE"
echo
echo "Log file saved to: $LOGFILE"
echo "############################################"
echo

cp $LOGFILE /mnt/boot/

############################################
# DONE
############################################
echo "Installation complete! You can now unmount and reboot."