#!/bin/bash
# build.sh — Build RPi Audio Mixer SD card image from scratch
#
# Produces a ready-to-flash Raspberry Pi 4B SD card image with:
#   - RPiOS Lite (64-bit, Bookworm) base
#   - Optional PREEMPT_RT kernel (cross-compiled on x86)
#   - JACK2, Carla, ALSA tools, Python deps
#   - Full mixer engine + web UI + Kivy touch UI
#   - First-boot setup wizard
#   - systemd services for auto-start
#
# Usage:
#   ./build.sh                          # Full build (default)
#   ./build.sh --skip-kernel            # Skip RT kernel build (use stock kernel)
#   ./build.sh --kernel-only            # Only build RT kernel, skip image
#   ./build.sh --output ./my-image.img  # Custom output path
#   ./build.sh --clean                  # Remove build artifacts
#
# Environment:
#   MIXER_VERSION      version tag (default: git describe or "dev")
#   MIXER_API_KEY      default API key (default: auto-generated)
#   MIXER_HOSTNAME     default hostname (default: pi-mixer)
#   MIXER_WIFI_SSID    pre-configure WiFi SSID (optional)
#   MIXER_WIFI_PASS    pre-configure WiFi password (optional)
#
# Requirements:
#   - Linux host (x86_64 or aarch64)
#   - sudo access
#   - ~10 GB free disk space
#   - Packages: qemu-user-static, wget, xz-utils, e2fsprogs, dosfstools
#   - Cross-compile (optional): gcc-aarch64-linux-gnu, libssl-dev:arm64
#
# Output:
#   build/out/rpi-audio-mixer-<version>.img     SD card image
#   build/out/rpi-audio-mixer-<version>.img.xz  Compressed image

set -euo pipefail

# ─── Configuration ───────────────────────────────────────────────────────────

SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
PROJECT_ROOT="$(dirname "$SCRIPT_DIR")"
BUILD_DIR="$SCRIPT_DIR"
OUT_DIR="$BUILD_DIR/out"
WORK_DIR="$BUILD_DIR/work"
DOWNLOAD_DIR="$BUILD_DIR/downloads"

MIXER_VERSION="${MIXER_VERSION:-$(git -C "$PROJECT_ROOT" describe --tags --always 2>/dev/null || echo 'dev')}"
MIXER_API_KEY="${MIXER_API_KEY:-mixer-$(head -c 12 /dev/urandom | base64 | tr -dc 'a-zA-Z0-9' | head -c 16)}"
MIXER_HOSTNAME="${MIXER_HOSTNAME:-pi-mixer}"
MIXER_WIFI_SSID="${MIXER_WIFI_SSID:-}"
MIXER_WIFI_PASS="${MIXER_WIFI_PASS:-}"

# RPiOS image
RPIOS_URL="${RPIOS_URL:-https://downloads.raspberrypi.com/raspios_lite_arm64/images/raspios_lite_arm64-2025-03-07/2025-03-07-raspios-bookworm-arm64-lite.img.xz}"
RPIOS_SHA256=""

# RPi kernel for RT build
RPI_KERNEL_REPO="${RPI_KERNEL_REPO:-https://github.com/raspberrypi/linux.git}"
RPI_KERNEL_BRANCH="${RPI_KERNEL_BRANCH:-rpi-6.12.y}"

# Image size (expandable on first boot)
IMAGE_SIZE_MB="${IMAGE_SIZE_MB:-8192}"        # 8 GB base image
BOOT_SIZE_MB="${BOOT_SIZE_MB:-256}"           # 256 MB FAT32 boot

# Build flags
CROSS_COMPILE="${CROSS_COMPILE:-aarch64-linux-gnu-}"
KERNEL_JOBS="${KERNEL_JOBS:-$(nproc 2>/dev/null || echo 4)}"
BUILD_KERNEL="${BUILD_KERNEL:-true}"

# QEMU static for chroot
QEMU_STATIC="/usr/bin/qemu-aarch64-static"

# ─── Colours ─────────────────────────────────────────────────────────────────

RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
BLUE='\033[0;34m'
BOLD='\033[1m'
NC='\033[0m'

