#!/usr/bin/env bash

#
# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
# SPDX-License-Identifier: Apache-2.0
#

set -e

function ensure_executable() {
  if ! command -v $1 &> /dev/null; then
    echo "$1 needs to be installed"
    exit
  fi
}

ensure_executable "docker-compose"
ensure_executable "tshark"
ensure_executable "gnuplot"

case "$(uname -s)" in
   Linux)
     sudo modprobe ip6table_filter
     ;;
   *)
     # TODO possibly add other things for Darwin
     ;;
esac

ROOT_DIR=$(realpath "$(dirname $0)/../../")
OUTPUT_DIR="${ROOT_DIR}/target/benchmark/results"
CLIENT_PCAP="sim/trace_node_left.pcap"
DOWNLOAD_MB=1000
UPLOAD_MB=0
DURATION=15
SCENARIO=${1:-"simple-p2p --delay=15ms --bandwidth=2Mbps --queue=25"}
NAME=${SCENARIO//[^a-zA-Z0-9]/-}

mkdir -p ${OUTPUT_DIR/NAME}
TMP=$(mktemp -d -t s2n-quic-bench-XXXXXXXXXX)

# Create the initial plot script for Gnuplot
cat <<EOT > "${TMP}/plotcmd.p"
    set grid
    set autoscale
    set title "${SCENARIO}"
    set timefmt "%H:%M:%S"
    set xdata time
    set xlabel "Time measurement"
    set ylabel "Megabits/second"
    set term svg size 900,480
    set output "${OUTPUT_DIR}/${NAME}.svg"
EOT

LOG_DIR="$OUTPUT_DIR/~logs/$NAME/s2n-quic"

if [ "$BUILD_S2N_QUIC" != false ]; then
    docker build . -f ./quic/s2n-quic-qns/etc/Dockerfile.build -t aws/s2n-quic
fi

ROOT_DIR="$ROOT_DIR" \
LOG_DIR="$LOG_DIR" \
DOWNLOAD_B=$(($DOWNLOAD_MB * 1000000)) \
UPLOAD_B=$(($UPLOAD_MB * 1000000)) \
DURATION="$DURATION" \
SCENARIO="$SCENARIO" \
docker-compose --file quic/s2n-quic-qns/benchmark/docker-compose.yml up --abort-on-container-exit --timeout 1 sim client server

tshark -r "$LOG_DIR/$CLIENT_PCAP" -t u -qz io,stat,1,"SUM(ip.len)"ip.len\&\&ip.dst==193.167.0.100 | awk -F '[\\|\s]' '{ print $2 $3}' > "${TMP}/s2n-quic.nsv"
# Find the time of the first real data transfer packet
QUIC_OFFSET=$(tshark -r "$LOG_DIR/$CLIENT_PCAP" -T fields -e frame.time_epoch -c 100 -Y "quic.header_form == 0" | head -n 1)
jq -r '.intervals[].sum | [(.start | strftime("%H:%M:%S")), .bits_per_second/1000000] | @tsv' "$LOG_DIR/client/perf_client.json" > "${TMP}/s2n-quic.tsv"

# Append to the plot script
cat <<EOT >> "${TMP}/plotcmd.p"
    plot "${TMP}/s2n-quic.nsv" using 1:(\$2/125000) title 's2n-quic Network' with lines lt 1, \
"${TMP}/s2n-quic.tsv" using (timecolumn(1) + "${QUIC_OFFSET}"):2 title 's2n-quic Streams' with lines lt 2,
EOT
# Remove trailing newline
truncate -s -1 "${TMP}/plotcmd.p"

if [ "$COMPARE_TO_MAIN" = true ] ; then
    LOG_DIR="$OUTPUT_DIR/~logs/$NAME/s2n-quic-main"

    ROOT_DIR="$ROOT_DIR" \
    LOG_DIR="$LOG_DIR" \
    DOWNLOAD_B=$(($DOWNLOAD_MB * 1000000)) \
    UPLOAD_B=$(($UPLOAD_MB * 1000000)) \
    DURATION="$DURATION" \
    SCENARIO="$SCENARIO" \
    docker-compose --file quic/s2n-quic-qns/benchmark/docker-compose.yml up --abort-on-container-exit --timeout 1 sim client server-main

    tshark -r "$LOG_DIR/$CLIENT_PCAP" -t u -qz io,stat,1,"SUM(ip.len)"ip.len\&\&ip.dst==193.167.0.100 | awk -F '[\\|\s]' '{ print $2 $3}' > "${TMP}/s2n-quic-main.nsv"
    # Find the time of the first real data transfer packet
    QUIC_MAIN_OFFSET=$(tshark -r "$LOG_DIR/$CLIENT_PCAP" -T fields -e frame.time_epoch -c 100 -Y "quic.header_form == 0" | head -n 1)
    jq -r '.intervals[].sum | [(.start | strftime("%H:%M:%S")), .bits_per_second/1000000] | @tsv' "$LOG_DIR/client/perf_client.json" > "${TMP}/s2n-quic-main.tsv"

    # Append to the plot script
    cat <<EOT >> "${TMP}/plotcmd.p"
"${TMP}/s2n-quic-main.nsv" using 1:(\$2/125000) title 's2n-quic main Network' with lines lt 3, \
"${TMP}/s2n-quic-main.tsv" using (timecolumn(1) + "${QUIC_MAIN_OFFSET}"):2 title 's2n-quic main Streams' with lines lt 4,
EOT

    # Remove trailing newline
    truncate -s -1 "${TMP}/plotcmd.p"
fi

LOG_DIR="$OUTPUT_DIR/~logs/$NAME/iperf"

# iperf3 occasionally crashes, so retry a couple times
MAX_ATTEMPTS=3
ATTEMPTS=1
until [ "$ATTEMPTS" -gt $MAX_ATTEMPTS ]
do
    echo "Attempt ${ATTEMPTS} of ${MAX_ATTEMPTS}"
    ROOT_DIR="$ROOT_DIR" \
    LOG_DIR="$LOG_DIR" \
    SCENARIO="$SCENARIO" \
    DURATION="$DURATION" \
    docker-compose --file quic/s2n-quic-qns/benchmark/docker-compose.yml up --abort-on-container-exit --timeout 1 --exit-code-from iperf_client sim iperf_client iperf_server && break
    ATTEMPTS=$((ATTEMPTS+1))
done

tshark -r "$LOG_DIR/$CLIENT_PCAP" -t u -qz io,stat,1,"SUM(ip.len)"ip.len\&\&ip.dst==193.167.0.90 | awk -F '[\\|\s]' '{ print $2 $3}' > "${TMP}/iperf.nsv"
# Find the time of the first real data transfer packet
IPERF_OFFSET=$(tshark -r "$LOG_DIR/$CLIENT_PCAP" -T fields -e frame.time_epoch -c 100 -Y "ip.len > 1000" | head -n 1)
jq -r '.intervals[].sum | [(.start | strftime("%H:%M:%S")), .bits_per_second/1000000] | @tsv' "$LOG_DIR/client/iperf.json" > "${TMP}/iperf.tsv"

# Append to the plot script
cat <<EOT >> "${TMP}/plotcmd.p"
"${TMP}/iperf.nsv" using 1:(\$2/125000) title 'iperf TCP Network' with lines lt 5, \
"${TMP}/iperf.tsv" using (timecolumn(1) + "${IPERF_OFFSET}"):2 title 'iperf TCP Streams' with lines lt 6
EOT

gnuplot "${TMP}/plotcmd.p"

rm -rf $TMP

echo "plot available in $OUTPUT_DIR/${NAME}.svg"