#!/usr/bin/env bash
# arch-autoinstall.sh — automated Arch Linux base installer
#
# If /answerfile.json exists (e.g. embedded in the ISO via build.sh --preconf),
# all prompts are answered from it. Missing fields fall back to interactive prompts.
#
# Answerfile fields: drive, kernel, keymap, hostname, username, encrypt, fido2_root,
#                    fido2_user, run_tui  (password always prompted interactively)

# -E: propagate ERR trap into functions and subshells; -e: abort on any error;
# -u: treat unset variables as errors; -o pipefail: catch failures inside pipes.
set -Eeuo pipefail

############################################
# LOGGING
############################################
LOGFILE="$HOME/arch-autoinstall.log"
{
    echo
    echo "############################################"
    echo " Arch Auto-Install Log - Started $(date)"
    echo "############################################"
    echo
} >> "$LOGFILE"
# Redirect all subsequent stdout and stderr to both the terminal and the log.
# `tee -a` appends so repeated runs accumulate without overwriting; process
# substitution keeps the descriptor open for the full lifetime of the script.
exec > >(tee -a "$LOGFILE") 2>&1

############################################
# Error handler — TUI prompt to send log via croc
############################################
error_handler() {
    # Capture exit code before any other command can overwrite $?.
    # $LINENO is passed as $1 from the trap, so we default to '?' if missing.
    local exit_code=$? line_num="${1:-?}"
    echo "" >> "$LOGFILE"
    echo "ERROR: installer failed at line $line_num (exit code $exit_code)" >> "$LOGFILE"

    # Prefer a dialog TUI if available; fall back to plain readline on headless installs.
    if command -v dialog &>/dev/null; then
        if dialog --clear --title " Installer Error " \
                  --yesno \
                  "Installation failed at line $line_num (exit code: $exit_code).\n\nSend the log to another system via croc for analysis?" \
                  9 62; then
            clear
            # croc may not be present in the base live ISO — install it on demand.
            if ! command -v croc &>/dev/null; then
                echo "Installing croc..."
                pacman -Sy --noconfirm croc 2>/dev/null || true
            fi
            if command -v croc &>/dev/null; then
                croc send "$LOGFILE"
            else
                echo "croc unavailable — log is at: $LOGFILE"
            fi
        else
            clear
            echo "Log saved to: $LOGFILE"
        fi
    else
        echo ""
        echo "Installation failed at line $line_num (exit code $exit_code)."
        read -rp "Send log via croc? [y/N]: " _croc_ans
        # ${_croc_ans,,} lowercases the answer so "Y", "y" both match.
        if [[ "${_croc_ans,,}" == "y" ]]; then
            command -v croc &>/dev/null || pacman -Sy --noconfirm croc 2>/dev/null || true
            croc send "$LOGFILE" || true
        else
            echo "Log saved to: $LOGFILE"
        fi
    fi
    exit "$exit_code"
}
# Fire error_handler on any ERR signal, passing the current line number.
# Single quotes prevent $LINENO from being evaluated at trap-definition time.
trap 'error_handler $LINENO' ERR

############################################
# ANSWERFILE
############################################
# Allow an external override via the environment; default to /answerfile.json
# which is where build.sh --preconf embeds the file in the live ISO.
ANSWERFILE="${ANSWERFILE:-/answerfile.json}"
AF_MODE=false
[[ -f "$ANSWERFILE" ]] && AF_MODE=true

af_get() {
    # af_get <jq-filter> [default]
    # Reads one field from the answerfile with a jq expression.
    # `// empty` returns empty string instead of "null" when the field is absent;
    # the outer `|| true` prevents a jq parse error from aborting the script.
    local val
    val=$(jq -r "${1} // empty" "$ANSWERFILE" 2>/dev/null || true)
    if [[ -z "$val" ]]; then printf '%s' "${2:-}"; else printf '%s' "$val"; fi
}

af_bool() {
    # Returns YES or NO from a JSON boolean field
    # `// false` provides a safe default when the key is absent entirely.
    local val; val=$(jq -r "${1} // false" "$ANSWERFILE" 2>/dev/null || true)
    [[ "$val" == "true" ]] && echo "YES" || echo "NO"
}