log()  { echo -e "${GREEN}[BUILD]${NC} $*"; }
warn() { echo -e "${YELLOW}[WARN]${NC}  $*"; }
err()  { echo -e "${RED}[ERROR]${NC} $*"; }
step() { echo -e "\n${BLUE}${BOLD}═══ $* ═══${NC}\n"; }

# ─── Usage ───────────────────────────────────────────────────────────────────

usage() {
    cat <<EOF
${BOLD}Usage:${NC} $0 [OPTIONS]

Build a ready-to-flash Raspberry Pi 4B SD card image for the RPi Audio Mixer.

${BOLD}Options:${NC}
  --skip-kernel       Skip PREEMPT_RT kernel build (use stock RPiOS kernel)
  --kernel-only       Only build the RT kernel, skip image creation
  --output PATH       Output image path (default: build/out/rpi-audio-mixer-<version>.img)
  --clean             Remove build artifacts (work + downloads + out)
  --image-size MB     Image size in MB (default: 8192)
  --version VERSION   Version tag for the image
  --help              Show this help

${BOLD}Examples:${NC}
  $0                                    # Full build with RT kernel
  $0 --skip-kernel                      # Quick build with stock kernel
  $0 --kernel-only                      # Only cross-compile the kernel
  $0 --output ./pi-mixer-v2.img         # Custom output path
  $0 --clean                            # Clean build artifacts

${BOLD}Environment:${NC}
  MIXER_VERSION       Version tag (default: git describe or "dev")
  MIXER_API_KEY       Default API key (auto-generated if unset)
  MIXER_HOSTNAME      Default hostname (default: pi-mixer)
  MIXER_WIFI_SSID     Pre-configure WiFi SSID (optional)
  MIXER_WIFI_PASS     Pre-configure WiFi password (optional)

${BOLD}Output:${NC}
  build/out/rpi-audio-mixer-<version>.img       Raw SD card image
  build/out/rpi-audio-mixer-<version>.img.xz    Compressed image
EOF
    exit 0
}

# ─── Argument parsing ────────────────────────────────────────────────────────

KERNEL_ONLY=false
OUTPUT_PATH=""
CLEAN=false

while [[ $# -gt 0 ]]; do
    case "$1" in
        --skip-kernel)   BUILD_KERNEL=false; shift ;;
        --kernel-only)   KERNEL_ONLY=true; shift ;;
        --output)        OUTPUT_PATH="$2"; shift 2 ;;
        --image-size)    IMAGE_SIZE_MB="$2"; shift 2 ;;
        --version)       MIXER_VERSION="$2"; shift 2 ;;
        --clean)         CLEAN=true; shift ;;
        --help)          usage ;;
        *)               err "Unknown option: $1"; usage ;;
    esac
done

if [[ -z "$OUTPUT_PATH" ]]; then
    OUTPUT_PATH="$OUT_DIR/rpi-audio-mixer-${MIXER_VERSION}.img"
fi

# ─── Clean ───────────────────────────────────────────────────────────────────

if $CLEAN; then
    step "Cleaning build artifacts"
    rm -rf "$WORK_DIR" "$DOWNLOAD_DIR" "$OUT_DIR"
    log "Build artifacts cleaned: $WORK_DIR $DOWNLOAD_DIR $OUT_DIR"
    exit 0
fi

# ─── Prerequisites ───────────────────────────────────────────────────────────

