#!/usr/bin/env bash
# Move a running RK1 root filesystem from eMMC to a specifically identified
# NVMe while retaining U-Boot and /boot on eMMC.

set -Eeuo pipefail
umask 022

TARGET=""
EXPECTED_SERIAL=""
BOOT_ENV=/boot/armbianEnv.txt
BOOT_ENV_BACKUP=/boot/armbianEnv.txt.pre-nvme
INSTALL_LOG=/var/log/armbian-install.log
INSTALL_LOG_BACKUP=/var/log/armbian-install.log.pre-nvme

usage() {
  cat <<'EOF'
Usage: sudo live-migrate-root-to-nvme.sh --target /dev/nvmeXnY --serial SERIAL

The target NVMe is erased. The running root must be on eMMC. U-Boot and /boot
remain on eMMC; only the root filesystem is moved to NVMe.
EOF
}

while (($#)); do
  case "$1" in
    --target) TARGET=${2:?}; shift 2 ;;
    --serial) EXPECTED_SERIAL=${2:?}; shift 2 ;;
    -h|--help) usage; exit 0 ;;
    *) echo "Unknown argument: $1" >&2; usage >&2; exit 2 ;;
  esac
done

[[ $EUID -eq 0 ]] || { echo "Run this script through sudo" >&2; exit 2; }
[[ -n "$TARGET" && -n "$EXPECTED_SERIAL" ]] || { usage >&2; exit 2; }

for command_name in armbian-install blkid chmod cmp cp debugfs dpkg dpkg-query \
  e2fsck findmnt flock grep lsof lsblk ps readlink sed swapon sync systemctl \
  udevadm umount xargs; do
  command -v "$command_name" >/dev/null || {
    echo "Missing required command: $command_name" >&2
    exit 1
  }
done

exec 9>/run/lock/rk1-nvme-migrate.lock
flock -n 9 || { echo "Another RK1 NVMe migration is already running" >&2; exit 1; }

TARGET=$(readlink -f -- "$TARGET")
[[ "$TARGET" =~ ^/dev/nvme[0-9]+n[0-9]+$ && -b "$TARGET" ]] || {
  echo "Refusing non-NVMe whole-disk target: $TARGET" >&2
  exit 1
}
[[ "$(lsblk -dnro TYPE "$TARGET")" == disk ]] || {
  echo "Target is not a whole disk: $TARGET" >&2
  exit 1
}
[[ "$(lsblk -dnro TRAN "$TARGET")" == nvme ]] || {
  echo "Target transport is not NVMe: $TARGET" >&2
  exit 1
}

actual_serial=$(lsblk -dnro SERIAL "$TARGET" | xargs)
actual_model=$(lsblk -dnro MODEL "$TARGET" | xargs)
[[ "$actual_serial" == "$EXPECTED_SERIAL" ]] || {
  echo "Target serial mismatch: expected '$EXPECTED_SERIAL', found '$actual_serial'" >&2
  exit 1
}

root_source=$(findmnt -no SOURCE --nofsroot /)
[[ "$root_source" =~ ^/dev/mmcblk[0-9]+p[0-9]+$ ]] || {
  echo "Running root is not on an eMMC partition: $root_source" >&2
  exit 1
}
root_disk="/dev/$(lsblk -no PKNAME "$root_source" | sed -n '1p')"
[[ -b "$root_disk" && "$root_disk" != "$TARGET" ]] || {
  echo "Could not safely distinguish root disk from target" >&2
  exit 1
}

target_mounts=$(lsblk -nrpo MOUNTPOINTS "$TARGET" | sed '/^[[:space:]]*$/d')
[[ -z "$target_mounts" ]] || {
  echo "Target or one of its partitions is mounted:" >&2
  printf '%s\n' "$target_mounts" >&2
  exit 1
}
while IFS= read -r swap_name; do
  [[ -z "$swap_name" ]] && continue
  if [[ "$swap_name" == "$TARGET" || "$swap_name" == "${TARGET}p"* ]]; then
    echo "Target contains active swap: $swap_name" >&2
    exit 1
  fi
done < <(swapon --noheadings --raw --show=NAME)

[[ -e /var/lib/rk1-media/provisioned ]] || {
  echo "RK1 media provisioning has not completed" >&2
  exit 1
}
[[ "$(dpkg-query -W -f='${Status}' rk1-media-stack 2>/dev/null)" == \
  "install ok installed" ]] || {
  echo "rk1-media-stack is not fully installed" >&2
  exit 1
}
dpkg_audit=$(dpkg --audit)
[[ -z "$dpkg_audit" ]] || {
  echo "dpkg reports an incomplete transaction:" >&2
  printf '%s\n' "$dpkg_audit" >&2
  exit 1
}

emmc_uuid=$(blkid -s UUID -o value "$root_source")
[[ -n "$emmc_uuid" ]] || { echo "Could not resolve eMMC UUID" >&2; exit 1; }
grep -Fxq "rootdev=UUID=$emmc_uuid" "$BOOT_ENV" || {
  echo "$BOOT_ENV does not currently point at the running eMMC root" >&2
  exit 1
}

