#!/usr/bin/env bash
#
# build_rpi_sd_no_chroot_auto.sh
#
# Fully automated Raspberry Pi SD card creation script WITHOUT chroot.
# When the Pi boots the first time, it installs snapclient automatically.
#
# WARNING: This script will format and overwrite $SD_DEVICE!
# Make sure you set $SD_DEVICE correctly.

set -e

#######################################
#            CONFIGURATION            #
#######################################
# Where to download Raspberry Pi OS Lite if not found
IMG_URL="https://downloads.raspberrypi.org/raspios_lite_armhf_latest"

# The exact .img filename you expect after extraction/unzipping
EXPECTED_IMG="rpi_os_lite_latest.img"
# Adjust to match the real image name. If the official download is .xz or .zip,
# you might need to adapt the script below.

# Your SD card device (e.g., /dev/sdb, /dev/mmcblk0)
SD_DEVICE="/dev/sda"

# Wi-Fi credentials
WIFI_SSID="Cloonar-Multimedia"
WIFI_PASS="K2MC28Zhk\$4zsx6Y"
WIFI_COUNTRY="AT"

# Snapserver IP or hostname
SNAPSERVER_HOST="snapcast.cloonar.com"

#######################################
#   FILE NAMES & PATHS ON THE PI      #
#######################################
# We'll copy a systemd service and script into the Pi's root filesystem.
FIRSTBOOT_SCRIPT="/usr/local/bin/firstboot.sh"
FIRSTBOOT_SERVICE="/etc/systemd/system/firstboot.service"

#######################################
#       DOWNLOAD IMAGE IF NEEDED      #
#######################################
download_image_if_needed() {
  if [[ ! -f "$EXPECTED_IMG" ]]; then
    echo ">>> $EXPECTED_IMG not found, downloading from $IMG_URL..."
    wget -O rpi_os_lite_latest.img.xz "$IMG_URL"
    echo ">>> Unzipping..."
    # unzip rpi_os_lite_latest.zip || true
    # If you get a .xz, un-comment:
    xz -dv *.xz

    if [[ ! -f "$EXPECTED_IMG" ]]; then
      echo "ERROR: Expected .img file '$EXPECTED_IMG' not found after extraction!"
      exit 1
    fi
  else
    echo ">>> Using existing image file: $EXPECTED_IMG"
  fi
}

#######################################
#          WRITE IMAGE TO SD          #
#######################################
write_image_to_sd() {
  echo ">>> Writing $EXPECTED_IMG to $SD_DEVICE..."
  umount "${SD_DEVICE}"* || true
  dd if="$EXPECTED_IMG" of="$SD_DEVICE" bs=4M conv=fsync status=progress
  sync
}

#######################################
#        MOUNT BOOT & ROOTFS          #
#######################################
mount_partitions() {
  if [[ $SD_DEVICE == *"mmcblk"* ]]; then
    BOOT_PART="${SD_DEVICE}p1"
    ROOT_PART="${SD_DEVICE}p2"
  else
    BOOT_PART="${SD_DEVICE}1"
    ROOT_PART="${SD_DEVICE}2"
  fi

  mkdir -p /mnt/rpi-boot
  mkdir -p /mnt/rpi-root

  echo ">>> Mounting boot partition ($BOOT_PART) -> /mnt/rpi-boot"
  mount "$BOOT_PART" /mnt/rpi-boot

  echo ">>> Mounting root partition ($ROOT_PART) -> /mnt/rpi-root"
  mount "$ROOT_PART" /mnt/rpi-root
}

unmount_partitions() {
  echo ">>> Unmounting partitions..."
  umount /mnt/rpi-boot || true
  umount /mnt/rpi-root || true
  sync
}

