#!/bin/bash

##############################################
#### Create Swap for all Linux Servers    ####
####                                      ####
#### Author: Phil Connor                  ####
#### Contact: pconnor@ara.com             ####
#### Version 3.51.20250729                ####
####                                      ####
#### Created 06/01/2023                   ####
##############################################
# v3.51 changes:
#   - Fixed: grep in pipeline crashes under set -euo pipefail when no matches found. Added || true guard
##############################################

# Exit on any error, undefined variables, and pipe failures
set -euo pipefail

# Script configuration constants
readonly SCRIPT_NAME="$(basename "$0")"
readonly SWAPFILE_PATH="/.swapfile"          # Standard location for swap file
readonly SWAPPINESS_VALUE=80                 # How aggressively to use swap (0-100)

# Logging function - outputs to stderr with script name prefix
log() {
    echo "[$SCRIPT_NAME] $*" >&2
}

# Error function - logs error message and exits with status 1
error() {
    log "ERROR: $*"
    exit 1
}

# Display usage information
usage() {
    cat <<EOF
Usage: $SCRIPT_NAME [-h|--help]

Creates and configures a swap file sized 1:1 with system RAM.
Handles creation, resizing, and removal of existing swap files.

Options:
  -h, --help    Show this help message and exit

Must be run as root.
EOF
    exit 0
}

# Clean up partial swap file on unexpected failure
cleanup_on_error() {
    log "Error detected, cleaning up partial swap file"
    swapoff "$SWAPFILE_PATH" 2>/dev/null || true
    rm -f "$SWAPFILE_PATH"
}

# Detect the operating system distribution (ubuntu, centos, etc.)
detect_os() {
    if command -v lsb_release >/dev/null 2>&1; then
        # Use lsb_release if available (most reliable)
        lsb_release -i | awk '{print $3}' | tr '[:upper:]' '[:lower:]'
    else
        # Fallback to parsing /etc/os-release
        # shellcheck source=/dev/null
        . /etc/os-release 2>/dev/null && echo "${ID:-unknown}" | tr '[:upper:]' '[:lower:]'
    fi
}

# Get total system memory in GB, rounded to nearest whole number
get_memory_gb() {
    local mem_kb
    # Extract memory from /proc/meminfo (in KB)
    mem_kb=$({ grep MemTotal /proc/meminfo || true; } | awk '{print $2}')
    
    if [[ -z "$mem_kb" || "$mem_kb" -eq 0 ]]; then
        error "Unable to determine system memory"
    fi
    
    local mem_gb
    # Convert KB to GB and round to nearest whole number
    mem_gb=$(awk "BEGIN {printf \"%.0f\", ($mem_kb/1024/1024)}")
    # Ensure minimum of 1GB to avoid division by zero issues
    [[ "$mem_gb" -eq 0 ]] && mem_gb=1
    
    echo "$mem_gb"
}

# Calculate swap size needed in MB (1:1 ratio with RAM)
get_swap_needed_mb() {
    local mem_gb="$1"
    echo $((mem_gb * 1024))
}

# Get the current swap file size in MB, or 0 if no swap file exists
get_current_swap_size() {
    if [[ -f "$SWAPFILE_PATH" ]]; then
        local size_bytes
        size_bytes=$(stat -c%s "$SWAPFILE_PATH" 2>/dev/null || echo 0)
        echo $((size_bytes / 1024 / 1024))
    else
        echo 0
    fi
}

# Check if our swap file is currently active
is_swap_active() {
    swapon --show=NAME --noheadings 2>/dev/null | grep -q "^${SWAPFILE_PATH}$"
}

# Check if there's enough disk space for the swap file (with 10% buffer)
check_disk_space() {
    local needed_mb="$1"
    local filesystem="/"
    
    log "Checking available disk space for ${needed_mb}MB swap file"
    
    local available_kb
    # Get available space in KB from df command
    available_kb=$(df --output=avail "$filesystem" | tail -n 1)
    local available_mb=$((available_kb / 1024))
    
    # Add 10% buffer for safety
    local required_mb=$((needed_mb + (needed_mb / 10)))
    
    if [[ "$available_mb" -lt "$required_mb" ]]; then
        error "Insufficient disk space. Need ${required_mb}MB (${needed_mb}MB + 10% buffer), but only ${available_mb}MB available on $filesystem"
    fi
    
    log "Disk space check passed: ${available_mb}MB available, ${required_mb}MB required"
}

# Verify script is running with root privileges
check_permissions() {
    if [[ $EUID -ne 0 ]]; then
        error "This script must be run as root! Login as root, or use sudo."
    fi
}

# Configure system swappiness (how aggressively to use swap)
setup_swappiness() {
    local sysconf="/etc/sysctl.conf"
    local procswap="/proc/sys/vm/swappiness"
    
    log "Configuring swappiness to $SWAPPINESS_VALUE"
    
    # If no swappiness setting exists, add it
    if ! grep -q "vm.swappiness" "$sysconf"; then
        echo "$SWAPPINESS_VALUE" > "$procswap"
        echo "vm.swappiness = $SWAPPINESS_VALUE" >> "$sysconf"
    # If setting exists but with different value, update it
    elif ! grep -q "vm.swappiness = $SWAPPINESS_VALUE" "$sysconf"; then
        sed -i "/vm.swappiness/d" "$sysconf"
        echo "$SWAPPINESS_VALUE" > "$procswap"
        echo "vm.swappiness = $SWAPPINESS_VALUE" >> "$sysconf"
    fi
}

