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

# ========= helpers =========
color() { echo -e "\033[1;32m$1\033[0m"; }
warn()  { echo -e "\033[1;33m$1\033[0m"; }
err()   { echo -e "\033[1;31m$1\033[0m"; }
need()  { command -v "$1" &>/dev/null || { err "Missing $1"; exit 1; }; }
to_mib(){ local g="$1"; echo $(( g * 1024 )); }

# ========= UI & preflight =========
if ! command -v whiptail &>/dev/null; then pacman -Sy --noconfirm libnewt; fi
if ! ping -c1 -W1 archlinux.org &>/dev/null; then
  warn "No ping to archlinux.org"
  whiptail --title "Network Warning" --yesno "No ping reply.\nContinue anyway?" 10 60 || exit 1
fi

whiptail --title "Arch Linux Installer" \
  --msgbox "Welcome! This installs Arch with BSPWM + LightDM.\nGTK greeter will be default; mini-greeter will be built from AUR and used if successful.\nYou'll pick partition sizes based on your disk." 12 68

# ========= disk choice & info =========
DISK=$(whiptail --title "Select Disk" --inputbox "Target disk (e.g. /dev/sda or /dev/nvme0n1):" 10 60 "/dev/sda" 3>&1 1>&2 2>&3) || exit 1
[[ -b "$DISK" ]] || { err "Device not found: $DISK"; exit 1; }
PSEP=""; [[ "$DISK" =~ [0-9]$ ]] && PSEP="p"

need lsblk; need blockdev
DISK_MODEL=$(lsblk -dno MODEL "$DISK" || true)
DISK_SIZE_BYTES=$(blockdev --getsize64 "$DISK")
DISK_SIZE_MIB=$(( DISK_SIZE_BYTES / 1024 / 1024 ))
DISK_SIZE_GIB=$(( DISK_SIZE_MIB / 1024 ))

whiptail --title "Disk Info" --msgbox \
"Disk:   $DISK
Model:  ${DISK_MODEL:-Unknown}
Size:   ${DISK_SIZE_GIB} GiB (${DISK_SIZE_MIB} MiB)

You will now choose partition sizes." 12 60

# ========= boot type =========
if [[ -d /sys/firmware/efi/efivars ]]; then BOOT_TYPE="UEFI"; else BOOT_TYPE="BIOS"; fi

# ========= system identity =========
HOSTNAME=$(whiptail --title "Hostname" --inputbox "Enter hostname:" 10 50 "archlinux" 3>&1 1>&2 2>&3)
USERNAME=$(whiptail --title "Username"  --inputbox "Enter username:"  10 50 "user"       3>&1 1>&2 2>&3)
USERPASS=$(whiptail --title "User Password" --passwordbox "Password for $USERNAME:" 10 50 3>&1 1>&2 2>&3)
ROOTPASS=$(whiptail --title "Root Password" --passwordbox "Root password:" 10 50 3>&1 1>&2 2>&3)

# ========= optional extras =========
EXTRAS=$(whiptail --title "Optional Packages" --checklist "Select any extras:" 15 62 6 \
  "alacritty" "GPU terminal" ON \
  "picom"     "Compositor"   ON \
  "neofetch"  "Sys info"     OFF \
  "git"       "VCS"          ON \
  "firefox"   "Browser"      OFF \
  "wireless_tools" "Legacy iw tools" OFF \
  3>&1 1>&2 2>&3) || true
EXTRAS=$(echo $EXTRAS | tr -d '"')

# ========= CPU / VM =========
CPU_VENDOR=$(lscpu | awk -F: '/Vendor ID/ {print $2}' | xargs || true)
grep -q hypervisor /proc/cpuinfo && VM=true || VM=false

# ========= partition sizes =========
RAM_GIB=$(awk '/MemTotal/ { printf "%.0f", $2/1024/1024 }' /proc/meminfo 2>/dev/null || echo 4)
DEF_SWAP_GIB=$(( RAM_GIB<8 ? RAM_GIB : 8 )); [[ $DEF_SWAP_GIB -lt 2 ]] && DEF_SWAP_GIB=2

