#!/usr/bin/env bash
# ──────────────────────────────────────────────────────────────────────────────
# Pi Multi-FX Pedal — Minimal SD Card Image Builder
#
# Takes a stripped DietPi installation (after dietpi-postinstall.sh + minimize.sh)
# and produces a minimal, flashable SD card image (compressed .img.gz).
#
# This script runs on the build host (x86 or RPi) and produces an image that
# can be written directly to an SD card for deployment.
#
# Usage:
#   sudo ./build-minimal-image.sh [--source-device /dev/sdX] [--output ./pedal-image.img.gz]
#   sudo ./build-minimal-image.sh --from-scratch   # Download DietPi + build from zero
#
# Options:
#   --source-device DEV   Build image FROM a running SD card (e.g. /dev/mmcblk0)
#   --from-scratch        Download a fresh DietPi image and provision it
#   --output PATH         Output image path (default: ./pi-multifx-pedal-<date>.img.gz)
#   --no-compress         Skip gzip compression (saves time for testing)
#   --help                Show this help
#
# Dependencies:
#   - pishrink (auto-installed if missing)
#   - rsync, fdisk, parted, losetup (standard Linux tools)
#   - wget, unzip (for --from-scratch)
#
# Target: < 500 MB compressed image, < 60 MB idle RAM on boot
# ──────────────────────────────────────────────────────────────────────────────

set -euo pipefail

SCRIPT_DIR="$(cd "$(dirname "$0")" && pwd)"
PROJECT_DIR="$(dirname "$SCRIPT_DIR")"
OUTPUT_DIR="$(pwd)"
COMPRESS=true
SOURCE_DEVICE=""
FROM_SCRATCH=false

# ── Colour helpers ───────────────────────────────────────────────────────────
info()  { printf "\e[34m[INFO] %s\e[0m\n" "$*"; }
ok()    { printf "\e[32m[ OK ] %s\e[0m\n" "$*"; }
warn()  { printf "\e[33m[WARN] %s\e[0m\n" "$*"; }
err()   { printf "\e[31m[FAIL] %s\e[0m\n" "$*"; }

# ── Parse args ───────────────────────────────────────────────────────────────
while [[ $# -gt 0 ]]; do
    case "$1" in
        --source-device)  SOURCE_DEVICE="$2"; shift 2 ;;
        --output)         OUTPUT_DIR="$(dirname "$2")"; shift 2 ;;
        --from-scratch)   FROM_SCRATCH=true; shift ;;
        --no-compress)    COMPRESS=false; shift ;;
        --help)           grep "^#" "$0" | grep -v "^#!/" | sed 's/^# \?//'; exit 0 ;;
        *)                err "Unknown option: $1"; exit 1 ;;
    esac
done

if [[ $EUID -ne 0 ]]; then
    err "This script must be run as root (sudo ./build-minimal-image.sh)"
    exit 1
fi

# ── Determine output filename ────────────────────────────────────────────────
DATE_TAG=$(date +%Y%m%d)
OUTPUT_IMG="${OUTPUT_DIR}/pi-multifx-pedal-${DATE_TAG}.img"
OUTPUT_FINAL="${OUTPUT_IMG}${COMPRESS:+".gz"}"

info "========== Pi Multi-FX Pedal — Image Builder =========="
echo ""

# ══════════════════════════════════════════════════════════════════════════════
# Mode 1: Build from scratch (download DietPi + provision)
# ══════════════════════════════════════════════════════════════════════════════

