#!/usr/bin/env bash
# archbaseos-guided-install.sh — guided (dialog-based) Arch Linux base installer
#
# If /answerfile.json exists (e.g. embedded via build.sh --preconf), all prompts
# are answered from it. Missing fields fall back to interactive prompts.

set -Eeuo pipefail

############################################
# LOGGING
############################################
LOGFILE="$HOME/archbaseos-guided-install.log"
{
    echo
    echo "############################################"
    echo " Arch Guided Install Log - Started $(date)"
    echo "############################################"
    echo
} >> "$LOGFILE"
exec > >(tee -a "$LOGFILE") 2>&1

############################################
# Error handler — TUI prompt to send log via croc
############################################
error_handler() {
    local exit_code=$? line_num="${1:-?}"
    echo "" >> "$LOGFILE"
    echo "ERROR: installer failed at line $line_num (exit code $exit_code)" >> "$LOGFILE"

    echo ""
    echo "Installation failed at line $line_num (exit code $exit_code)."
    read -rp "Send log via croc? [y/N]: " _croc_ans
    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
    exit "$exit_code"
}
trap 'error_handler $LINENO' ERR

############################################
# 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..."
}

# Returns the correct partition device for a given drive and partition number.
# NVMe and eMMC use a 'p' separator (e.g. /dev/nvme0n1p1), others don't.
part() { [[ "$1" == *nvme* || "$1" == *mmcblk* ]] && echo "${1}p${2}" || echo "${1}${2}"; }

############################################
# ANSWERFILE
############################################
ANSWERFILE="${ANSWERFILE:-/answerfile.json}"
AF_MODE=false
[[ -f "$ANSWERFILE" ]] && AF_MODE=true

