#!/usr/bin/env bash
set -euo pipefail

# Onefinity CNC Controller - SD Card Backup & Restore
#
# Backs up the Raspberry Pi's SD card over SSH as a compressed image.
# Compression runs on the local machine (fast), raw bytes stream from the Pi.
#
# Usage:
#   ./onefinity-backup.sh backup              # backup with defaults
#   ./onefinity-backup.sh backup -o myfile.gz  # custom output file
#   ./onefinity-backup.sh restore image.gz     # restore to SD card
#   ./onefinity-backup.sh verify image.gz      # verify image integrity
#
# Environment:
#   ONEFINITY_HOST  - Pi IP/hostname    (default: 10.1.10.55)
#   ONEFINITY_USER  - SSH user          (default: bbmc)
#   ONEFINITY_PASS  - sudo password     (default: onefinity)

HOST="${ONEFINITY_HOST:-10.1.10.55}"
USER="${ONEFINITY_USER:-bbmc}"
PASS="${ONEFINITY_PASS:-onefinity}"
BACKUP_DIR="$(cd "$(dirname "$0")" && pwd)"
DEVICE="/dev/mmcblk0"

ssh_cmd() {
    ssh -o ConnectTimeout=10 -o StrictHostKeyChecking=no -o LogLevel=ERROR "$USER@$HOST" "$@"
}

sudo_ssh() {
    ssh_cmd "echo '$PASS' | sudo -S bash -c '$1' 2>/dev/null"
}

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

# ── Backup ──────────────────────────────────────────────────────────────────

do_backup() {
    local outfile=""

    while [[ $# -gt 0 ]]; do
        case "$1" in
            -o|--output) outfile="$2"; shift 2 ;;
            *) die "Unknown option: $1" ;;
        esac
    done

    if [[ -z "$outfile" ]]; then
        outfile="$BACKUP_DIR/onefinity-$(date +%Y%m%d-%H%M).img.gz"
    fi

    echo "╔══════════════════════════════════════════════════════╗"
    echo "║  Onefinity CNC Controller - SD Card Backup          ║"
    echo "╚══════════════════════════════════════════════════════╝"
    echo ""
    echo "  Host:   $USER@$HOST"
    echo "  Device: $DEVICE"
    echo "  Output: $outfile"
    echo ""

    # Check connectivity
    echo "→ Checking SSH connection..."
    ssh_cmd 'hostname' >/dev/null 2>&1 || die "Cannot SSH to $USER@$HOST"

    # Get card size
    local card_bytes
    card_bytes=$(sudo_ssh "blockdev --getsize64 $DEVICE")
    local card_gb=$(echo "scale=1; $card_bytes / 1073741824" | bc)
    echo "  SD card: ${card_gb}GB ($card_bytes bytes)"
    echo ""

    # Check for enough local disk space (compressed is ~4% of raw)
    local avail_bytes
    avail_bytes=$(df -P "$(dirname "$outfile")" | tail -1 | awk '{print $4 * 1024}')
    local need_bytes=$((card_bytes / 10))  # conservative: assume 10% compressed
    if (( avail_bytes < need_bytes )); then
        die "Not enough local disk space. Need ~$(echo "scale=1; $need_bytes/1073741824" | bc)GB, have $(echo "scale=1; $avail_bytes/1073741824" | bc)GB"
    fi

    # Stream raw dd from Pi, compress locally with gzip
    # The Pi's SD card reads at ~20MB/s which is the bottleneck.
    # Compressing locally on a fast machine is much better than on the ARM.
    echo "→ Streaming SD card image (this takes ~20-50 minutes)..."
    echo "  Pi: dd → SSH → local gzip → $outfile"
    echo ""

    local start_time=$SECONDS
    local tmpfile="${outfile}.partial"

    ssh_cmd "echo '$PASS' | sudo -S dd if=$DEVICE bs=4M 2>/dev/null" 2>/dev/null \
        | gzip -1 > "$tmpfile" &
    local pid=$!

    # Progress monitor
    while kill -0 $pid 2>/dev/null; do
        sleep 15
        if [[ -f "$tmpfile" ]]; then
            local size_h
            size_h=$(ls -lh "$tmpfile" 2>/dev/null | awk '{print $5}')
            local elapsed=$(( SECONDS - start_time ))
            local min=$(( elapsed / 60 ))
            local sec=$(( elapsed % 60 ))
            printf "\r  %dm%02ds elapsed — %s compressed" "$min" "$sec" "$size_h"
        fi
    done

    wait $pid
    local exit_code=$?
    echo ""

    if [[ $exit_code -ne 0 ]]; then
        rm -f "$tmpfile"
        die "Backup failed (exit code $exit_code)"
    fi

    mv "$tmpfile" "$outfile"

    local elapsed=$(( SECONDS - start_time ))
    local final_size
    final_size=$(ls -lh "$outfile" | awk '{print $5}')

    echo ""
    echo "→ Verifying image integrity..."
    if gzip -t "$outfile" 2>/dev/null; then
        echo "  ✓ gzip integrity OK"
    else
        die "Image file is corrupt!"
    fi

    # Verify full size by counting decompressed bytes
    local actual_bytes
    actual_bytes=$(gzip -d -c "$outfile" | wc -c | tr -d ' ')
    if [[ "$actual_bytes" -eq "$card_bytes" ]]; then
        echo "  ✓ Size matches: $actual_bytes bytes (full ${card_gb}GB card)"
    else
        echo "  ⚠ Size mismatch: expected $card_bytes, got $actual_bytes"
    fi

    echo ""
    echo "╔══════════════════════════════════════════════════════╗"
    echo "  ✓ Backup complete"
    echo "  File: $outfile"
    echo "  Size: $final_size compressed (${card_gb}GB raw)"
    echo "  Time: $(( elapsed / 60 ))m $(( elapsed % 60 ))s"
    echo "╚══════════════════════════════════════════════════════╝"
}

