#!/usr/bin/env bash
# Script to build a Debian 13 (Trixie) amd64 persistent USB image with X11, ratpoison,
# auto-login, and autostarting TreeSheets and Impala, from a host running Debian 11 arm64.
# 
# **IMPORTANT:** Review and adjust configuration variables below (image size, usernames, etc.)
# before running. Run this script as root on a Debian-based arm64 host.
# It will create an image file and set up a chroot with a Debian amd64 system.
# No changes are made to the host system aside from installing required packages (if needed).
# 
# The script will prompt for confirmation before overwriting any existing image file or writing to a USB device.
# All major actions are logged to the console for transparency.

set -euo pipefail

#############################
# Configuration Variables
#############################
DEBIAN_RELEASE="trixie"            # Target Debian release name
TARGET_ARCH="amd64"               # Target architecture for the USB system
IMAGE_SIZE="4G"                   # Size of the image file to create (e.g., "4G" for 4 GiB)
IMAGE_NAME="debian13-${TARGET_ARCH}-usb.img"  # Name of the image file to create
BUILD_DIR="$PWD/debian_usb_build" # Working directory for mounts and temporary files (created if not exists)
USERNAME="user"                   # Default username for auto-login
USERPASSWORD="password"           # Password for the default user (change or leave as desired)
HOSTNAME="debian-usb"             # Hostname for the new system
TIMEZONE="Etc/UTC"                # Timezone for the new system (can be changed to user's timezone)
LOCALE="en_US.UTF-8"              # Locale to generate for the new system

# Package lists for installation in the target system:
# Base system and utilities
BASE_PACKAGES="systemd-sysv,systemd,locales,tzdata,dialog"  # core packages (some are normally included by debootstrap second stage)
# Desktop/X11 and user applications
X11_PACKAGES="xserver-xorg-core,xserver-xorg-video-fbdev,xserver-xorg-video-vesa,xinit,xterm,ratpoison"
APP_PACKAGES="treesheets"         # TreeSheets is in Debian repo. Impala might need manual installation (not in Debian).
NETWORK_PACKAGES="iwd"            # Use iwd for Wi-Fi (Impala requires iwd). Alternatively, could include dhclient or systemd-networkd if needed.
# Bootloaders and kernel
BOOT_PACKAGES="grub-efi-amd64,linux-image-amd64,extlinux,syslinux-common" 
# Note: extlinux (Syslinux for ext filesystems) and syslinux-common provide BIOS boot support; grub-efi-amd64 for UEFI.

#############################
# Script Setup and Functions
#############################

# Ensure the working directory exists
mkdir -p "${BUILD_DIR}"
# Define mount points relative to working directory
ROOTFS_DIR="${BUILD_DIR}/rootfs"    # Mount point for the ext4 root filesystem
EFI_DIR="${BUILD_DIR}/efiboot"      # Mount point for the FAT32 EFI/boot partition

# Trap to cleanup mounts/loop on exit or error
cleanup() {
    echo "Cleaning up: Unmounting and detaching loop devices..."
    # Try to unmount in reverse order of mounting
    umount -lf "${ROOTFS_DIR}/boot/efi" 2>/dev/null || true
    umount -lf "${EFI_DIR}" 2>/dev/null || true
    umount -lf "${ROOTFS_DIR}/dev/pts" 2>/dev/null || true
    umount -lf "${ROOTFS_DIR}/dev" 2>/dev/null || true
    umount -lf "${ROOTFS_DIR}/proc" 2>/dev/null || true
    umount -lf "${ROOTFS_DIR}/sys" 2>/dev/null || true
    umount -lf "${ROOTFS_DIR}" 2>/dev/null || true
    if losetup -a | grep -q "$IMAGE_NAME"; then
        # Detach all loop devices associated with our image
        LOOP_DEV="$(losetup -j "${BUILD_DIR}/${IMAGE_NAME}" | cut -d: -f1 || true)"
        if [ -n "${LOOP_DEV}" ]; then
            losetup -d "${LOOP_DEV}" 2>/dev/null || true
        fi
    fi
}
trap cleanup EXIT

# Confirm function for dangerous operations
confirm() {
    local prompt="$1"
    read -r -p "$prompt " response
    case "$response" in
        [Yy][Ee][Ss]|[Yy]) true ;;
        *) false ;;
    esac
}

# Require root privileges
if [[ $EUID -ne 0 ]]; then
    echo "This script must be run as root. Exiting."
    exit 1
fi