af_get() {
    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() {
    local val; val=$(jq -r "${1} // false" "$ANSWERFILE" 2>/dev/null || true)
    [[ "$val" == "true" ]] && echo "YES" || echo "NO"
}

get_mac_suffix() {
    local mac
    mac=$(ip link show 2>/dev/null \
        | awk '/^[0-9]+: [^l][^o]/{iface=1} iface && /link\/ether/{print $2; iface=0; exit}')
    printf '%s' "${mac//:/}"
}

if $AF_MODE; then
    echo "== Arch Linux Guided Installer (answerfile mode) =="
    command -v jq &>/dev/null || pacman -Sy --noconfirm jq
else
    echo "== Arch Linux FIDO2-Ready Installer =="
fi

############################################
# NETWORK CHECK
############################################
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

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

KEYMAPS=(
    "us|English US"
    "de|German"
)

############################################
# User input
############################################
if $AF_MODE; then
    KERNEL=$(af_get '.kernel' 'linux')
    RAW_HOSTNAME=$(af_get '.hostname' '')
    if [[ -n "$RAW_HOSTNAME" ]]; then
        HOSTNAME="${RAW_HOSTNAME}-$(get_mac_suffix)"
    else
        HOSTNAME="arch"
    fi
    USERNAME=$(af_get '.username' '')
    ENCRYPT_DISK=$(af_bool '.encrypt')
    ENABLE_FIDO_ROOT=$(af_bool '.fido2_root')
    ENABLE_FIDO_USER=$(af_bool '.fido2_user')
    RUN_TUI=$(af_bool '.run_tui')
    KEYMAP=$(af_get '.keymap' 'us')
    echo "Kernel:   $KERNEL / Hostname: $HOSTNAME / Username: $USERNAME"
    echo "Encrypt:  $ENCRYPT_DISK / FIDO2 root: $ENABLE_FIDO_ROOT / FIDO2 user: $ENABLE_FIDO_USER"
    echo "Keymap:   $KEYMAP"
else
    KERNEL=$(ask "Kernel (linux, linux-lts, linux-zen)")
    HOSTNAME=$(ask "Hostname")
    USERNAME=$(ask "Username")
    read -rp "Enable disk encryption? (YES/NO): " ENCRYPT_DISK
    ENABLE_FIDO_ROOT="NO"
    if [[ "$ENCRYPT_DISK" == "YES" ]]; then
        read -rp "Enable FIDO2 for unlocking root? (YES/NO): " ENABLE_FIDO_ROOT
    fi
    read -rp "Enable FIDO2 for user login? (YES/NO): " ENABLE_FIDO_USER

    echo ""
    echo "Select keyboard layout for installed system:"
    for i in "${!KEYMAPS[@]}"; do
        _km_code="${KEYMAPS[$i]%%|*}"
        _km_name="${KEYMAPS[$i]##*|}"
        printf "  %d) %-14s (%s)\n" $((i+1)) "$_km_name" "$_km_code"
    done
    read -rp "Choice [1]: " _KM_CHOICE
    _KM_CHOICE=$(( ${_KM_CHOICE:-1} - 1 ))
    if (( _KM_CHOICE >= 0 && _KM_CHOICE < ${#KEYMAPS[@]} )); then
        KEYMAP="${KEYMAPS[$_KM_CHOICE]%%|*}"
    else
        KEYMAP="${KEYMAPS[0]%%|*}"
    fi
fi

read -rp "Password for $USERNAME: " USERPASS
[[ -z "$USERPASS" ]] && { echo "Error: password cannot be empty."; exit 1; }

lsblk

if $AF_MODE && [[ -n "$(af_get '.drive')" ]]; then
    DRIVE=$(af_get '.drive')
    echo "Drive (from answerfile): $DRIVE"
    echo "WARNING: All data on $DRIVE will be erased. Proceeding in 5 seconds..."
    sleep 5
else
    DRIVE=$(ask "Enter install drive (e.g., /dev/sda)")
    confirm "$DRIVE" || exit 1
fi

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

############################################
# 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=$(part "$DRIVE" 1)
ROOT_PART=$(part "$DRIVE" 2)
SWAP_PART=$(part "$DRIVE" 3)

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

############################################
# Encryption (optional)
############################################
LUKS_BACKUP_KEY=""

if [[ "$ENCRYPT_DISK" == "YES" ]]; then
    echo "Formatting LUKS2 root..."
    cryptsetup luksFormat "$ROOT_PART" --type luks2
    cryptsetup open "$ROOT_PART" cryptroot

    # ── Auto-generate backup LUKS key ─────────────────────────────────────────
    LUKS_BACKUP_KEY=$(mktemp /tmp/luks-backup-key.XXXXXX)
    dd if=/dev/urandom bs=64 count=1 2>/dev/null | base64 -w0 > "$LUKS_BACKUP_KEY"
    echo "Enrolling auto-generated backup LUKS key..."
    cryptsetup luksAddKey "$ROOT_PART" "$LUKS_BACKUP_KEY"

    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

    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

else
    echo "Skipping encryption — formatting root directly."

    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 "$EFI_PART" /mnt/boot

# Place backup key inside the new system (readable only by root, inside LUKS container)
if [[ -n "$LUKS_BACKUP_KEY" ]]; then
    install -m 400 "$LUKS_BACKUP_KEY" /mnt/_LUKS_BACKUP_KEY
    rm -f "$LUKS_BACKUP_KEY"
    echo "Backup LUKS key written to /_LUKS_BACKUP_KEY in new system."
fi

############################################
# Base System Install
############################################
GPU_INFO=$(lspci | grep -E "VGA|3D" || true)
GPU_PKGS=""

if echo "$GPU_INFO" | grep -qi nvidia; then
    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

# shellcheck disable=SC2086
pacstrap /mnt \
    base base-devel "$KERNEL" linux-firmware vim zsh git networkmanager grub efibootmgr \
    btrfs-progs cryptsetup libfido2 pam-u2f sudo less jq $GPU_PKGS

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

############################################
# COPY ANSWERFILE INTO NEW SYSTEM
############################################
if $AF_MODE; then
    install -m 644 "$ANSWERFILE" /mnt/answerfile.json
fi

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

arch-chroot /mnt /usr/bin/env \
  HOSTNAME="$HOSTNAME" \
  USERNAME="$USERNAME" \
  USERPASS="$USERPASS" \
  ENCRYPT_DISK="$ENCRYPT_DISK" \
  ENABLE_FIDO_ROOT="$ENABLE_FIDO_ROOT" \
  ENABLE_FIDO_USER="$ENABLE_FIDO_USER" \
  ROOT_UUID="$ROOT_UUID" \
  ROOT_PART="$ROOT_PART" \
  KEYMAP="$KEYMAP" \
  /bin/bash <<'CHROOT_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

echo "KEYMAP=${KEYMAP}" > /etc/vconsole.conf

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

echo "$HOSTNAME" > /etc/hostname
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 failed — clone manually after first boot."

mkdir -p /etc/skel/{Desktop,Documents,Downloads,Music,Pictures,Public,Templates,Videos}

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

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

# Initramfs
if [[ "$ENCRYPT_DISK" == "YES" && "$ENABLE_FIDO_ROOT" == "YES" ]]; then
    sed -i 's/^HOOKS=.*/HOOKS=(base udev systemd autodetect microcode modconf kms consolefont block sd-encrypt 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 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

mkinitcpio -P

# GRUB
if [[ "$ENCRYPT_DISK" == "YES" ]]; then
    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
else
    GRUB_CMDLINE="root=UUID=${ROOT_UUID} rootflags=subvol=@"
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 — directory + PAM only; key enrollment happens outside chroot
if [[ "$ENABLE_FIDO_USER" == "YES" ]]; then
    mkdir -p "/home/$USERNAME/.config/Yubico"
    chown "$USERNAME:$USERNAME" "/home/$USERNAME/.config/Yubico"
    echo "auth required pam_u2f.so cue" >> /etc/pam.d/system-auth
fi

CHROOT_EOF

# pamu2fcfg must run outside arch-chroot: inside the chroot the host's udev manages
# /dev/hidraw* permissions and the new user has no access to the device.
if [[ "$ENABLE_FIDO_USER" == "YES" ]]; then
    echo "Enrolling FIDO2 key for user login (outside chroot)..."
    U2F_KEYFILE="/mnt/home/${USERNAME}/.config/Yubico/u2f_keys"
    mkdir -p "/mnt/home/${USERNAME}/.config/Yubico"
    pamu2fcfg -u "$USERNAME" -o "pam://$HOSTNAME" -i "pam://$HOSTNAME" > "$U2F_KEYFILE"
    _NEWUID=$(arch-chroot /mnt id -u "$USERNAME" 2>/dev/null || echo "1000")
    _NEWGID=$(arch-chroot /mnt id -g "$USERNAME" 2>/dev/null || echo "1000")
    chown -R "$_NEWUID:$_NEWGID" "/mnt/home/${USERNAME}/.config/Yubico"
    chmod 600 "$U2F_KEYFILE"
    echo "FIDO2 key enrolled for $USERNAME."
fi

############################################
# DOTFILES SETUP (in-chroot, optional)
############################################
if $AF_MODE; then
    _DO_TUI="${RUN_TUI}"
else
    read -rp "Run dotfiles TUI setup inside chroot now? [YES/no]: " _TUI_IN
    _TUI_IN="${_TUI_IN:-YES}"
    [[ "${_TUI_IN^^}" == "YES" ]] && _DO_TUI="YES" || _DO_TUI="NO"
fi

if [[ "${_DO_TUI^^}" == "YES" ]]; then
    echo "${USERNAME} ALL=(ALL) NOPASSWD: ALL" \
        | arch-chroot /mnt tee /etc/sudoers.d/99-setup-nopasswd > /dev/null

    echo "Running simple-install.sh as ${USERNAME} inside chroot..."
    arch-chroot /mnt runuser -u "${USERNAME}" -- \
        bash "/home/${USERNAME}/Dotfiles/setup/simple-install.sh" \
        || echo "Warning: tui-install exited with errors — check ~/dotfiles-install.log in the new system."

    arch-chroot /mnt rm -f /etc/sudoers.d/99-setup-nopasswd
fi

# Remove answerfile from new system after setup completes
if $AF_MODE && [[ -f /mnt/answerfile.json ]]; then
    rm -f /mnt/answerfile.json
fi

cp "$LOGFILE" /mnt/boot/ 2>/dev/null || true

echo "Installation complete! Log saved to /mnt/boot/$(basename "$LOGFILE")"
echo "  umount -R /mnt && reboot"
if [[ "${_DO_TUI^^}" != "YES" ]]; then
    echo
    echo "After first boot, login as ${USERNAME} and run:"
    echo "  ~/Dotfiles/setup/simple-install.sh"
fi
if [[ "$ENCRYPT_DISK" == "YES" ]]; then
    echo
    echo "LUKS backup key stored at /_LUKS_BACKUP_KEY in the new system."
    echo "Keep this file secure — it can unlock the root partition."
fi