# Set up automated cache clearing cron job (every 5 minutes)
setup_cache_clearing() {
    local os="$1"
    local ctab
    
    # Different crontab locations for different distributions
    if [[ "$os" == "ubuntu" ]]; then
        ctab="/var/spool/cron/crontabs/root"
    else 
        ctab="/var/spool/cron/root"
    fi
    
    log "Setting up cache clearing cron job"
    
    # Remove any existing cache clearing jobs that use 'echo 3' (more aggressive)
    if crontab -l 2>/dev/null | grep -q '/usr/bin/sync; echo 3'; then
        sed -i "/\/usr\/bin\/sync.*echo 3/d" "$ctab" 2>/dev/null || true
    fi
    
    # Add cache clearing job if it doesn't exist (echo 1 = page cache only)
    if ! crontab -l 2>/dev/null | grep -q '/usr/bin/sync; echo 1'; then
        (crontab -u root -l 2>/dev/null; echo "*/5 * * * * /usr/bin/sync; echo 1 > /proc/sys/vm/drop_caches") | crontab -u root -
    fi
}

# Remove existing swap file and clean up fstab entries
remove_swap() {
    local backup_time
    
    # Create timestamp for backup file
    backup_time=$(date +%y-%m-%d--%H-%M-%S)
    
    log "Removing existing swap file: $SWAPFILE_PATH"
    
    # Disable swap file (ignore errors if already disabled)
    swapoff "$SWAPFILE_PATH" 2>/dev/null || true
    
    # Backup fstab before modifying
    cp /etc/fstab "/etc/fstab.$backup_time"
    
    # Remove swap entries from fstab
    sed -i "\|${SWAPFILE_PATH}|d" /etc/fstab
    
    # Delete the swap file
    rm -f "$SWAPFILE_PATH"
}

# Create and configure a new swap file
create_swap() {
    local swap_mb="$1"
    
    if [[ "$swap_mb" -eq 0 ]]; then
        error "Cannot create swap: swap size cannot be 0 MB"
    fi
    
    log "Creating swap file of size ${swap_mb}MB at $SWAPFILE_PATH"
    
    # Set trap to clean up partial swap file on failure
    trap cleanup_on_error ERR
    
    # Create swap file using dd with progress display (oflag=direct avoids polluting page cache)
    dd if=/dev/zero of="$SWAPFILE_PATH" bs=1M count="$swap_mb" oflag=direct status=progress
    
    # Set proper permissions (only root can read/write)
    chmod 600 "$SWAPFILE_PATH"
    
    # Format the file as swap space
    mkswap "$SWAPFILE_PATH"
    
    # Enable the swap file
    swapon "$SWAPFILE_PATH"
    
    # Add to fstab for persistent mounting if not already present
    if ! grep -q "$SWAPFILE_PATH" /etc/fstab; then
        echo "$SWAPFILE_PATH swap swap defaults 0 0" >> /etc/fstab
    fi
    
    # Clear the error trap now that swap is fully created
    trap - ERR
    
    log "Swap file created and enabled successfully"
}

# Main function - orchestrates the entire swap setup process
main() {
    # Handle --help flag
    if [[ "${1:-}" == "-h" || "${1:-}" == "--help" ]]; then
        usage
    fi
    
    # Ensure script is run with root privileges
    check_permissions
    
    # Detect operating system for distribution-specific configurations
    local os
    os=$(detect_os)
    
    # Get system memory information
    local mem_gb
    mem_gb=$(get_memory_gb)
    
    # Calculate required swap size
    local needed_mb
    needed_mb=$(get_swap_needed_mb "$mem_gb")
    
    # Check current swap configuration
    local current_size
    current_size=$(get_current_swap_size)
    
    # Configure system settings
    setup_swappiness
    setup_cache_clearing "$os"
    
    # If swap file exists at the correct size and is active, nothing to do
    if [[ "$current_size" -eq "$needed_mb" ]] && is_swap_active; then
        log "Swap size is already correct and active"
        log "Swap setup completed successfully"
        return 0
    fi
    
    # If swap file exists but wrong size, remove first so disk space check is accurate
    if [[ "$current_size" -ne 0 && "$needed_mb" -ne "$current_size" ]]; then
        remove_swap
    fi
    
    # Verify system has enough disk space (after potential removal)
    if [[ "$needed_mb" -ne "$current_size" ]]; then
        check_disk_space "$needed_mb"
        create_swap "$needed_mb"
    else
        # File is the right size but not active, re-enable it
        log "Swap file exists at correct size but is not active, enabling"
        chmod 600 "$SWAPFILE_PATH"
        mkswap "$SWAPFILE_PATH"
        swapon "$SWAPFILE_PATH"
        if ! grep -q "$SWAPFILE_PATH" /etc/fstab; then
            echo "$SWAPFILE_PATH swap swap defaults 0 0" >> /etc/fstab
        fi
    fi
    
    log "Swap setup completed successfully"
}

# Execute main function with all script arguments
main "$@"