#!/bin/bash

#NVIDIA Jetson Nano TX1 TX2 TX2i
# Create a swap file and set up permissions
# If a parameter is passed, it should be the place to create the swapfile

FSTAB="/etc/fstab"
SWAPDIRECTORY="$PWD"
SWAPSIZE=8
NAME_SWAP="swapfile"
AUTOMOUNT="N"
SWAP_OFF="N"
SWAP_STATUS="N"


function usage
{
    echo "usage: createSwapFile [[[-d directory ] [-s size] -a] | [-h] | [--off]]"
    echo "  -d | --dir    <directoryname> Directory to place swapfile"
    echo "  -n | --name   <swapname> Name swap file"
    echo "  -s | --size   <gigabytes>"
    echo "  -a | --auto   Enable swap on boot in $FSTAB "
    echo "  -t | --status Check if the swap is currently active"
    echo "  --off         Switch off the swap"
    echo "  -h | --help   This message"
}

while [ "$1" != "" ]; do
    case $1 in
        -d | --dir )            shift
                                SWAPDIRECTORY=$1
                                ;;
        -n | --name )           shift
                                NAME_SWAP=$1
                                ;;
        -s | --size )           shift 
                                SWAPSIZE=$1
                                ;;
        -a | --auto )           AUTOMOUNT="Y"
                                ;;
        -t | --status )         SWAP_STATUS="Y"
                                ;;
        --off )                 SWAP_OFF="Y"
                                ;;
        -h | --help )           usage
                                exit
                                ;;
        * )                     usage
                                exit 1
    esac
    shift
done

SWAPLOCATION="$SWAPDIRECTORY/$NAME_SWAP"

if [ "$SWAP_STATUS" = "Y" ]; then
    line_found=$(swapon --show=NAME | grep -x -n "$SWAPLOCATION" | cut -d ':' -f1)
    if [ -z "$line_found" ] ; then
        exit 0
    fi
    # Extract the line where is swap
    counter=1
    while read -r line
    do
        if [[ "$counter" = "$line_found" ]]; then
            echo "$line"
            exit 0
        fi
        counter=$((counter+1))
    done < <(swapon -s)
    exit 0
fi

if [[ $EUID != 0 ]]; then
    tput setaf 1
    echo "Launch with sudo!"
    tput sgr0
    
    usage
    
    exit 1
fi

# Switch off the swapfile
if [ "$SWAP_OFF" = "Y" ]; then
    # Remove from FSTAB if exist
    pwdesc=$(echo $SWAPLOCATION | sed 's_/_\\/_g')
    if grep -Fxq "$SWAPLOCATION none swap sw 0 0" $FSTAB ; then
        echo "Remove swapfile from $FSTAB"
        sed -i "/$pwdesc none swap sw 0 0/d" $FSTAB
    fi
    if [ -f "$SWAPLOCATION" ]; then
        echo "Remove swapfile in $SWAPDIRECTORY"
        sudo swapoff "$SWAPLOCATION"
        sudo rm "$SWAPLOCATION"
        exit 0
    else
        echo "Swap file does not exist in $SWAPDIRECTORY"
        exit 1
    fi
fi

echo "Creating Swapfile at: " $SWAPDIRECTORY
echo "Swapfile Size: " $SWAPSIZE"G"
echo "Automount: " $AUTOMOUNT

#Create a swapfile for Ubuntu at the current directory location
fallocate -l $SWAPSIZE"G" $SWAPLOCATION
cd $SWAPDIRECTORY
#List out the file
ls -lh $SWAPLOCATION
# Change permissions so that only root can use it
sudo chmod 600 $SWAPLOCATION
#List out the file
ls -lh $SWAPLOCATION
#Set up the Linux swap area
sudo mkswap $SWAPLOCATION
#Now start using the swapfile
sudo swapon $SWAPLOCATION
#Show that it's now being used
swapon -s

if [ "$AUTOMOUNT" = "Y" ]; then
    if ! grep -Fxq "$SWAPLOCATION none swap sw 0 0" $FSTAB ; then
        echo "Modifying $FSTAB to enable on boot"    
        sudo sh -c 'echo "'$SWAPLOCATION' none swap sw 0 0" >> '$FSTAB''
        echo "Added in $FSTAB \"$SWAPLOCATION none swap sw 0 0\""
    else
        echo "Swap already in $FSTAB"
    fi
fi
