#!/bin/bash
# enable-hibernation.sh — one-time setup for laptop hibernation on btrfs.
# Idempotent: re-runs are no-ops once everything is in place. Run with sudo.
# After successful run, reboot before testing `systemctl hibernate`.

set -euo pipefail

if [ "$EUID" -ne 0 ]; then
    echo "Run as root: sudo bash $0" >&2
    exit 1
fi

SWAPFILE=/swapfile
SWAPSIZE=16g
ROOT_PARTUUID=$(findmnt -no PARTUUID /)
[ -n "$ROOT_PARTUUID" ] || { echo "ERROR: could not determine root PARTUUID via findmnt." >&2; exit 1; }

echo "==> Step 1/6: ensure swap file exists at $SWAPFILE (size $SWAPSIZE)"
if [ ! -f "$SWAPFILE" ]; then
    btrfs filesystem mkswapfile --size "$SWAPSIZE" "$SWAPFILE"
    echo "    created $SWAPFILE"
else
    actual_bytes=$(stat -c %s "$SWAPFILE")
    min_bytes=$((14 * 1024 * 1024 * 1024))
    if [ "$actual_bytes" -lt "$min_bytes" ]; then
        echo "ERROR: $SWAPFILE exists but is smaller than 14 GiB (got $actual_bytes bytes)." >&2
        echo "       Remove it manually and re-run." >&2
        exit 1
    fi
    echo "    already exists ($((actual_bytes / 1024 / 1024 / 1024)) GiB) — skipping creation"
fi

echo "==> Step 2/6: activate swap and persist in fstab"
if ! swapon --show | grep -q "^$SWAPFILE "; then
    swapon "$SWAPFILE"
    echo "    swapon $SWAPFILE"
else
    echo "    already active — skipping swapon"
fi
if ! grep -qE "^$SWAPFILE\s" /etc/fstab; then
    printf '%s\tnone\tswap\tdefaults\t0 0\n' "$SWAPFILE" >> /etc/fstab
    echo "    appended to /etc/fstab"
else
    echo "    /etc/fstab already references $SWAPFILE — skipping"
fi

echo "==> Step 3/6: compute resume params"
RESUME_OFFSET=$(btrfs inspect-internal map-swapfile -r "$SWAPFILE")
echo "    resume=PARTUUID=$ROOT_PARTUUID"
echo "    resume_offset=$RESUME_OFFSET"

echo "==> Step 4/6: ensure 'resume' hook in /etc/mkinitcpio.conf"
if grep -qE "^HOOKS=\(.*\bresume\b.*\)" /etc/mkinitcpio.conf; then
    echo "    resume hook already present — skipping"
else
    sed -i -E 's/(^HOOKS=\(.*\bblock\b)/\1 resume/' /etc/mkinitcpio.conf
    if ! grep -qE "^HOOKS=\(.*\bresume\b.*\)" /etc/mkinitcpio.conf; then
        echo "ERROR: failed to insert 'resume' hook. Check /etc/mkinitcpio.conf manually." >&2
        exit 1
    fi
    echo "    inserted 'resume' after 'block'"
fi

echo "==> Step 5/6: append resume params to /etc/kernel/cmdline"
if grep -q "resume=" /etc/kernel/cmdline; then
    echo "    cmdline already contains resume= — skipping"
else
    sed -i "1 s|\$| resume=PARTUUID=$ROOT_PARTUUID resume_offset=$RESUME_OFFSET|" /etc/kernel/cmdline
    echo "    appended to /etc/kernel/cmdline"
fi

echo "==> Step 6/6: regenerate UKI"
mkinitcpio -P

echo
echo "Done. Reboot to activate hibernation:"
echo "    sudo reboot"
echo
echo "After reboot, test with: systemctl hibernate"