#!/bin/bash

SWAPPATH=/
SWAPNAME="swapfile"
TOTALRAM=$(free -m | gawk  '/Mem:/{print $2}')
SWAPPINESS="10"
SYSCTLCONF="sysctl.conf"
SYSCTLDIR="sysctl.conf.d"
VERSION="0.1.0"
BUILDDATE="2018-06-05"

# Colors
RED='\033[0;31m'
LRED="\033[1;31m"
BLUE="\033[0;34m"
LBLUE="\033[1;34m"
GREEN="\033[0;32m"
LGREEN="\033[1;32m"
YELLOW="\033[1;33m"
CYAN="\033[0;36m"
LCYAN="\033[1;36m"
PURPLE="\033[0;35m"
LPURPLE="\033[1;35m"
BWHITE="\e[1m"
NC='\033[0m' # No Color

if [ ! -f /usr/bin/gawk ];then
    echo -e "\n${LRED}Please install 'gawk' with your package manager${NC}\n"
    exit 0
fi
function autoSwap(){
        if [[ $EUID -ne 0 ]]; then
            printf "${LRED}This script must be run as root${NC}\n" 1>&2
            exit 1
        fi

        echo -e "\n${LGREEN}----------------------------------------------------------------${NC}"
        echo -e "\n${LCYAN}Total Physical RAM is: ${TOTALRAM} MB\n${NC}"
        echo -e "${LPURPLE}Creating a ${TOTALRAM} MB file in:${SWAPPATH}..\n${NC}"
        echo -e "${YELLOW}Enter to Accept | Ctrl+c to cancel\n${NC}"
        echo -e "${LGREEN}------------------------------------------------------------------${NC}"
        read

        echo -e "${LCYAN}\nCreating file ${SWAPPATH}${SWAPNAME}...\n${NC}"
        dd if=/dev/zero of=${SWAPPATH}${SWAPNAME} count=${TOTALRAM} bs=1MiB
        ls -lh ${SWAPPATH}${SWAPNAME}

        sleep 1

        echo -e "\n${LCYAN}Setting permission of ${SWAPPATH}${SWAPNAME} to 600...\n${NC}"
        chmod 600 ${SWAPPATH}${SWAPNAME}
        ls -lh ${SWAPPATH}${SWAPNAME}

        sleep 1
        
        echo -e -n "\n${YELLOW}Do you want to mark ${SWAPPATH}${SWAPNAME} as swap now?(Y/n): ${NC}"
        read markswap

        if [ "${markswap}" = "n" ] || [ "${markswap}" = "N" ];then
            echo "\n${LCYAN}Not making ${SWAPPATH}/${SWAPFILE} as swap...${NC}\n"
        fi

        echo -e "\n${LCYAN}Marking ${SWAPPATH}${SWAPNAME} as swap...\n\n${NC}"
        mkswap ${SWAPPATH}${SWAPNAME}

        echo -e "${LCYAN}Enabling Swap file for use...${NC}\n"
        swapon ${SWAPPATH}${SWAPNAME}

        echo -e -n "\n${YELLOW}Do you want to add swapfile to fstab?(Y/n): ${NC}"
        read addfstab

        if [ "${addfstab}" = "n" ];then
            echo -e "\n${LCYAN} Not adding ${SWAPPATH}${SWAPNAME} to fstab${NC}\n"
        fi

        echo -e "\n${LCYAN}Creating backup of /etc/fstab ...\n${NC}"
        cp -v /etc/fstab /etc/fstab.bak

        echo -e "${LCYAN}Adding ${SWAPPATH}${SWAPNAME} to /etc/fstab ..\n${NC}"
        echo "${SWAPPATH}${SWAPNAME} none swap sw 0 0" | sudo tee -a /etc/fstab

        CHECKFSTAB=$(cat /etc/fstab | grep "swap" | gawk  '/swap/{print $1}')

        if [ "${CHECKFSTAB}" != "/swapfile" ];then
            echo -e "${LRED}Swap line does not exist in /etc/fstab, please manually check this.${NC}\n"

        else
            echo -e "${LGREEN}Swap successfully added to /etc/fstab! Reboot to verify success...${NC}\n"

            echo -e "${YELLOW}Do you want to reboot your system now?(n/Y): ${NC}"
            read rebootCheck

            if [ "${rebootCheck}" = "y" ] || [ "${rebootCheck}" = "Y" ];then
                echo -e "${LRED}Rebooting system in 3 seconds...${NC}\n"

                sleep 3

                sudo reboot
            else
                echo -e "\n${LGREEN}Swapfile setup complete, setting swappiness level...${NC}\n"
            fi
        fi
}