while :; do
  if [[ "$BOOT_TYPE" == "UEFI" ]]; then
    ESP_GIB=$(whiptail --title "UEFI ESP" --inputbox "ESP size GiB [min 0.3, typical 0.5]:" 10 60 "0.5" 3>&1 1>&2 2>&3) || exit 1
    SWAP_GIB=$(whiptail --title "Swap" --inputbox "Swap size GiB (0 for none):" 10 60 "$DEF_SWAP_GIB" 3>&1 1>&2 2>&3) || exit 1
    ROOT_GIB=$(whiptail --title "Root" --inputbox "Root size GiB (or ALL):" 10 60 "ALL" 3>&1 1>&2 2>&3) || exit 1

    local_used_mib=1
    if [[ "$ESP_GIB" =~ ^[0-9]+(\.[0-9]+)?$ ]]; then
      ESP_MIB=$(printf "%.0f" "$(awk "BEGIN{print $ESP_GIB*1024}")")
      [[ $ESP_MIB -lt 300 ]] && { err "ESP must be >= 0.3 GiB"; continue; }
    else err "ESP must be a number"; continue; fi
    local_used_mib=$(( local_used_mib + ESP_MIB ))

    if [[ "$SWAP_GIB" =~ ^[0-9]+$ ]]; then SWAP_MIB=$(to_mib "$SWAP_GIB"); else err "Swap must be integer GiB"; continue; fi
    local_used_mib=$(( local_used_mib + SWAP_MIB ))

    if [[ "$ROOT_GIB" =~ ^[Aa][Ll][Ll]$ ]]; then
      ROOT_MIB=$(( DISK_SIZE_MIB - local_used_mib ))
      [[ $ROOT_MIB -lt 1024 ]] && { err "Not enough space for root"; continue; }
      USE_ALL_ROOT=true
    elif [[ "$ROOT_GIB" =~ ^[0-9]+$ ]]; then
      ROOT_MIB=$(to_mib "$ROOT_GIB")
      total_mib=$(( local_used_mib + ROOT_MIB ))
      [[ $total_mib -le $DISK_SIZE_MIB ]] || { err "Total exceeds disk (${DISK_SIZE_GIB} GiB)"; continue; }
      USE_ALL_ROOT=false
    else err "Root must be integer GiB or ALL"; continue; fi

    break

  else
    SWAP_GIB=$(whiptail --title "Swap" --inputbox "Swap size GiB (0 for none):" 10 60 "$DEF_SWAP_GIB" 3>&1 1>&2 2>&3) || exit 1
    ROOT_GIB=$(whiptail --title "Root" --inputbox "Root size GiB (or ALL):" 10 60 "ALL" 3>&1 1>&2 2>&3) || exit 1

    ROOT_START_MIB=1
    if [[ "$SWAP_GIB" =~ ^[0-9]+$ ]]; then SWAP_MIB=$(to_mib "$SWAP_GIB"); else err "Swap must be integer GiB"; continue; fi

    if [[ "$ROOT_GIB" =~ ^[Aa][Ll][Ll]$ ]]; then
      ROOT_MIB=$(( DISK_SIZE_MIB - ROOT_START_MIB - SWAP_MIB ))
      [[ $ROOT_MIB -lt 1024 ]] && { err "Not enough space for root"; continue; }
      USE_ALL_ROOT=true
    elif [[ "$ROOT_GIB" =~ ^[0-9]+$ ]]; then
      ROOT_MIB=$(to_mib "$ROOT_GIB")
      total_mib=$(( ROOT_START_MIB + ROOT_MIB + SWAP_MIB ))
      [[ $total_mib -le $DISK_SIZE_MIB ]] || { err "Total exceeds disk (${DISK_SIZE_GIB} GiB)"; continue; }
      USE_ALL_ROOT=false
    else err "Root must be integer GiB or ALL"; continue; fi

    break
  fi
done

# ========= confirm destructive =========
whiptail --title "Partitioning" --yesno "WIPE and partition $DISK for $BOOT_TYPE?\nProceed?" 10 60 || exit 1

# ========= partition =========
color "Partitioning $DISK ..."
sgdisk --zap-all "$DISK" || true

if [[ "$BOOT_TYPE" == "UEFI" ]]; then
  parted -s "$DISK" mklabel gpt

  ESP_START=1
  ESP_END=$(( ESP_START + ESP_MIB ))
  if [[ $SWAP_MIB -gt 0 ]]; then SWAP_START=$ESP_END; SWAP_END=$(( SWAP_START + SWAP_MIB )); fi
  ROOT_START=${SWAP_END:-$ESP_END}
  ROOT_END_STR=$([[ "${USE_ALL_ROOT:-false}" == true ]] && echo "100%" || echo "$(( ROOT_START + ROOT_MIB ))MiB")

  parted -s "$DISK" mkpart ESP fat32 "${ESP_START}MiB" "${ESP_END}MiB"
  parted -s "$DISK" set 1 esp on
  if [[ $SWAP_MIB -gt 0 ]]; then
    parted -s "$DISK" mkpart primary linux-swap "${SWAP_START}MiB" "${SWAP_END}MiB"
    SWAP_PART="${DISK}${PSEP}2"; ROOT_INDEX=3
  else
    ROOT_INDEX=2
  fi
  parted -s "$DISK" mkpart primary ext4 "${ROOT_START}MiB" "$ROOT_END_STR"

  EFI_PART="${DISK}${PSEP}1"
  ROOT_PART="${DISK}${PSEP}${ROOT_INDEX}"

