#!/bin/bash
#
# Pre-systemd-swap setup script
# Prepares btrfs subvolumes and disables existing swap partitions
#

set -euo pipefail

# =============================================================================
# Helper Functions
# =============================================================================

die() {
    echo "ERROR: $1" >&2
    exit 1
}

log_info() {
    echo "INFO: $1"
}

# =============================================================================
# Configuration Parsing
# =============================================================================

CONFIG_FILE="/etc/systemd/swap.conf"

if [[ ! -f "$CONFIG_FILE" ]]; then
    die "Configuration file not found: $CONFIG_FILE"
fi

# Extract swapfc_path folder name (second component of path).
# grep returns exit 1 when no match; || true prevents set -e from aborting.
swapfcfolder=$(grep "^swapfc_path=" "$CONFIG_FILE" | cut -d'/' -f2 || true)

if [[ -z "$swapfcfolder" ]]; then
    # Default path mirrors the default in config.rs
    swapfcfolder="swapfile"
fi

# =============================================================================
# System Detection
# =============================================================================

# Get root partition device
rootpartition=$(df --output=source / | tail -n 1)

if [[ -z "$rootpartition" ]]; then
    die "Could not determine root partition"
fi

# Get filesystem type
fstype=$(lsblk -fno FSTYPE "$rootpartition" 2>/dev/null || echo "unknown")

# Check if path is a btrfs subvolume and get top level id
toplevelid=""
if [[ "$fstype" == "btrfs" ]] && [[ -d "/$swapfcfolder" ]]; then
    toplevelid=$(btrfs subvolume show "/$swapfcfolder" 2>/dev/null | grep -i "top level id" | awk '{print $4}' || echo "")
fi

# =============================================================================
# Btrfs Subvolume Setup
# =============================================================================

if [[ "$toplevelid" != "5" ]] && [[ "$fstype" == "btrfs" ]]; then
    log_info "Setting up btrfs subvolume for swap files"

    # Mount root partition on /mnt
    if ! mount "$rootpartition" /mnt; then
        die "Failed to mount $rootpartition on /mnt"
    fi

    # Create subvolume if it doesn't exist
    if [[ ! -e "/mnt/@${swapfcfolder}" ]]; then
        log_info "Creating subvolume @$swapfcfolder"
        if ! btrfs subvolume create "/mnt/@$swapfcfolder"; then
            umount /mnt
            die "Failed to create btrfs subvolume"
        fi
    fi

    # Add to fstab if not already present
    if ! grep -q "@${swapfcfolder}" /etc/fstab; then
        log_info "Adding subvolume to /etc/fstab"
        uuid=$(blkid -s UUID -o value "$rootpartition")
        if [[ -n "$uuid" ]]; then
            echo "UUID=$uuid /$swapfcfolder btrfs subvol=/@$swapfcfolder,defaults,noatime 0 0" >>/etc/fstab
        else
            die "Could not get UUID for $rootpartition"
        fi
    fi

    # Unmount root partition
    umount /mnt

    # Reload systemd to pick up fstab changes
    systemctl daemon-reload

    # Create mount point if needed
    if [[ ! -e "/$swapfcfolder" ]]; then
        mkdir -p "/$swapfcfolder"
    fi

    # Mount the subvolume
    if ! mount "/$swapfcfolder"; then
        die "Failed to mount /$swapfcfolder"
    fi

    log_info "Btrfs subvolume setup complete"
fi

# =============================================================================
# Disable Existing Swap Partitions
# =============================================================================

# Find swap partitions (type 8x devices with SWAP)
while IFS= read -r partswap; do
    if [[ -n "$partswap" ]]; then
        log_info "Disabling swap partition: /dev/$partswap"
        swapoff "/dev/$partswap" 2>/dev/null || true
    fi
done < <(lsblk -lno NAME,FSTYPE | awk '$2 == "swap" {print $1}')

# =============================================================================
# Dynamic Kernel Tuning
# =============================================================================

# Calculate min_free_kbytes: 3% of RAM, capped at 512MB (524288 KB).
# Gives kswapd enough headroom to reclaim before OOM.
total_ram_kb=$(awk '/^MemTotal:/ {print $2}' /proc/meminfo)
if [[ -n "$total_ram_kb" ]]; then
    min_free=$((total_ram_kb * 3 / 100))
    if [[ "$min_free" -gt 524288 ]]; then
        min_free=524288
    fi
    log_info "Setting vm.min_free_kbytes = $min_free (3% of ${total_ram_kb}KB)"
    echo "$min_free" > /proc/sys/vm/min_free_kbytes
fi

log_info "Pre-systemd-swap setup complete"
