#!/usr/bin/env bash
# Shrink a container's root disk by copying it into a smaller volume. Run as root
# on the Proxmox node. Without --yes it only measures and prints the plan.
# Spec: docs/plans/lxc-shrink.md (section B).
set -euo pipefail

CONF_DIR=${PCT_CONF_DIR:-/etc/pve/lxc}
BACKUP_DIR=${PCT_SHRINK_BACKUP_DIR:-/root}
GIB=1073741824
ctid="" conf="" mounted=0 new_volid=""

die() { echo "pct-shrink-rootfs: $*" >&2; exit 1; }

usage() { die "usage: pct-shrink-rootfs [--yes] <vmid> <new-size-GiB>"; }

conf_get() { awk -v k="$1" 'index($0, k ": ") == 1 { print substr($0, length(k) + 3); exit }' "$conf"; }

refuse_if() { if grep -Eq "$1" "$conf"; then die "$2"; fi; }

# The value of the unusedK line that holds exactly this volid (no regex on the volid).
unused_key_of() {
  awk -v v="$1" '/^unused[0-9]+: / { if (substr($0, index($0, ": ") + 2) == v) { print substr($0, 1, index($0, ":") - 1); exit } }' "$conf"
}

size_bytes() {
  [[ $1 =~ ^([0-9]+)([KMGT]?)$ ]] || die "cannot read the current rootfs size '$1'"
  local m=1
  case ${BASH_REMATCH[2]} in K) m=1024 ;; M) m=$((1024 ** 2)) ;; G) m=$GIB ;; T) m=$((1024 ** 4)) ;; *) ;; esac
  echo $((BASH_REMATCH[1] * m))
}

manual_cleanup() {
  cat >&2 <<EOT
pct-shrink-rootfs: 'pct unmount $ctid' failed. Nothing else was touched. Clean up by hand:
  pct unmount $ctid
  pct unlock $ctid    # if the lock remains
EOT
}

unmount_or_die() {
  mounted=0
  pct unmount "$ctid" || { manual_cleanup; exit 1; }
}

on_exit() {
  local rc=$?
  trap - EXIT
  if ((mounted)); then unmount_or_die; fi
  if ((rc != 0)) && [[ -n $new_volid ]]; then
    echo "pct-shrink-rootfs: the new volume $new_volid is still in the config as mp9;" \
      "remove it with 'pct set $ctid --delete mp9' and then the unusedK that holds it." >&2
  fi
  exit "$rc"
}

# Runs inside flock. Writes like PVE's file_set_contents: "<file>.tmp.<pid>" in the same
# pmxcfs directory, then rename. Re-checks both volids under the lock before writing.
swap_config() {
  local conf=$1 old=$2 new=$3 size=$4 backup=$5 tmp=$1.tmp.$$
  cp "$conf" "$backup"
  if ! awk -v old="$old" -v new="$new" -v size="$size" '
    { key = substr($0, 1, index($0, ":") - 1); val = substr($0, index($0, ": ") + 2) }
    key ~ /^unused[0-9]+$/ { used[substr(key, 7) + 0] = 1 }
    key == "mp9" { split(val, m, ","); ok_new = (m[1] == new); next }
    key == "rootfs" {
      n = split(val, a, ",")
      ok_old = (a[1] == old)
      out = new; seen = 0
      for (i = 2; i <= n; i++) {
        if (substr(a[i], 1, 5) == "size=") { out = out ",size=" size; seen = 1 }
        else out = out "," a[i]
      }
      if (!seen) out = out ",size=" size
      print "rootfs: " out; next
    }
    { print }
    END { if (!ok_old || !ok_new) exit 3; k = 0; while (k in used) k++; print "unused" k ": " old }
  ' "$conf" > "$tmp"; then
    rm -f "$tmp"
    echo "pct-shrink-rootfs: config changed during the copy (rootfs or mp9 differs); not swapped" >&2
    return 1
  fi
  mv -f "$tmp" "$conf"
}
export -f swap_config

free_new_volume() {
  local key vol=$new_volid
  pct set "$ctid" --delete mp9
  new_volid=""  # no longer mp9: the on_exit hint would be wrong
  key=$(unused_key_of "$vol")
  [[ -n $key ]] || die "cannot find $vol among the unused volumes; delete it by hand"
  pct set "$ctid" --delete "$key"
}

