#!/bin/bash

function grn() {
    echo -ne "\033[32;1m" "$@" "\033[m"
}
function gr() {
    echo -e "\033[32;1m" "$@" "\033[m"
}

function yl() {
    echo -e "\033[33;1m" "$@" "\033[m"
}


HN="$(hostname)"

# Let the SSH server run without systemd
mkdir /var/run/sshd

# Start the SSH server
/usr/sbin/sshd -E /var/log/auth.log
ssh-keygen -f ~/.ssh/id_rsa -N "" &>/dev/null
cat ~/.ssh/id_rsa.pub >> ~/.ssh/authorized_keys
if [ -e ~/.ssh/config ];
then
    yl "SSH config already existed, not modifying it";
else
    cat << EOF > ~/.ssh/config
Host $HN
    StrictHostKeyChecking no
    Port 2222
EOF
fi

export HOME=/root

FILEPATH="$(firewheel config get minimega.files_dir)"
mkdir -p $FILEPATH

# Verifying extra minimega environment variables
if [[ -n "${MM_FILEPATH}" ]]; then
    mkdir -p "${MM_FILEPATH}"
    firewheel config set -s minimega.files_dir "${MM_FILEPATH}"
fi

if [[ -n "${MM_BASE}" ]]; then
    firewheel config set -s minimega.base_dir "${MM_BASE}"
fi

grn "Starting minimega...       "
/start-minimega.sh &>/dev/null
gr "done"

grn "Setting up FIREWHEEL...    "
mkdir -p /tmp/firewheel
firewheel config set -s cluster.control "$HN" &>/dev/null
firewheel config set -s cluster.compute "" &>/dev/null
firewheel config set -s cluster.compute "$HN" &>/dev/null
firewheel config set -s discovery.hostname localhost &>/dev/null
firewheel config set -s minimega.use_gre True

firewheel init &>/dev/null
firewheel sync &>/dev/null

gr "done"

grn "Starting FIREWHEEL...      "
firewheel restart hard > /dev/null 2>&1
gr "done"

source /fwpy/bin/activate

if tty -s; then
    if [ "$#" -eq 0 ]; then
        tmux
    else
        # If a command was passed to the container, execute it
        "$@"
    fi
else
    # If it's not a TTY, it's likely a CI/CD job, in which case that job will
    # keep the container open after the end of this script. If there's ever a
    # use case for keeping the container alive when not in a TTY, this could be
    # changed
    yl "Not a TTY, exiting..."
    exit 0
fi

grn "Stopping experiment...     "
ct=0
while ps aux | grep -e tmux -e minimega -e miniweb -e discovery -e fwpy -e qemu | grep -v grep >/dev/null; do
    ps aux | grep -e tmux -e minimega -e miniweb -e discovery -e fwpy -e qemu | grep -v grep | awk '{print $2}' | xargs kill 2>/dev/null || true
    ((++ct))
    if [ "$ct" -gt 5 ]; then
        break
    fi
    sleep 0.5
done
sync
sleep 2
gr "done"

grn "Cleanup...                 "
ct=0
until rm -rf /tmp/*; do
    ((++ct))
    if [ "$ct" -gt 5 ]; then
        break
    fi
    sleep 0.5
done
gr "done"

gr "Container shutdown."
