#!/usr/bin/env bash
# ============================================================================
#  ArxOS VM Setup  --  one script that turns any Linux PC into a ready
#  QEMU/KVM host and boots ArxOS 0.0.1 in a virtual machine.
#
#  For people who have NEVER used Linux virtualization before.
#  It installs everything, fixes every permission and network gotcha,
#  finds (or downloads) the ArxOS ISO, asks you how big the VM should be,
#  and starts it -- with nothing left to fix after a reboot.
#
#  Every block below is commented so you can see exactly what it does.
#  Nothing destructive touches your real disk: it only creates ONE VM file.
# ============================================================================

set -uo pipefail            # -u: error on unset vars · -o pipefail: catch errors in pipes
                            # (we deliberately do NOT use -e so the menu can recover from errors)

# ---- pretty output helpers (colours degrade gracefully on dumb terminals) --
if [ -t 1 ]; then
  B=$'\e[1m'; DIM=$'\e[2m'; R=$'\e[0m'; G=$'\e[38;5;173m'; OK=$'\e[32m'; NO=$'\e[31m'; Y=$'\e[33m'
else
  B=""; DIM=""; R=""; G=""; OK=""; NO=""; Y=""
fi
say()  { printf '%s\n' "$*"; }                       # plain line
head() { printf '\n%s%s%s\n' "$B$G" "$*" "$R"; }     # section heading
good() { printf '%s✓%s %s\n' "$OK" "$R" "$*"; }      # success line
warn() { printf '%s!%s %s\n' "$Y" "$R" "$*"; }       # warning line
bad()  { printf '%s✗%s %s\n' "$NO" "$R" "$*"; }      # error line
ask()  { printf '%s%s%s ' "$B" "$*" "$R"; }          # a prompt (no newline)

# ---- run a command with sudo, showing the user what will run --------------
# Beginners are told exactly which privileged command runs and why.
SUDO=""; [ "$(id -u)" -ne 0 ] && SUDO="sudo"          # only prepend sudo if not already root
run() { printf '%s$ %s%s\n' "$DIM" "$*" "$R"; "$@"; } # echo the command, then run it

# ---- global state ----------------------------------------------------------
VM_NAME="arxos-vm"                                   # libvirt domain name we create
POOL_DIR="/var/lib/libvirt/images"                   # libvirt's own storage folder (correct perms, no /home headaches)
ISO_PATH=""                                          # filled in by the ISO step
VM_RAM=""; VM_CPUS=""; VM_DISK=""                    # filled in by the sizing step
R2_URL="https://pub-d5ff3efb1d204998aa120ded02d070b9.r2.dev/arxos-0.0.1.iso"
R2_SHA="18d87568b4e2cf4d2e83bb0052c93a5d76b12708bd942e86029ea612c1f2d44d"

# ===========================================================================
#  STEP 0  --  detect the distro so we use the right package manager
# ===========================================================================
detect_distro() {
  # /etc/os-release is the standard file every modern Linux ships that names the distro.
  . /etc/os-release 2>/dev/null || true
  DISTRO_ID="${ID:-unknown}"; DISTRO_LIKE="${ID_LIKE:-}"
  # Decide which package manager + package names to use.
  if   command -v pacman  >/dev/null 2>&1; then PKG="pacman"
  elif command -v apt     >/dev/null 2>&1; then PKG="apt"
  elif command -v dnf     >/dev/null 2>&1; then PKG="dnf"
  elif command -v zypper  >/dev/null 2>&1; then PKG="zypper"
  else PKG=""; fi
}

