#!/usr/bin/env bash

set -e

if [ "$UID" -ne 0 ]; then
    echo "Script must be run as root"
    exit 1
fi

EXTRA_TUNING=()
FIRST_CPU=0

export EXP_ROOT="$PWD"

source ./configs/base.sh
source ./configs/nic.sh

if [ -e "./configs/$(uname -r).sh" ]; then
    echo "Loading kernel-specific config for $(uname -r)"
    source ./configs/$(uname -r).sh
fi

while [[ $# -gt 0 ]]; do
    case $1 in
        -e)
            EXPERIMENT="$2"
            shift
            shift
            ;;
        --extra-tuning)
            EXTRA_TUNING+=("$2")
            shift
            shift
            ;;
        -t|--threads)
            THREADS="$2"
            shift
            shift
            ;;
        --first-cpu)
            FIRST_CPU=$2
            shift
            shift
            ;;
        --fstack)
            FSTACK=true
            USERSTACK=true
            shift
            ;;
        --caladan)
            CALADAN=true
            USERSTACK=true
            shift
            ;;
        -f|--force)
            FORCE_REMOVE=true
            shift
            ;;
        --aux)
            RUN_AUX=true
            shift
            ;;
        --)
            shift
            break
            ;;
    esac
done

if [ -z "$EXPERIMENT" ]; then
    echo "Must specify an experiment to run"
    exit 1
fi

if [ ! -f "./experiments/$EXPERIMENT.sh" ]; then
     echo "$EXPERIMENT does not exist"
     exit 1
fi

apply_tunings() {
    source ./tunings/base.sh

    if [ "$USERSTACK" != "true" ]; then
        source ./tunings/normalize_queues.sh
        if [ "${#EXTRA_TUNING[@]}" -gt 0 ]; then
            for tuning in "${EXTRA_TUNING[@]}"; do
                source ./tunings/"$tuning".sh
            done
        fi
    fi
}

apply_tunings

# Initialization for DPDK
if [ "$USERSTACK" == "true" ]; then
    # mlx4 supports bifurcation, and we only need to work around a kernel driver bug by re-loading the module
    modprobe -r mlx4_ib mlx4_en mlx4_core
    modprobe mlx4_en mlx4_ib 
    echo 1024 > /sys/kernel/mm/hugepages/hugepages-2048kB/nr_hugepages
    echo 1024 > /sys/devices/system/node/node0/hugepages/hugepages-2048kB/nr_hugepages
    echo 1024 > /sys/devices/system/node/node1/hugepages/hugepages-2048kB/nr_hugepages
    mkdir /mnt/huge || true
    mount -t hugetlbfs nodev /mnt/huge || true
fi

if [ "$CALADAN" == "true" ]; then
    pkill iokerneld || true
    pkill memcached || true
    if [ "$THREADS" -le 1 ]; then
        echo "Must allocate one thread for IOKernel"
        exit 1
    fi
    CALADAN_CONFIG=/tmp/caladan.config
    sed "s/{threads}/$((THREADS - 1))/g" ./caladan.config > $CALADAN_CONFIG
    $CALADAN_SRC/scripts/setup_machine.sh || true
    sleep 5
    echo "Starting up Caladan iokernel"
    $CALADAN_SRC/iokerneld simple noht &
    CALADAN_PID=$!
    sleep 5
fi

source "./experiments/$EXPERIMENT.sh"

pre_start $@

EXPERIMENT_NAME="$EXPERIMENT.$(uname -r)"
if [ "${#EXTRA_TUNING[@]}" -gt 0 ]; then
    for tuning in "${EXTRA_TUNING[@]}"; do
        EXPERIMENT_NAME="$EXPERIMENT_NAME.$tuning"
    done
fi
if [ $FIRST_CPU -gt 0 ]; then
    EXPERIMENT_NAME="$EXPERIMENT_NAME.first$FIRST_CPU"
fi
if [ "$FSTACK" == "true" ]; then
    EXPERIMENT_NAME="$EXPERIMENT_NAME.fstack"
fi
if [ "$CALADAN" == "true" ]; then
    EXPERIMENT_NAME="$EXPERIMENT_NAME.caladan"
fi
if [ ! -z "$EXPERIMENT_NAME_EXT" ]; then
    EXPERIMENT_NAME="$EXPERIMENT_NAME.$EXPERIMENT_NAME_EXT"
fi

DATA_OUT="$EXP_ROOT/data/$EXPERIMENT_NAME"

if [ -d "$DATA_OUT" ] && [ "$RUN_AUX" == "true" ]; then
    AUX_NUM=0
    if [ -f "$DATA_OUT/.last_exp_num" ]; then
        AUX_NUM="$(cat "$DATA_OUT/.last_exp_num")"
    fi
    EXP_AUX_SUFFIX="_aux$AUX_NUM"
fi

if [ -d "$DATA_OUT" ] && [ -z "$EXP_AUX_SUFFIX" ]; then
    if [ "$FORCE_REMOVE" == "true" ]; then
        rm -rf "$DATA_OUT"
    else
        echo "Experiment directory $DATA_OUT already exists. Remove it manually or use --force. To run auxillary experiments for multiple sets of data, use --aux"
        exit 1
    fi
fi
mkdir -p "$DATA_OUT"

run_exp

post_start

if [ "$CALADAN" == "true" ]; then
    kill -KILL "$CALADAN_PID"
fi

if [ ! -z "$EXP_AUX_SUFFIX" ]; then
    echo "$((AUX_NUM + 1))" > "$DATA_OUT/.last_exp_num"
fi