#!/bin/sh
set -eu

### User-configurable vars
IMG_DIR="./out"
CONFIG_PART_SUFFIX="s1"           # config partition after flashing
ENV_SEARCH_DIR="./out"            # optional; leave empty to use selected image dir

MNT="/tmp/monok8s-usb-config.$$"

cleanup() {
    if mount | grep -q "on ${MNT} "; then
        diskutil unmount "$MNT" >/dev/null 2>&1 || true
    fi
    rmdir "$MNT" >/dev/null 2>&1 || true
}

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

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

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

require_cmd() {
    command -v "$1" >/dev/null 2>&1 || die "Missing required command: $1"
}

wait_for_disk() {
    dev="$1"
    timeout="${2:-15}"
    i=0
    while [ "$i" -lt "$timeout" ]; do
        [ -e "$dev" ] && return 0
        i=$((i + 1))
        sleep 1
    done
    return 1
}

wait_for_mountable_partition() {
    dev="$1"
    timeout="${2:-20}"
    i=0
    while [ "$i" -lt "$timeout" ]; do
        if [ -e "$dev" ] && diskutil info "$dev" >/dev/null 2>&1; then
            return 0
        fi
        i=$((i + 1))
        sleep 1
    done
    return 1
}

select_image() {
    dir="$1"

    [ -d "$dir" ] || die "Image directory not found: $dir"

    OLD_IFS="${IFS}"
    IFS='
'
    set -- $(find "$dir" -maxdepth 1 -type f -name '*.img.gz' | sort)
    IFS="${OLD_IFS}"

    [ "$#" -gt 0 ] || die "No .img.gz images found in $dir"

    echo "Available images:"
    n=1
    for img in "$@"; do
        size="$(stat -f%z "$img" 2>/dev/null || echo "?")"
        echo "  [$n] $(basename "$img") (${size} bytes)"
        n=$((n + 1))
    done

    echo
    printf "Select image number: "
    read -r choice

    case "$choice" in
        ''|*[!0-9]*)
            die "Invalid selection: $choice"
            ;;
    esac

    n=1
    for img in "$@"; do
        if [ "$n" -eq "$choice" ]; then
            SELECTED_IMG="$img"
            return 0
        fi
        n=$((n + 1))
    done

    die "Selection out of range: $choice"
}

list_candidate_disks() {
    found=0
    for dev in /dev/disk*; do
        base="$(basename "$dev")"

        case "$base" in
            disk[0-9])
                ;;
            *)
                continue
                ;;
        esac

        info="$(diskutil list "$dev" 2>/dev/null || true)"
        [ -n "$info" ] || continue

        # Skip Apple disks
        echo "$info" | grep -q "APFS" && continue

        info="$(diskutil info "$dev" 2>/dev/null || true)"
        found=1

        media_name="$(echo "$info" | sed -n 's/^ *Device \/ Media Name: *//p' | head -n1)"
        protocol="$(echo "$info" | sed -n 's/^ *Protocol: *//p' | head -n1)"
        size="$(echo "$info" | sed -n 's/^ *Disk Size: *//p' | head -n1)"

        echo "$dev|$media_name|$protocol|$size"
    done

    [ "$found" -eq 1 ] || return 1
    return 0
}

select_disk() {
    candidates="$(list_candidate_disks || true)"

    if [ -z "$candidates" ]; then
        warn "No obvious external candidate disks found."
        printf "Enter target disk manually (example: /dev/disk7): "
        read -r TARGET_DISK
        [ -n "$TARGET_DISK" ] || die "No disk entered"
        return 0
    fi

    echo "Candidate target disks:"
    n=1
    echo "$candidates" | while IFS='|' read -r dev media protocol size; do
        echo "  [$n] $dev - $media - $protocol - $size"
        n=$((n + 1))
    done

    echo
    printf "Select disk number (or type a full /dev/diskN path): "
    read -r choice

    case "$choice" in
        /dev/disk[0-9]*)
            TARGET_DISK="$choice"
            return 0
            ;;
        ''|*[!0-9]*)
            die "Invalid selection: $choice"
            ;;
    esac

    n=1
    OLD_IFS="${IFS}"
    IFS='