# ===========================================================================
#  STEP 1  --  install the full QEMU/KVM stack (packages differ per distro)
# ===========================================================================
install_stack() {
  head "Step 1 · Install QEMU/KVM + libvirt + virt-manager + UEFI + TPM"
  # Check for hardware virtualization first -- without it a VM would be painfully slow.
  if ! grep -Eq 'vmx|svm' /proc/cpuinfo; then
    warn "Your CPU does not report hardware virtualization (VT-x/AMD-V)."
    warn "Turn on 'Virtualization' / 'SVM' / 'VT-x' in your BIOS/UEFI, then rerun."
  else
    good "CPU virtualization (KVM) is available."
  fi

  case "$PKG" in
    pacman)  # Arch / EndeavourOS / Manjaro / ArxOS
      run $SUDO pacman -S --needed --noconfirm \
        qemu-full libvirt virt-manager dnsmasq edk2-ovmf swtpm ;;
    apt)     # Ubuntu / Debian / Mint / Pop!_OS
      run $SUDO apt update
      run $SUDO apt install -y \
        qemu-system-x86 qemu-utils libvirt-daemon-system virt-manager \
        dnsmasq-base ovmf swtpm ;;
    dnf)     # Fedora / RHEL / Nobara
      run $SUDO dnf install -y \
        @virtualization qemu-kvm libvirt virt-manager edk2-ovmf swtpm ;;
    zypper)  # openSUSE
      run $SUDO zypper install -y \
        qemu-kvm libvirt virt-manager ovmf swtpm ;;
    *)
      bad "No supported package manager found (pacman/apt/dnf/zypper)."
      warn "Install QEMU, libvirt, virt-manager, OVMF and swtpm by hand, then rerun."
      return 1 ;;
  esac
  good "Packages installed."
}

# ===========================================================================
#  STEP 2  --  make it 'just work' forever: service, groups, network, perms
#             (this is the part beginners always get stuck on)
# ===========================================================================
enable_everything() {
  head "Step 2 · Enable the service, your permissions, and networking"

  # 2a. Start the libvirt daemon now AND on every future boot, so you never
  #     have to start it by hand again.  Some distros use a socket, some a service.
  run $SUDO systemctl enable --now libvirtd.service 2>/dev/null \
    || run $SUDO systemctl enable --now libvirtd.socket 2>/dev/null || true
  good "libvirt daemon is running and will auto-start on every boot."

  # 2b. Add YOU to the 'libvirt' and 'kvm' groups so you can manage VMs
  #     without typing sudo for every action.  (Takes effect after re-login.)
  for grp in libvirt kvm; do
    if getent group "$grp" >/dev/null 2>&1; then
      run $SUDO usermod -aG "$grp" "$USER"
    fi
  done
  good "Added '$USER' to the libvirt and kvm groups."

  # 2c. firewalld ships a special 'libvirt' zone, but only picks it up after a
  #     reload -- skipping this is the #1 cause of 'network default failed to start'.
  if systemctl is-active --quiet firewalld 2>/dev/null; then
    run $SUDO firewall-cmd --reload || true
    good "Reloaded firewalld so it knows libvirt's network zone."
  fi

  # 2d. Turn on the built-in NAT network called 'default'.  This gives the VM
  #     internet through your PC with zero manual bridge setup.  We start it now
  #     AND mark it autostart so it comes back automatically after every reboot.
  if ! $SUDO virsh net-info default >/dev/null 2>&1; then
    # On a few minimal installs the default network isn't predefined -- create it.
    cat > /tmp/arxos-default-net.xml <<'NETXML'
<network>
  <name>default</name>
  <forward mode='nat'/>
  <bridge name='virbr0' stp='on' delay='0'/>
  <ip address='192.168.122.1' netmask='255.255.255.0'>
    <dhcp><range start='192.168.122.2' end='192.168.122.254'/></dhcp>
  </ip>
</network>
NETXML
    run $SUDO virsh net-define /tmp/arxos-default-net.xml || true
  fi
  run $SUDO virsh net-start default 2>/dev/null || true      # start it now (ok if already running)
  run $SUDO virsh net-autostart default 2>/dev/null || true  # make it come up on every boot
  good "NAT networking is on and set to auto-start after reboot."

  # 2e. Storage permissions: we keep the VM's disk in libvirt's own folder
  #     ($POOL_DIR), which already has the right owner (the 'qemu' user can read
  #     it).  This avoids the classic 'Permission denied' you hit when a VM disk
  #     lives under /home.  Make sure the default storage pool exists + autostarts.
  run $SUDO mkdir -p "$POOL_DIR"
  if ! $SUDO virsh pool-info default >/dev/null 2>&1; then
    run $SUDO virsh pool-define-as default dir --target "$POOL_DIR" || true
    run $SUDO virsh pool-build default 2>/dev/null || true
  fi
  run $SUDO virsh pool-start default 2>/dev/null || true
  run $SUDO virsh pool-autostart default 2>/dev/null || true
  good "Storage is in $POOL_DIR with correct permissions (no post-reboot fixes)."
}

