#!/usr/bin/env bash

# Script to create a read-only Raspberry Pi OS Lite image with Snapcast client.
# Requires sudo privileges for many operations.
# Ensure you have all dependencies from shell.nix (e.g., qemu-arm-static, parted, etc.)

set -euo pipefail # Exit on error, undefined variable, or pipe failure

# --- Configuration & Defaults ---
RASPI_OS_LITE_URL="https://downloads.raspberrypi.com/raspios_lite_arm64/images/raspios_lite_arm64-2025-05-07/2025-05-06-raspios-bookworm-arm64-lite.img.xz"
# Check for the latest image URL from https://www.raspberrypi.com/software/operating-systems/
# The script assumes an .img.xz file. If it's .zip, adjust extraction.

WORK_DIR="rpi_build_temp"
# QEMU is no longer used in this script.

# --- Helper Functions ---
info() { echo -e "\033[0;32m[INFO]\033[0m $1"; }
warn() { echo -e "\033[0;33m[WARN]\033[0m $1"; }
error() { echo -e "\033[0;31m[ERROR]\033[0m $1" >&2; exit 1; }
cleanup_exit() {
    warn "Cleaning up and exiting..."
    local rootfs_mount_point="${PWD}/${WORK_DIR}/rootfs" # PWD should be the script's initial dir, WORK_DIR is relative

    # Attempt to unmount in reverse order of mounting
    if mount | grep -q "${rootfs_mount_point}/boot"; then
        info "Cleanup: Unmounting ${rootfs_mount_point}/boot..."
        sudo umount "${rootfs_mount_point}/boot" || sudo umount -l "${rootfs_mount_point}/boot" || warn "Failed to unmount ${rootfs_mount_point}/boot during cleanup."
    fi
    if mount | grep -q "${rootfs_mount_point}"; then # Check rootfs itself
        info "Cleanup: Unmounting ${rootfs_mount_point}..."
        sudo umount "${rootfs_mount_point}" || sudo umount -l "${rootfs_mount_point}" || warn "Failed to unmount ${rootfs_mount_point} during cleanup."
    fi
    if [ -n "${LOOP_DEV:-}" ] && losetup -a | grep -q "${LOOP_DEV}"; then
        info "Cleanup: Detaching loop device ${LOOP_DEV}..."
        sudo losetup -d "${LOOP_DEV}" || warn "Failed to detach ${LOOP_DEV} during cleanup."
    fi
    # sudo rm -rf "${WORK_DIR}" # Optional: clean work dir on error
    exit 1
}
trap cleanup_exit ERR INT TERM

# --- Argument Parsing ---
DEVICE_TYPE=""
HOSTNAME_PI=""
CONFIG_FILE="./config.txt"
OUTPUT_IMAGE_FILE=""

usage() {
    echo "Usage: $0 -d <device_type> -n <hostname> [-c <config_file>] [-o <output_image>]"
    echo "  -d: Device type (rpizero2w | rpi4)"
    echo "  -n: Desired hostname for the Raspberry Pi"
    echo "  -c: Path to config.txt (default: ./config.txt)"
    echo "  -o: Output image file name (default: snapcast-client-<device_type>-<hostname>.img)"
    exit 1
}

while getopts "d:n:c:o:h" opt; do
    case ${opt} in
        d) DEVICE_TYPE="${OPTARG}";;
        n) HOSTNAME_PI="${OPTARG}";;
        c) CONFIG_FILE="${OPTARG}";;
        o) OUTPUT_IMAGE_FILE="${OPTARG}";;
        h) usage;;
        *) usage;;
    esac
done

if [ -z "${DEVICE_TYPE}" ]; then
    error "Mandatory argument -d <device_type> is missing. Use -h for help."
fi
if [ -z "${HOSTNAME_PI}" ]; then
    error "Mandatory argument -n <hostname> is missing. Use -h for help."
