#!/usr/bin/env bash

# Copyright 2016 Midokura SARL
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

# Useful common functions are defined here.  Some of them were shamelessly
# stolen from devstack.

# Save trace setting
XTRACE=$(set +o | grep xtrace)
set +o xtrace

# Control Functions
# -----------------

# Prints line number and "message" in warning format
# warn $LINENO "message"
function warn {
    local exitcode=$?
    local xtrace=$(set +o | grep xtrace)
    set +o xtrace
    local msg="[WARNING] ${BASH_SOURCE[2]}:$1 $2"
    echo $msg 1>&2;
    if [[ -n ${LOGDIR} ]]; then
        echo $msg >> "${LOGDIR}/error.log"
    fi
    $xtrace
    return $exitcode
}

# Prints backtrace info
# filename:lineno:function
# backtrace level
function backtrace {
    local level=$1
    local deep=$((${#BASH_SOURCE[@]} - 1))
    echo "[Call Trace]"
    while [ $level -le $deep ]; do
        echo "${BASH_SOURCE[$deep]}:${BASH_LINENO[$deep-1]}:${FUNCNAME[$deep-1]}"
        deep=$((deep - 1))
    done
}

# Prints line number and "message" in error format
# err $LINENO "message"
function err {
    local exitcode=$?
    local xtrace=$(set +o | grep xtrace)
    set +o xtrace
    local msg="[ERROR] ${BASH_SOURCE[2]}:$1 $2"
    echo $msg 1>&2;
    if [[ -n ${LOGDIR} ]]; then
        echo $msg >> "${LOGDIR}/error.log"
    fi
    $xtrace
    return $exitcode
}

# Prints line number and "message" then exits
# die $LINENO "message"
function die {
    local exitcode=$?
    set +o xtrace
    local line=$1; shift
    if [ $exitcode == 0 ]; then
        exitcode=1
    fi
    backtrace 2
    err $line "$*"
    # Give buffers a second to flush
    sleep 1
    exit $exitcode
}

# Checks an environment variable is not set or has length 0 OR if the
# exit code is non-zero and prints "message" and exits
# NOTE: env-var is the variable name without a '$'
# die_if_not_set $LINENO env-var "message"
function die_if_not_set {
    local exitcode=$?
    local xtrace=$(set +o | grep xtrace)
    set +o xtrace
    local line=$1; shift
    local evar=$1; shift
    if ! is_set $evar || [ $exitcode != 0 ]; then
        die $line "$*"
    fi
    $xtrace
}

# Test if the named environment variable is set and not zero length
# is_set env-var
function is_set {
    local var=\$"$1"
    eval "[ -n \"$var\" ]" # For ex.: sh -c "[ -n \"$var\" ]" would be better, but several exercises depends on this
}


# Package Functions
# -----------------

# Wrapper for ``apt-get``
# apt_get operation package [package ...]
function apt_get {
    sudo DEBIAN_FRONTEND=noninteractive \
        apt-get --option "Dpkg::Options::=--force-confnew" --assume-yes "$@"
}

# Update package repository
# Uses globals ``NO_UPDATE_REPOS``, ``REPOS_UPDATED``, ``RETRY_UPDATE``
# install_package package [package ...]
function update_package_repo {
    NO_UPDATE_REPOS=${NO_UPDATE_REPOS:-False}
    REPOS_UPDATED=${REPOS_UPDATED:-False}

    if [[ "$NO_UPDATE_REPOS" = "True" ]]; then
        return 0
    fi

    local xtrace=$(set +o | grep xtrace)
    set +o xtrace
    if [[ "$REPOS_UPDATED" != "True" ]]; then
        # if there are transient errors pulling the updates, that's fine.
        # It may be secondary repositories that we don't really care about.
        apt_get update  || /bin/true
        REPOS_UPDATED=True
    fi
    $xtrace
}

# Package installer
# install_package package [package ...]
function install_package {
    update_package_repo
    apt_get install --no-install-recommends "$@"
}

# Function to tell if a package is installed
# is_package_installed package [package ...]
function is_package_installed {
    dpkg -s "$@" > /dev/null 2> /dev/null
}


# System Functions
# -----------------

# Service wrapper to stop services
# stop_service service-name
# Checks if kmod is loaded
function is_kmod_loaded {
    lsmod | grep -w $1 >& /dev/null
}


# Process Functions
# -----------------

function is_screen_running {
    type -p screen > /dev/null && screen -ls | egrep -q "[0-9]\.$1"
}

function create_screen {
    local name=$1
    screen -d -m -S $name -t shell -s /bin/bash
    sleep 1

    # Set a reasonable status bar
    SCREEN_HARDSTATUS='%{= .} %-Lw%{= .}%> %n%f %t*%{= .}%+Lw%< %-=%{g}(%{d}%H/%l%{g})'
    screen -r $name -X hardstatus alwayslastline "$SCREEN_HARDSTATUS"
    screen -r $name -X setenv PROMPT_COMMAND /bin/true
}


# Helper to launch a process in a named screen
# Uses globals ``CURRENT_LOG_TIME``, ``LOGDIR``,
# ``SERVICE_DIR``
# screen_process name "command-line"
# Run a command in a shell in a screen window
function screen_process {
    local name=$1
    local command="$2"

    SERVICE_DIR=${SERVICE_DIR:-/tmp/status}
    mkdir -p ${SERVICE_DIR}/${SCREEN_NAME}

    # Append the process to the screen rc file
    screen_rc ${SCREEN_NAME} "$name" "$command"

    screen -S ${SCREEN_NAME} -X screen -t $name

    if [[ -n ${LOGDIR} ]]; then
        screen -S ${SCREEN_NAME} -p $name -X logfile ${LOGDIR}/${name}.log.${CURRENT_LOG_TIME}
        screen -S ${SCREEN_NAME} -p $name -X log on
        ln -sf ${LOGDIR}/${name}.log.${CURRENT_LOG_TIME} ${LOGDIR}/${name}.log
    fi

    # sleep to allow bash to be ready to be send the command - we are
    # creating a new window in screen and then sends characters, so if
    # bash isn't running by the time we send the command, nothing happens
    sleep 3

    NL=`echo -ne '\015'`
    screen -S ${SCREEN_NAME} -p $name -X stuff "$command & echo \$! >$SERVICE_DIR/${SCREEN_NAME}/${name}.pid; fg || echo \"$name failed to start\" | tee \"$SERVICE_DIR/${SCREEN_NAME}/${name}.failure\"$NL"
}

# _run_process() is designed to be backgrounded by run_process() to simulate a
# fork.  It includes the dirty work of closing extra filehandles and preparing log
# files to produce the same logs as screen_it().  The log filename is derived
# from the service name.
# _run_process service "command-line"
function _run_process {
    local service=$1
    local command="$2"

    # Undo logging redirections and close the extra descriptors
    exec 1>&3
    exec 2>&3
    exec 3>&-
    exec 6>&-

    local real_logfile="${LOGDIR}/${service}.log.${CURRENT_LOG_TIME}"
    if [[ -n ${LOGDIR} ]]; then
        exec 1>&"$real_logfile" 2>&1
        ln -sf "$real_logfile" ${LOGDIR}/${service}.log

        # Hack to get stdout from the Python interpreter for the logs.
        export PYTHONUNBUFFERED=1
    fi

    SERVICE_DIR=${SERVICE_DIR:-/tmp/status}
    mkdir -p ${SERVICE_DIR}/${SCREEN_NAME}
    setsid $command & echo $! > ${SERVICE_DIR}/${SCREEN_NAME}/${service}.pid

    # Just silently exit this process
    exit 0
}

# Run a single service under screen or directly
# If the command includes shell metachatacters (;<>*) it must be run using a shell
# run_process service "command-line"
function run_process {
    local service=$1
    local command="$2"

    if [[ "$USE_SCREEN" = "True" ]]; then
        screen_process "$service" "cd $TOP_DIR && $command"
    else
        # Spawn directly without screen
        cd $TOP_DIR
        _run_process "$service" "$command" &
        cd -
    fi
}

# Screen rc file builder
# Uses globals ``SCREENRC``
# screen_rc service "command-line"
function screen_rc {
    local screen=$1
    SCREENRC=$DEVMIDO_DIR/$screen-screenrc
    if [[ ! -e $SCREENRC ]]; then
        # Name the screen session
        echo "sessionname $screen" > $SCREENRC
        # Set a reasonable statusbar
        echo "hardstatus alwayslastline '$SCREEN_HARDSTATUS'" >> $SCREENRC
        # Some distributions override PROMPT_COMMAND for the screen terminal type - turn that off
        echo "setenv PROMPT_COMMAND /bin/true" >> $SCREENRC
        echo "screen -t shell bash" >> $SCREENRC
    fi
    # If this service doesn't already exist in the screenrc file
    if ! grep $1 $SCREENRC 2>&1 > /dev/null; then
        NL=`echo -ne '\015'`
        echo "screen -t $1 bash" >> $SCREENRC
        echo "stuff \"$2$NL\"" >> $SCREENRC

        if [[ -n ${LOGDIR} ]]; then
            echo "logfile ${LOGDIR}/${1}.log.${CURRENT_LOG_TIME}" >>$SCREENRC
            echo "log on" >>$SCREENRC
        fi
    fi
}

# Set an option in an INI file
# iniset config-file section option value
function iniset {
    local xtrace=$(set +o | grep xtrace)
    set +o xtrace
    local file=$1
    local section=$2
    local option=$3
    local value=$4

    [[ -z $section || -z $option ]] && return

    if ! grep -q "^\[$section\]" "$file" 2>/dev/null; then
        # Add section at the end
        echo -e "\n[$section]" >>"$file"
    fi
    if ! ini_has_option "$file" "$section" "$option"; then
        # Add it
        sed -i -e "/^\[$section\]/ a\\
$option = $value
" "$file"
    else
        local sep=$(echo -ne "\x01")
        # Replace it
        sed -i -e '/^\['${section}'\]/,/^\[.*\]/ s'${sep}'^\('${option}'[ \t]*=[ \t]*\).*$'${sep}'\1'"${value}"${sep} "$file"
    fi
    $xtrace
}

# Determinate is the given option present in the INI file
# ini_has_option config-file section option
function ini_has_option {
    local xtrace=$(set +o | grep xtrace)
    set +o xtrace
    local file=$1
    local section=$2
    local option=$3
    local line

    line=$(sed -ne "/^\[$section\]/,/^\[.*\]/ { /^$option[ \t]*=/ p; }" "$file")
    $xtrace
    [ -n "$line" ]
}

function stop_process {
    local service=$1

    # Kill via pid if we have one available
    pkill -F $SERVICE_DIR/$SCREEN_NAME/$service.pid
    rm $SERVICE_DIR/$SCREEN_NAME/$service.pid
    screen -S $SCREEN_NAME -p $service -X kill
}


# MN Functions
# ------------

# Wrapper for mn-conf command
# Uses globals ``ZOOKEEPER_HOSTS``
function configure_mn {
    local value="$2"

    # quote with "" only when necessary.  we don't always quote because
    # mn-conf complains for quoted booleans.  eg. "false"
    if [[ "${value}" =~ ":" || "${value}" = "" ]]; then
        value="\"${value}\""
    fi

    # In some commands, mn-conf creates a local file, which requires root
    # access.  For simplicity, always call mn-conf with root for now.
    echo $1 : "${value}" | MIDO_ZOOKEEPER_HOSTS="$ZOOKEEPER_HOSTS" sudo mn-conf set
}

# Restore xtrace
$XTRACE
