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

# -------- configurable (change if your root is elsewhere) -------
PART=/dev/mmcblk1p1
DISK=/dev/mmcblk1
TMPROOT=/mnt/newroot
SAFETY_MIB=16         # extra headroom added to partition after shrink
TMPFS_SIZE=256M       # enough to hold binaries + libs
# ---------------------------------------------------------------

require() { command -v "$1" >/dev/null 2>&1 || { echo "Missing required tool: $1"; exit 1; }; }

# Tools we will need in the RAM-root
TOOLS=(
  /bin/busybox
  /bin/sh
  /bin/mount
  /bin/umount
  /usr/sbin/chroot
  /usr/sbin/pivot_root
  /sbin/resize2fs
  /sbin/e2fsck
  /sbin/tune2fs
  /sbin/blkid
  /sbin/partprobe
  /sbin/lsblk
  /sbin/udevadm
  /usr/sbin/parted
)

echo "[*] Preflight checks…"
for t in "${TOOLS[@]}"; do
  require "$(basename "$t")" || true
  [[ -x "$t" ]] || { echo "Required binary not found or not executable: $t"; exit 1; }
done

[[ -b "$PART" ]] || { echo "Partition not found: $PART"; exit 1; }
[[ -b "$DISK" ]] || { echo "Disk not found: $DISK"; exit 1; }

ROOTDEV=$(findmnt -no SOURCE /) || true
if [[ "$ROOTDEV" != "$PART" ]]; then
  echo "Root appears to be '$ROOTDEV', not '$PART'. Edit PART/DISK at top of script."
  exit 1
fi

echo "[*] Mount tmpfs at $TMPROOT…"
mkdir -p "$TMPROOT"
mount -t tmpfs -o size="$TMPFS_SIZE",mode=0755 tmpfs "$TMPROOT"

echo "[*] Create minimal filesystem layout…"
for d in bin sbin usr/bin usr/sbin lib lib64 etc proc sys dev run old_root; do
  mkdir -p "$TMPROOT/$d"
done

# Helper: copy binary + all shared libs it needs
copy_with_libs() {
  local src="$1"
  local dst="$TMPROOT$src"
  local dstdir
  dstdir=$(dirname "$dst")
  mkdir -p "$dstdir"
  install -m 0755 "$src" "$dst"

  # ldd may be a shell script; use it anyway to parse deps
  # Ignore 'not a dynamic executable' for static binaries
  local libs
  libs=$(ldd "$src" 2>/dev/null | awk '/=>/ {print $3} /^[[:space:]]*\/.*/ {print $1}' | sort -u || true)
  for lib in $libs; do
    [[ -f "$lib" ]] || continue
    local libdst="$TMPROOT$lib"
    mkdir -p "$(dirname "$libdst")"
    install -m 0644 "$lib" "$libdst"
  done

  # Also copy the dynamic linker if shown explicitly
  local ldso
  ldso=$(ldd "$src" 2>/dev/null | awk '/ld-linux|ld-musl|ld64/ {print $1}' || true)
  if [[ -n "${ldso:-}" && -f "$ldso" ]]; then
    mkdir -p "$TMPROOT$(dirname "$ldso")"
    install -m 0755 "$ldso" "$TMPROOT$ldso"
  fi
}

echo "[*] Copy required tools and libraries into tmpfs…"
for bin in "${TOOLS[@]}"; do
  copy_with_libs "$bin"
done

# Provide busybox applets fallback inside tmpfs
if [[ -x /bin/busybox ]]; then
  for app in sh mount umount ls ln cat echo grep awk sed cp mv rm mkdir dmesg sync reboot; do
    [[ -e "$TMPROOT/bin/$app" ]] || ln -sf /bin/busybox "$TMPROOT/bin/$app"
  done
fi

# Minimal /etc for chroot comfort
echo "root:x:0:0:root:/root:/bin/sh" > "$TMPROOT/etc/passwd" || true

echo "[*] Bind-mount system psuedo-fs into tmpfs…"
mount --bind /proc "$TMPROOT/proc"
mount --bind /sys  "$TMPROOT/sys"
mount --bind /dev  "$TMPROOT/dev"
mount --bind /run  "$TMPROOT/run"

echo "[*] Make mounts private to allow pivot_root cleanly…"
mount --make-rprivate /

echo "[*] Pivot into tmpfs root…"
cd "$TMPROOT"
pivot_root . old_root