fi
if [ "${DEVICE_TYPE}" != "rpizero2w" ] && [ "${DEVICE_TYPE}" != "rpi4" ]; then
    error "Invalid device type: '${DEVICE_TYPE}'. Must be 'rpizero2w' or 'rpi4'."
fi
if [ ! -f "${CONFIG_FILE}" ]; then
    error "Config file not found: ${CONFIG_FILE}"
fi
if [ -z "${OUTPUT_IMAGE_FILE}" ]; then
    OUTPUT_IMAGE_FILE="snapcast-client-${DEVICE_TYPE}-${HOSTNAME_PI}.img"
fi

info "Starting Raspberry Pi Image Builder..."
info "Device Type: ${DEVICE_TYPE}"
info "Hostname: ${HOSTNAME_PI}"
info "Config File: ${CONFIG_FILE}"
info "Output Image: ${OUTPUT_IMAGE_FILE}"

# --- Load Configuration ---
source "${CONFIG_FILE}"
if [ -z "${WIFI_SSID:-}" ] || [ -z "${WIFI_PSK:-}" ] || [ -z "${SNAPCAST_SERVER:-}" ]; then
    error "WIFI_SSID, WIFI_PSK, or SNAPCAST_SERVER not set in config file."
fi

# --- Prepare Workspace ---
sudo rm -rf "${WORK_DIR}"
mkdir -p "${WORK_DIR}"
cd "${WORK_DIR}"

# --- 1. Base Image Acquisition ---
IMG_XZ_NAME=$(basename "${RASPI_OS_LITE_URL}")
IMG_NAME="${IMG_XZ_NAME%.xz}"

# Check for uncompressed image first
if [ -f "${IMG_NAME}" ]; then
    info "Using existing uncompressed image: ${IMG_NAME}"
# Else, check for compressed image
elif [ -f "${IMG_XZ_NAME}" ]; then
    info "Found existing compressed image: ${IMG_XZ_NAME}. Extracting..."
    xz -d -k "${IMG_XZ_NAME}" # -k to keep the original .xz file
    if [ ! -f "${IMG_NAME}" ]; then # Double check extraction
        error "Failed to extract ${IMG_XZ_NAME} to ${IMG_NAME}"
    fi
    info "Extraction complete: ${IMG_NAME}"
# Else, download and extract
else
    info "Downloading Raspberry Pi OS Lite image: ${IMG_XZ_NAME}..."
    wget -q --show-progress -O "${IMG_XZ_NAME}" "${RASPI_OS_LITE_URL}"
    info "Extracting image..."
    xz -d -k "${IMG_XZ_NAME}" # -k to keep the original .xz file
    if [ ! -f "${IMG_NAME}" ]; then # Double check extraction
        error "Failed to extract ${IMG_XZ_NAME} to ${IMG_NAME}"
    fi
    info "Extraction complete: ${IMG_NAME}"
fi

# Always work on a copy for the output image
info "Copying ${IMG_NAME} to ${OUTPUT_IMAGE_FILE}..."
cp "${IMG_NAME}" "${OUTPUT_IMAGE_FILE}"

# --- 2. Mount Image Partitions ---
info "Setting up loop device for ${OUTPUT_IMAGE_FILE}..."
LOOP_DEV=$(sudo losetup -Pf --show "${OUTPUT_IMAGE_FILE}")
if [ -z "${LOOP_DEV}" ]; then error "Failed to setup loop device."; fi
info "Loop device: ${LOOP_DEV}"

# Wait for device nodes to be created
sleep 2
sudo partprobe "${LOOP_DEV}"
sleep 2

BOOT_PART="${LOOP_DEV}p1"
ROOT_PART="${LOOP_DEV}p2"