if [[ "$FROM_SCRATCH" == true ]]; then
    info "Mode: from-scratch — downloading DietPi image..."

    WORKDIR="$(mktemp -d -t pedal-image-XXXXXX)"
    trap 'rm -rf "$WORKDIR"' EXIT

    # ── Download latest DietPi for RPi 4 ──────────────────────────────
    DIETPI_URL="https://dietpi.com/downloads/images/DietPi_RPi-ARMv8-Bookworm.img.xz"
    DIETPI_IMG="${WORKDIR}/DietPi_RPi-ARMv8-Bookworm.img"
    MOUNT_DIR="${WORKDIR}/mnt"
    BOOT_DIR="${WORKDIR}/boot"

    if [[ ! -f "$DIETPI_IMG" ]]; then
        info "Downloading DietPi for RPi 4B (ARMv8 Bookworm)..."
        wget -q --show-progress "$DIETPI_URL" -O "${DIETPI_IMG}.xz"
        info "Extracting..."
        unxz -f "${DIETPI_IMG}.xz"
        ok "DietPi image ready: $(ls -lh "$DIETPI_IMG" | awk '{print $5}')"
    fi

    # ── Expand the image to 4 GB (room for Python deps + NAM models) ──
    info "Expanding image to 4 GB..."
    ORIG_SIZE=$(stat -c%s "$DIETPI_IMG" 2>/dev/null || stat -f%z "$DIETPI_IMG" 2>/dev/null)
    dd if=/dev/zero bs=1M count=$((4096 - ORIG_SIZE / 1048576)) >> "$DIETPI_IMG" 2>/dev/null

    # Fix partition table — expand last partition to fill new space
    PART_START=$(fdisk -l "$DIETPI_IMG" 2>/dev/null | grep "^${DIETPI_IMG}" | tail -1 | awk '{print $2}')
    if [[ -z "$PART_START" ]]; then
        # Try parsed differently
        PART_INFO=$(fdisk -l "$DIETPI_IMG" 2>/dev/null | tail -2 | head -1)
        PART_START=$(echo "$PART_INFO" | awk '{print $2}')
    fi

    if [[ -n "$PART_START" ]]; then
        # fdisk: delete last partition, recreate with same start, larger end
        printf "d\nn\np\n\n${PART_START}\n\nw\n" | fdisk "$DIETPI_IMG" > /dev/null 2>&1 || true
        ok "Partition table expanded"
    fi

    # ── Mount image ──────────────────────────────────────────────────────
    mkdir -p "$MOUNT_DIR" "$BOOT_DIR"
    LOOP_DEV=$(losetup --find --show --partscan "$DIETPI_IMG")
    ok "Loop device: $LOOP_DEV"

    # Wait for partitions to appear
    sleep 2

    # Mount root partition (usually p2 for DietPi) and boot (p1)
    ROOT_PART="${LOOP_DEV}p2"
    BOOT_PART="${LOOP_DEV}p1"

    if [[ ! -b "$ROOT_PART" ]]; then
        ROOT_PART="${LOOP_DEV}p2"
        BOOT_PART="${LOOP_DEV}p1"
    fi

    mount "$ROOT_PART" "$MOUNT_DIR"
    mount "$BOOT_PART" "$BOOT_DIR"
    ok "Image mounted at $MOUNT_DIR"

    # ── Copy project scripts into the image ────────────────────────────
    info "Staging project files..."
    mkdir -p "$MOUNT_DIR/opt/pi-multifx-pedal/scripts" \
             "$MOUNT_DIR/opt/pi-multifx-pedal/src" \
             "$MOUNT_DIR/opt/pi-multifx-pedal/docs" \
             "$MOUNT_DIR/opt/pi-multifx-pedal/presets"

    # Use a bind mount so we access the project root cleanly
    # (copying via rsync preserves scripts for chroot execution)
    rsync -a --delete \
        --exclude='.git' \
        --exclude='__pycache__' \
        --exclude='*.pyc' \
        --exclude='.venv' \
        --exclude='venv' \
        --exclude='hardware/gerber' \
        "$PROJECT_DIR/" "$MOUNT_DIR/opt/pi-multifx-pedal/"

    ok "Project files staged"

    # ── Chroot and run postinstall ────────────────────────────────────
    info "Running postinstall inside chroot..."

    # Bind mount required filesystems for chroot
    mount --bind /dev "$MOUNT_DIR/dev"
    mount --bind /proc "$MOUNT_DIR/proc"
    mount --bind /sys "$MOUNT_DIR/sys"

    # Copy DNS config for apt inside chroot
    cp /etc/resolv.conf "$MOUNT_DIR/etc/resolv.conf" 2>/dev/null || true

    # Run postinstall
    chroot "$MOUNT_DIR" /bin/bash -c "
        cd /opt/pi-multifx-pedal
        bash scripts/dietpi-postinstall.sh --unattended
    " || warn "Postinstall exited with non-zero (some steps may need reboot inside chroot)"

    # Run minimize
    chroot "$MOUNT_DIR" /bin/bash -c "
        cd /opt/pi-multifx-pedal
        bash scripts/minimize.sh --aggressive
    " || true

    # ── Cleanup chroot ────────────────────────────────────────────────
    umount "$MOUNT_DIR/dev" 2>/dev/null || true
    umount "$MOUNT_DIR/proc" 2>/dev/null || true
    umount "$MOUNT_DIR/sys" 2>/dev/null || true
    rm -f "$MOUNT_DIR/etc/resolv.conf"

    # ── Clean apt caches inside image ─────────────────────────────────
    chroot "$MOUNT_DIR" apt-get clean -qq 2>/dev/null || true
    rm -rf "$MOUNT_DIR/var/cache/apt/archives/"* 2>/dev/null || true
    rm -rf "$MOUNT_DIR/var/log/"*.log 2>/dev/null || true
    rm -rf "$MOUNT_DIR/tmp/"* 2>/dev/null || true

    # ── Sync and unmount ──────────────────────────────────────────────
    sync
    umount "$MOUNT_DIR" 2>/dev/null || true
    umount "$BOOT_DIR" 2>/dev/null || true
    losetup -d "$LOOP_DEV" 2>/dev/null || true

    info "Image built from scratch: $DIETPI_IMG"
    cp "$DIETPI_IMG" "$OUTPUT_IMG"
    ok "Copied to $OUTPUT_IMG"