# We are now inside the new root. Paths from here are relative to $TMPROOT.
# Create a function to run commands safely here.
inram() { "$@"; }

echo "[*] Remount pseudo-fs inside new root…"
inram mount -t proc proc /proc || true
inram mount -t sysfs sysfs /sys || true
# /dev and /run already bind-mounted

echo "[*] Stop services that could hold the old root (best-effort)…"
inram /bin/busybox ps aux 2>/dev/null || true
# Not trying to stop everything; pivot isolates us sufficiently.

echo "[*] Unmount old_root pseudo-fs to free the device…"
inram umount -l /old_root/proc || true
inram umount -l /old_root/sys  || true
inram umount -l /old_root/run  || true
# Keep /old_root/dev mounted for block device access until we’re done:
# Instead, we already have /dev in new root, so we can unmount it too:
inram umount -l /old_root/dev || true

echo "[*] Sync and ensure $PART is not mounted…"
inram sync
# In case anything still had it, attempt lazy unmount of previous root:
inram umount -lf /old_root || true

# ----------------- Shrink filesystem to minimum -----------------
echo "[*] Checking filesystem integrity on $PART…"
inram e2fsck -f -y "$PART"

echo "[*] Querying minimum filesystem size (in 4K blocks by default)…"
MINBLOCKS_LINE=$(inram resize2fs -P "$PART" 2>/dev/null | tail -n1)
# Expected: "Estimated minimum size of the filesystem: N"
MINBLOCKS=$(echo "$MINBLOCKS_LINE" | awk -F': ' '{print $2}')
if ! [[ "$MINBLOCKS" =~ ^[0-9]+$ ]]; then
  echo "Failed to determine minimum filesystem size. Got: $MINBLOCKS_LINE"
  inram reboot -f
fi
echo "[*] Minimum ext size: ${MINBLOCKS} blocks"

# Determine ext block size (usually 4096)
BLKSZ=$(inram tune2fs -l "$PART" | awk -F': ' '/Block size:/ {print $2}')
if ! [[ "$BLKSZ" =~ ^[0-9]+$ ]]; then BLKSZ=4096; fi
echo "[*] Filesystem block size: ${BLKSZ} bytes"

# Shrink filesystem to minimum (exact blocks)
echo "[*] Shrinking filesystem to ${MINBLOCKS} blocks…"
inram resize2fs -p "$PART" "${MINBLOCKS}"

# ----------------- Resize partition to fit FS (+safety) --------
echo "[*] Reading current partition start (MiB)…"
# parted machine-readable: number:start:end:size:type,name,flags...
LINE=$(inram parted -m "$DISK" unit MiB print | awk -F: -v p="$PART" '
  BEGIN{partnum=""}
  /^1:/ {line=$0}
  END{print line}')
if [[ -z "$LINE" ]]; then
  echo "Could not read partition info for $DISK"
  inram reboot -f
fi

START_MIB=$(echo "$LINE" | awk -F: '{print $2}' | sed 's/MiB//')
if ! [[ "$START_MIB" =~ ^[0-9.]+$ ]]; then
  echo "Failed to parse start MiB from: $LINE"
  inram reboot -f
fi
echo "[*] Partition start: ${START_MIB} MiB"

# Compute needed size in MiB = ceil(MINBLOCKS * BLKSZ / MiB) + SAFETY
need_bytes=$(( MINBLOCKS * BLKSZ ))
MiB=1048576
NEED_MIB=$(( (need_bytes + MiB - 1) / MiB ))
NEW_END_MIB=$(python3 - <<PY
import math
start=${START_MIB}
need=${NEED_MIB}
safety=${SAFETY_MIB}
print(int(math.ceil(start + need + safety)))
PY
)
echo "[*] Will set new partition end to: ${NEW_END_MIB} MiB"

echo "[*] Resizing partition boundary…"
inram parted -s "$DISK" unit MiB resizepart 1 "${NEW_END_MIB}"

echo "[*] Inform kernel of partition changes…"
inram partprobe "$DISK" || true
inram udevadm settle || true
sleep 2

# ----------------- Expand FS to fill (tiny grow to edge) -------
echo "[*] Final filesystem check…"
inram e2fsck -f -y "$PART"

echo "[*] Expanding filesystem to fill the resized partition…"
inram resize2fs -p "$PART"

echo "[*] Done. Syncing and rebooting…"
inram sync
inram reboot -f

# End of script