#######################################
#   PREPARE BOOT PARTITION (NO CHROOT)
#######################################
enable_ssh_wifi_hifiberry() {
  echo ">>> Enabling SSH..."
  touch /mnt/rpi-boot/ssh

  echo ">>> Creating wpa_supplicant.conf..."
  cat << EOF > /mnt/rpi-boot/wpa_supplicant.conf
country=$WIFI_COUNTRY
ctrl_interface=DIR=/var/run/wpa_supplicant GROUP=netdev
update_config=1

network={
    ssid="$WIFI_SSID"
    psk="$WIFI_PASS"
}
EOF

  echo ">>> Enabling HiFiBerry Amp+ overlay in config.txt..."
  cat << EOF >> /mnt/rpi-boot/config.txt

# --- HiFiBerry Amp+ overlay ---
dtoverlay=hifiberry-amp
dtparam=audio=off
EOF
}

#######################################
#   INSTALL FIRST-BOOT SERVICE/SCRIPT
#   (PLACED INTO THE Pi ROOTFS)
#######################################
install_firstboot_autosetup() {
  echo ">>> Creating first-boot script ($FIRSTBOOT_SCRIPT) in Pi root filesystem..."

  cat << 'FOO' > /mnt/rpi-root$FIRSTBOOT_SCRIPT
#!/usr/bin/env bash
set -e

echo ">>> [FIRST BOOT] Installing snapclient..."

# Update system
apt-get update
apt-get upgrade -y

# Install snapclient
apt-get install -y snapclient

# Configure snapclient
cat << SCLIENT > /etc/default/snapclient
SNAPCLIENT_OPTS="--host SNAPSERVER_PLACEHOLDER"
SCLIENT

systemctl enable snapclient
systemctl start snapclient

echo ">>> [FIRST BOOT] Snapclient installation complete."

# Disable this service so it doesn't run again
systemctl disable firstboot.service

# Optional: remove script & service
rm -f /usr/local/bin/firstboot.sh
rm -f /etc/systemd/system/firstboot.service
FOO

  # Replace placeholder with actual snapserver host
  sed -i "s|SNAPSERVER_PLACEHOLDER|$SNAPSERVER_HOST|g" /mnt/rpi-root$FIRSTBOOT_SCRIPT

  # Make it executable
  chmod +x /mnt/rpi-root$FIRSTBOOT_SCRIPT


  echo ">>> Creating first-boot systemd service ($FIRSTBOOT_SERVICE)..."
  cat << EOF > /mnt/rpi-root$FIRSTBOOT_SERVICE
[Unit]
Description=Run once at first boot to install snapclient
After=network-online.target
Wants=network-online.target

[Service]
Type=oneshot
ExecStart=$FIRSTBOOT_SCRIPT
RemainAfterExit=true

[Install]
WantedBy=multi-user.target
EOF

  echo ">>> Enabling firstboot.service so it runs automatically..."
  # We can manually create the symlink to "enable" it:
  mkdir -p /mnt/rpi-root/etc/systemd/system/multi-user.target.wants
  ln -s $FIRSTBOOT_SERVICE /mnt/rpi-root/etc/systemd/system/multi-user.target.wants/firstboot.service
}


#######################################
#                MAIN                 #
#######################################
main() {
  # 1) Download image if needed
  download_image_if_needed

  # 2) Write image to SD
  write_image_to_sd

  # 3) Mount partitions
  mount_partitions

  # 4) Enable SSH, Wi-Fi, HiFiBerry in /boot
  enable_ssh_wifi_hifiberry

  # 5) Place the "first boot" script & service in / (root partition)
  install_firstboot_autosetup

  # 6) Unmount
  unmount_partitions

  echo "=============================================================="
  echo "Done! The SD card is fully configured for auto snapclient setup."
  echo "Insert it into your Pi, power on, and wait a couple of minutes."
  echo ""
  echo "What happens on first boot:"
  echo "  1. Pi connects to Wi-Fi, SSH is enabled"
  echo "  2. 'firstboot.service' runs once, installs snapclient, configures it"
  echo "  3. The service disables itself and removes the script"
  echo ""
  echo "No user intervention is required. Enjoy!"
  echo "=============================================================="
}

main "$@"