#!/bin/sh
# Keep cmux-tui's durable state on a reserved filesystem inside the existing VM
# disk. The image is fully allocated during the bake, so general writes cannot
# consume the daemon's state budget and leave SQLite/WAL writes with no room.
# When the state filesystem is low, this helper grows it in bounded steps while
# preserving root-filesystem headroom. The VM's provisioned disk size is
# unchanged.

set -eu

STATE_IMAGE=/var/lib/cmux/cmux-tui-state.ext4
STATE_LOCK=/var/lib/cmux/cmux-tui-state.lock
STATE_RESERVATION_MARKER=/etc/cmux/cmux-tui-state-reservation
STATE_SIZE_BYTES=1073741824
STATE_MAX_BYTES=8589934592
STATE_MIN_FREE_KB=131072
ROOT_MIN_FREE_KB=1048576

usage() {
  echo "usage: $0 ensure|create HOME USER" >&2
  exit 2
}

[ "$#" = 3 ] || usage
action=$1
home=$2
user=$3
state="$home/.local/state/cmux-tui"

case "$home" in
  /?*) ;;
  *) echo "cmux-tui-state-mount: HOME must be absolute" >&2; exit 2 ;;
esac
case "$user" in
  ""|*[!A-Za-z0-9._-]*) echo "cmux-tui-state-mount: invalid USER" >&2; exit 2 ;;
esac

is_mounted() {
  mountpoint -q "$state" 2>/dev/null
}

free_kb() {
  df -Pk "$1" | awk 'NR == 2 { print $4; exit }'
}

grow_if_needed() {
  state_free_kb=$(free_kb "$state")
  [ -n "$state_free_kb" ] || { echo "cmux-tui-state-mount: cannot read state filesystem free space" >&2; return 1; }
  [ "$state_free_kb" -ge "$STATE_MIN_FREE_KB" ] && return 0

  command -v findmnt >/dev/null 2>&1 || { echo "cmux-tui-state-mount: findmnt is required for state growth" >&2; return 1; }
  command -v losetup >/dev/null 2>&1 || { echo "cmux-tui-state-mount: losetup is required for state growth" >&2; return 1; }
  command -v resize2fs >/dev/null 2>&1 || { echo "cmux-tui-state-mount: resize2fs is required for state growth" >&2; return 1; }

  # A previous grow may have allocated the larger backing file and then
  # failed before resize2fs completed. Finish that resize before allocating
  # another step, so retries never double-count unused backing blocks.
  loop_device=$(findmnt -n -o SOURCE --target "$state")
  case "$loop_device" in
    /dev/loop*) ;;
    *) echo "cmux-tui-state-mount: state mount is not backed by a loop device" >&2; return 1 ;;
  esac
  losetup -c "$loop_device"
  resize2fs "$loop_device" >/dev/null
  state_free_kb=$(free_kb "$state")
  [ -n "$state_free_kb" ] || { echo "cmux-tui-state-mount: cannot reread state filesystem free space" >&2; return 1; }
  [ "$state_free_kb" -ge "$STATE_MIN_FREE_KB" ] && return 0

  current_bytes=$(stat -c %s "$STATE_IMAGE")
  [ "$current_bytes" -lt "$STATE_MAX_BYTES" ] || {
    echo "cmux-tui-state-mount: state filesystem reached its ${STATE_MAX_BYTES}-byte cap" >&2
    return 1
  }
  next_bytes=$((current_bytes * 2))
  [ "$next_bytes" -le "$STATE_MAX_BYTES" ] || next_bytes=$STATE_MAX_BYTES
  growth_kb=$(((next_bytes - current_bytes + 1023) / 1024))
  root_free_kb=$(free_kb "$STATE_IMAGE")
  [ -n "$root_free_kb" ] || { echo "cmux-tui-state-mount: cannot read root filesystem free space" >&2; return 1; }
  if [ "$root_free_kb" -lt $((growth_kb + ROOT_MIN_FREE_KB)) ]; then
    echo "cmux-tui-state-mount: insufficient root headroom to grow state image (free=${root_free_kb}KiB required=$((growth_kb + ROOT_MIN_FREE_KB))KiB)" >&2
    return 1
  fi

  command -v fallocate >/dev/null 2>&1 || { echo "cmux-tui-state-mount: fallocate is required for state growth" >&2; return 1; }

  echo "cmux-tui-state-mount: growing state image from ${current_bytes} to ${next_bytes} bytes" >&2
  (umask 077; fallocate -l "$next_bytes" "$STATE_IMAGE")
  losetup -c "$loop_device"
  resize2fs "$loop_device" >/dev/null
}

