#!/bin/bash -e
#
#  Copyright (C) 2015 ScyllaDB

GREEN='\033[0;32m'
NO_COLOR='\033[0m' # No Color

if [ "`id -u`" -ne 0 ]; then
    echo "Requires root permission."
    exit 1
fi

print_usage() {
    echo "scylla_setup --disks /dev/hda,/dev/hdb... --nic eth0 --ntp-domain centos --ami --setup-nic --developer-mode --no-kernel-check --no-verify-package --no-enable-service --no-selinux-setup --no-bootparam-setup --no-ntp-setup --no-raid-setup --no-coredump-setup --no-sysconfig-setup"
    echo "  --disks			specify disks for RAID"
    echo "  --nic				specify NIC"
    echo "  --ntp-domain			specify NTP domain"
    echo "  --ami				setup AMI instance"
    echo "  --setup-nic				optimize NIC queue"
    echo "  --developer-mode			enable developer mode"
    echo "  --no-kernel-check           skip kernel version check"
    echo "  --no-verify-package          skip verifying packages"
    echo "  --no-enable-service		skip enabling service"
    echo "  --no-selinux-setup		skip selinux setup"
    echo "  --no-bootparam-setup		skip bootparam setup"
    echo "  --no-ntp-setup		skip ntp setup"
    echo "  --no-raid-setup		skip raid setup"
    echo "  --no-coredump-setup		skip coredump setup"
    echo "  --no-sysconfig-setup		skip sysconfig setup"
    echo "  --no-io-setup		skip IO configuration setup"
    echo "  --no-version-check		skip daily version check"
    exit 1
}

interactive_ask_service() {
    echo $1
    echo $2
    while true; do
        if [ "$3" == "yes" ]; then
        prompt="[YES/no]"
        elif [ "$3" == "no" ]; then
        prompt="[yes/NO]"
        else
        prompt="[yes/no]"
        fi
        result=""
        while [ x == x"$result" ]; do
            read -p $prompt ans
            if [ x == x"$ans" ]; then
                result="$3"
            else
                result="$ans"
            fi
        done
        case $(echo $result | tr '[:upper:]' '[:lower:]') in
            "y" | "yes")
                return 1
                ;;
            "n" | "no")
                return 0
                ;;
        esac
    done
}

verify_package() {
    if [ "$ID" = "ubuntu" ]; then
        dpkg -s $1 > /dev/null 2>&1 &&:
    else
        rpm -q $1 > /dev/null 2>&1 &&:
    fi
    if [ $? -eq 1 ]; then
        echo "$1 package is not installed."
        exit 1
    fi
}

AMI=0
SET_NIC=0
DEV_MODE=0
KERNEL_CHECK=1
VERIFY_PACKAGE=1
ENABLE_SERVICE=1
ENABLE_CHECK_VERSION=1
SELINUX_SETUP=1
BOOTPARAM_SETUP=1
NTP_SETUP=1
RAID_SETUP=1
COREDUMP_SETUP=1
SYSCONFIG_SETUP=1
IO_SETUP=1

if [ $# -ne 0 ]; then
    INTERACTIVE=0
else
    INTERACTIVE=1
fi

while [ $# -gt 0 ]; do
    case "$1" in
        "--disks")
            DISKS="$2"
            shift 2
            ;;
        "--nic")
            NIC="$2"
            shift 2
            ;;
        "--ntp-domain")
            NTP_DOMAIN="$2"
            shift 2
            ;;
        "--ami")
            AMI=1
            shift 1
            ;;
        "--setup-nic")
            SET_NIC=1
            shift 1
            ;;
        "--developer-mode")
            DEV_MODE=1
            shift 1
            ;;
        "--no-kernel-check")
            KERNEL_CHECK=0
            shift 1
            ;;
        "--no-verify-package")
            VERIFY_PACKAGE=0
            shift 1
            ;;
        "--no-enable-service")
            ENABLE_SERVICE=0
            shift 1
            ;;
        "--no-version-check")
            ENABLE_CHECK_VERSION=0
            shift 1
            ;;
        "--no-selinux-setup")
            SELINUX_SETUP=0
            shift 1
            ;;
        "--no-bootparam-setup")
            BOOTPARAM_SETUP=0
            shift 1
            ;;
        "--no-ntp-setup")
            NTP_SETUP=0
            shift 1
            ;;
        "--no-raid-setup")
            RAID_SETUP=0
            shift 1
            ;;
        "--no-coredump-setup")
            COREDUMP_SETUP=0
            shift 1
            ;;
        "--no-sysconfig-setup")
            SYSCONFIG_SETUP=0
            shift 1
            ;;
        "--no-io-setup")
            IO_SETUP=0
            shift 1
            ;;
        "-h" | "--help")
            print_usage
            shift 1
            ;;
    esac