get_mac_suffix() {
    # Produce a compact MAC address string (no colons) from the first non-loopback
    # interface to create a unique per-machine hostname suffix at deploy time.
    # The awk pattern skips "lo" by requiring the 3rd char to differ from 'o'.
    local mac
    mac=$(ip link show 2>/dev/null \
        | awk '/^[0-9]+: [^l][^o]/{iface=1} iface && /link\/ether/{print $2; iface=0; exit}')
    # Remove colons via parameter substitution: ${var//pattern/replacement}.
    printf '%s' "${mac//:/}"
}

# Map drive paths to their correct partition suffix.
# NVMe (/dev/nvme0n1) and eMMC (/dev/mmcblk0) use a 'p' before the part number;
# conventional block devices (/dev/sda) do not.
part() { [[ "$1" == *nvme* || "$1" == *mmcblk* ]] && echo "${1}p${2}" || echo "${1}${2}"; }

if $AF_MODE; then
    echo "Answerfile detected: $ANSWERFILE"
    # jq is required to parse the answerfile; it may not be on the base live ISO.
    command -v jq &>/dev/null || pacman -Sy --noconfirm jq
fi

############################################
# NETWORK CHECK
############################################
# Send one ICMP packet (-c1) with a 3-second timeout (-W3) to Cloudflare DNS.
# In AF_MODE we warn but continue — packages in pacstrap may already be cached.
if ! ping -c1 -W3 1.1.1.1 &>/dev/null; then
    if $AF_MODE; then
        echo "Warning: no internet connection detected — continuing in answerfile mode."
    else
        echo "No internet connection detected."
        echo "  Wired: ensure the cable is plugged in."
        echo "  WiFi:  switch to another TTY (Alt+F2) and run: iwctl"
        echo ""
        read -rp "Press Enter once connected (or Ctrl+C to abort)..."
        if ! ping -c1 -W3 1.1.1.1 &>/dev/null; then
            echo "Warning: still offline. Packages cannot be downloaded without network."
            read -rp "Continue anyway? [y/N]: " _net_ans
            [[ "${_net_ans,,}" == "y" ]] || { echo "Aborted — no network."; exit 1; }
        fi
    fi
fi

############################################
# KEYMAP
# To add more layouts: append "code|Display Name" to KEYMAPS
############################################
# Each entry is "keymap-code|Human-readable name"; the pipe acts as a cheap delimiter
# so both pieces of data live in a single array element without needing two arrays.
KEYMAPS=(
    "us|English US"
    "de|German"
)

if $AF_MODE; then
    LIVE_KEYMAP=$(af_get '.keymap' 'us')
