#!/bin/bash

# omarchy:summary=Set up hibernation with swap and boot resume configuration
# omarchy:requires-sudo=true
# omarchy:args=[--force] [--no-rebuild]

FORCE=false
NO_REBUILD=false
for arg in "$@"; do
  case "$arg" in
    --force) FORCE=true ;;
    --no-rebuild) NO_REBUILD=true ;;
  esac
done

if [[ ! -f /sys/power/image_size ]]; then
  echo -e "Hibernation is not supported on your system" >&2
  exit 0
fi

# When --no-rebuild is set, the caller is responsible for the UKI rebuild
# (e.g. running before limine-mkinitcpio-hook is installed during initial
# install), so we only require limine-mkinitcpio when we'd invoke it ourselves.
if ! $NO_REBUILD && ! command -v limine-mkinitcpio &>/dev/null; then
  echo "Skipping hibernation setup (requires Limine bootloader)"
  exit 0
fi

MKINITCPIO_CONF="/etc/mkinitcpio.conf.d/omarchy_resume.conf"
SWAP_FILE="/swap/swapfile"
RESUME_DROP_IN="/etc/limine-entry-tool.d/resume.conf"

# Check if hibernation is already configured
if [[ -f $MKINITCPIO_CONF ]] && grep -q "^HOOKS+=(resume)$" "$MKINITCPIO_CONF"; then
  # Fix empty resume_offset if btrfs map-swapfile failed during initial setup
  if [[ -f $RESUME_DROP_IN ]] && grep -q 'resume_offset="$' "$RESUME_DROP_IN" && [[ -f $SWAP_FILE ]]; then
    RESUME_OFFSET=$(sudo btrfs inspect-internal map-swapfile -r "$SWAP_FILE" 2>/dev/null)
    if [[ -n $RESUME_OFFSET ]]; then
      echo "Fixing empty resume_offset ($RESUME_OFFSET)"
      sudo sed -i "s/resume_offset=\"$/resume_offset=$RESUME_OFFSET\"/" "$RESUME_DROP_IN"
      sudo sed -i "s/resume_offset=\"$/resume_offset=$RESUME_OFFSET\"/" /etc/default/limine
      $NO_REBUILD || sudo limine-mkinitcpio
    fi
  fi
  echo "Hibernation is already set up"
  exit 0
fi

if ! $FORCE; then
  MEM_TOTAL_HUMAN=$(free --human | awk '/Mem/ {print $2}')
  if ! gum confirm "Use $MEM_TOTAL_HUMAN on boot drive to make hibernation available?"; then
    exit 0
  fi
fi

SWAP_SUBVOLUME="/swap"

# Create btrfs subvolume for swap
if ! sudo btrfs subvolume show "$SWAP_SUBVOLUME" &>/dev/null; then
  echo "Creating Btrfs subvolume"
  sudo btrfs subvolume create "$SWAP_SUBVOLUME"
  sudo chattr +C "$SWAP_SUBVOLUME"
fi

# Create swapfile
if ! sudo swaplabel "$SWAP_FILE" &>/dev/null; then
  echo "Creating swapfile in Btrfs subvolume"
  MEM_TOTAL_KB="$(awk '/MemTotal/ {print $2}' /proc/meminfo)k"
  sudo btrfs filesystem mkswapfile -s "$MEM_TOTAL_KB" "$SWAP_FILE"
fi

# Add swapfile to fstab
if ! grep -Fq "$SWAP_FILE" /etc/fstab; then
  echo "Adding swapfile to /etc/fstab"
  sudo cp -a /etc/fstab "/etc/fstab.$(date +%Y%m%d%H%M%S).back"
  printf "\n# Btrfs swapfile for system hibernation\n%s none swap defaults,pri=0 0 0\n" "$SWAP_FILE" | sudo tee -a /etc/fstab >/dev/null
fi

# Enable swap
if ! swapon --show | grep -q "$SWAP_FILE"; then
  echo "Enabling swap on $SWAP_FILE"
  sudo swapon -p 0 "$SWAP_FILE"
fi

# Add resume hook to mkinitcpio
sudo mkdir -p /etc/mkinitcpio.conf.d
echo "Adding resume hook to $MKINITCPIO_CONF"
echo "HOOKS+=(resume)" | sudo tee "$MKINITCPIO_CONF" >/dev/null

# Ensure keyboard backlight doesn't prevent sleep
sudo cp -p "$OMARCHY_PATH/default/systemd/system-sleep/keyboard-backlight" /usr/lib/systemd/system-sleep/

# Add resume= kernel parameters so the initramfs resume hook knows where to find the
# hibernation image. Without these, resume happens late (after GPU drivers load) and fails.
if [[ ! -f $RESUME_DROP_IN ]]; then
  echo "Adding resume kernel parameters"
  sudo swapon -p 0 "$SWAP_FILE" 2>/dev/null
  RESUME_DEVICE=$(findmnt -no SOURCE -T "$SWAP_FILE" | sed 's/\[.*\]//')
  RESUME_OFFSET=$(sudo btrfs inspect-internal map-swapfile -r "$SWAP_FILE")
  if [[ -n $RESUME_OFFSET ]]; then
    sudo mkdir -p /etc/limine-entry-tool.d
    echo "KERNEL_CMDLINE[default]+=\" resume=$RESUME_DEVICE resume_offset=$RESUME_OFFSET\"" | sudo tee "$RESUME_DROP_IN" >/dev/null
    sudo tee -a /etc/default/limine < "$RESUME_DROP_IN" >/dev/null
  else
    echo "Warning: Could not determine resume offset for $SWAP_FILE" >&2
  fi
fi

# Use ACPI alarm for RTC wakeup on s2idle systems (needed for suspend-then-hibernate)
if grep -q "\[s2idle\]" /sys/power/mem_sleep 2>/dev/null; then
  LIMINE_DROP_IN="/etc/limine-entry-tool.d/rtc-alarm.conf"
  if [[ ! -f $LIMINE_DROP_IN ]]; then
    echo "Enabling ACPI RTC alarm for s2idle suspend"
    sudo mkdir -p /etc/limine-entry-tool.d
    echo 'KERNEL_CMDLINE[default]+=" rtc_cmos.use_acpi_alarm=1"' | sudo tee "$LIMINE_DROP_IN" >/dev/null
    sudo tee -a /etc/default/limine < "$LIMINE_DROP_IN" >/dev/null
  fi
fi

if ! $NO_REBUILD; then
  # limine-mkinitcpio rebuilds initramfs/UKI for all kernels and updates the
  # /boot/limine.conf entries via limine-entry-tool. The limine bootloader
  # binary on the ESP doesn't change here, so we don't need limine-update
  # (which would also re-deploy the binary and rebuild a second time).
  echo "Regenerating initramfs..."
  sudo limine-mkinitcpio
  echo
fi

if ! $FORCE && ! $NO_REBUILD && gum confirm "Reboot to enable hibernation?"; then
  omarchy-system-reboot
fi