# ══════════════════════════════════════════════════════════════════════════════
# Mode 2: Build from a running SD card device (e.g. /dev/mmcblk0)
# ══════════════════════════════════════════════════════════════════════════════

elif [[ -n "$SOURCE_DEVICE" ]]; then
    info "Mode: from-device — imaging $SOURCE_DEVICE..."

    if [[ ! -b "$SOURCE_DEVICE" ]]; then
        err "Not a block device: $SOURCE_DEVICE"
        exit 1
    fi

    # ── Verify we're not imaging our own boot device ───────────────────
    ROOT_DEV=$(findmnt -n -o SOURCE / 2>/dev/null | sed 's/[0-9]*$//' || true)
    if [[ "$SOURCE_DEVICE" == "$ROOT_DEV" ]]; then
        err "Refusing to image the device we're booted from ($SOURCE_DEVICE)"
        err "Boot from a different SD card or use --from-scratch"
        exit 1
    fi

    info "Reading $SOURCE_DEVICE..."
    dd if="$SOURCE_DEVICE" of="$OUTPUT_IMG" bs=4M status=progress
    sync
    ok "Raw image created: $OUTPUT_IMG"

    # ── Run pishrink to shrink the image ───────────────────────────────
    info "Shrinking image with PiShrink..."

    if ! command -v pishrink &>/dev/null; then
        info "Installing PiShrink..."
        wget -q https://raw.githubusercontent.com/Drewsif/PiShrink/master/pishrink.sh -O /tmp/pishrink.sh
        chmod +x /tmp/pishrink.sh
        PISHRINK="/tmp/pishrink.sh"
    else
        PISHRINK="pishrink"
    fi

    $PISHRINK -v "$OUTPUT_IMG" -p 2>/dev/null || {
        warn "PiShrink encountered issues, trying without verbose..."
        $PISHRINK "$OUTPUT_IMG" || warn "PiShrink failed — image is raw (unshrunk)"
    }

    ok "Image shrunk"

# ══════════════════════════════════════════════════════════════════════════════
# Mode 3: No source specified — build from current system
# ══════════════════════════════════════════════════════════════════════════════

else
    warn "No --source-device or --from-scratch specified."
    warn ""
    warn "If you want to build from your CURRENT running system:"
    warn "  sudo ./build-minimal-image.sh --source-device /dev/mmcblk0"
    warn ""
    warn "If you want to build from a fresh DietPi download:"
    warn "  sudo ./build-minimal-image.sh --from-scratch"
    echo ""
    err "Choose one of the above."
    exit 1
fi

# ══════════════════════════════════════════════════════════════════════════════
# Final compression
# ══════════════════════════════════════════════════════════════════════════════

if [[ "$COMPRESS" == true ]] && [[ -f "$OUTPUT_IMG" ]]; then
    info "Compressing image (gzip -9)..."
    gzip -9 -f "$OUTPUT_IMG"
    FINAL_FILE="${OUTPUT_IMG}.gz"
    ok "Compressed: $FINAL_FILE"
elif [[ -f "$OUTPUT_IMG" ]]; then
    FINAL_FILE="$OUTPUT_IMG"
fi

# ── Summary ──────────────────────────────────────────────────────────────────
echo ""
info "========== Build Complete =========="
echo ""
echo "  Output:  $FINAL_FILE"
if [[ -f "$FINAL_FILE" ]]; then
    SIZE=$(stat -c%s "$FINAL_FILE" 2>/dev/null || stat -f%z "$FINAL_FILE" 2>/dev/null)
    echo "  Size:    $(numfmt --to=iec-i "$SIZE") ($(numfmt --to=si "$SIZE"))"
fi
echo ""
echo "  Write to SD card:"
echo "    gunzip -c $FINAL_FILE | sudo dd of=/dev/mmcblk0 bs=4M status=progress"
echo "    sync"
echo ""
echo "  First-boot instructions:"
echo "    1. Insert SD card, power on"
echo "    2. The pedal service starts automatically (multi-fx-pedal.target)"
echo "    3. Check status: systemctl status pi-multifx-pedal.service"
echo "    4. Verify JACK:   jack_wait -c"
echo ""