#!/bin/bash
# bee-remount-medium — find and remount the live ISO medium to /run/live/medium
#
# Run this after reconnecting the ISO source disc (USB/CD) if the live medium
# was lost and /run/live/medium/live/filesystem.squashfs is missing.
#
# Usage: bee-remount-medium [--wait]
#   --wait  keep retrying every 5 seconds until the medium is found (useful
#           while physically reconnecting the device)

set -euo pipefail

MEDIUM_DIR="/run/live/medium"
SQUASHFS_REL="live/filesystem.squashfs"
WAIT_MODE=0

for arg in "$@"; do
    case "$arg" in
        --wait|-w) WAIT_MODE=1 ;;
        --help|-h)
            echo "Usage: bee-remount-medium [--wait]"
            echo "  Finds and remounts the live ISO medium to $MEDIUM_DIR"
            echo "  --wait  retry every 5 s until a medium with squashfs is found"
            exit 0 ;;
    esac
done

log() { echo "[$(date +%H:%M:%S)] $*"; }
die() { log "ERROR: $*" >&2; exit 1; }

# Return all candidate block devices (optical + removable USB mass storage)
find_candidates() {
    # CD/DVD drives
    for dev in /dev/sr* /dev/scd*; do
        [ -b "$dev" ] && echo "$dev"
    done
    # USB/removable disks and partitions
    for dev in /dev/sd* /dev/vd*; do
        [ -b "$dev" ] || continue
        # Only whole disks or partitions — skip the same device we are running from
        local removable
        local base
        base=$(basename "$dev")
        removable=$(cat "/sys/block/${base%%[0-9]*}/removable" 2>/dev/null || echo 0)
        [ "$removable" = "1" ] && echo "$dev"
    done
}

# Try to mount $1 to $MEDIUM_DIR and check for squashfs
try_mount() {
    local dev="$1"
    local tmpdir
    tmpdir=$(mktemp -d /tmp/bee-probe-XXXXXX)
    if mount -o ro "$dev" "$tmpdir" 2>/dev/null; then
        if [ -f "${tmpdir}/${SQUASHFS_REL}" ]; then
            # Unmount probe mount and mount properly onto live path
            umount "$tmpdir" 2>/dev/null || true
            rmdir "$tmpdir"  2>/dev/null || true
            # Unmount whatever is currently on MEDIUM_DIR (may be empty/stale)
            umount "$MEDIUM_DIR" 2>/dev/null || true
            mkdir -p "$MEDIUM_DIR"
            if mount -o ro "$dev" "$MEDIUM_DIR"; then
                log "Mounted $dev on $MEDIUM_DIR"
                return 0
            else
                log "Mount of $dev on $MEDIUM_DIR failed"
                return 1
            fi
        fi
        umount "$tmpdir" 2>/dev/null || true
    fi
    rmdir "$tmpdir" 2>/dev/null || true
    return 1
}

attempt() {
    log "Scanning for ISO medium..."
    for dev in $(find_candidates); do
        log "  Trying $dev ..."
        if try_mount "$dev"; then
            local sq="${MEDIUM_DIR}/${SQUASHFS_REL}"
            log "SUCCESS: squashfs available at $sq ($(du -sh "$sq" | cut -f1))"
            return 0
        fi
    done
    return 1
}

if [ "$WAIT_MODE" = "1" ]; then
    log "Waiting for live medium (press Ctrl+C to abort)..."
    while true; do
        if attempt; then
            exit 0
        fi
        log "  Not found — retrying in 5 s (reconnect the disc now)"
        sleep 5
    done
else
    attempt || die "No ISO medium with ${SQUASHFS_REL} found. Reconnect the disc and re-run, or use --wait."
fi