# ── Restore ─────────────────────────────────────────────────────────────────

do_restore() {
    local imgfile="$1"
    local target="${2:-}"

    [[ -f "$imgfile" ]] || die "Image file not found: $imgfile"

    echo "╔══════════════════════════════════════════════════════╗"
    echo "║  Onefinity CNC Controller - SD Card Restore         ║"
    echo "╚══════════════════════════════════════════════════════╝"
    echo ""

    if [[ -n "$target" ]]; then
        # ── Local restore: write to a local SD card device ──
        [[ -b "$target" ]] || die "$target is not a block device"

        local target_bytes
        target_bytes=$(diskutil info -plist "$target" 2>/dev/null \
            | plutil -extract TotalSize raw - 2>/dev/null \
            || blockdev --getsize64 "$target" 2>/dev/null \
            || echo 0)

        echo "  Image:  $imgfile"
        echo "  Target: $target ($(echo "scale=1; $target_bytes/1073741824" | bc)GB)"
        echo ""
        echo "  ⚠  THIS WILL ERASE ALL DATA ON $target"
        echo ""
        read -rp "  Type YES to continue: " confirm
        [[ "$confirm" == "YES" ]] || die "Aborted"

        echo ""
        echo "→ Unmounting target..."
        diskutil unmountDisk "$target" 2>/dev/null || true

        echo "→ Writing image to $target..."
        local raw_target
        raw_target=$(echo "$target" | sed 's|/dev/disk|/dev/rdisk|')
        gzip -d -c "$imgfile" | sudo dd of="$raw_target" bs=4M status=progress
        sync

        echo ""
        echo "  ✓ Restore complete. Safe to eject $target."

    else
        # ── Remote restore: write back to Pi over SSH ──
        echo "  Image:  $imgfile"
        echo "  Target: $USER@$HOST:$DEVICE"
        echo ""
        echo "  ⚠  THIS WILL ERASE THE PI'S SD CARD"
        echo "  ⚠  The Pi must be booted from USB/network, not the SD card"
        echo ""
        read -rp "  Type YES to continue: " confirm
        [[ "$confirm" == "YES" ]] || die "Aborted"

        echo ""
        echo "→ Writing image to $HOST:$DEVICE..."
        gzip -d -c "$imgfile" \
            | ssh_cmd "echo '$PASS' | sudo -S dd of=$DEVICE bs=4M 2>/dev/null"

        echo ""
        echo "  ✓ Remote restore complete."
    fi
}

# ── Verify ──────────────────────────────────────────────────────────────────

do_verify() {
    local imgfile="$1"
    [[ -f "$imgfile" ]] || die "Image file not found: $imgfile"

    echo "Verifying: $imgfile"
    echo ""

    local compressed_size
    compressed_size=$(ls -lh "$imgfile" | awk '{print $5}')
    echo "  Compressed size: $compressed_size"

    echo "  Checking gzip integrity..."
    if gzip -t "$imgfile" 2>/dev/null; then
        echo "  ✓ gzip OK"
    else
        die "gzip integrity check FAILED"
    fi

    echo "  Counting uncompressed bytes..."
    local raw_bytes
    raw_bytes=$(gzip -d -c "$imgfile" | wc -c | tr -d ' ')
    local raw_gb=$(echo "scale=1; $raw_bytes / 1073741824" | bc)
    echo "  ✓ Uncompressed size: ${raw_gb}GB ($raw_bytes bytes)"

    echo "  Checking partition table..."
    gzip -d -c "$imgfile" 2>/dev/null | head -c 512 | xxd | head -4 || true

    echo ""
    echo "  ✓ Image looks valid"
}

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

usage() {
    cat <<EOF
Usage: $(basename "$0") <command> [options]

Commands:
  backup  [-o file.img.gz]           Backup SD card from Pi over SSH
  restore <image.gz> [/dev/diskN]    Restore image to local SD card or remote Pi
  verify  <image.gz>                 Verify image integrity

Environment variables:
  ONEFINITY_HOST  Pi address    (default: 10.1.10.55)
  ONEFINITY_USER  SSH user      (default: bbmc)
  ONEFINITY_PASS  sudo password (default: onefinity)

Examples:
  $(basename "$0") backup
  $(basename "$0") backup -o /tmp/mybackup.img.gz
  $(basename "$0") restore backup/onefinity-20260430.img.gz /dev/disk4
  $(basename "$0") verify backup/onefinity-20260430.img.gz
EOF
    exit 1
}

[[ $# -ge 1 ]] || usage

case "$1" in
    backup)  shift; do_backup "$@" ;;
    restore) shift; [[ $# -ge 1 ]] || usage; do_restore "$@" ;;
    verify)  shift; [[ $# -ge 1 ]] || usage; do_verify "$@" ;;
    *)       usage ;;
esac