ensure_locked() {
  # Older images do not carry a reservation. Keep their existing layout
  # working until they are replaced by an image that does.
  [ -f "$STATE_IMAGE" ] || return 0
  mkdir -p "$state"
  if is_mounted; then
    grow_if_needed
    return 0
  fi
  mount -o loop "$STATE_IMAGE" "$state"
  chown "$user:$user" "$state"
  chmod 0700 "$state"
  grow_if_needed
}

create_locked() {
  command -v fallocate >/dev/null 2>&1 || { echo "cmux-tui-state-mount: fallocate is required" >&2; return 1; }
  command -v mkfs.ext4 >/dev/null 2>&1 || { echo "cmux-tui-state-mount: mkfs.ext4 is required" >&2; return 1; }
  command -v mount >/dev/null 2>&1 || { echo "cmux-tui-state-mount: mount is required" >&2; return 1; }
  mkdir -p "$home/.local/state" /var/lib/cmux
  chmod 0700 /var/lib/cmux

  if [ -e "$STATE_IMAGE" ]; then
    [ -f "$STATE_IMAGE" ] || { echo "cmux-tui-state-mount: state image is not a regular file" >&2; return 1; }
    if [ -f "$STATE_RESERVATION_MARKER" ]; then
      ensure_locked
      return 0
    fi
    # An image without the marker is an interrupted or pre-marker creation.
    # Remove it so a retry can seed and restore the root-backed state again.
    if is_mounted; then
      command -v umount >/dev/null 2>&1 || { echo "cmux-tui-state-mount: umount is required" >&2; return 1; }
      umount "$state"
    fi
    rm -f "$STATE_IMAGE"
  fi

  seed=$(mktemp -d /tmp/cmux-tui-state-seed.XXXXXX)
  created_image=0
  cleanup() {
    if [ "$created_image" = 1 ]; then
      if is_mounted; then umount "$state" 2>/dev/null || true; fi
      [ -f "$STATE_RESERVATION_MARKER" ] || rm -f "$STATE_IMAGE" 2>/dev/null || true
    fi
    rm -rf "$seed"
  }
  trap cleanup EXIT HUP INT TERM
  if [ -d "$state" ]; then
    cp -a "$state/." "$seed/"
  fi

  # fallocate, rather than truncate, reserves real blocks on the unchanged
  # root filesystem before the image is formatted and mounted.
  (umask 077; fallocate -l "$STATE_SIZE_BYTES" "$STATE_IMAGE")
  created_image=1
  mkfs.ext4 -F -m 0 -L cmux-tui-state "$STATE_IMAGE" >/dev/null
  ensure_locked
  cp -a "$seed/." "$state/"
  chown -R "$user:$user" "$state"
  created_image=0
  trap - EXIT HUP INT TERM
  cleanup
}

with_state_lock() {
  mkdir -p /var/lib/cmux
  command -v flock >/dev/null 2>&1 || { echo "cmux-tui-state-mount: flock is required" >&2; return 1; }
  exec 9>"$STATE_LOCK"
  flock -x 9
  "$@"
}

ensure() { with_state_lock ensure_locked; }
create() { with_state_lock create_locked; }

case "$action" in
  ensure) ensure ;;
  create) create ;;
  *) usage ;;
esac