function setSwappinessLvl(){
    if [[ $EUID -ne 0 ]]; then
        printf "${LRED}This script must be run as root${NC}\n" 1>&2
        exit 1
    fi

    echo -e -n "\n${YELLOW}Enter a number for swappiness (0-100)[default=10]: ${NC}"
    read setSwappiness

    if [ -z "${setSwappiness}" ];then
        echo -e "\n${LCYAN}Setting default (10)\n${NC}"

        if [ -d "/etc/${SYSCTLDIR}" ];then

            SYSCTLCONF="swappiness.conf"
            touch /etc/${SYSCTLDIR}/${SYSCTLCONF}
            echo "vm.swappiness=${SWAPPINESS}" | sudo tee -a /etc/${SYSCTLDIR}/${SYSCTLCONF}

        else

            echo "vm.swappiness=${SWAPPINESS}" | sudo tee -a /etc/${SYSCTLCONF}
        fi
    

    elif [[ ! "${setSwappiness}" =~ ^-?[0-9]+$ ]];then
        echo -e "\n${LRED}${setSwappiness} is not a number\n${NC}"
        exit 0
    
    elif [ "${setSwappiness}" -lt "0" ];then
        echo -e "\n${LRED}Please input a number greater than 0 or less than 100\n${NC}"
        exit 0
    
    elif [ "${setSwappiness}" -gt "100" ];then
        echo -e -n "\n${LRED}Please input a number greater than 0 or less than 100\n${NC}"
        exit 0

    else
        if [ -d "/etc/${SYSCTLDIR}" ];then

            SYSCTLCONF="swappiness.conf"
            USINGSYSCTLDIR="true"
            touch /etc/${SYSCTLDIR}/${SYSCTLCONF}
            echo "vm.swappiness=${setSwappiness}" > /etc/${SYSCTLDIR}/${SYSCTLCONF}

        else

            echo "vm.swappiness=${setSwappiness}" >> /etc/${SYSCTLCONF}
        fi


        if [ "${USINGSYSCTLDIR}" = "true" ];then

            SWAPPINESSSET=$(cat /etc/${SYSCTLDIR}/${SYSCTLCONF} | grep vm.swappiness=${setSwappiness})

        else

            SWAPPINESSSET=$(cat /etc/${SYSCTLCONF} | grep vm.swappiness=${setSwappiness})

        fi

        if [[ ${SWAPPINESSSET} = "vm.swappiness=${setSwappiness}" ]] ;then
            echo -e "${LGREEN}Swappiness value set successfully!\n${NC}"
        else
            echo -e "${LRED}Swappiness not set sucessfully...\n${NC}"
        fi
    fi
}

liveSwappiness(){
    CURRENTSWAPPINESS=$(cat /proc/sys/vm/swappiness)
    echo -e "\n${LCYAN}Current swappiness value: (reboot to apply any changes you set):${NC}\n"
    echo -e "${LGREEN}${CURRENTSWAPPINESS}${NC}\n"
}

function help(){
    echo -e "\nUsage: $0 [-a] [--auto] [-s] [--swappiness]\n\n"
    echo -e "-a | --auto                Automatically create and size swapfile based on physical memory\n"
    echo -e "-ss | --set-swappiness     Set the swappiness value to minimize swap usage\n\n"
    echo -e "-ls | --live-swappiness    View current swappiness level"
}

function version(){
    echo -e "Version ${VERSION}, Build Date: ${BUILDDATE}\n"
}

case "${1}" in

    -a|--auto)
        autoSwap

        while true;do
            echo -e "\n${YELLOW}Do you want to set the swappiness level?(y/n):${NC} "
            read swappinessLevel

            if [ "${swappinessLevel}" = "y" ];then
                setSwappinessLvl
                break && exit 0;

            elif [ "${swappinessLevel}" = "n" ];then
                break && exit 0;

            else
                echo -e "\n${LRED}Please enter 'y' or 'n'${NC}\n"
            fi
        done 
    ;;

    -v|version|--version)
        version
    ;;

    -ss|--set-swappiness)
        setSwappinessLvl
    ;;

    -ls|--live-swappiness)
        liveSwappiness
    ;;

    *)
        help
    ;;
esac