#!/bin/bash

CLIENT_VM="client"
CLIENT_IP="10.6.6.10"
SERVER_VM="server"
SERVER_IP="10.6.6.20"
USER="ubuntu"
PASSWORD="ubuntu"

SIZE="10"
FILENAME="${SIZE}MB.file"

TCPDUMP_DURATION=300
TCP_CC_COMBINATIONS=("reno reno" "reno bbr" "bbr reno" "bbr bbr")

if [ $# -eq 0 ]; then
    echo "Specify output directory for dumps. Ex. ./run_benchmarks dumps/profile1"
    exit
else
    OUTPUT_DIR="$1"
    echo "Using specified output directory: $OUTPUT_DIR"
fi

for cmd in sshpass tcpdump; do
    if ! command -v $cmd &> /dev/null; then
        echo "Error: $cmd is not installed. Please install it first."
        exit 1
    fi
done

echo "Available network interfaces:"
interfaces=($(ip -o link show | awk -F': ' '{print $2}'))
PS3="Select interface number: "
select interface in "${interfaces[@]}"; do
    if [[ -n "$interface" ]]; then
        echo "Selected interface: $interface"
        break
    else
        echo "Invalid selection. Try again."
    fi
done

mkdir -p "$OUTPUT_DIR"

run_vm() {
    vm=$1
    cmd=$2
    sshpass -p "$PASSWORD" ssh -o StrictHostKeyChecking=no "$USER@$vm" "$cmd"
}

set_tcp_cc() {
    vm=$1
    cc=$2
    run_vm "$vm" "sudo sysctl -w net.ipv4.tcp_congestion_control=$cc"
}

run_benchmark() {
    protocol=$1
    client_cc=$2
    server_cc=$3
    suffix=$4

    if [[ "$suffix" == "tcp" ]]; then
        base_name="${protocol}_${client_cc}_${server_cc}"
    else
        base_name="${protocol}"
    fi

    pcap_file="${OUTPUT_DIR}/${base_name}.pcap"

    case $protocol in
        http1)
            filter="tcp and port 80 and host ${CLIENT_IP} and host ${SERVER_IP}" ;;
        http2)
            filter="tcp and port 81 and host ${CLIENT_IP} and host ${SERVER_IP}" ;;
        http3)
            filter="udp and port 443 and host ${CLIENT_IP} and host ${SERVER_IP}" ;;
        tftp)
            filter="udp and host ${CLIENT_IP} and host ${SERVER_IP}" ;;
    esac

    echo "Starting $protocol capture with filter: $filter"

    timeout $TCPDUMP_DURATION tcpdump -i "$interface" -w "$pcap_file" "$filter" >/dev/null 2>&1 &
    tcpdump_pid=$!
    sleep 2

    case $protocol in
        http1)
            run_vm "$CLIENT_VM" "curl -s -o /dev/null http://$SERVER_VM:80/$FILENAME" ;;
        http2)
            run_vm "$CLIENT_VM" "curl -s -o /dev/null --http2-prior-knowledge http://$SERVER_VM:81/$FILENAME" ;;
        http3)
            run_vm "$CLIENT_VM" "./curl -s -o /dev/null --http3-only --insecure https://$SERVER_VM:443/$FILENAME" ;;
        tftp)
            run_vm "$CLIENT_VM" "rm -f $FILENAME; tftp $SERVER_VM -c get /var/www/$FILENAME" ;;
    esac

    sleep 4
    kill $tcpdump_pid 2>/dev/null
    wait $tcpdump_pid 2>/dev/null

    echo "Capture saved to $pcap_file"
}

for combination in "${TCP_CC_COMBINATIONS[@]}"; do
    read client_cc server_cc <<< "$combination"
    echo -e "\nRunning TCP benchmarks with ($client_cc/$server_cc) congestion control"

    set_tcp_cc "$CLIENT_VM" "$client_cc"
    set_tcp_cc "$SERVER_VM" "$server_cc"

    run_benchmark "http1" "$client_cc" "$server_cc" "tcp"
    run_benchmark "http2" "$client_cc" "$server_cc" "tcp"
done

echo -e "\nRunning UDP benchmarks"
run_benchmark "http3" "" "" "udp"
run_benchmark "tftp" "" "" "udp"

echo -e "\nAll benchmarks have been run, dumps have been stored in $OUTPUT_DIR"