else
  parted -s "$DISK" mklabel msdos
  ROOT_START=1
  if [[ "${USE_ALL_ROOT:-false}" == true ]]; then
    if [[ $SWAP_MIB -gt 0 ]]; then
      ROOT_END=$(( DISK_SIZE_MIB - SWAP_MIB ))
      ROOT_END_STR="${ROOT_END}MiB"
      SWAP_START=$ROOT_END; SWAP_END_STR="100%"
    else
      ROOT_END_STR="100%"
    fi
  else
    ROOT_END_STR="$(( ROOT_START + ROOT_MIB ))MiB"
    if [[ $SWAP_MIB -gt 0 ]]; then SWAP_START=$(( ROOT_START + ROOT_MIB )); SWAP_END_STR="100%"; fi
  fi

  parted -s "$DISK" mkpart primary ext4 "${ROOT_START}MiB" "$ROOT_END_STR"
  parted -s "$DISK" set 1 boot on
  ROOT_PART="${DISK}${PSEP}1"
  if [[ $SWAP_MIB -gt 0 ]]; then
    parted -s "$DISK" mkpart primary linux-swap "${SWAP_START}MiB" "$SWAP_END_STR"
    SWAP_PART="${DISK}${PSEP}2"
  fi
fi

# ========= format & mount =========
color "Formatting ..."
if [[ "$BOOT_TYPE" == "UEFI" ]]; then mkfs.fat -F32 "$EFI_PART"; fi
mkfs.ext4 -F "$ROOT_PART"
if [[ "${SWAP_PART:-}" ]]; then mkswap "$SWAP_PART" && swapon "$SWAP_PART"; fi

mount "$ROOT_PART" /mnt
mkdir -p /mnt/boot
if [[ "$BOOT_TYPE" == "UEFI" ]]; then mkdir -p /mnt/boot/efi && mount "$EFI_PART" /mnt/boot/efi; fi

# ========= packages =========
PKGS="base linux linux-firmware nano vim networkmanager \
xorg-server xorg-xinit xorg-xrandr xorg-xsetroot xorg-xprop xorg-xset xterm mesa \
bspwm sxhkd rofi lightdm lightdm-gtk-greeter ttf-dejavu iwd grub sudo accountsservice"

if ! $VM; then
  [[ "$CPU_VENDOR" == "GenuineIntel" ]] && PKGS+=" intel-ucode"
  [[ "$CPU_VENDOR" == "AuthenticAMD" ]] && PKGS+=" amd-ucode"
fi
PKGS+=" $EXTRAS"

color "Installing base system ..."

if lspci | grep -qi virtualbox; then
  PKGS+=" virtualbox-guest-utils"
fi

pacstrap -K /mnt $PKGS --noconfirm
genfstab -U /mnt >> /mnt/etc/fstab

# ========= chroot config =========
arch-chroot /mnt /usr/bin/env \
  HOSTNAME="$HOSTNAME" USERNAME="$USERNAME" USERPASS="$USERPASS" ROOTPASS="$ROOTPASS" \
  BOOT_TYPE="$BOOT_TYPE" DISK="$DISK" bash -s <<'EOFCHROOT'
set -euo pipefail

# ---- basics ----
echo "$HOSTNAME" > /etc/hostname
cat >/etc/hosts <<HOSTS
127.0.0.1   localhost
::1         localhost
127.0.1.1   $HOSTNAME.localdomain $HOSTNAME
HOSTS

ln -sf /usr/share/zoneinfo/Australia/Melbourne /etc/localtime
hwclock --systohc
sed -i 's/^#en_US.UTF-8/en_US.UTF-8/' /etc/locale.gen
locale-gen
echo "LANG=en_US.UTF-8" > /etc/locale.conf

# ensure bspwm session entry
install -d -m 0755 /usr/share/xsessions
cat >/usr/share/xsessions/bspwm.desktop <<'DESK'
[Desktop Entry]
Name=bspwm
Comment=Binary space partitioning window manager
Exec=bspwm
TryExec=bspwm
Type=Application
DesktopNames=bspwm
DESK