# ===========================================================================
#  STEP 3  --  find the ArxOS ISO fast (or download it)
# ===========================================================================
find_iso() {
  head "Step 3 · Locate the ArxOS ISO"
  say "Scanning your computer for an ArxOS image (this is quick)..."

  # Fast path first: look in the folders people actually keep downloads/USBs,
  # before doing an exhaustive search.  -maxdepth keeps each scan snappy.
  local hits=() quick
  for quick in "$HOME/Downloads" "$HOME" /run/media /media /mnt "$POOL_DIR" "$(pwd)"; do
    [ -d "$quick" ] || continue
    while IFS= read -r f; do hits+=("$f"); done \
      < <(find "$quick" -maxdepth 4 -iname 'arxos*.iso' -type f 2>/dev/null)
  done

  # If nothing turned up, do ONE full-system sweep, skipping virtual filesystems
  # (/proc /sys /dev) so it stays fast and never hangs.
  if [ "${#hits[@]}" -eq 0 ]; then
    say "Not in the usual places -- doing a full search..."
    while IFS= read -r f; do hits+=("$f"); done \
      < <(find / -xdev -iname 'arxos*.iso' -type f 2>/dev/null)
  fi

  # De-duplicate the list.
  mapfile -t hits < <(printf '%s\n' "${hits[@]}" | awk 'NF' | sort -u)

  if [ "${#hits[@]}" -gt 0 ]; then
    say "Found:"
    local i=1; for f in "${hits[@]}"; do printf '  %s%d)%s %s\n' "$B" "$i" "$R" "$f"; i=$((i+1)); done
    printf '  %s%d)%s Download a fresh copy instead\n' "$B" "$i" "$R"
    ask "Pick a number (or 'b' to go back):"; read -r pick
    [ "$pick" = "b" ] && return 2
    if [ "$pick" -ge 1 ] 2>/dev/null && [ "$pick" -le "${#hits[@]}" ]; then
      ISO_PATH="${hits[$((pick-1))]}"; good "Using $ISO_PATH"; return 0
    fi
  else
    warn "No ArxOS ISO found on this computer."
  fi

  # Download branch (either chosen above or because nothing was found).
  ask "Download ArxOS 0.0.1 (~4.4 GB) now? [Y/n]:"; read -r dl
  case "${dl:-y}" in
    n|N) warn "Skipped. Put an arxos*.iso on this PC and rerun Step 3."; return 2 ;;
  esac
  ISO_PATH="$POOL_DIR/arxos-0.0.1.iso"
  # curl -L follows redirects, -C - resumes a half-finished download, --progress-bar shows a bar.
  run $SUDO curl -L -C - --progress-bar -o "$ISO_PATH" "$R2_URL" || { bad "Download failed."; return 1; }
  # Verify the download is intact by comparing its SHA256 to the official value.
  say "Verifying the download is intact (SHA256)..."
  if echo "$R2_SHA  $ISO_PATH" | $SUDO sha256sum -c - >/dev/null 2>&1; then
    good "Checksum matches -- the ISO is genuine and complete."
  else
    bad "Checksum MISMATCH -- the download is corrupt. Delete it and rerun Step 3."
    return 1
  fi
  good "Ready: $ISO_PATH"
}