'
    for line in $candidates; do
        if [ "$n" -eq "$choice" ]; then
            TARGET_DISK="$(echo "$line" | cut -d'|' -f1)"
            IFS="${OLD_IFS}"
            return 0
        fi
        n=$((n + 1))
    done
    IFS="${OLD_IFS}"

    die "Selection out of range: $choice"
}

### Sanity checks
require_cmd diskutil
require_cmd gunzip
require_cmd dd
require_cmd sync
require_cmd mount
require_cmd find
require_cmd cp
require_cmd sed
require_cmd stat
require_cmd pv
require_cmd cut
require_cmd grep
require_cmd head

select_image "$IMG_DIR"
IMG="$SELECTED_IMG"

select_disk

TARGET_RAW_DISK="/dev/r$(basename "$TARGET_DISK")"
TARGET_BASENAME="$(basename "$TARGET_DISK")"

### Derived vars
TARGET_CONFIG_PART="/dev/${TARGET_BASENAME}${CONFIG_PART_SUFFIX}"
IMG_DIR_ABS="$(cd "$(dirname "$IMG")" && pwd)"
SEARCH_DIR="${ENV_SEARCH_DIR:-$IMG_DIR_ABS}"

[ -f "$IMG" ] || die "Image not found: $IMG"
IMG_SIZE_BYTES="$(stat -f%z "$IMG")"

[ -b "$TARGET_DISK" ] || die "Target disk not found: $TARGET_DISK"
[ "$(id -u)" -eq 0 ] || die "Run this script with sudo"

case "$TARGET_DISK" in
    /dev/disk[0-9]*)
        ;;
    *)
        die "TARGET_DISK must look like /dev/diskN"
        ;;
esac

mkdir -p "$MNT"
trap cleanup EXIT

log "Image: $IMG"
log "Compressed image size: $IMG_SIZE_BYTES bytes"
log "Target disk: $TARGET_DISK"
log "Raw target disk: $TARGET_RAW_DISK"
log "Expected config partition after flash: $TARGET_CONFIG_PART"

log "Target disk info:"
diskutil info "$TARGET_DISK" || die "Unable to inspect target disk"

echo
printf "Type 'YES' to continue: "
read -r CONFIRM
[ "$CONFIRM" = "YES" ] || die "Aborted."

log "Unmounting target disk..."
diskutil unmountDisk force "$TARGET_DISK" || die "Failed to unmount $TARGET_DISK"

log "Flashing image with progress..."
pv -s "$IMG_SIZE_BYTES" "$IMG" \
    | gunzip \
    | dd of="$TARGET_RAW_DISK" bs=4m

log "Syncing writes..."
sync

log "Ejecting and re-reading disk..."
diskutil eject "$TARGET_DISK" >/dev/null 2>&1 || true
sleep 2

log "Waiting for disk to come back..."
if ! wait_for_disk "$TARGET_DISK" 15; then
    warn "Disk node did not reappear quickly. You may need to replug the drive."
fi

log "Disk layout after flashing:"
diskutil list "$TARGET_DISK" || warn "Could not list flashed disk yet"

ENV_FILES="$(find "$SEARCH_DIR" -maxdepth 1 -type f -name '*.env' | sort || true)"

if [ -n "$ENV_FILES" ]; then
    log "Found env files:"
    echo "$ENV_FILES" | sed 's/^/  /'

    log "Waiting for config partition..."
    wait_for_mountable_partition "$TARGET_CONFIG_PART" 20 || die "Config partition not ready: $TARGET_CONFIG_PART"

    log "Mounting flashed config partition..."
    diskutil mount -mountPoint "$MNT" "$TARGET_CONFIG_PART" >/dev/null || die "Failed to mount $TARGET_CONFIG_PART"

    log "Copying env files to config partition..."
    echo "$ENV_FILES" | while IFS= read -r f; do
        [ -n "$f" ] || continue
        base="$(basename "$f")"
        cp "$f" "$MNT/$base"
        echo "  copied: $base"
    done

    log "Syncing copied config..."
    sync

    log "Config partition contents:"
    ls -la "$MNT" || true

    log "Unmounting config partition..."
    diskutil unmount "$MNT" >/dev/null || warn "Unmount failed for $MNT"
else
    log "No .env files found in $SEARCH_DIR"
fi

log "Done."