if [[ -e "$BOOT_ENV_BACKUP" ]]; then
  cmp --silent "$BOOT_ENV" "$BOOT_ENV_BACKUP" || {
    echo "Existing boot-environment backup differs; refusing to overwrite it" >&2
    exit 1
  }
else
  cp -a -- "$BOOT_ENV" "$BOOT_ENV_BACKUP"
fi

for apt_unit in apt-daily.service apt-daily-upgrade.service; do
  apt_state=$(systemctl is-active "$apt_unit" 2>/dev/null || true)
  case "$apt_state" in
    inactive|failed|unknown) ;;
    *) echo "Refusing migration while $apt_unit is $apt_state" >&2; exit 1 ;;
  esac
done

provision_state=$(systemctl is-active rk1-media-provision.service 2>/dev/null || true)
case "$provision_state" in
  inactive|failed|unknown|active) ;;
  *) echo "Provisioning is not quiescent: $provision_state" >&2; exit 1 ;;
esac

echo "Confirmed source root: $root_source (UUID=$emmc_uuid)"
echo "ERASING target: $TARGET | $actual_model | serial $actual_serial"

apt_daily_timer_was_active=0
apt_upgrade_timer_was_active=0
unattended_was_active=0
kodi_was_active=0
systemctl is-active --quiet apt-daily.timer && apt_daily_timer_was_active=1
systemctl is-active --quiet apt-daily-upgrade.timer && apt_upgrade_timer_was_active=1
systemctl is-active --quiet unattended-upgrades.service && unattended_was_active=1
systemctl is-active --quiet kodi-rk.service && kodi_was_active=1

restore_services=0
restore_boot_env=0
cleanup() {
  local rc=$?
  trap - EXIT
  if ((restore_boot_env)); then
    echo "Migration did not validate; restoring the eMMC boot environment" >&2
    if ! cp -a -- "$BOOT_ENV_BACKUP" "$BOOT_ENV" || \
      ! cmp --silent "$BOOT_ENV_BACKUP" "$BOOT_ENV"; then
      echo "CRITICAL: could not restore and verify $BOOT_ENV" >&2
      rc=1
    fi
    if ! sync; then
      echo "CRITICAL: writeback failed while restoring $BOOT_ENV" >&2
      rc=1
    fi
  fi
  if ((restore_services)); then
    ((apt_daily_timer_was_active)) && systemctl start apt-daily.timer 2>/dev/null || true
    ((apt_upgrade_timer_was_active)) && systemctl start apt-daily-upgrade.timer 2>/dev/null || true
    ((unattended_was_active)) && systemctl start unattended-upgrades.service 2>/dev/null || true
    ((kodi_was_active)) && systemctl start kodi-rk.service 2>/dev/null || true
  fi
  exit "$rc"
}
trap cleanup EXIT

restore_services=1
systemctl stop apt-daily.timer apt-daily-upgrade.timer 2>/dev/null || true
systemctl stop apt-daily.service apt-daily-upgrade.service 2>/dev/null || true
for apt_unit in apt-daily.service apt-daily-upgrade.service; do
  apt_state=$(systemctl is-active "$apt_unit" 2>/dev/null || true)
  case "$apt_state" in
    inactive|failed|unknown) ;;
    *) echo "Package service did not stop cleanly: $apt_unit is $apt_state" >&2; exit 1 ;;
  esac
done
systemctl stop rk1-media-provision.service
! systemctl is-active --quiet rk1-media-provision.service || {
  echo "Provisioning service did not stop cleanly" >&2
  exit 1
}
systemctl stop unattended-upgrades.service 2>/dev/null || true
if systemctl list-unit-files kodi-rk.service --no-legend 2>/dev/null | grep -q '^kodi-rk\.service'; then
  systemctl stop kodi-rk.service
  ! systemctl is-active --quiet kodi-rk.service || {
    echo "Kodi did not stop cleanly" >&2
    exit 1
  }
fi

package_processes=$(ps -eo comm= | grep -E \
  '^(apt|apt-get|dpkg|unattended-upgr)$' || true)
[[ -z "$package_processes" ]] || {
  echo "Package-management process remains active:" >&2
  printf '%s\n' "$package_processes" >&2
  exit 1
}
lock_holders=$(lsof -t /var/lib/dpkg/lock /var/lib/dpkg/lock-frontend \
  /var/cache/apt/archives/lock 2>/dev/null || true)
[[ -z "$lock_holders" ]] || {
  echo "A package-management lock is still held by PID(s): $lock_holders" >&2
  exit 1
}
dpkg_audit=$(dpkg --audit)
[[ -z "$dpkg_audit" ]] || {
  echo "dpkg became inconsistent before migration:" >&2
  printf '%s\n' "$dpkg_audit" >&2
  exit 1
}

