#!/bin/bash
###############################################################################
# setup.sh - Interactive Debian VM bootstrap wizard                          #
# Jordan-friendly edition (whiptail UI, aliases, neofetch, CIFS, etc.)       #
###############################################################################
set -euo pipefail

# =============================================================================
# Default SSH key (paste your key between the quotes)
# =============================================================================
DEFAULT_SSH_KEY="ssh-rsa AAAAB3NzaC1yc2EAAAADAQABAAABAQCoxUu/nC1C03LvxIhCMzyMu7CAfIp9+Rbt4vmx8q3ER1EPP2K53fnjUmOijC4YY2jgPHHXEoTgC6Rlcrl3eYFoqbhRc4nweN6Z3LXRghmfNXVmMRSouXEMWhxhPjk9r+w9+3E9+6p9X9YtQu+u76ArWcY9MgvD6Awvo66hSFgkzeXzgCcKkTdMkSOUwuHfm8Ja9TzSIUfnB6SAiKWLIejDntYJHSKoqsSzsovYRUc/W+al09MfIMWwN9vJwk7WM7O3E+YPL5Zcpmr4jaoFULf6hWtgn688nDU+4V0POIzRNnk4EPH5qo+AmSL7MwQ0Bh7z5EgiAJiAryrT/GnU41w7 rsa-key-20240415"

###############################################################################
# Ensure running as root                                                     #
###############################################################################
if [[ $EUID -ne 0 ]]; then
    if command -v whiptail &>/dev/null; then
        whiptail --title "Error" --msgbox "This script must be run as root." 10 60
    else
        echo "Error: This script must be run as root." >&2
    fi
    exit 1
fi

###############################################################################
# Logging & helpers                                                           #
###############################################################################
LOGFILE="/var/log/freshbox.log"
exec &> >(tee -a "$LOGFILE")

# Colors
RED=$'\e[31m'; YLW=$'\e[33m'; CLR=$'\e[0m'
# Error exit
die() { echo -e "${RED}[ERROR] $*${CLR}" >&2; exit 1; }
# Install packages quietly
ensure_pkgs() { apt-get -qq update; apt-get -y install "$@"; }
# Message dialog or stdout
msg() { if command -v whiptail &>/dev/null; then whiptail --title "Setup" --msgbox "$1" 10 60; else echo -e "${YLW}$1${CLR}"; fi; }
# Yes/No dialog (returns success on Yes)
confirm() { if command -v whiptail &>/dev/null; then whiptail --yesno "$1" 10 60; else return 0; fi; }
# Gauge progress for long tasks
gauge_run() {
    local MESSAGE=$1; shift
    (
        # Run the command quietly, hide stdout/stderr, background it
        "$@" > /dev/null 2>&1 & CMD_PID=$!
        while kill -0 "$CMD_PID" 2>/dev/null; do
            echo "XXX"; echo 50; echo "$MESSAGE"; echo "XXX"
            sleep 1
        done
    ) | whiptail --gauge "$MESSAGE" 6 60 0
}

###############################################################################
# Ensure whiptail exists                                                     #
###############################################################################
if ! command -v whiptail &>/dev/null; then
    echo "Installing whiptail …"
    ensure_pkgs whiptail
fi

###############################################################################
# Welcome                                                                    #
###############################################################################
whiptail --title "VM Setup Wizard" --yesno "\
Welcome!\n\nThis wizard will update the system, configure sudo, hostname,\naliases, neofetch, optional tools, CIFS mounts, and SSH keys.\n\nContinue?" 12 70 || exit 0

###############################################################################
# Core packages & sudo                                                       #
###############################################################################
msg "Installing core tools…"
ensure_pkgs sudo curl gnupg lsb-release nfs-common