# ---- user + sudo ----
useradd -m -G wheel -s /bin/bash "$USERNAME"
echo "$USERNAME:$USERPASS" | chpasswd
echo "root:$ROOTPASS" | chpasswd
install -d -m 0755 /etc/sudoers.d
printf '%%wheel ALL=(ALL:ALL) ALL\n' > /etc/sudoers.d/01-wheel
chmod 0440 /etc/sudoers.d/01-wheel

# ---- LightDM default to GTK greeter ----
install -d -m 0755 /etc/lightdm
cat >/etc/lightdm/lightdm.conf <<'CONF'
[LightDM]
logind-check-graphical=true

[Seat:*]
user-session=bspwm
greeter-session=lightdm-gtk-greeter
greeter-allow-guest=false
CONF

# ---- Greeter toggle helper ----
cat >/usr/local/bin/switch-greeter <<'SH'
#!/usr/bin/env bash
set -e
choice="$1"
case "$choice" in
  gtk)
    sudo sed -i 's/^greeter-session=.*/greeter-session=lightdm-gtk-greeter/' /etc/lightdm/lightdm.conf
    ;;
  mini)
    if command -v lightdm-mini-greeter >/dev/null; then
      sudo sed -i 's/^greeter-session=.*/greeter-session=lightdm-mini-greeter/' /etc/lightdm/lightdm.conf
    else
      echo "lightdm-mini-greeter is not installed."
      exit 1
    fi
    ;;
  *)
    echo "Usage: switch-greeter [gtk|mini]"
    exit 1
    ;;
esac
sudo systemctl restart lightdm
SH
chmod +x /usr/local/bin/switch-greeter

# drop a user-owned copy in HOME for convenience
cp /usr/local/bin/switch-greeter /home/$USERNAME/switch-greeter
chown $USERNAME:$USERNAME /home/$USERNAME/switch-greeter


cat >/etc/lightdm/lightdm-mini-greeter.conf <<MINI
[greeter]
user=$USERNAME
password-hint-enabled=false
show-password-label=false
MINI

# ---- Try AUR mini-greeter ----
pacman -S --needed --noconfirm base-devel git
id builder &>/dev/null || useradd -m builder
echo 'builder ALL=(ALL) NOPASSWD: ALL' >/etc/sudoers.d/99_builder
chmod 0440 /etc/sudoers.d/99_builder

MINI_OK=false
sudo -u builder bash -lc '
  set -e
  rm -rf ~/aur && mkdir -p ~/aur && cd ~/aur
  git clone https://aur.archlinux.org/lightdm-mini-greeter.git
  cd lightdm-mini-greeter
  if makepkg -si --noconfirm --needed; then
    exit 0
  else
    exit 1
  fi
' && MINI_OK=true || MINI_OK=false

rm -f /etc/sudoers.d/99_builder
userdel -r builder || true

# Keep GTK as default; user can switch later with switch-greeter

# ---- BSPWM defaults ----
mkdir -p /home/$USERNAME/.config/{bspwm,sxhkd}
if [[ -f /usr/share/doc/bspwm/examples/bspwmrc ]]; then
  cp /usr/share/doc/bspwm/examples/bspwmrc /home/$USERNAME/.config/bspwm/
else
  cat >/home/$USERNAME/.config/bspwm/bspwmrc <<'BRC'
#!/usr/bin/env bash
sxhkd &
bspc monitor -d I II III IV V VI VII VIII IX X
BRC
fi
if [[ -f /usr/share/doc/bspwm/examples/sxhkdrc ]]; then
  cp /usr/share/doc/bspwm/examples/sxhkdrc /home/$USERNAME/.config/sxhkd/
else
  cat >/home/$USERNAME/.config/sxhkd/sxhkdrc <<'SRC'
super + Return
  alacritty || xterm || xterm
super + q
  bspc node -c
SRC
fi
chmod +x /home/$USERNAME/.config/bspwm/bspwmrc
chown -R $USERNAME:$USERNAME /home/$USERNAME/.config

# ---- services & bootloader ----
systemctl enable NetworkManager
systemctl set-default graphical.target
systemctl enable lightdm

if [[ "$BOOT_TYPE" == "UEFI" ]]; then
  pacman -S --noconfirm --needed efibootmgr
  grub-install --target=x86_64-efi --efi-directory=/boot/efi --bootloader-id=GRUB
else
  grub-install --target=i386-pc "$DISK"
fi
grub-mkconfig -o /boot/grub/grub.cfg
systemctl enable vboxservice || true

EOFCHROOT

whiptail --title "Installation Complete" \
  --msgbox "Done! Reboot into your new Arch + BSPWM desktop.\nGTK greeter by default; mini-greeter if built successfully." 10 70