# Check and install required host tools
REQUIRED_TOOLS=(debootstrap qemu-user-static parted losetup mkfs.vfat mkfs.ext4)
MISSING_TOOLS=()
for tool in "${REQUIRED_TOOLS[@]}"; do
    if ! command -v "$tool" &>/dev/null; then
        MISSING_TOOLS+=("$tool")
    fi
done
if (( ${#MISSING_TOOLS[@]} > 0 )); then
    echo "Installing missing host tools: ${MISSING_TOOLS[*]}..."
    apt-get update && apt-get install -y "${MISSING_TOOLS[@]}"
fi

#############################
# Create and Partition Image
#############################
IMAGE_PATH="${BUILD_DIR}/${IMAGE_NAME}"
if [[ -f "$IMAGE_PATH" ]]; then
    echo "Image file $IMAGE_PATH already exists."
    if ! confirm "Overwrite existing image file? [yes/NO]"; then
        echo "Aborting to avoid overwriting existing image."
        exit 1
    fi
    rm -f "$IMAGE_PATH"
fi

echo ">>> Creating blank image file of size $IMAGE_SIZE at $IMAGE_PATH..."
# Create an empty file of the specified size
truncate -s "$IMAGE_SIZE" "$IMAGE_PATH"

echo ">>> Partitioning image file..."
# Create partition table: Partition 1 = FAT32 (for EFI & Syslinux), Partition 2 = ext4 (root)
parted -s "$IMAGE_PATH" mklabel msdos \
    mkpart primary fat32 1MiB 300MiB \
    mkpart primary ext4 300MiB 100% \
    set 1 boot on

# Set up loop device with partitions
LOOP_DEV="$(losetup -f --show -P "$IMAGE_PATH")"
echo "    Loop device $LOOP_DEV created for $IMAGE_PATH"
# The loop device now has partitions accessible as ${LOOP_DEV}p1 and p2
LOOP_P1="${LOOP_DEV}p1"
LOOP_P2="${LOOP_DEV}p2"

echo ">>> Creating filesystems..."
# Format the EFI/FAT32 partition (partition 1)
mkfs.vfat -F 32 -n EFI "$LOOP_P1"
# Format the root ext4 partition (partition 2) with a label
mkfs.ext4 -L rootfs "$LOOP_P2"

# Create mount points
mkdir -p "$ROOTFS_DIR" "$EFI_DIR"

echo ">>> Mounting image partitions..."
mount "$LOOP_P2" "$ROOTFS_DIR"
mkdir -p "${ROOTFS_DIR}/boot/efi"
mount "$LOOP_P1" "$ROOTFS_DIR/boot/efi"
# Also mount the EFI partition separately if needed (not strictly necessary since it's at rootfs/boot/efi)
mount "$LOOP_P1" "$EFI_DIR"

#############################
# Debootstrap: Base System
#############################
echo ">>> Bootstrapping Debian $DEBIAN_RELEASE ($TARGET_ARCH)..."
# First stage debootstrap (download and extract base system)
debootstrap --arch="$TARGET_ARCH" --foreign "$DEBIAN_RELEASE" "$ROOTFS_DIR" http://deb.debian.org/debian

# Enable QEMU for chroot (copy qemu static binary into the new system)
echo ">>> Copying QEMU static binary for $TARGET_ARCH into chroot..."
cp "$(which qemu-${TARGET_ARCH}-static)" "${ROOTFS_DIR}/usr/bin/"

# Prepare essential mount points for chroot environment
echo ">>> Mounting special filesystems for chroot..."
mount -t proc proc "${ROOTFS_DIR}/proc"
mount -t sysfs sys "${ROOTFS_DIR}/sys"
mount -o bind /dev "${ROOTFS_DIR}/dev"
mount -o bind /dev/pts "${ROOTFS_DIR}/dev/pts"
# Use a tmpfs for /run inside chroot to avoid interference with host /run
mount -t tmpfs tmpfs "${ROOTFS_DIR}/run"
mkdir -p "${ROOTFS_DIR}/run/lock"  # for any lock files

# Second stage debootstrap (configure base system inside chroot)
echo ">>> Running debootstrap second-stage in chroot..."
chroot "$ROOTFS_DIR" /debootstrap/debootstrap --second-stage

# Basic system configuration: hostname, hosts, timezone, locale
echo ">>> Configuring base system (hostname, timezone, locale)..."
echo "$HOSTNAME" > "${ROOTFS_DIR}/etc/hostname"
# Set up /etc/hosts with minimal entries
cat > "${ROOTFS_DIR}/etc/hosts" <<EOF
127.0.0.1   localhost
127.0.1.1   ${HOSTNAME}
EOF

# Timezone
echo "$TIMEZONE" > "${ROOTFS_DIR}/etc/timezone"
ln -sf "/usr/share/zoneinfo/$TIMEZONE" "${ROOTFS_DIR}/etc/localtime"

# Locale (generate specified locale)
chroot "$ROOTFS_DIR" bash -c "echo '$LOCALE UTF-8' > /etc/locale.gen"
chroot "$ROOTFS_DIR" locale-gen

# Set default LANG
echo "LANG=$LOCALE" > "${ROOTFS_DIR}/etc/default/locale"

#############################
# Install Packages in Chroot
#############################
echo ">>> Installing required packages in the target system..."
# Configure apt sources (use default deb.debian.org for stable)
cat > "${ROOTFS_DIR}/etc/apt/sources.list" <<EOF
deb http://deb.debian.org/debian $DEBIAN_RELEASE main contrib non-free non-free-firmware
deb http://deb.debian.org/debian $DEBIAN_RELEASE-updates main contrib non-free non-free-firmware
deb http://security.debian.org/debian-security $DEBIAN_RELEASE-security main contrib non-free non-free-firmware
EOF

# Update apt cache and install packages
chroot "$ROOTFS_DIR" apt-get update
# Use apt-get in one command to install all desired packages
chroot "$ROOTFS_DIR" apt-get install -y --no-install-recommends \
    $BASE_PACKAGES,$X11_PACKAGES,$APP_PACKAGES,$NETWORK_PACKAGES,$BOOT_PACKAGES

# Set the system's timezone and reconfigure tzdata (non-interactively)
chroot "$ROOTFS_DIR" bash -c "DEBIAN_FRONTEND=noninteractive dpkg-reconfigure tzdata"

#############################
# User Setup and Autologin
#############################
echo ">>> Setting up default user and auto-login..."
# Create the user with home directory and add to groups (sudo,netdev,audio,video)
chroot "$ROOTFS_DIR" useradd -m -s /bin/bash "$USERNAME"
chroot "$ROOTFS_DIR" bash -c "echo '${USERNAME}:${USERPASSWORD}' | chpasswd"
# Set root password (optional: here we set same as user, or leave locked by not setting)
chroot "$ROOTFS_DIR" bash -c "echo 'root:${USERPASSWORD}' | chpasswd"
# Add user to necessary groups
chroot "$ROOTFS_DIR" usermod -aG sudo,netdev,audio,video "$USERNAME"

# Configure autologin on tty1 via systemd getty override
AUTOLOGIN_CONF="${ROOTFS_DIR}/etc/systemd/system/getty@tty1.service.d"
mkdir -p "$AUTOLOGIN_CONF"
cat > "$AUTOLOGIN_CONF/autologin.conf" <<EOF
[Service]
ExecStart=
ExecStart=-/sbin/agetty --autologin $USERNAME --noclear %I 38400 linux
EOF

# Set up startx on login (for console auto-login sessions)
# Add a command to .bash_profile to launch X only for the autologin on tty1
USER_HOME="${ROOTFS_DIR}/home/${USERNAME}"
cat >> "${USER_HOME}/.bash_profile" <<'EOF'
# If logging in on tty1, start X automatically
if [[ -z $DISPLAY && $(tty) == "/dev/tty1" ]]; then
    startx -- -nocursor
    logout
fi
EOF
chroot "$ROOTFS_DIR" chown "$USERNAME:$USERNAME" "/home/$USERNAME/.bash_profile"

#############################
# X11 Autostart (ratpoison + apps)
#############################
echo ">>> Configuring X11 session (ratpoison) and application autostart..."
# Create an .xinitrc for the user to start Ratpoison
cat > "${USER_HOME}/.xinitrc" <<'EOF'
#!/bin/bash
# .xinitrc: run Ratpoison window manager
exec ratpoison
EOF
chroot "$ROOTFS_DIR" chown "$USERNAME:$USERNAME" "/home/$USERNAME/.xinitrc"
chroot "$ROOTFS_DIR" chmod +x "/home/$USERNAME/.xinitrc"

# Configure Ratpoison autostart: .ratpoisonrc to launch TreeSheets and Impala at startup
cat > "${USER_HOME}/.ratpoisonrc" <<EOF
# Disable startup message
startup_message off
# Set a blank cursor (useful if -nocursor used for X)
exec xsetroot -cursor_name left_ptr
# Autostart applications:
exec treesheets      # launch TreeSheets GUI on start
exec xterm -e impala # open an xterm and run impala TUI inside it
EOF
# Note: impala is not an official Debian package. Ensure the 'impala' binary is installed in the system or adjust this line.
chroot "$ROOTFS_DIR" chown "$USERNAME:$USERNAME" "/home/$USERNAME/.ratpoisonrc"

#############################
# Bootloader Setup (Syslinux & GRUB)
#############################
echo ">>> Installing and configuring bootloaders (Syslinux for BIOS, GRUB for UEFI)..."
# 1. EXTLINUX (Syslinux) for BIOS boot:
# Create extlinux directory
chroot "$ROOTFS_DIR" mkdir -p /boot/extlinux
# Copy Syslinux BIOS modules to /boot/extlinux
chroot "$ROOTFS_DIR" cp -r /usr/lib/syslinux/modules/bios/* /boot/extlinux/ 2>/dev/null || true
# Install extlinux bootloader on the ext4 partition
chroot "$ROOTFS_DIR" extlinux --install /boot/extlinux

# Create extlinux configuration file
ROOT_UUID=$(blkid -s UUID -o value "$LOOP_P2")  # get UUID of root partition
cat > "${ROOTFS_DIR}/boot/extlinux/extlinux.conf" <<EOF
DEFAULT linux
LABEL linux
    LINUX ../vmlinuz
    INITRD ../initrd.img
    APPEND root=UUID=${ROOT_UUID} ro quiet
EOF
# The vmlinuz and initrd.img symlinks point to the latest kernel and initrd in /boot.
# We use '../' because extlinux directory is /boot/extlinux, going up to /boot for the files.

# Ensure the extlinux config and modules are owned by root (should already be)
chroot "$ROOTFS_DIR" chown -R root:root /boot/extlinux

# 2. GRUB for UEFI boot:
# Install GRUB EFI (this was already installed via apt in BOOT_PACKAGES)
# Perform grub-install targeting x86_64 EFI, pointing to the mounted EFI partition.
chroot "$ROOTFS_DIR" grub-install --target=x86_64-efi --efi-directory=/boot/efi --bootloader-id=Debian --removable --no-nvram
# Generate GRUB configuration file
chroot "$ROOTFS_DIR" update-grub

#############################
# Finalization
#############################
# Remove the QEMU static binary from the target system (not needed on real x86_64 hardware)
rm -f "${ROOTFS_DIR}/usr/bin/qemu-${TARGET_ARCH}-static"

# Sync data to disk
sync

echo ">>> Unmounting chroot filesystems..."
# These will be also handled by the trap on exit, but we unmount explicitly here for clarity:
umount "${ROOTFS_DIR}/proc" || true
umount "${ROOTFS_DIR}/sys" || true
umount "${ROOTFS_DIR}/dev/pts" || true
umount "${ROOTFS_DIR}/dev" || true
umount "${ROOTFS_DIR}/run" || true
umount "${ROOTFS_DIR}/boot/efi" || true
umount "${EFI_DIR}" || true
umount "${ROOTFS_DIR}" || true

# Write Syslinux MBR boot code to the image (for BIOS boot).
# Use dd to write the first 440 bytes from Syslinux's mbr.bin to the image's MBR.
echo ">>> Writing Syslinux MBR boot code to image..."
dd if="${ROOTFS_DIR}/usr/lib/SYSLINUX/mbr.bin" of="$LOOP_DEV" bs=440 count=1 conv=notrunc

# Detach loop device
losetup -d "$LOOP_DEV"

echo ">>> Debian USB image creation completed successfully!"
echo "Image file: $IMAGE_PATH"
# Offer to write image to a USB device
if confirm "Write the image to a USB drive now? (This will destroy contents on the target drive) [yes/NO]"; then
    read -rp "Enter the device path for the USB (e.g., /dev/sdX): " USBDEV
    if [[ -n "$USBDEV" ]]; then
        echo "WARNING: About to overwrite $USBDEV with the image. This will erase all data on $USBDEV."
        if confirm "Are you absolutely sure? Type 'yes' to continue: "; then
            echo ">>> Writing image to $USBDEV ... (this may take a while)"
            dd if="$IMAGE_PATH" of="$USBDEV" bs=4M status=progress conv=fsync
            echo ">>> Syncing data to $USBDEV..."
            sync
            echo "Image written to $USBDEV successfully. You can now boot the USB drive on an x86_64 machine."
        else
            echo "Skipped writing image to USB."
        fi
    fi
else
    echo "Image creation complete. You can write $IMAGE_PATH to a USB device later using 'dd' or a similar tool."
fi