mkdir -p rootfs
info "Mounting root partition (${ROOT_PART}) to rootfs/..."
sudo mount "${ROOT_PART}" rootfs
info "Mounting boot partition (${BOOT_PART}) to rootfs/boot/..."
# Note: Newer RPi OS might use /boot/firmware. Adjust if needed.
# Check if /boot/firmware exists, if so, use it.
# For simplicity, this script assumes /boot. If issues, this is a place to check.
# A more robust check: BOOT_MOUNT_POINT="rootfs/boot"; if [ -d "rootfs/boot/firmware" ]; then BOOT_MOUNT_POINT="rootfs/boot/firmware"; fi
# sudo mount "${BOOT_PART}" "${BOOT_MOUNT_POINT}"
sudo mount "${BOOT_PART}" rootfs/boot 

# --- 3. System Configuration (Directly on Mounted Rootfs) ---
# QEMU and chroot are no longer used. All operations are on the mounted rootfs.

if ! command -v dpkg-deb &> /dev/null; then
    error "dpkg-deb command not found. Please ensure dpkg is installed and in your PATH (e.g., via Nix shell)."
fi

info "Setting hostname to ${HOSTNAME_PI} on rootfs..."
sudo sh -c "echo '${HOSTNAME_PI}' > rootfs/etc/hostname"
sudo sed -i "s/127.0.1.1.*raspberrypi/127.0.1.1\t${HOSTNAME_PI}/g" rootfs/etc/hosts
sudo sed -i "s/raspberrypi/${HOSTNAME_PI}/g" rootfs/etc/hosts # Also replace other occurrences

info "Configuring Wi-Fi (wpa_supplicant) on rootfs..."
sudo sh -c "cat > rootfs/boot/wpa_supplicant.conf <<WPA
ctrl_interface=DIR=/var/run/wpa_supplicant GROUP=netdev
update_config=1
country=AT # Set your country code

network={
    ssid=\"${WIFI_SSID}\"
    psk=\"${WIFI_PSK}\"
    scan_ssid=1
}
WPA"
sudo chmod 600 rootfs/boot/wpa_supplicant.conf

info "Downloading and installing Snapcast client from .deb..."
SNAPCLIENT_DEB_URL="https://github.com/badaix/snapcast/releases/download/v0.31.0/snapclient_0.31.0-1_arm64_bookworm_with-pulse.deb"
SNAPCLIENT_DEB_NAME=$(basename "${SNAPCLIENT_DEB_URL}")

if [ ! -f "${SNAPCLIENT_DEB_NAME}" ]; then
    wget -q --show-progress -O "${SNAPCLIENT_DEB_NAME}" "${SNAPCLIENT_DEB_URL}"
else
    info "Using existing ${SNAPCLIENT_DEB_NAME}"
fi

mkdir -p snapclient_deb_extract
info "Extracting ${SNAPCLIENT_DEB_NAME}..."
dpkg-deb -x "${SNAPCLIENT_DEB_NAME}" snapclient_deb_extract
info "Copying extracted files to rootfs using rsync..."
# Use rsync to handle merging and symlinks like /lib -> /usr/lib correctly
# The trailing slash on snapclient_deb_extract/ is important for rsync
if ! command -v rsync &> /dev/null; then
    error "rsync command not found. Please ensure rsync is installed and in your PATH (e.g., via Nix shell)."
fi

if ! command -v openssl &> /dev/null; then
    error "openssl command not found. Please ensure openssl is installed and in your PATH (e.g., via Nix shell)."
fi
sudo rsync -aK --chown=root:root snapclient_deb_extract/ rootfs/
rm -rf snapclient_deb_extract "${SNAPCLIENT_DEB_NAME}"
info "Snapclient files installed."

info "Attempting to create 'snapclient' user and group..."
SNAPCLIENT_UID=987 # Choose an appropriate UID/GID
SNAPCLIENT_GID=987
SNAPCLIENT_USER="snapclient"
SNAPCLIENT_GROUP="snapclient"
SNAPCLIENT_HOME="/var/lib/snapclient" # A typical home for system users, though not strictly needed if nologin
SNAPCLIENT_SHELL="/usr/sbin/nologin"

