#!/bin/bash
# Volt Network Benchmark - Throughput Tests
# Tests TCP/UDP throughput using iperf3

set -e

SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"

# Parse arguments
SERVER_IP="${1:?Usage: $0 <server-ip> [backend-name] [duration]}"
BACKEND="${2:-unknown}"
DURATION="${3:-30}"

# Setup results directory
TIMESTAMP=$(date +%Y-%m-%d_%H%M%S)
RESULTS_DIR="${SCRIPT_DIR}/results/${BACKEND}/${TIMESTAMP}"
mkdir -p "$RESULTS_DIR"

echo "=== Volt Throughput Benchmark ==="
echo "Server: $SERVER_IP"
echo "Backend: $BACKEND"
echo "Duration: ${DURATION}s per test"
echo "Results: $RESULTS_DIR"
echo ""

# Function to run iperf3 test
run_iperf3() {
    local test_name="$1"
    local extra_args="$2"
    local output_file="${RESULTS_DIR}/${test_name}.json"
    
    echo "[$(date +%H:%M:%S)] Running: $test_name"
    
    if iperf3 -c "$SERVER_IP" -t "$DURATION" $extra_args -J > "$output_file" 2>&1; then
        # Extract key metrics
        if [ -f "$output_file" ] && command -v jq &> /dev/null; then
            local bps=$(jq -r '.end.sum_sent.bits_per_second // .end.sum.bits_per_second // 0' "$output_file" 2>/dev/null)
            local gbps=$(echo "scale=2; $bps / 1000000000" | bc 2>/dev/null || echo "N/A")
            echo "  → ${gbps} Gbps"
        else
            echo "  → Complete (see JSON for results)"
        fi
    else
        echo "  → FAILED"
        return 1
    fi
}

# Verify connectivity
echo "[$(date +%H:%M:%S)] Verifying connectivity to $SERVER_IP:5201..."
if ! timeout 5 bash -c "echo > /dev/tcp/$SERVER_IP/5201" 2>/dev/null; then
    echo "ERROR: Cannot connect to iperf3 server at $SERVER_IP:5201"
    echo "Ensure iperf3 -s is running on the server"
    exit 1
fi
echo "  → Connected"
echo ""

# Record system info
echo "=== System Info ===" > "${RESULTS_DIR}/system-info.txt"
echo "Date: $(date)" >> "${RESULTS_DIR}/system-info.txt"
echo "Kernel: $(uname -r)" >> "${RESULTS_DIR}/system-info.txt"
echo "Backend: $BACKEND" >> "${RESULTS_DIR}/system-info.txt"
ip addr show 2>/dev/null | grep -E "inet |mtu" >> "${RESULTS_DIR}/system-info.txt" || true
echo "" >> "${RESULTS_DIR}/system-info.txt"

# TCP Tests
echo "--- TCP Throughput Tests ---"
echo ""

# Single stream TCP
run_iperf3 "tcp-single" ""

# Wait between tests
sleep 2

# Multi-stream TCP (8 parallel)
run_iperf3 "tcp-multi-8" "-P 8"

sleep 2

# Reverse direction (download)
run_iperf3 "tcp-reverse" "-R"

sleep 2

# UDP Tests
echo ""
echo "--- UDP Throughput Tests ---"
echo ""

# UDP maximum bandwidth (let iperf3 find the limit)
run_iperf3 "udp-max" "-u -b 0"

sleep 2

# UDP at specific rates for comparison
for rate in 1G 5G 10G; do
    run_iperf3 "udp-${rate}" "-u -b ${rate}"
    sleep 2
done

# Generate summary
echo ""
echo "=== Summary ==="
SUMMARY_FILE="${RESULTS_DIR}/throughput-summary.txt"
{
    echo "Volt Throughput Benchmark Results"
    echo "======================================"
    echo "Backend: $BACKEND"
    echo "Server: $SERVER_IP"
    echo "Date: $(date)"
    echo "Duration: ${DURATION}s per test"
    echo ""
    echo "Results:"
    echo "--------"
    
    for json_file in "${RESULTS_DIR}"/*.json; do
        if [ -f "$json_file" ] && command -v jq &> /dev/null; then
            test_name=$(basename "$json_file" .json)
            
            # Try to extract metrics based on test type
            if [[ "$test_name" == udp-* ]]; then
                bps=$(jq -r '.end.sum.bits_per_second // 0' "$json_file" 2>/dev/null)
                loss=$(jq -r '.end.sum.lost_percent // 0' "$json_file" 2>/dev/null)
                gbps=$(echo "scale=2; $bps / 1000000000" | bc 2>/dev/null || echo "N/A")
                printf "%-20s %8s Gbps  (loss: %.2f%%)\n" "$test_name:" "$gbps" "$loss"
            else
                bps=$(jq -r '.end.sum_sent.bits_per_second // 0' "$json_file" 2>/dev/null)
                gbps=$(echo "scale=2; $bps / 1000000000" | bc 2>/dev/null || echo "N/A")
                printf "%-20s %8s Gbps\n" "$test_name:" "$gbps"
            fi
        fi
    done
} | tee "$SUMMARY_FILE"

echo ""
echo "Full results saved to: $RESULTS_DIR"
echo "JSON files available for detailed analysis"