else
    echo "Select keyboard layout:"
    for i in "${!KEYMAPS[@]}"; do
        # %%|* strips everything from the first '|' onward → the keymap code.
        _km_code="${KEYMAPS[$i]%%|*}"
        # ##*| strips everything up to and including the last '|' → the display name.
        _km_name="${KEYMAPS[$i]##*|}"
        printf "  %d) %-14s (%s)\n" $((i+1)) "$_km_name" "$_km_code"
    done
    read -rp "Choice [1]: " _KM_IDX
    # Convert 1-based user input to 0-based array index; default to 1 if blank.
    _KM_IDX=$(( ${_KM_IDX:-1} - 1 ))
    if (( _KM_IDX >= 0 && _KM_IDX < ${#KEYMAPS[@]} )); then
        LIVE_KEYMAP="${KEYMAPS[$_KM_IDX]%%|*}"
    else
        # Out-of-range input silently falls back to the first entry (us).
        LIVE_KEYMAP="${KEYMAPS[0]%%|*}"
    fi
fi
# Apply the chosen layout to the live environment immediately so subsequent
# interactive prompts can be typed with the correct physical keyboard.
loadkeys "$LIVE_KEYMAP"
# Also export as KEYMAP so the chroot heredoc can write /etc/vconsole.conf.
KEYMAP="$LIVE_KEYMAP"

############################################
# SAFETY WARNING
############################################
if $AF_MODE; then
    echo "WARNING: Automated install — all data on $(af_get '.drive' '/dev/?') will be ERASED."
    # Give the operator a 5-second window to abort with Ctrl-C before destructive ops begin.
    echo "Proceeding in 5 seconds... (Ctrl-C to abort)"
    sleep 5
else
    echo "WARNING: This will ERASE ALL DATA on the selected drive!"
    read -rp "Type 'YES' to continue: " confirm
    # Require the exact string "YES" (case-sensitive) to prevent accidental wipes.
    [[ "$confirm" == "YES" ]] || { echo "Aborted."; exit 1; }
fi

############################################
# REQUIRED PACKAGES FOR INSTALL ENVIRONMENT
############################################
# Install tools needed by this script into the live environment before partitioning:
# parted — disk partitioning; cryptsetup — LUKS; libfido2/pam-u2f — FIDO2 enrollment.
pacman -Sy --noconfirm parted cryptsetup libfido2 pam-u2f

############################################
# DRIVE SELECTION
############################################
# Print current block device layout so the operator can confirm the target drive.
lsblk
if $AF_MODE && [[ -n "$(af_get '.drive')" ]]; then
    DRIVE=$(af_get '.drive')
    echo "Drive (from answerfile): $DRIVE"
else
    read -rp "Enter target drive (e.g., /dev/sda): " DRIVE
fi

############################################
# USER INPUT
############################################
if $AF_MODE; then
    KERNEL=$(af_get '.kernel' 'linux')
    RAW_HOSTNAME=$(af_get '.hostname' '')
    # Append MAC suffix to make the hostname unique when the same answerfile is
    # deployed to multiple machines (e.g. a lab rollout or reinstall fleet).
    if [[ -n "$RAW_HOSTNAME" ]]; then
        HOSTNAME="${RAW_HOSTNAME}-$(get_mac_suffix)"
    else
        HOSTNAME="arch"
    fi
    USERNAME=$(af_get '.username' '')
    ENCRYPT_DISK=$(af_bool '.encrypt')
    FIDO_ROOT=$(af_bool '.fido2_root')
    FIDO_USER=$(af_bool '.fido2_user')
    RUN_TUI=$(af_bool '.run_tui')
    echo "Kernel:   $KERNEL"
    echo "Hostname: $HOSTNAME"
    echo "Username: $USERNAME"
    echo "Encrypt:  $ENCRYPT_DISK  / FIDO2 root: $FIDO_ROOT  / FIDO2 user: $FIDO_USER"
else
    read -rp "Enter kernel package (e.g., linux, linux-lts): " KERNEL
    read -rp "Enter hostname: " HOSTNAME
    read -rp "Enter username: " USERNAME
    read -rp "Enable disk encryption? (YES/NO): " ENCRYPT_DISK
    # Default FIDO2 root to NO; only ask when encryption is active because FIDO2
    # unlocking only makes sense on an encrypted root partition.
    FIDO_ROOT="NO"
    if [[ "$ENCRYPT_DISK" == "YES" ]]; then
        read -rp "Enable FIDO2 unlocking for root partition? (YES/NO): " FIDO_ROOT
    fi
    read -rp "Enable FIDO2 authentication for user login? (YES/NO): " FIDO_USER
fi

# Password always prompted — never stored in answerfile
read -rp "Enter password for $USERNAME: " USERPASS
[[ -z "$USERPASS" ]] && { echo "Error: password cannot be empty."; exit 1; }

# In interactive mode, decide now whether to run the dotfiles TUI inside chroot.
# AF_MODE already has RUN_TUI set from the answerfile.
if ! $AF_MODE; then
    read -rp "Run dotfiles TUI setup inside chroot now? [YES/no]: " _RUN_TUI_IN
    # Default to YES if the user presses Enter without typing anything.
    _RUN_TUI_IN="${_RUN_TUI_IN:-YES}"
    [[ "${_RUN_TUI_IN^^}" == "YES" ]] && RUN_TUI="YES" || RUN_TUI="NO"
fi

############################################
# RAM / PARTITION SIZING
############################################
# Read installed RAM in GiB; used to size both the swap and to calculate remaining
# root space.  `--giga` prints in GiB units; awk extracts the total column.
RAM_GB=$(free --giga | awk '/^Mem:/ {print $2}')
# 15 GiB boot partition: large enough for multiple kernels, initramfs, and UKIs.
BOOT_SIZE=15GiB
SWAP_SIZE="${RAM_GB}GiB"

# lsblk -b reports raw bytes; divide down to GiB for integer arithmetic.
DISK_SIZE=$(lsblk -b -dn -o SIZE "$DRIVE")
DISK_GIB=$((DISK_SIZE / 1024 / 1024 / 1024))
# Root gets whatever is left after reserving 15 GiB for boot and RAM_GB for swap.
ROOT_GIB=$((DISK_GIB - RAM_GB - 15))
if (( ROOT_GIB < 8 )); then
    echo "ERROR: Not enough disk space. Root would be only ${ROOT_GIB}GiB (need ≥8GiB)."
    exit 1
fi

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

############################################
# PARTITION DISK
############################################
# --script suppresses interactive prompts; mklabel gpt wipes any existing table.
# Partition 1: ESP at 1MiB–15GiB, flagged boot/esp so firmware can find it.
# Partition 2: ROOT immediately after the ESP.
# Partition 3: SWAP fills the rest (100%).
# Starting at 1MiB aligns the first partition to a 1 MiB boundary, which is
# optimal for both SSD erase blocks and spinner track alignment.
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%

# Resolve partition device paths via the `part` helper (handles NVMe 'p' suffix).
BOOT_PART=$(part "$DRIVE" 1)
ROOT_PART=$(part "$DRIVE" 2)
SWAP_PART=$(part "$DRIVE" 3)

############################################
# FORMAT BOOT + SWAP
############################################
# FAT32 is required for the EFI System Partition; -F32 sets the FAT variant.
mkfs.fat -F32 "$BOOT_PART"
# mkswap writes the swap header; swapon activates it for use by pacstrap.
mkswap "$SWAP_PART"
swapon "$SWAP_PART"

############################################
# ENCRYPTION (OPTIONAL)
############################################
LUKS_BACKUP_KEY=""   # path to key file, set only when encryption is active

if [[ "$ENCRYPT_DISK" == "YES" ]]; then
    echo "Encrypting root partition..."
    # -v (verbose) shows progress; luksFormat initialises a LUKS1 container by default
    # and prompts for the primary passphrase interactively.
    cryptsetup -v luksFormat "$ROOT_PART"
    # Open the container and expose it as /dev/mapper/cryptroot for formatting.
    cryptsetup open "$ROOT_PART" cryptroot

    # ── Auto-generate backup LUKS key ──────────────────────────────────────────
    # A random key is enrolled as a second LUKS slot so recovery is possible
    # without the primary passphrase. It is written to /_LUKS_BACKUP_KEY in the
    # new system (inside the encrypted container) where only root can read it.
    LUKS_BACKUP_KEY=$(mktemp /tmp/luks-backup-key.XXXXXX)
    # Read 64 bytes from /dev/urandom and base64-encode them (-w0 disables line
    # wrapping) to produce a printable key with high entropy.
    dd if=/dev/urandom bs=64 count=1 2>/dev/null | base64 -w0 > "$LUKS_BACKUP_KEY"
    echo "Enrolling auto-generated backup LUKS key..."
    # luksAddKey reads the existing passphrase interactively to authorise adding
    # the new key file into a free LUKS slot.
    cryptsetup luksAddKey "$ROOT_PART" "$LUKS_BACKUP_KEY"

    # ── Optional FIDO2 enrollment ─────────────────────────────────────────────
    if [[ "$FIDO_ROOT" == "YES" ]]; then
        echo "Insert FIDO2 key for LUKS and touch when prompted..."
        # --fido2-device=auto finds any connected FIDO2 token automatically.
        # --fido2-with-client-pin=no skips PIN verification; the hardware
        # user-presence tap alone is sufficient for this setup.
        systemd-cryptenroll "$ROOT_PART" --fido2-device=auto --fido2-with-client-pin=no
    fi

    ############################################
    # BTRFS ON ENCRYPTED ROOT
    ############################################
    # Format the decrypted mapper device, not the raw partition.
    mkfs.btrfs /dev/mapper/cryptroot
    # Mount flat (no subvolume) first so we can create the subvolume layout.
    mount /dev/mapper/cryptroot /mnt
    # @ is the conventional name for the root subvolume in btrfs-on-Arch setups;
    # @home separates user data so it can be snapshotted or rolled back independently.
    btrfs subvolume create /mnt/@
    btrfs subvolume create /mnt/@home
    # Unmount the flat view before remounting through the subvolumes.
    umount /mnt

    # Mount each subvolume at its final path for pacstrap to populate.
    mount -o subvol=@ /dev/mapper/cryptroot /mnt
    mkdir -p /mnt/home
    mount -o subvol=@home /dev/mapper/cryptroot /mnt/home

else
    echo "Skipping encryption — formatting root directly."

    ############################################
    # BTRFS ON UNENCRYPTED ROOT
    ############################################
    # Same subvolume layout as the encrypted branch for consistency.
    mkfs.btrfs "$ROOT_PART"
    mount "$ROOT_PART" /mnt
    btrfs subvolume create /mnt/@
    btrfs subvolume create /mnt/@home
    umount /mnt

    mount -o subvol=@ "$ROOT_PART" /mnt
    mkdir -p /mnt/home
    mount -o subvol=@home "$ROOT_PART" /mnt/home
fi

mkdir -p /mnt/boot
# Mount the ESP so GRUB and the kernel can be installed inside the new root tree.
mount "$BOOT_PART" /mnt/boot

# Place backup key inside the new system (only accessible when disk is unlocked)
if [[ -n "$LUKS_BACKUP_KEY" ]]; then
    # `install -m 400` creates the destination file with mode 0400 (root read-only)
    # in one atomic step, avoiding an intermediate world-readable state.
    install -m 400 "$LUKS_BACKUP_KEY" /mnt/_LUKS_BACKUP_KEY
    # Remove the temp file from the live environment immediately after copying.
    rm -f "$LUKS_BACKUP_KEY"
    echo "Backup LUKS key written to /_LUKS_BACKUP_KEY in new system."
fi

############################################
# GPU DETECTION
############################################
# lspci lists PCI devices; filter for VGA and 3D controllers to identify the GPU.
# `|| true` prevents set -e from aborting if no GPU line is found.
GPU_INFO=$(lspci | grep -E "VGA|3D" || true)
GPU_PKGS=""
if echo "$GPU_INFO" | grep -qi "NVIDIA"; then
    GPU_PKGS="nvidia nvidia-utils"
elif echo "$GPU_INFO" | grep -qi "AMD"; then
    # xf86-video-amdgpu provides the Xorg DDX driver; the kernel DRM driver ships in linux-firmware.
    GPU_PKGS="xf86-video-amdgpu"
elif echo "$GPU_INFO" | grep -qi "Intel"; then
    GPU_PKGS="xf86-video-intel"
fi
echo "Detected GPU: ${GPU_INFO:-none}"

############################################
# BASE INSTALL
############################################
# pacstrap installs packages into /mnt using pacman's keyring (-K initialises it).
# GPU_PKGS is intentionally unquoted so it expands to multiple words (or nothing).
# shellcheck disable=SC2086
pacstrap -K /mnt base base-devel "$KERNEL" linux-firmware vim bash zsh git less btop fastfetch \
    networkmanager grub cryptsetup libfido2 pam-u2f efibootmgr sudo btrfs-progs lvm2 jq $GPU_PKGS

############################################
# FSTAB
############################################
# genfstab -U uses UUIDs (rather than device names or labels) so the fstab
# remains valid even if drive letter assignments change after a hardware swap.
genfstab -U /mnt >> /mnt/etc/fstab

############################################
# COPY ANSWERFILE INTO NEW SYSTEM
############################################
# Make the answerfile available inside the new system so the dotfiles TUI
# installer can read the same selections without re-prompting the user.
if $AF_MODE; then
    install -m 644 "$ANSWERFILE" /mnt/answerfile.json
fi

############################################
# PASS VARIABLES INTO CHROOT
############################################
# Variables must be exported so the arch-chroot heredoc inherits them.
# The heredoc uses single-quote delimiter ('CHROOT_EOF') which prevents expansion
# outside the chroot, so env vars must be in the environment, not inline.
export HOSTNAME USERNAME USERPASS ROOT_PART KERNEL FIDO_ROOT FIDO_USER ENCRYPT_DISK KEYMAP

############################################
# CHROOT CONFIGURATION
############################################
# arch-chroot binds /proc, /sys, /dev and runs commands inside the new root.
# Single-quoted 'CHROOT_EOF' prevents the outer shell from expanding variables;
# the inner shell resolves them from the exported environment above.
arch-chroot /mnt /bin/bash <<'CHROOT_EOF'
set -euo pipefail

# Locale
# Uncomment (create) the en_US.UTF-8 entry so locale-gen can compile it.
echo "en_US.UTF-8 UTF-8" > /etc/locale.gen
locale-gen
echo "LANG=en_US.UTF-8" > /etc/locale.conf
# vconsole.conf sets the keymap for the virtual console (TTY) on every boot.
echo "KEYMAP=${KEYMAP}" > /etc/vconsole.conf

# Time / hostname
# Symlink the timezone file so /etc/localtime always points at the correct zoneinfo.
ln -sf /usr/share/zoneinfo/Europe/Vienna /etc/localtime
# --systohc writes the current system time to the hardware clock (RTC).
hwclock --systohc
echo "$HOSTNAME" > /etc/hostname

# NetworkManager
# Enable at boot so the system has networking on first login without manual setup.
systemctl enable NetworkManager

# Populate /etc/skel before user creation so useradd -m copies everything
echo "Cloning dotfiles into /etc/skel..."
git clone https://git.abdelbaki.eu/The_miro/Dotfiles.git /etc/skel/Dotfiles \
    || echo "Warning: dotfiles clone into skel failed — will fall back to direct clone."
# Seed /etc/skel with base shell dotfiles from the repo clone
if [[ -d /etc/skel/Dotfiles ]]; then
    D=/etc/skel/Dotfiles
    # Copy individually; `2>/dev/null || true` silently skips any file that is absent
    # in this repo revision rather than aborting the install.
    cp "$D/.zshrc"  /etc/skel/.zshrc  2>/dev/null || true
    cp "$D/.bashrc" /etc/skel/.bashrc 2>/dev/null || true
    cp "$D/.vimrc"  /etc/skel/.vimrc  2>/dev/null || true
fi

# User
# -m: create home dir; -G wheel: add to sudoers group; -s /bin/zsh: default shell.
useradd -m -G wheel -s /bin/zsh "$USERNAME"
# chpasswd reads "user:pass" from stdin to set the password non-interactively.
echo "$USERNAME:$USERPASS" | chpasswd
# Grant wheel group full sudo access (ALL:ALL covers any user/group runas context).
echo "%wheel ALL=(ALL:ALL) ALL" >> /etc/sudoers

###################################################
# INITRAMFS CONFIG
###################################################
# The HOOKS line controls which initramfs modules are compiled in.
# Three variants depending on encryption and FIDO2 choices:
# 1. FIDO2 root unlock: needs `systemd` + `sd-encrypt` for systemd-cryptsetup.
# 2. Password-only LUKS: uses classic `encrypt` hook (no systemd dependency).
# 3. Unencrypted: minimal hook set — no encrypt hook needed.
if [[ "$ENCRYPT_DISK" == "YES" && "$FIDO_ROOT" == "YES" ]]; then
    sed -i 's/^HOOKS=.*/HOOKS=(base udev systemd autodetect microcode modconf kms consolefont block sd-encrypt lvm2 btrfs filesystems keyboard keymap fsck)/' /etc/mkinitcpio.conf
elif [[ "$ENCRYPT_DISK" == "YES" ]]; then
    sed -i 's/^HOOKS=.*/HOOKS=(base udev autodetect microcode modconf kms consolefont block encrypt lvm2 btrfs filesystems keyboard keymap fsck)/' /etc/mkinitcpio.conf
else
    sed -i 's/^HOOKS=.*/HOOKS=(base udev autodetect microcode modconf kms consolefont block btrfs filesystems keyboard fsck)/' /etc/mkinitcpio.conf
fi

# Regenerate the initramfs for the selected kernel preset.
mkinitcpio -p "$KERNEL"

###################################################
# GRUB CONFIG
###################################################
# Read the raw partition UUID; used in kernel cmdline for stable device identification.
UUID=$(blkid -s UUID -o value "$ROOT_PART")
if [[ "$ENCRYPT_DISK" == "YES" ]]; then
    if [[ "$FIDO_ROOT" == "YES" ]]; then
        # systemd-cryptsetup syntax: rd.luks.name=<UUID>=<name>.
        # The FIDO2 device is auto-detected by systemd-cryptenroll at boot.
        KERNEL_CMD="rd.luks.name=${UUID}=cryptroot root=/dev/mapper/cryptroot"
    else
        # Classic initramfs encrypt hook syntax: cryptdevice=UUID=<UUID>:<name>.
        KERNEL_CMD="cryptdevice=UUID=${UUID}:cryptroot root=/dev/mapper/cryptroot"
    fi
else
    # Unencrypted btrfs: root= by UUID and rootflags to select the @ subvolume.
    KERNEL_CMD="root=UUID=${UUID} rootflags=subvol=@"
fi

# Inject the kernel command line into /etc/default/grub using sed.
# The `|` delimiter avoids conflicts with `/` in device paths.
sed -i "s|^GRUB_CMDLINE_LINUX=.*|GRUB_CMDLINE_LINUX=\"$KERNEL_CMD\"|" /etc/default/grub

# Install GRUB to the EFI partition; --bootloader-id sets the NVRAM entry name.
grub-install --target=x86_64-efi --efi-directory=/boot --bootloader-id=M-Archy-GRUB
# Generate grub.cfg from /etc/default/grub and detected kernels.
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..."
    # pamu2fcfg registers the FIDO2 key and writes the credential to u2f_keys.
    # Running under sudo -u ensures the file is generated with the correct UID context.
    sudo -u "$USERNAME" pamu2fcfg -u "$USERNAME" > "/home/$USERNAME/.config/Yubico/u2f_keys"
    chown "$USERNAME":"$USERNAME" "/home/$USERNAME/.config/Yubico/u2f_keys"
    # Append the PAM rule so pam_u2f is required for system-local-login sessions.
    echo "auth required pam_u2f.so" >> /etc/pam.d/system-local-login
fi

###################################################
# CLONE DOTFILES
###################################################
# Fall back to a direct clone if /etc/skel/Dotfiles was not copied into home
# (happens when the skel clone failed or useradd did not copy skel).
echo "Cloning dotfiles..."
git clone https://git.abdelbaki.eu/The_miro/Dotfiles.git "/home/$USERNAME/Dotfiles" \
    && chown -R "$USERNAME":"$USERNAME" "/home/$USERNAME/Dotfiles" \
    || echo "Warning: dotfiles clone failed — clone manually after first boot."

CHROOT_EOF

############################################
# DOTFILES TUI SETUP (in-chroot, optional)
############################################
if [[ "${RUN_TUI^^}" == "YES" ]]; then
    # Grant passwordless sudo temporarily so the TUI installer can call pacman/yay
    # without needing a password inside the chroot (the real sudoers is already set).
    # The file is removed immediately after the TUI exits.
    echo "${USERNAME} ALL=(ALL) NOPASSWD: ALL" \
        | arch-chroot /mnt tee /etc/sudoers.d/99-setup-nopasswd > /dev/null

    echo "Running tui-install.sh as ${USERNAME} inside chroot..."
    # `runuser -u` switches to the unprivileged user inside the chroot so that
    # AUR helpers and dotfiles are owned/built by the correct UID.
    arch-chroot /mnt runuser -u "${USERNAME}" -- \
        bash "/home/${USERNAME}/Dotfiles/setup/tui-install.sh" \
        || echo "Warning: tui-install exited with errors — check ~/dotfiles-install.log in the new system."

    # Remove the temporary no-password sudoers drop-in after setup is done.
    arch-chroot /mnt rm -f /etc/sudoers.d/99-setup-nopasswd
fi

# Remove answerfile from new system after setup is complete (contains sensitive paths/config)
if $AF_MODE && [[ -f /mnt/answerfile.json ]]; then
    rm -f /mnt/answerfile.json
fi

############################################
# SUMMARY
############################################
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:-none}"
echo
echo "Disk encryption:   $ENCRYPT_DISK"
echo "FIDO2 root unlock: $FIDO_ROOT"
echo "FIDO2 user login:  $FIDO_USER"
[[ "$ENCRYPT_DISK" == "YES" ]] && echo "LUKS backup key:   /_LUKS_BACKUP_KEY (in new system)"
echo
echo "Boot size:         $BOOT_SIZE"
echo "Root size:         ${ROOT_GIB}GiB"
echo "Swap size:         $SWAP_SIZE"
echo
echo "Log file:          $LOGFILE"
echo "############################################"
echo

# Copy install log into /boot so it survives and is accessible before login.
# `2>/dev/null || true` silently ignores failures (e.g. /mnt/boot not writable).
cp "$LOGFILE" /mnt/boot/ 2>/dev/null || true

echo "Installation complete! Unmount and reboot:"
echo "  umount -R /mnt && reboot"
if [[ "${RUN_TUI^^}" != "YES" ]]; then
    echo
    echo "After first boot, login as $USERNAME and run:"
    echo "  ~/Dotfiles/setup/tui-install.sh"
fi