#!/usr/bin/env bash
# Mount a Raspberry Pi OS .img or .img.xz, edit cloud-init NoCloud files on the boot
# partition, then unmount and (if .img.xz) recompress to a NEW file with current
# date/time prefix. Original image is never overwritten. Requires sudo for loop/mount.
#
# Usage:
#   ./edit-cloudinit-on-image.sh <path-to-image.img.xz>
#   ./edit-cloudinit-on-image.sh <path-to-image.img>
#
# Options:
#   --no-recompress   If image was .img.xz, leave decompressed .img and do not overwrite
#   --replace-with-repo  Copy user-data, meta-data, network-config from repo before editing
#
# Example:
#   ./edit-cloudinit-on-image.sh /path/to/gnss-bootstrap-20260223-215010.img.xz
#
# The original image is preserved; output uses a new timestamp (e.g. gnss-bootstrap-20250305-143022.img.xz).

set -e

REPO_ROOT="$(cd "$(dirname "$0")/../.." && pwd)"
CLOUDINIT_SRC="$REPO_ROOT/emmc-provisioning/cloud-init"
NO_RECOMPRESS=""
REPLACE_WITH_REPO=""

while [[ $# -gt 0 ]]; do
  case "$1" in
    --no-recompress)   NO_RECOMPRESS=1; shift ;;
    --replace-with-repo) REPLACE_WITH_REPO=1; shift ;;
    -*) echo "Unknown option: $1"; exit 1 ;;
    *)  break ;;
  esac
done

IMAGE_IN="$1"
[[ -n "$IMAGE_IN" && -f "$IMAGE_IN" ]] || {
  echo "Usage: $0 [--no-recompress] [--replace-with-repo] <path-to-image.img.xz|.img>"
  echo "Example: $0 gnss-bootstrap-20260223-215010.img.xz"
  exit 1
}

IMAGE_IN="$(realpath "$IMAGE_IN")"
WORK_DIR=""
IMG_FILE=""
ORIGINAL_XZ=""
cleanup() {
  if [[ -n "$MNT" && -d "$MNT" ]]; then
    sudo umount "$MNT" 2>/dev/null || true
  fi
  if [[ -n "$LOOP" && -b "$LOOP" ]]; then
    sudo losetup -d "$LOOP" 2>/dev/null || true
  fi
  if [[ -n "$WORK_DIR" && -d "$WORK_DIR" ]]; then
    if [[ -n "$NO_RECOMPRESS" && -n "$ORIGINAL_XZ" ]]; then
      echo "Work dir left at: $WORK_DIR"
    else
      rm -rf "$WORK_DIR"
    fi
  fi
}
trap cleanup EXIT

if [[ "$IMAGE_IN" == *.img.xz ]]; then
  ORIGINAL_XZ="$IMAGE_IN"
  echo "Decompressing $(basename "$ORIGINAL_XZ") (this may take a minute and needs ~3GB free)…"
  WORK_DIR=$(mktemp -d -p "${TMPDIR:-/tmp}" edit-cloudinit.XXXXXX)
  IMG_FILE="$WORK_DIR/image.img"
  xz -T 0 -d -k -f -c "$ORIGINAL_XZ" > "$IMG_FILE"
else
  WORK_DIR=$(mktemp -d -p "${TMPDIR:-/tmp}" edit-cloudinit.XXXXXX)
  IMG_FILE="$IMAGE_IN"
fi

echo "Attaching image and mounting boot partition…"
LOOP=$(sudo losetup -f --show -P "$IMG_FILE")
# Force kernel to scan partition table so /dev/loopNp1, p2 etc. appear
sudo partprobe "$LOOP" 2>/dev/null || true
sleep 0.5
boot_part="${LOOP}p1"
[[ -b "$boot_part" ]] || boot_part="${LOOP}p2"
[[ -b "$boot_part" ]] || {
  echo "Boot partition not found on image. Partitions on image:"
  ls -la "${LOOP}"p* 2>/dev/null || true
  sudo fdisk -l "$IMG_FILE" 2>/dev/null || true
  exit 1
}

MNT="$WORK_DIR/mnt"
mkdir -p "$MNT"
sudo mount "$boot_part" "$MNT"

if [[ -n "$REPLACE_WITH_REPO" && -d "$CLOUDINIT_SRC" ]]; then
  echo "Copying cloud-init files from repo into boot partition…"
  for f in user-data meta-data network-config; do
    if [[ -f "$CLOUDINIT_SRC/$f" ]]; then
      sudo cp "$CLOUDINIT_SRC/$f" "$MNT/$f"
    elif [[ -f "$CLOUDINIT_SRC/$f.bootstrap" ]]; then
      sudo cp "$CLOUDINIT_SRC/$f.bootstrap" "$MNT/$f"
    fi
  done
fi

echo ""
echo "Boot partition is mounted at: $MNT"
echo "Cloud-init files to edit:"
echo "  $MNT/user-data"
echo "  $MNT/meta-data"
echo "  $MNT/network-config"
echo ""
EDITOR="${EDITOR:-nano}"
echo "Opening editor ($EDITOR). Save and exit when done."
read -r -p "Press Enter to open $EDITOR on these files…"
sudo "$EDITOR" "$MNT/user-data" "$MNT/meta-data" "$MNT/network-config"

echo "Unmounting…"
sudo umount "$MNT"
MNT=""
sudo losetup -d "$LOOP"
LOOP=""

if [[ -n "$ORIGINAL_XZ" && -z "$NO_RECOMPRESS" ]]; then
  ORIG_DIR="$(dirname "$ORIGINAL_XZ")"
  ORIG_BASE="$(basename "$ORIGINAL_XZ" .img.xz)"
  if [[ "$ORIG_BASE" =~ ^(.+)-[0-9]{8}-[0-9]{6}$ ]]; then
    BASE_NAME="${BASH_REMATCH[1]}"
  else
    BASE_NAME="$ORIG_BASE"
  fi
  TIMESTAMP=$(date +%Y%m%d-%H%M%S)
  NEW_IMG_XZ="$ORIG_DIR/${BASE_NAME}-${TIMESTAMP}.img.xz"
  echo "Recompressing to $(basename "$NEW_IMG_XZ")…"
  xz -T 0 -z -f -k "$IMG_FILE"
  mv -f "${IMG_FILE}.xz" "$NEW_IMG_XZ"
  echo "Done. New image (original unchanged): $NEW_IMG_XZ"
  rm -rf "$WORK_DIR"
  WORK_DIR=""
elif [[ -n "$NO_RECOMPRESS" && -n "$ORIGINAL_XZ" ]]; then
  echo "Left decompressed image at: $IMG_FILE"
  echo "Recompress manually: xz -z -k \"$IMG_FILE\""
fi