# Create group if it doesn't exist
if ! sudo grep -q "^${SNAPCLIENT_GROUP}:" rootfs/etc/group; then
    info "Creating group '${SNAPCLIENT_GROUP}' (${SNAPCLIENT_GID}) in rootfs/etc/group"
    sudo sh -c "echo '${SNAPCLIENT_GROUP}:x:${SNAPCLIENT_GID}:' >> rootfs/etc/group"
else
    info "Group '${SNAPCLIENT_GROUP}' already exists in rootfs/etc/group."
fi

# Create user if it doesn't exist
if ! sudo grep -q "^${SNAPCLIENT_USER}:" rootfs/etc/passwd; then
    info "Creating user '${SNAPCLIENT_USER}' (${SNAPCLIENT_UID}) in rootfs/etc/passwd"
    sudo sh -c "echo '${SNAPCLIENT_USER}:x:${SNAPCLIENT_UID}:${SNAPCLIENT_GID}:${SNAPCLIENT_USER} system user:${SNAPCLIENT_HOME}:${SNAPCLIENT_SHELL}' >> rootfs/etc/passwd"
    info "Creating basic shadow entry for '${SNAPCLIENT_USER}' (account locked)"
    # '!' in password field locks the account. '*' also works.
    sudo sh -c "echo '${SNAPCLIENT_USER}:!:19700:0:99999:7:::' >> rootfs/etc/shadow"
    # Create home directory if it doesn't exist and set permissions
    sudo mkdir -p "rootfs${SNAPCLIENT_HOME}"
    sudo chown "${SNAPCLIENT_UID}:${SNAPCLIENT_GID}" "rootfs${SNAPCLIENT_HOME}"
    sudo chmod 700 "rootfs${SNAPCLIENT_HOME}"
else
    info "User '${SNAPCLIENT_USER}' already exists in rootfs/etc/passwd."
fi
# Remove previous warnings
# warn "The snapclient user and group were NOT automatically created."
# warn "Ensure 'snapclient' user/group exist on target or adjust service file."

info "Creating Snapcast systemd service file on rootfs..."
# Note: SNAPCAST_SERVER_IP is from the config file via 'source "${CONFIG_FILE}"'
sudo sh -c "cat > rootfs/etc/systemd/system/snapclient.service <<SERVICE
[Unit]
Description=Snapcast client
After=network-online.target sound.target
Wants=network-online.target

[Service]
ExecStart=/usr/bin/snapclient -h ${SNAPCAST_SERVER} --player pulse # Or alsa, adjust player if needed
Restart=always
User=${SNAPCLIENT_USER}
Group=${SNAPCLIENT_GROUP}
# User and group should now be created by this script.

[Install]
WantedBy=multi-user.target
SERVICE"

info "Enabling Snapcast systemd service on rootfs..."
sudo mkdir -p rootfs/etc/systemd/system/multi-user.target.wants
sudo ln -sf ../../../../lib/systemd/system/snapclient.service rootfs/etc/systemd/system/multi-user.target.wants/snapclient.service
# Note: The above symlink path assumes snapclient.service from the .deb is installed to /lib/systemd/system/snapclient.service
# If dpkg-deb -x places it in /usr/lib/systemd/system, adjust the symlink source.
# A common location for package-installed units is /lib/systemd/system.
# If the .deb actually places it in /etc/systemd/system, then the symlink would be:
# sudo ln -sf ../snapclient.service rootfs/etc/systemd/system/multi-user.target.wants/snapclient.service
# Let's assume the .deb installs it to /usr/lib/systemd/system or /lib/systemd/system.
# The .deb extraction copies to rootfs/, so if the .deb has ./usr/lib/systemd/system/snapclient.service,
# it will be at rootfs/usr/lib/systemd/system/snapclient.service.
# The service file we created is at rootfs/etc/systemd/system/snapclient.service.
# Systemd prefers /etc/systemd/system over /lib/systemd/system.
# So, if we create it in /etc/systemd/system, that should be fine.
# The symlink should point to the file in /etc/systemd/system if we create it there.