done

if [ $INTERACTIVE -eq 0 ] && [ $RAID_SETUP -eq 1 ] && [ "$DISKS" = "" ]; then
    print_usage
fi
if [ $INTERACTIVE -eq 0 ] && [ $SYSCONFIG_SETUP -eq 1 ] && [ "$NIC" = "" ]; then
    print_usage
fi

. /etc/os-release

printf "${GREEN}Skip any of the following steps by answering 'no'${NO_COLOR}\n"

if [ $INTERACTIVE -eq 1 ]; then
    interactive_ask_service "Do you want to run kernel version check?" "Answer yes to have this script verify that the currently installed kernel is qualified to run Scylla; answer no to skip this check." "yes" &&:
    KERNEL_CHECK=$?
fi
if [ $KERNEL_CHECK -eq 1 ]; then
    /usr/lib/scylla/scylla_kernel_check
fi

if [ $INTERACTIVE -eq 1 ]; then
    interactive_ask_service "Do you verify ScyllaDB packages installed?" "Answer yes to have this script check that ScyllaDB is already installed; answer no to skip this check." "yes" &&:
    VERIFY_PACKAGE=$?
fi

if [ $VERIFY_PACKAGE -eq 1 ]; then
    verify_package scylla-jmx
    verify_package scylla-tools
fi

if [ $INTERACTIVE -eq 1 ]; then
    interactive_ask_service "Do you want to enable ScyllaDB services?" "Answer yes to automatically start Scylla when the node boots; answer no to skip this step." "yes" &&:
    ENABLE_SERVICE=$?
    if [ $ENABLE_SERVICE -eq 1 ] && [ ! -f /etc/scylla.d/housekeeping.cfg ]; then
        interactive_ask_service "Do you want to enable ScyllaDB version check?" "Answer yes to automatically start Scylla-housekeeping that check for newer version, when the node boots; answer no to skip this step." "yes" &&:
        ENABLE_CHECK_VERSION=$?
    fi
fi

if [ $ENABLE_SERVICE -eq 1 ]; then
    if [ "$ID" = "fedora" ] || [ "$ID" = "centos" ]; then
        systemctl enable scylla-server.service
        systemctl enable collectd.service
        if [ $ENABLE_CHECK_VERSION -eq 1 ]; then
            systemctl unmask scylla-housekeeping.timer
        else
            systemctl mask scylla-housekeeping.timer
            systemctl stop scylla-housekeeping.timer || true
        fi
    fi
    if [ $ENABLE_CHECK_VERSION -eq 1 ]; then
        if [ ! -f /etc/scylla.d/housekeeping.cfg ]; then
           printf "[housekeeping]\ncheck-version: True\n" > /etc/scylla.d/housekeeping.cfg
        fi
    else
        if [ ! -f /etc/scylla.d/housekeeping.cfg ]; then
           printf "[housekeeping]\ncheck-version: False\n" > /etc/scylla.d/housekeeping.cfg
        fi
    fi
fi

if [ $INTERACTIVE -eq 1 ]; then
    interactive_ask_service "Do you want to disable SELinux?" "Answer yes to disable SELinux; this improves performance; answer no to keep it activated." "yes" &&:
    SELINUX_SETUP=$?
fi
if [ $SELINUX_SETUP -eq 1 ]; then
    /usr/lib/scylla/scylla_selinux_setup
fi

if [ $INTERACTIVE -eq 1 ]; then
    interactive_ask_service "Do you want to setup bootloader options?" "Anwser yes to enable hugepages at boot time; this improve performance. Answer no to skip this step." "yes" &&:
    BOOTPARAM_SETUP=$?
fi
if [ $BOOTPARAM_SETUP -eq 1 ]; then
    /usr/lib/scylla/scylla_bootparam_setup
fi

if [ $INTERACTIVE -eq 1 ]; then
    interactive_ask_service "Do you want to setup NTP?" "Answer yes to enable time synchronization at boot time. This keeps time right on the node. Answer no to do nothing." "yes" &&:
    NTP_SETUP=$?
fi
if [ $NTP_SETUP -eq 1 ]; then
    if [ "$NTP_DOMAIN" != "" ]; then
        /usr/lib/scylla/scylla_ntp_setup --subdomain $NTP_DOMAIN
    else
        /usr/lib/scylla/scylla_ntp_setup
    fi
fi

