#!/bin/sh
set -eu

MNT=/mnt
ROOT_MNT=/mnt-root
USB_DEV=/dev/sda
EMMC_DEV=/dev/mmcblk0
EMMC_ROOT_PART=/dev/mmcblk0p1
CONFIG_DIR=/opt/monok8s/config

DEFAULT_GLOB="${MNT}/${BUILD_TAG}_RELEASE_IMAGE"

cleanup() {
    umount "$ROOT_MNT" 2>/dev/null || true
    umount "$MNT" 2>/dev/null || true
}

log() {
    echo "[*] $*"
}

warn() {
    echo "[!] $*" >&2
}

die() {
    echo "[!] ERROR: $*" >&2
    exit 1
}

reread_partitions() {
    log "Re-reading partition table..."
    blockdev --rereadpt "$EMMC_DEV" 2>/dev/null || true
    partprobe "$EMMC_DEV" 2>/dev/null || true
    partx -u "$EMMC_DEV" 2>/dev/null || true
    mdev -s 2>/dev/null || true
}

get_sectors() {
    devname="${1##*/}"
    cat "/sys/class/block/$devname/size" 2>/dev/null || echo 0
}

wait_for_partition_ready() {
    part="$1"
    timeout="${2:-15}"

    i=0
    while [ "$i" -lt "$timeout" ]; do
        if [ -b "$part" ]; then
            sectors="$(get_sectors "$part")"
            if [ "${sectors:-0}" -gt 0 ]; then
                log "Partition $part is present with $sectors sectors"
                return 0
            fi
        fi
        i=$((i + 1))
        sleep 1
    done

    return 1
}

mkdir -p "$MNT" "$ROOT_MNT"
trap cleanup EXIT

log "Mounting USB..."
mount "$USB_DEV" "$MNT"

# Find a default image (first match)
DEFAULT_IMG="$(ls $DEFAULT_GLOB 2>/dev/null | head -n1 || true)"

echo
if [ -n "$DEFAULT_IMG" ]; then
    echo "Default image found:"
    echo "  $DEFAULT_IMG"
else
    echo "No default image found under $DEFAULT_GLOB"
fi
echo

printf "Image Location (%s): " "${DEFAULT_IMG:-/mnt/...}"
read -r IMG

if [ -z "$IMG" ]; then
    IMG="$DEFAULT_IMG"
fi

[ -n "$IMG" ] || die "No image selected"
[ -f "$IMG" ] || die "Invalid image: $IMG"

echo
echo "Looking for env files on USB..."
ENV_FILES="$(find "$MNT" -maxdepth 1 -type f -name '*.env' | sort || true)"

if [ -n "$ENV_FILES" ]; then
    echo "Found env files:"
    echo "$ENV_FILES" | sed 's/^/  /'
else
    echo "No *.env files found in $MNT"
fi
echo

echo "About to write:"
echo "  Image:  $IMG"
echo "  Target: $EMMC_DEV"
echo

printf "Type 'YES' to continue: "
read -r CONFIRM

[ "$CONFIRM" = "YES" ] || die "Aborted."

log "Flashing image to eMMC..."
gunzip -c "$IMG" > "$EMMC_DEV"

log "Syncing flash write..."
sync

reread_partitions

log "Waiting for partition nodes..."
if ! wait_for_partition_ready "$EMMC_ROOT_PART" 15; then
    warn "Initial partition scan did not settle; retrying after another reread..."
    reread_partitions
    sleep 2
    wait_for_partition_ready "$EMMC_ROOT_PART" 15 || die "root partition not ready: $EMMC_ROOT_PART"
fi

log "Partition sizes:"
echo "  $(basename "$EMMC_DEV"):      $(get_sectors "$EMMC_DEV") sectors"
echo "  $(basename "$EMMC_ROOT_PART"): $(get_sectors "$EMMC_ROOT_PART") sectors"

if [ -n "$ENV_FILES" ]; then
    log "Mounting flashed root partition..."
    mount "$EMMC_ROOT_PART" "$ROOT_MNT" || {
        warn "Mount failed; retrying once after partition reread..."
        reread_partitions
        sleep 2
        mount "$EMMC_ROOT_PART" "$ROOT_MNT"
    }

    log "Creating config dir: $CONFIG_DIR"
    mkdir -p "$ROOT_MNT$CONFIG_DIR"

    log "Copying env files..."
    for f in $ENV_FILES; do
        base="$(basename "$f")"
        cp "$f" "$ROOT_MNT$CONFIG_DIR/$base"
        echo "  copied: $base"
    done

    log "Syncing copied config..."
    sync

    umount "$ROOT_MNT"
fi

log "Done."
log "Rebooting..."
reboot -f