info "Enabling dhcpcd systemd service on rootfs..."
# Ensure the multi-user.target.wants directory exists (already created for snapclient)
# Standard path for dhcpcd.service in base RPi OS images is /lib/systemd/system/dhcpcd.service or /usr/lib/systemd/system/dhcpcd.service
# The symlink needs to point from /etc/systemd/system/multi-user.target.wants/ to that location.
if [ -f "rootfs/lib/systemd/system/dhcpcd.service" ]; then
    sudo ln -sf ../../../../lib/systemd/system/dhcpcd.service rootfs/etc/systemd/system/multi-user.target.wants/dhcpcd.service
    info "dhcpcd.service enabled (symlink created from /lib)."
elif [ -f "rootfs/usr/lib/systemd/system/dhcpcd.service" ]; then
    sudo ln -sf ../../../../usr/lib/systemd/system/dhcpcd.service rootfs/etc/systemd/system/multi-user.target.wants/dhcpcd.service
    info "dhcpcd.service enabled (symlink created from /usr/lib)."
else
    warn "dhcpcd.service file not found in rootfs/lib/systemd/system/ or rootfs/usr/lib/systemd/system/. Cannot enable dhcpcd."
fi

# Let's ensure the symlink points to the one we just created:
sudo rm -f rootfs/etc/systemd/system/multi-user.target.wants/snapclient.service # remove if exists
sudo ln -s ../snapclient.service rootfs/etc/systemd/system/multi-user.target.wants/snapclient.service
info "Snapclient service enabled (symlink created)."


if [ "${DEVICE_TYPE}" == "rpizero2w" ]; then
    info "Applying HifiBerry DAC+ overlay for Raspberry Pi Zero 2 W on rootfs..."
    sudo sh -c "echo 'dtoverlay=hifiberry-dacplus' >> rootfs/boot/config.txt"
fi

info "Configuring for read-only filesystem on rootfs..."
# 1. Modify /etc/fstab
sudo sed -i -E 's/(\s+\/\s+ext4\s+)(defaults,noatime)(\s+0\s+1)/\1ro,defaults,noatime\3/' rootfs/etc/fstab
sudo sed -i -E 's/(\s+\/boot\s+vfat\s+)(defaults)(\s+0\s+2)/\1ro,defaults,nofail\3/' rootfs/etc/fstab

# 2. Add fastboot and ro to /boot/cmdline.txt
BOOT_CMDLINE_FILE="rootfs/boot/cmdline.txt"
if [ -f "${BOOT_CMDLINE_FILE}" ]; then
    if ! sudo grep -q "fastboot" "${BOOT_CMDLINE_FILE}"; then
        sudo sed -i '1 s/$/ fastboot/' "${BOOT_CMDLINE_FILE}"
    fi
    if ! sudo grep -q " ro" "${BOOT_CMDLINE_FILE}"; then # space before ro is important
        sudo sed -i '1 s/$/ ro/' "${BOOT_CMDLINE_FILE}" # Add ro if not present
    fi
else
    warn "${BOOT_CMDLINE_FILE} not found. Skipping cmdline modifications."
fi

info "Creating userconf.txt for predefined user setup..."
DEFAULT_USER="rpiuser"
DEFAULT_PASS="raspberry"
# Check if openssl is available (already done above, but good to be mindful here)
ENCRYPTED_PASS=$(echo "${DEFAULT_PASS}" | openssl passwd -6 -stdin)
if [ -z "${ENCRYPTED_PASS}" ]; then
    error "Failed to encrypt password using openssl."