if [[ -s "$INSTALL_LOG" ]]; then
  [[ ! -e "$INSTALL_LOG_BACKUP" ]] || {
    echo "Installer log backup already exists: $INSTALL_LOG_BACKUP" >&2
    exit 1
  }
  cp -a -- "$INSTALL_LOG" "$INSTALL_LOG_BACKUP"
fi
: >"$INSTALL_LOG"
chmod 0600 "$INSTALL_LOG"

sync
restore_boot_env=1
armbian-install --target "$TARGET" --boot sd --fs ext4 --yes

target_partition="${TARGET}p1"
udevadm settle
[[ -b "$target_partition" ]] || {
  echo "Installer did not create $target_partition" >&2
  exit 1
}

mapfile -t post_install_mounts < <(findmnt -rn -S "$target_partition" -o TARGET)
for mountpoint in "${post_install_mounts[@]}"; do
  [[ "$mountpoint" == /mnt/armbian-install.* ]] || {
    echo "Target remains mounted at unexpected path: $mountpoint" >&2
    exit 1
  }
  umount "$mountpoint" || {
    echo "Could not unmount installer path: $mountpoint" >&2
    exit 1
  }
done
findmnt -rn -S "$target_partition" >/dev/null && {
  echo "Target partition remains mounted after installer teardown" >&2
  exit 1
}

grep -Fq "scenario: sd install to $TARGET completed" "$INSTALL_LOG" || {
  echo "Installer did not record successful sd-mode completion" >&2
  exit 1
}

nvme_uuid=$(blkid -s UUID -o value "$target_partition")
[[ -n "$nvme_uuid" && "$nvme_uuid" != "$emmc_uuid" ]] || {
  echo "NVMe UUID is missing or duplicates the eMMC UUID" >&2
  exit 1
}
mapfile -t uuid_devices < <(blkid -t "UUID=$nvme_uuid" -o device)
[[ ${#uuid_devices[@]} -eq 1 && "${uuid_devices[0]}" == "$target_partition" ]] || {
  echo "NVMe filesystem UUID is not unique" >&2
  printf '%s\n' "${uuid_devices[@]}" >&2
  exit 1
}
grep -Fxq "rootdev=UUID=$nvme_uuid" "$BOOT_ENV" || {
  echo "eMMC boot environment does not point to the new NVMe root" >&2
  exit 1
}
grep -Fxq 'rootfstype=ext4' "$BOOT_ENV" || {
  echo "eMMC boot environment does not select ext4 for the new root" >&2
  exit 1
}

e2fsck -fn "$target_partition"
target_fstab=$(debugfs -R 'cat /etc/fstab' "$target_partition" 2>/dev/null)
root_mount_count=$(grep -Ec '^[^#[:space:]][^[:space:]]*[[:space:]]+/[[:space:]]+' \
  <<<"$target_fstab" || true)
[[ "$root_mount_count" -eq 1 ]] || {
  echo "Target fstab does not contain exactly one root mount" >&2
  exit 1
}
grep -Eq "^UUID=${nvme_uuid}[[:space:]]+/[[:space:]]+ext4" <<<"$target_fstab" || {
  echo "Target fstab does not mount the NVMe UUID at /" >&2
  exit 1
}
grep -Eq "^UUID=${emmc_uuid}[[:space:]]+/media/boot-media[[:space:]]+ext4" \
  <<<"$target_fstab" || {
  echo "Target fstab does not mount the eMMC boot medium" >&2
  exit 1
}
grep -Eq '^/media/boot-media/boot[[:space:]]+/boot[[:space:]]+none[[:space:]]+bind' \
  <<<"$target_fstab" || {
  echo "Target fstab does not bind the eMMC /boot directory" >&2
  exit 1
}

while read -r fstab_source _; do
  case "$fstab_source" in
    UUID=*)
      fstab_uuid=${fstab_source#UUID=}
      mapfile -t fstab_devices < <(blkid -t "UUID=$fstab_uuid" -o device)
      [[ ${#fstab_devices[@]} -eq 1 ]] || {
        echo "Target fstab UUID does not resolve uniquely: $fstab_uuid" >&2
        exit 1
      }
      ;;
    PARTUUID=*)
      fstab_partuuid=${fstab_source#PARTUUID=}
      mapfile -t fstab_devices < <(blkid -t "PARTUUID=$fstab_partuuid" -o device)
      [[ ${#fstab_devices[@]} -eq 1 ]] || {
        echo "Target fstab PARTUUID does not resolve uniquely: $fstab_partuuid" >&2
        exit 1
      }
      ;;
  esac
done <<<"$target_fstab"

sync
restore_boot_env=0
restore_services=0
echo "MIGRATION_VALIDATED"
echo "NVMe root: $target_partition (UUID=$nvme_uuid)"
echo "eMMC boot fallback: $BOOT_ENV_BACKUP"
echo "REBOOT_REQUIRED"