# ===========================================================================
#  STEP 4  --  size the VM based on YOUR hardware
# ===========================================================================
size_vm() {
  head "Step 4 · Choose how much of your PC to give the VM"

  # Read total RAM (kB) from /proc/meminfo and total logical CPUs from nproc.
  local ram_kb ram_mb host_cpus
  ram_kb=$(awk '/MemTotal/{print $2}' /proc/meminfo)
  ram_mb=$(( ram_kb / 1024 ))
  host_cpus=$(nproc)

  # Suggest sane defaults: half the RAM (min 4 GB), half the cores (min 2).
  local sug_ram=$(( ram_mb / 2 )); [ "$sug_ram" -lt 4096 ] && sug_ram=4096
  local max_ram=$(( ram_mb - 2048 ))                 # always leave 2 GB for the host
  [ "$sug_ram" -gt "$max_ram" ] && sug_ram=$max_ram
  local sug_cpu=$(( host_cpus / 2 )); [ "$sug_cpu" -lt 2 ] && sug_cpu=2

  say "Your PC has ${B}$(( ram_mb/1024 )) GB RAM${R} and ${B}${host_cpus} CPU cores${R}."
  say ""

  # --- RAM ---
  ask "RAM for the VM in MB [suggested ${sug_ram}, max ${max_ram}] (or 'b' to go back):"; read -r a
  [ "$a" = "b" ] && return 2
  VM_RAM="${a:-$sug_ram}"
  # keep it inside safe bounds so the host never starves
  [ "$VM_RAM" -lt 2048 ] 2>/dev/null && VM_RAM=2048
  [ "$VM_RAM" -gt "$max_ram" ] 2>/dev/null && VM_RAM=$max_ram

  # --- CPUs ---
  ask "CPU cores for the VM [suggested ${sug_cpu}, max ${host_cpus}]:"; read -r c
  VM_CPUS="${c:-$sug_cpu}"
  [ "$VM_CPUS" -lt 1 ] 2>/dev/null && VM_CPUS=1
  [ "$VM_CPUS" -gt "$host_cpus" ] 2>/dev/null && VM_CPUS=$host_cpus

  # --- Disk ---
  ask "Virtual disk size in GB [suggested 40, minimum 25]:"; read -r d
  VM_DISK="${d:-40}"
  [ "$VM_DISK" -lt 25 ] 2>/dev/null && VM_DISK=25

  good "VM will get: ${VM_RAM} MB RAM · ${VM_CPUS} cores · ${VM_DISK} GB disk (thin, grows as used)."
}

# ===========================================================================
#  STEP 5  --  create the disk, define the VM, and start it
# ===========================================================================
create_vm() {
  head "Step 5 · Build and launch the ArxOS VM"

  # If a VM with this name already exists, offer to replace it (so reruns are safe).
  if $SUDO virsh dominfo "$VM_NAME" >/dev/null 2>&1; then
    warn "A VM named '$VM_NAME' already exists."
    ask "Delete it and rebuild? [y/N] (or 'b' to go back):"; read -r rr
    [ "$rr" = "b" ] && return 2
    case "$rr" in
      y|Y) run $SUDO virsh destroy "$VM_NAME" 2>/dev/null || true
           run $SUDO virsh undefine "$VM_NAME" --nvram 2>/dev/null || true ;;
      *)   warn "Keeping the existing VM. Open 'virt-manager' to use it."; return 0 ;;
    esac
  fi

  # 5a. Create the thin virtual disk in libvirt's storage folder.
  #     qcow2 grows on demand, so a 40 GB disk uses almost no space until filled.
  local disk="$POOL_DIR/${VM_NAME}.qcow2"
  run $SUDO qemu-img create -f qcow2 "$disk" "${VM_DISK}G"

  # 5b. virt-install writes the whole VM definition for us -- UEFI firmware,
  #     an emulated TPM 2.0 (needed for ArxOS DEATHSTROKE), NAT networking,
  #     and the ISO as a boot CD.  --noautoconsole returns control immediately.
  run $SUDO virt-install \
    --name "$VM_NAME" \
    --memory "$VM_RAM" \
    --vcpus "$VM_CPUS" \
    --cpu host-passthrough \
    --machine q35 \
    --boot uefi \
    --tpm backend.type=emulator,backend.version=2.0,model=tpm-crb \
    --disk path="$disk",format=qcow2,bus=virtio \
    --cdrom "$ISO_PATH" \
    --os-variant archlinux \
    --network network=default,model=virtio \
    --graphics spice \
    --video virtio \
    --noautoconsole

  good "The ArxOS VM is created and starting."
  say ""
  say "Open it now with:  ${B}virt-manager${R}   (double-click '${VM_NAME}')"
  say "Inside the VM, double-click ${B}Install...${R} on the desktop to install ArxOS."
  say ""
  warn "If group changes just applied, log out and back in once so 'virsh' works without sudo."
}