if [ $INTERACTIVE -eq 1 ]; then
    interactive_ask_service "Do you want to setup RAID and XFS?" "It is recommended to use RAID0 and XFS for Scylla data. If you select yes, you will be prompt to choose which disks to use for Scylla data. Selected disks will be formatted in the process." "yes" &&:
    RAID_SETUP=$?
    if [ $RAID_SETUP -eq 1 ]; then
        echo "Please select disks from following list: "
        while true; do
            blkid -c /dev/null|cut -f 1 -d ' '|sed s/://g|grep -v loop|while read dev
            do
                count_raw=$(grep $dev /proc/mounts|wc -l)
                count_pvs=0
                if [ -f /usr/sbin/pvs ]; then
                    count_pvs=$(pvs|grep $dev|wc -l)
                fi
                count_swap=$(swapon --show |grep `realpath $dev`|wc -l)
                if [ $count_raw -eq 0 -a $count_pvs -eq 0 -a $count_swap -eq 0 ]; then
                    echo -n "$dev "
                fi
            done

            echo "type 'done' to finish selection. selected: $DISKS"
            echo -n "> "
            read dsk
            if [ "$dsk" = "done" ]; then
                if [ "$DISKS" = "" ]; then
                    continue
                fi
                break
            fi
            if [ "$dsk" = "" ]; then
                continue
            fi
            if [ -b $dsk ]; then
                if [ "$DISKS" = "" ]; then
                    DISKS=$dsk
                else
                    DISKS="$DISKS,$dsk"
                fi
            else
                echo "$dsk not found"
            fi
        done
    fi
fi
if [ $RAID_SETUP -eq 1 ]; then
    /usr/lib/scylla/scylla_raid_setup --disks $DISKS --update-fstab
fi

if [ $INTERACTIVE -eq 1 ]; then
    interactive_ask_service "Do you want to setup coredump?" "Answer yes to enable core dumps; this allow to do post mortem analysis of Scylla state after a crash. Answer no to do nothing." "yes" &&:
    COREDUMP_SETUP=$?
fi
if [ $COREDUMP_SETUP -eq 1 ]; then
    if [ "$DISKS" != "" ]; then
        /usr/lib/scylla/scylla_coredump_setup --dump-to-raiddir
    else
        /usr/lib/scylla/scylla_coredump_setup
    fi
fi

if [ $INTERACTIVE -eq 1 ]; then
    interactive_ask_service "Do you want to setup sysconfig?" "Answer yes to do system wide configuration customized for Scylla. Answer no to do nothing." "yes" &&:
    SYSCONFIG_SETUP=$?
    if [ $SYSCONFIG_SETUP -eq 1 ]; then
        NICS=$(for i in /sys/class/net/*;do nic=`basename $i`; if [ "$nic" != "lo" ]; then echo $nic; fi; done)
        NR_NICS=`echo $NICS|wc -w`
        if [ $NR_NICS -eq 0 ]; then
            echo "NIC not found."
            exit 1
        elif [ $NR_NICS -eq 1 ]; then
            NIC=$NICS
        else
            echo "Please select NIC from following list: "
            while true; do
                echo $NICS
                echo -n "> "
                read NIC
                if [ -e /sys/class/net/$NIC ]; then
                    break
                fi
            done
        fi
        interactive_ask_service "Do you want to optimize NIC queue settings?" "Answer yes to enable network card optimization; this improve performance. Answer no to skip this optimization." "yes" &&:
        SET_NIC=$?
    fi
fi
if [ $SYSCONFIG_SETUP -eq 1 ]; then
    SETUP_ARGS=
    if [ $SET_NIC -eq 1 ]; then
        SETUP_ARGS="--setup-nic"
    fi
    /usr/lib/scylla/scylla_sysconfig_setup --nic $NIC $SETUP_ARGS
fi

if [ $INTERACTIVE -eq 1 ]; then
    interactive_ask_service "Do you want to setup IO configuration?" "Answer yes to let iotune study what are your disks IO profile and adapt Scylla to it. Answer no to skip this action." "yes" &&:
    IO_SETUP=$?
fi
if [ $IO_SETUP -eq 1 ]; then
    /usr/lib/scylla/scylla_io_setup
fi

if [ $DEV_MODE -eq 1 ]; then
    /usr/lib/scylla/scylla_dev_mode_setup --developer-mode 1
fi

echo "ScyllaDB setup finished."
if [ $BOOTPARAM_SETUP -eq 1 ]; then
    echo "Please restart machine before using ScyllaDB, since you have changed "
    echo "boot parameters."
fi