fi
sudo sh -c "echo '${DEFAULT_USER}:${ENCRYPTED_PASS}' > rootfs/boot/userconf.txt"
sudo chmod 600 rootfs/boot/userconf.txt # Set appropriate permissions
info "userconf.txt created with user '${DEFAULT_USER}'."

info "Attempting to disable unnecessary write-heavy services on rootfs..."
# This is a best-effort attempt by removing common symlinks.
# The exact paths might vary or services might not be installed.
DISABLED_SERVICES_COUNT=0
declare -a services_to_disable=(
    "apt-daily.timer"
    "apt-daily-upgrade.timer"
    "man-db.timer"
    "dphys-swapfile.service" # RPi OS specific swap service
    "logrotate.timer"
    "motd-news.timer"
)
declare -a common_wants_dirs=(
    "multi-user.target.wants"
    "timers.target.wants"
    "sysinit.target.wants"
    # Add other .wants directories if known
)
for service in "${services_to_disable[@]}"; do
    for wants_dir in "${common_wants_dirs[@]}"; do
        link_path="rootfs/etc/systemd/system/${wants_dir}/${service}"
        if [ -L "${link_path}" ]; then
            info "Disabling ${service} by removing symlink ${link_path}"
            sudo rm -f "${link_path}"
            DISABLED_SERVICES_COUNT=$((DISABLED_SERVICES_COUNT + 1))
        fi
    done
    # Also check for service files directly in /lib/systemd/system and mask them if we want to be more aggressive
    # For now, just removing symlinks from /etc/systemd/system is less intrusive.
done
info "Attempted to disable ${DISABLED_SERVICES_COUNT} services by removing symlinks."
warn "Service disabling is best-effort. Review target system services."

# No apt cache to clean as we didn't use apt

info "System configuration on rootfs complete."

# --- 5. Cleanup & Unmount ---
# Pseudo-filesystems and QEMU are no longer used, so their cleanup is removed.

info "Unmounting partitions..."
# Unmount boot first, then root
# BOOT_MOUNT_POINT="rootfs/boot"; if [ -d "rootfs/boot/firmware" ]; then BOOT_MOUNT_POINT="rootfs/boot/firmware"; fi
# sudo umount "${BOOT_MOUNT_POINT}"
sudo umount rootfs/boot
sudo umount rootfs

info "Detaching loop device ${LOOP_DEV}..."
sudo losetup -d "${LOOP_DEV}"
unset LOOP_DEV # Important for trap

rmdir rootfs
cd .. # Back to original directory

# --- 6. Shrink Image (Optional) ---
info "Image shrinking (optional step)..."
info "If you want to shrink the image, you can use a tool like PiShrink."
info "Example: sudo pishrink.sh ${WORK_DIR}/${OUTPUT_IMAGE_FILE}"
# Check if pishrink is available and executable
# if [ -x "./pishrink.sh" ]; then
#   info "Running PiShrink..."
#   sudo ./pishrink.sh "${WORK_DIR}/${OUTPUT_IMAGE_FILE}"
# else
#   warn "PiShrink script (pishrink.sh) not found or not executable in current directory. Skipping shrink."
# fi


# --- 7. Final Output ---
FINAL_IMAGE_PATH="${WORK_DIR}/${OUTPUT_IMAGE_FILE}"
info "---------------------------------------------------------------------"
info "Raspberry Pi image created successfully!"
info "Output image: ${FINAL_IMAGE_PATH}"
info "Device Type: ${DEVICE_TYPE}"
info "Hostname: ${HOSTNAME_PI}"
info "Wi-Fi SSID: ${WIFI_SSID}"
info "Snapcast Server: ${SNAPCAST_SERVER}"
info ""
info "To write to an SD card (e.g., /dev/sdX - BE VERY CAREFUL):"
info "  sudo dd bs=4M if=${FINAL_IMAGE_PATH} of=/dev/sdX status=progress conv=fsync"
info "---------------------------------------------------------------------"

exit 0