# ===========================================================================
#  MAIN MENU  --  run the steps in order, or jump to any one; 'b' goes back
# ===========================================================================
main_menu() {
  detect_distro
  while true; do
    head "ArxOS VM Setup  ·  ${DISTRO_ID} detected  ·  package manager: ${PKG:-none}"
    say "  ${B}1)${R} Install QEMU/KVM + tools        (Step 1)"
    say "  ${B}2)${R} Enable service, perms, network  (Step 2)"
    say "  ${B}3)${R} Find or download the ArxOS ISO  (Step 3)"
    say "  ${B}4)${R} Choose VM size (RAM/CPU/disk)   (Step 4)"
    say "  ${B}5)${R} Create and start the VM         (Step 5)"
    say "  ${B}a)${R} Do everything, in order (recommended for first-timers)"
    say "  ${B}q)${R} Quit"
    ask "Choose:"; read -r choice
    case "$choice" in
      1) install_stack ;;
      2) enable_everything ;;
      3) find_iso ;;
      4) size_vm ;;
      5) [ -z "$ISO_PATH" ] && { warn "Do Step 3 (find ISO) first."; continue; }
         [ -z "$VM_RAM" ]  && { warn "Do Step 4 (VM size) first."; continue; }
         create_vm ;;
      a|A)
         # Run the whole pipeline. Any step returning 2 means the user chose 'back'.
         install_stack        || { warn "Step 1 had an issue -- review above."; continue; }
         enable_everything    || { warn "Step 2 had an issue -- review above."; continue; }
         find_iso;   rc=$?; [ "$rc" = 2 ] && continue; [ "$rc" != 0 ] && continue
         size_vm;    rc=$?; [ "$rc" = 2 ] && continue
         create_vm
         ;;
      q|Q) say "Bye. Run this script again any time."; exit 0 ;;
      b|B) : ;;   # already at the top menu; 'back' just redraws
      *)   warn "Type 1-5, a, or q." ;;
    esac
  done
}

# Friendly banner, then start.
clear 2>/dev/null || true
say "${B}${G}"
say "   █████╗ ██████╗ ██╗  ██╗ ██████╗ ███████╗   VM Setup"
say "  ██╔══██╗██╔══██╗╚██╗██╔╝██╔═══██╗██╔════╝"
say "  ███████║██████╔╝ ╚███╔╝ ██║   ██║███████╗   QEMU/KVM, done for you"
say "  ██╔══██║██╔══██╗ ██╔██╗ ██║   ██║╚════██║"
say "  ██║  ██║██║  ██║██╔╝ ██╗╚██████╔╝███████║   Stingray Labs"
say "  ╚═╝  ╚═╝╚═╝  ╚═╝╚═╝  ╚═╝ ╚═════╝ ╚══════╝"
say "${R}"
say "This will install virtualization, fix all the fiddly bits, and boot ArxOS in a VM."
say "It only creates ONE VM file. Your real system and files are never touched."
say ""
main_menu