check_deps() {
    local missing=()

    command -v wget >/dev/null 2>&1 || missing+=(wget)
    command -v xz >/dev/null 2>&1 || missing+=(xz-utils)
    command -v mkfs.vfat >/dev/null 2>&1 || missing+=(dosfstools)
    command -v mkfs.ext4 >/dev/null 2>&1 || missing+=(e2fsprogs)

    if [[ ! -f "$QEMU_STATIC" ]]; then
        missing+=("qemu-user-static (sudo apt install qemu-user-static)")
    fi

    if $BUILD_KERNEL; then
        command -v "${CROSS_COMPILE}gcc" >/dev/null 2>&1 || missing+=("gcc-aarch64-linux-gnu (sudo apt install gcc-aarch64-linux-gnu)")
    fi

    if [[ ${#missing[@]} -gt 0 ]]; then
        err "Missing dependencies:"
        for m in "${missing[@]}"; do
            echo "  - $m"
        done
        exit 1
    fi
}

# ─── Directory setup ─────────────────────────────────────────────────────────

setup_dirs() {
    mkdir -p "$OUT_DIR" "$WORK_DIR" "$DOWNLOAD_DIR"
    # Pre-create the kernel build workdir
    if $BUILD_KERNEL; then
        mkdir -p "$WORK_DIR/kernel"
    fi
}

# ─── Step 1: Download RPiOS Lite image ────────────────────────────────────────

download_base_image() {
    local img_xz="$DOWNLOAD_DIR/$(basename "$RPIOS_URL")"
    local img="${img_xz%.xz}"

    if [[ -f "$img" ]]; then
        log "Base image already downloaded: $img"
        echo "$img"
        return
    fi

    step "Step 1/7: Downloading RPiOS Lite (64-bit, Bookworm)"
    log "URL: $RPIOS_URL"
    wget -nc -O "$img_xz" "$RPIOS_URL" || {
        err "Failed to download RPiOS base image. Check your internet connection."
        err "You can also download manually and place it at: $img_xz"
        exit 1
    }

    log "Decompressing..."
    xz -d -k "$img_xz"  # -k keeps the .xz

    echo "$img"
}

# ─── Step 1.5 (optional): Build PREEMPT_RT kernel ────────────────────────────

build_rt_kernel() {
    local kernel_dir="$WORK_DIR/kernel/linux-rpi"

    step "Step 2/7: Building PREEMPT_RT kernel (cross-compiled on $(uname -m))"

    # Clone if not already present
    if [[ ! -d "$kernel_dir" ]]; then
        log "Cloning Raspberry Pi kernel (branch: $RPI_KERNEL_BRANCH)..."
        git clone --depth=1 -b "$RPI_KERNEL_BRANCH" "$RPI_KERNEL_REPO" "$kernel_dir"
    else
        log "Kernel source already present at $kernel_dir"
    fi

    cd "$kernel_dir"

    # Apply PREEMPT_RT config
    if [[ ! -f ".config" ]] || ! grep -q "CONFIG_PREEMPT_RT=y" .config 2>/dev/null; then
        log "Configuring kernel with PREEMPT_RT..."
        make ARCH=arm64 CROSS_COMPILE="$CROSS_COMPILE" bcm2711_defconfig

        # Enable RT configs
        scripts/config -e CONFIG_PREEMPT_RT
        scripts/config -d CONFIG_PREEMPT
        scripts/config -e CONFIG_HIGH_RES_TIMERS
        scripts/config -e CONFIG_HZ_1000
        scripts/config -e CONFIG_NO_HZ_FULL
        scripts/config -e CONFIG_RCU_NOCB_CPU
        scripts/config -e CONFIG_CPU_FREQ_DEFAULT_GOV_PERFORMANCE
        scripts/config -e CONFIG_CPU_FREQ_GOV_PERFORMANCE
        scripts/config -e CONFIG_IRQ_FORCED_THREADING
        scripts/config --set-val CONFIG_HZ 1000

        # USB audio (build as module)
        scripts/config -m CONFIG_SND_USB_AUDIO

        make ARCH=arm64 CROSS_COMPILE="$CROSS_COMPILE" olddefconfig
    else
        log "Kernel .config already has PREEMPT_RT enabled"
    fi

    log "Building kernel (jobs: $KERNEL_JOBS)..."
    make -j "$KERNEL_JOBS" ARCH=arm64 CROSS_COMPILE="$CROSS_COMPILE" \
        Image modules dtbs 2>&1 | tail -20

    log "Building modules..."
    make -j "$KERNEL_JOBS" ARCH=arm64 CROSS_COMPILE="$CROSS_COMPILE" modules

    # Install modules to staging
    local mod_staging="$WORK_DIR/kernel/modules"
    rm -rf "$mod_staging"
    mkdir -p "$mod_staging"
    make ARCH=arm64 CROSS_COMPILE="$CROSS_COMPILE" \
        INSTALL_MOD_PATH="$mod_staging" modules_install

    log "Kernel build complete"
    log "  Image:  $kernel_dir/arch/arm64/boot/Image"
    log "  DTBs:   $kernel_dir/arch/arm64/boot/dts/broadcom/*.dtb"
    log "  Modules: $mod_staging/lib/modules/"

    cd "$PROJECT_ROOT"
}

# ─── Step 2: Create image file ───────────────────────────────────────────────

create_image() {
    local img_file="$1"

    step "Step 3/7: Creating SD card image (${IMAGE_SIZE_MB}MB)"

    rm -f "$img_file"
    dd if=/dev/zero of="$img_file" bs=1M count="$IMAGE_SIZE_MB" status=progress 2>/dev/null || \
        dd if=/dev/zero of="$img_file" bs=1M count="$IMAGE_SIZE_MB"

    # Partition table (MBR)
    # Part 1: FAT32 boot, 256MB
    # Part 2: ext4 root, 6GB
    # Part 3: ext4 /data, remaining
    local root_size_mb=6144
    local data_start_mb=$((BOOT_SIZE_MB + root_size_mb))

    log "Partitioning image..."
    cat <<FDISK | fdisk "$img_file" >/dev/null 2>&1
o
n
p
1

+${BOOT_SIZE_MB}M
t
c
n
p
2

+${root_size_mb}M
n
p
3


w
FDISK

    # Mount via loopback
    local loopdev
    loopdev=$(sudo losetup -f --show -P "$img_file")
    log "Loopback device: $loopdev"

    # Format partitions
    log "Formatting partitions..."
    sudo mkfs.vfat -F 32 -n BOOT "${loopdev}p1" >/dev/null 2>&1
    sudo mkfs.ext4 -F -L rootfs "${loopdev}p2" >/dev/null 2>&1
    sudo mkfs.ext4 -F -L data "${loopdev}p3" >/dev/null 2>&1

    echo "$loopdev"
}

# ─── Step 3: Mount image partitions ──────────────────────────────────────────

mount_image() {
    local loopdev="$1"
    local mount_point="$2"

    mkdir -p "$mount_point"/{boot,data}

    sudo mount "${loopdev}p2" "$mount_point"
    sudo mount "${loopdev}p1" "$mount_point/boot"
    sudo mount "${loopdev}p3" "$mount_point/data"

    log "Mounted image at $mount_point"
}

unmount_image() {
    local mount_point="$1"
    local loopdev="$2"

    log "Unmounting image..."
    sudo umount -R "$mount_point" 2>/dev/null || true
    sudo losetup -d "$loopdev" 2>/dev/null || true
}

# ─── Step 4: Populate root filesystem ────────────────────────────────────────

populate_rootfs() {
    local base_img="$1"
    local mount_point="$2"

    step "Step 4/7: Populating root filesystem"

    # Mount base image to copy from it
    local base_mount="$WORK_DIR/base-mount"
    mkdir -p "$base_mount"
    local base_loop
    base_loop=$(sudo losetup -f --show -P "$base_img")
    sudo mount "${base_loop}p2" "$base_mount" 2>/dev/null || \
        sudo mount "${base_loop}" "$base_mount"  # maybe not partitioned

    log "Copying base OS files... (this takes a few minutes)"
    sudo rsync -a --info=progress2 "$base_mount/" "$mount_point/" 2>/dev/null || \
        sudo rsync -a "$base_mount/" "$mount_point/"

    sudo umount "$base_mount" 2>/dev/null || true
    sudo losetup -d "$base_loop" 2>/dev/null || true
    rmdir "$base_mount"

    # Set up first-boot expand script
    log "Setting up first-boot root expansion..."
    sudo mkdir -p "$mount_point/etc/init.d"
}

# ─── Step 5: Configure the system (chroot) ───────────────────────────────────

configure_system() {
    local mount_point="$1"

    step "Step 5/7: Configuring system (chroot)"

    # Copy QEMU for chroot
    if [[ "$(uname -m)" != "aarch64" ]]; then
        sudo cp "$QEMU_STATIC" "$mount_point/usr/bin/"
    fi

    # Prepare chroot environment
    sudo mount -t proc proc "$mount_point/proc" || true
    sudo mount -t sysfs sys "$mount_point/sys" || true
    sudo mount -t devtmpfs dev "$mount_point/dev" || true
    sudo mount -t devpts devpts "$mount_point/dev/pts" || true
    # /run may already be tmpfs under chroot
    sudo mount -t tmpfs tmpfs "$mount_point/run" 2>/dev/null || true

    # Copy the configuration script into the image
    local config_script="$mount_point/tmp/configure-system.sh"
    sudo cp "$BUILD_DIR/configure-system.sh" "$config_script"
    sudo chmod +x "$config_script"

    # Copy project source
    log "Copying project source into image..."
    sudo mkdir -p "$mount_point/opt/rpi-mixer"
    sudo rsync -a --exclude='.git' --exclude='build/work' --exclude='build/downloads' \
        --exclude='build/out' --exclude='__pycache__' --exclude='*.pyc' \
        --exclude='.pytest_cache' --exclude='venv' --exclude='.venv' \
        "$PROJECT_ROOT/" "$mount_point/opt/rpi-mixer/"

    # Run configuration inside chroot
    log "Running system configuration inside chroot..."
    sudo chroot "$mount_point" /bin/bash /tmp/configure-system.sh || {
        warn "Chroot configuration had non-zero exit. Continuing..."
    }

    # Clean up
    sudo rm -f "$mount_point/tmp/configure-system.sh"
    sudo rm -f "$mount_point/usr/bin/qemu-aarch64-static"

    # Unmount chroot mounts
    sudo umount "$mount_point/proc" 2>/dev/null || true
    sudo umount "$mount_point/sys" 2>/dev/null || true
    sudo umount "$mount_point/dev/pts" 2>/dev/null || true
    sudo umount "$mount_point/dev" 2>/dev/null || true
    sudo umount "$mount_point/run" 2>/dev/null || true
}

# ─── Step 6: Install kernel (if RT built) ────────────────────────────────────

install_rt_kernel() {
    local mount_point="$1"

    step "Step 6/7: Installing PREEMPT_RT kernel"

    local kernel_dir="$WORK_DIR/kernel/linux-rpi"
    local mod_staging="$WORK_DIR/kernel/modules"

    # Copy kernel image
    log "Installing kernel image..."
    sudo cp "$kernel_dir/arch/arm64/boot/Image" \
        "$mount_point/boot/kernel8-rt.img"

    # Copy DTBs
    log "Installing Device Tree Blobs..."
    sudo cp "$kernel_dir/arch/arm64/boot/dts/broadcom/"*.dtb \
        "$mount_point/boot/"

    # Copy overlays
    if [[ -d "$kernel_dir/arch/arm64/boot/dts/overlays" ]]; then
        sudo mkdir -p "$mount_point/boot/overlays"
        sudo cp "$kernel_dir/arch/arm64/boot/dts/overlays/"*.dtbo \
            "$mount_point/boot/overlays/" 2>/dev/null || true
    fi

    # Install kernel modules
    log "Installing kernel modules..."
    local modules_ver
    modules_ver=$(ls "$mod_staging/lib/modules/" | head -1)
    sudo cp -r "$mod_staging/lib/modules/$modules_ver" \
        "$mount_point/lib/modules/"

    # Update config.txt to use RT kernel
    log "Updating /boot/config.txt..."
    sudo tee -a "$mount_point/boot/config.txt" >/dev/null <<'EOF'

# PREEMPT_RT kernel
kernel=kernel8-rt.img

EOF

    # Update cmdline.txt for RT
    log "Updating /boot/cmdline.txt..."
    local existing_cmdline
    existing_cmdline=$(sudo cat "$mount_point/boot/cmdline.txt" 2>/dev/null || echo "")
    # Remove console=serial0 if present (quieter boot)
    existing_cmdline="${existing_cmdline//console=serial0,115200/}"
    sudo tee "$mount_point/boot/cmdline.txt" >/dev/null <<CMDLINE
${existing_cmdline} quiet loglevel=3 threadirqs nohz_full=1-3 rcu_nocbs=1-3 isolcpus=1-3 usbcore.autosuspend=-1 fsck.mode=skip
CMDLINE
}

# ─── Step 7: Finalize and compress ───────────────────────────────────────────

finalize_image() {
    local mount_point="$1"
    local loopdev="$2"
    local img_file="$3"

    step "Step 7/7: Finalizing image"

    # Unmount
    unmount_image "$mount_point" "$loopdev"

    # Compress
    log "Compressing image..."
    xz -3 -f -T0 -v "$img_file"
    log "Compressed image: ${img_file}.xz"

    # Report sizes
    local raw_size compressed_size
    raw_size=$(du -h "$img_file" | cut -f1)
    compressed_size=$(du -h "${img_file}.xz" | cut -f1)

    echo ""
    echo -e "${GREEN}${BOLD}╔══════════════════════════════════════════╗${NC}"
    echo -e "${GREEN}${BOLD}║   Build Complete                          ║${NC}"
    echo -e "${GREEN}${BOLD}╚══════════════════════════════════════════╝${NC}"
    echo ""
    echo "  Version:  $MIXER_VERSION"
    echo "  Image:    $img_file ($raw_size)"
    echo "  Compressed: ${img_file}.xz ($compressed_size)"
    echo ""
    if $BUILD_KERNEL; then
        echo "  Kernel:   PREEMPT_RT (cross-compiled)"
    else
        echo "  Kernel:   Stock RPiOS (PREEMPT_RT not included)"
    fi
    echo ""
    echo -e "${BOLD}To flash to SD card:${NC}"
    echo "  sudo dd if=$img_file of=/dev/sdX bs=4M status=progress conv=fsync"
    echo ""
    echo -e "${BOLD}First boot:${NC}"
    echo "  - The setup wizard will guide you through configuration"
    echo "  - Default hostname: $MIXER_HOSTNAME"
    echo "  - Default API key:  $MIXER_API_KEY"
    echo "  - Web UI: http://<pi-ip>:8080"
    echo "  - SSH: pi@<pi-ip> (password: raspberry, CHANGE ON FIRST BOOT)"
    echo ""
}

# ─── Main ────────────────────────────────────────────────────────────────────

main() {
    echo -e "${BLUE}${BOLD}"
    echo "╔══════════════════════════════════════════╗"
    echo "║   RPi Audio Mixer — Image Builder        ║"
    echo "║   Version: $MIXER_VERSION"
    echo "╚══════════════════════════════════════════╝"
    echo -e "${NC}"

    check_deps
    setup_dirs

    # ── Kernel-only mode ──
    if $KERNEL_ONLY; then
        if $BUILD_KERNEL; then
            build_rt_kernel
        else
            err "--kernel-only requires kernel build (remove --skip-kernel)"
            exit 1
        fi
        log "Kernel build complete. Output in: $WORK_DIR/kernel/"
        log "  Image: $WORK_DIR/kernel/linux-rpi/arch/arm64/boot/Image"
        log "  Modules: $WORK_DIR/kernel/modules/lib/modules/"
        exit 0
    fi

    # ── Build RT kernel (if enabled) ──
    if $BUILD_KERNEL; then
        build_rt_kernel
    else
        log "Skipping PREEMPT_RT kernel build (--skip-kernel)"
    fi

    # ── Download base image ──
    local base_img
    base_img=$(download_base_image)

    # ── Create new image ──
    local loopdev
    loopdev=$(create_image "$OUTPUT_PATH")

    local mount_point="$WORK_DIR/mnt"
    mount_image "$loopdev" "$mount_point"

    # ── Populate from base ──
    populate_rootfs "$base_img" "$mount_point"

    # ── Configure system ──
    configure_system "$mount_point"

    # ── Install RT kernel if built ──
    if $BUILD_KERNEL; then
        install_rt_kernel "$mount_point"
    fi

    # ── Finalize ──
    finalize_image "$mount_point" "$loopdev" "$OUTPUT_PATH"
}

main "$@"