main() {
  local yes=0 size root rootfs old_volid storage cur_bytes used iused out mp9 new_size backup key
  while [[ ${1:-} == --* ]]; do
    [[ $1 == --yes ]] || usage
    yes=1
    shift
  done
  [[ $# -eq 2 ]] || usage
  ctid=$1 size=$2
  [[ $ctid =~ ^[0-9]+$ ]] || die "vmid must be a number"
  [[ $size =~ ^[1-9][0-9]{0,8}$ ]] || die "new size must be a positive integer (GiB)"

  [[ $(id -u) -eq 0 ]] || die "must run as root on the node"
  command -v pct > /dev/null || die "pct not found; run this on a Proxmox node"
  conf=$CONF_DIR/$ctid.conf
  [[ -f $conf ]] || die "no config $conf"
  [[ $(pct status "$ctid") == "status: stopped" ]] || die "container $ctid must be stopped"
  refuse_if '^\[' "config has a snapshot or pending section"
  refuse_if '^lock:' "container is locked"
  refuse_if '^template: 1$' "container is a template"
  refuse_if '^protection: 1$' "container has protection set"
  refuse_if '^(mp[0-9]+|lxc\.mount\.entry):' "container has a mount point (mpN or lxc.mount.entry); not supported"
  out=$(ha-manager status)
  if grep -Eq "(^| )ct:$ctid( |\$)" <<< "$out"; then die "container is HA-managed"; fi
  out=$(pvesr list)
  if grep -Eq "^$ctid-" <<< "$out"; then die "container has a replication job"; fi

  rootfs=$(conf_get rootfs)
  old_volid=${rootfs%%,*}
  [[ $old_volid == *:* && $old_volid != /* ]] || die "rootfs is not a volume: $old_volid"
  storage=${old_volid%%:*}
  [[ $rootfs =~ (^|,)size=([^,]+) ]] || die "rootfs has no size= option"
  cur_bytes=$(size_bytes "${BASH_REMATCH[2]}")
  ((size * GIB < cur_bytes)) || die "new size must be smaller than the current rootfs (${BASH_REMATCH[2]})"

  trap on_exit EXIT
  root=/var/lib/lxc/$ctid/rootfs
  mounted=1
  pct mount "$ctid"
  used=$(df -B1 --output=used "$root" | tail -n1 | tr -d ' ')
  iused=$(df -i --output=iused "$root" | tail -n1 | tr -d ' ')
  unmount_or_die
  [[ $used =~ ^[0-9]+$ && $iused =~ ^[0-9]+$ ]] || die "cannot read df output"

  echo "Container $ctid, rootfs $old_volid ($storage): $used bytes, $iused inodes used."
  echo "Config dir: $CONF_DIR; backup dir: $BACKUP_DIR"
  # 5 % root reserve + ~2 % ext4 metadata; mkfs.ext4 default is one inode per 16 KiB.
  ((used * 1100 <= size * GIB * 930)) || die "data too full for ${size}G: need used * 1.1 <= 93 % of the disk"
  ((iused * 11 <= size * 65536 * 10)) || die "too many inodes for ${size}G: $iused used, ${size}G gives $((size * 65536))"

  if ((!yes)); then
    echo "Plan: allocate ${size}G on $storage, rsync the root disk into it, swap it into rootfs."
    echo "The old volume is kept as unusedK. Re-run with --yes to do it."
    return 0
  fi

  pct set "$ctid" --mp9 "$storage:$size,mp=/.shrink-new"
  mp9=$(conf_get mp9)
  [[ -n $mp9 ]] || die "mp9 not found after pct set"
  new_volid=${mp9%%,*}
  [[ ,$mp9, =~ ,size=([^,]+), ]] || die "mp9 has no size= option: $mp9"
  new_size=${BASH_REMATCH[1]}

  mounted=1
  pct mount "$ctid"
  if ! rsync -aHAX --numeric-ids --whole-file --sparse --one-file-system \
    --exclude=/.shrink-new "$root/" "$root/.shrink-new/"; then
    unmount_or_die
    free_new_volume
    die "rsync failed; the new volume was removed and the config is unchanged"
  fi
  unmount_or_die

  backup=$BACKUP_DIR/pct-shrink-$ctid-$(date +%Y%m%d-%H%M%S).conf
  flock "/run/lock/lxc/pve-config-$ctid.lock" \
    bash -c 'set -euo pipefail; swap_config "$@"' _ "$conf" "$old_volid" "$new_volid" "$new_size" "$backup"
  new_volid=""

  key=$(unused_key_of "$old_volid")
  cat <<EOT
Done. rootfs is now on the new ${new_size} volume; config backup: $backup
The old root disk is kept as '$key: $old_volid'.
Start the container and check it. Then free the old volume:
  pct set $ctid --delete $key
rsync does not copy chattr flags (+i, +a); re-apply them in the container if you used any.
EOT
}

if [[ "${BASH_SOURCE[0]}" == "$0" ]]; then
  main "$@"
fi