# Choose sudo users
mapfile -t USERS < <(awk -F: '$3>=1000 && $1!="nobody"{print $1}' /etc/passwd)
SEL_SUDO=""
if ((${#USERS[@]})); then
    OPTS=(); for u in "${USERS[@]}"; do OPTS+=("$u" "" OFF); done
    SEL_SUDO=$(whiptail --title "Sudo Access" --checklist \
        "Select users to grant passwordless sudo:" 15 60 6 \
        "${OPTS[@]}" 3>&1 1>&2 2>&3) || true
    for u in $SEL_SUDO; do usermod -aG sudo "${u//\"/}"; done
fi

###############################################################################
# System upgrade                                                             #
###############################################################################
gauge_run "Applying system updates…" apt-get -y dist-upgrade

###############################################################################
# Hostname                                                                   #
###############################################################################
MACS=$(ip -brief link | awk '$1!~"lo"{print $1": "$3}')
NEW_HOST=$(whiptail --title "Hostname" --inputbox "\
Current MAC addresses:\n$MACS\n\nEnter new hostname:" 15 70 "$(hostname)" 3>&1 1>&2 2>&3) || die "Hostname required"
OLD_HOST=$(hostname)
[[ -n $NEW_HOST ]] || die "Hostname cannot be blank"
sed -i "s/$OLD_HOST/$NEW_HOST/g" /etc/hosts
hostnamectl set-hostname "$NEW_HOST"
msg "Hostname set to $NEW_HOST"

###############################################################################
# SSH keys for root                                                           #
###############################################################################
install -d -m 700 /root/.ssh
chmod 700 /root/.ssh
test -n "$DEFAULT_SSH_KEY" && {
    grep -qxF "$DEFAULT_SSH_KEY" /root/.ssh/authorized_keys 2>/dev/null || \
        echo "$DEFAULT_SSH_KEY" >> /root/.ssh/authorized_keys
    chmod 600 /root/.ssh/authorized_keys
    msg "Default SSH key added for root"
}
if confirm "Add additional SSH public key for root?"; then
    KEY=$(whiptail --title "SSH Key (root)" --inputbox "Paste your public key for root:" 12 70 3>&1 1>&2 2>&3) || die "Cancelled"
    grep -qxF "$KEY" /root/.ssh/authorized_keys 2>/dev/null || echo "$KEY" >> /root/.ssh/authorized_keys
    chmod 600 /root/.ssh/authorized_keys
    msg "Additional SSH key added for root"
fi

###############################################################################
# SSH keys for user jordanwages                                              #
###############################################################################
USER_NAME="jordanwages"
USER_HOME="/home/$USER_NAME"
install -d -m 700 "$USER_HOME/.ssh"
chown $USER_NAME:$USER_NAME "$USER_HOME/.ssh"
test -n "$DEFAULT_SSH_KEY" && {
    grep -qxF "$DEFAULT_SSH_KEY" "$USER_HOME/.ssh/authorized_keys" 2>/dev/null || \
        echo "$DEFAULT_SSH_KEY" >> "$USER_HOME/.ssh/authorized_keys"
    chown $USER_NAME:$USER_NAME "$USER_HOME/.ssh/authorized_keys"
    chmod 600 "$USER_HOME/.ssh/authorized_keys"
    msg "Default SSH key added for $USER_NAME"
}
if confirm "Add additional SSH public key for $USER_NAME?"; then
    KEY=$(whiptail --title "SSH Key ($USER_NAME)" --inputbox "Paste public key for $USER_NAME:" 12 70 3>&1 1>&2 2>&3) || die "Cancelled"
    grep -qxF "$KEY" "$USER_HOME/.ssh/authorized_keys" 2>/dev/null || echo "$KEY" >> "$USER_HOME/.ssh/authorized_keys"
    chown $USER_NAME:$USER_NAME "$USER_HOME/.ssh/authorized_keys"
    chmod 600 "$USER_HOME/.ssh/authorized_keys"
    msg "Additional SSH key added for $USER_NAME"
fi

###############################################################################
# Optional utilities                                                         #
###############################################################################
TOOLS=(htop curl wget jq git tree)
OPTS=(); for t in "${TOOLS[@]}"; do OPTS+=("$t" "" OFF); done
SEL_RAW=$(whiptail --title "Extra Utilities" --checklist \
    "Select additional packages to install:" 15 60 8 \
    "${OPTS[@]}" 3>&1 1>&2 2>&3) || true
IFS=' ' read -r -a SEL_PKGS <<< "${SEL_RAW//\"/}"
ensure_pkgs bat ncdu neofetch cifs-utils "${SEL_PKGS[@]}"
# Aliases
echo "alias cat='batcat --paging=never'" >/etc/profile.d/10-bat_alias.sh
echo "alias du='ncdu'" >/etc/profile.d/10-ncdu_alias.sh
chmod +x /etc/profile.d/10-*alias.sh
echo 'neofetch' >/etc/profile.d/90-neofetch.sh
chmod +x /etc/profile.d/90-neofetch.sh

###############################################################################
# NFS mounts                                                                 #
###############################################################################
NFS_HOSTS=(jimmu keiko keitai)
OPTS=(); for h in "${NFS_HOSTS[@]}"; do OPTS+=("$h" "" OFF); done
SEL_HOSTS=$(whiptail --title "NFS Mounts" --checklist \
    "Select NAS shares to mount under /media:" 15 60 6 \
    "${OPTS[@]}" 3>&1 1>&2 2>&3) || true

mkdir -p /media
for host in "${NFS_HOSTS[@]}"; do
    TEMPLATE="${host}.wageshouse:/Data /media/${host} nfs4 _netdev,x-systemd.automount,noatime 0 0"
    if [[ $SEL_HOSTS == *\"$host\"* ]]; then
        mkdir -p "/media/$host"
        grep -q "${host}.wageshouse" /etc/fstab && \
            sed -i "#${host}\.wageshouse#d" /etc/fstab
        echo "$TEMPLATE" >> /etc/fstab
    else
        grep -q "${host}.wageshouse" /etc/fstab || echo "# $TEMPLATE" >> /etc/fstab
    fi
done

if ! mount -a 2>>"$LOGFILE"; then
    msg "⚠️  Some NFS mounts failed. Re-run the script to check configuration."
fi

###############################################################################
# Summary                                                                    #
###############################################################################
msg "\
Setup complete!\n\n• Hostname:  $NEW_HOST\n• Sudoers:   ${SEL_SUDO//\"/ }\n• Packages:  bat ncdu neofetch ${SEL_PKGS[*]}\n• CIFS:      ${SEL_HOSTS//\"/ }\n\nLog saved to $LOGFILE\n\nEnjoy your new Debian VM!"