#!/bin/bash
###############################################################################
# iperf-full-matrix.sh
#
# Two test suites:
#   A. Management network (vmbr0 / VLAN1): all-pairs single-stream TCP, 10s
#   B. Storage network (VLAN1000): tsys6+tsys7 → tsys4+tsys5, stress test
#
# Output: returned-logs/iperf/
###############################################################################
set -uo pipefail

SSH_OPTS=(-o BatchMode=yes -o ConnectTimeout=8 -o ServerAliveInterval=10 -o StrictHostKeyChecking=accept-new)
LOG_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)/returned-logs/iperf"
mkdir -p "$LOG_DIR"

HOSTS=(pfv-tsys1 pfv-tsys3 pfv-tsys4 pfv-tsys5 pfv-tsys6 pfv-tsys7)

# Storage IPs (known, static on VLAN1000)
declare -A SIP
SIP[pfv-tsys1]="10.100.100.1"
SIP[pfv-tsys3]="10.100.100.3"
SIP[pfv-tsys4]="10.100.100.4"
SIP[pfv-tsys5]="10.100.100.5"
SIP[pfv-tsys6]="10.100.100.6"
SIP[pfv-tsys7]="10.100.100.7"

# ===========================================================================
# STEP 0: Discover management IPs (vmbr0)
# ===========================================================================
echo "==================================================================="
echo " STEP 0: Discover management network IPs (vmbr0)"
echo "==================================================================="
declare -A MIP
for host in "${HOSTS[@]}"; do
    if ! ssh "${SSH_OPTS[@]}" "root@$host" 'echo ok' >/dev/null 2>&1; then
        echo "  [$host] UNREACHABLE"
        continue
    fi
    mip=$(ssh "${SSH_OPTS[@]}" "root@$host" 'ip -o -4 addr show dev vmbr0 2>/dev/null | awk "{print \$4}" | cut -d/ -f1 | head -1')
    if [ -n "$mip" ]; then
        MIP[$host]="$mip"
        echo "  [$host] vmbr0 = $mip"
    else
        echo "  [$host] no vmbr0 IPv4 — skipping"
    fi
done

# ===========================================================================
# STEP 1: Ensure iperf3 installed on all hosts
# ===========================================================================
echo ""
echo "==================================================================="
echo " STEP 1: Ensure iperf3 installed"
echo "==================================================================="
for host in "${HOSTS[@]}"; do
    [ -z "${MIP[$host]:-}" ] && continue
    if ! ssh "${SSH_OPTS[@]}" "root@$host" 'command -v iperf3 >/dev/null 2>&1' 2>/dev/null; then
        echo -n "  [$host] installing iperf3... "
        ssh "${SSH_OPTS[@]}" "root@$host" \
            'DEBIAN_FRONTEND=noninteractive apt-get update -qq >/dev/null 2>&1 && \
             DEBIAN_FRONTEND=noninteractive apt-get install -y -qq iperf3 >/dev/null 2>&1 && echo OK || echo FAILED'
    else
        echo "  [$host] iperf3 already installed"
    fi
done

# Helper: start iperf3 server in one-shot mode bound to a specific IP
start_server() {
    local host="$1" ip="$2"
    ssh "${SSH_OPTS[@]}" "root@$host" \
        "pkill -x iperf3 2>/dev/null; nohup iperf3 -s -1 -B ${ip} >/dev/null 2>&1 &" 2>/dev/null
    sleep 1
}

# Helper: run iperf3 client test, save output, extract result
run_test() {
    local client="$1" server="$2" sip="$3" cip="$4" label="$5" logfile="$6"
    shift 6
    local extra="$*"

    echo -n "  [$label] ... "
    {
        echo "=== iperf3: $label ==="
        echo "Client: $client ($cip)  →  Server: $server ($sip)"
        echo "Date: $(date -u +%Y-%m-%dT%H:%M:%SZ)"
        echo "Args: $extra"
        echo ""
        ssh "${SSH_OPTS[@]}" "root@$client" "iperf3 -c ${sip} -B ${cip} ${extra}" 2>&1
        echo ""
        echo "=== END ==="
    } > "$logfile" 2>&1

    # Extract result
    sum=$(grep '\[SUM\].*sender$' "$logfile" | tail -1)
    if [ -n "$sum" ]; then
        bitrate=$(echo "$sum" | awk '{print $6, $7}')
        retrans=$(echo "$sum" | awk '{print $8}')
    else
        single=$(grep 'sender$' "$logfile" | tail -1)
        bitrate=$(echo "$single" | awk '{print $7, $8}')
        retrans=$(echo "$single" | awk '{print $9}')
    fi
    echo "${bitrate:-?} (retrans: ${retrans:-?})"
}

# ===========================================================================
# SUITE A: Management network (vmbr0) — all pairs, single-stream TCP, 10s
# ===========================================================================
echo ""
echo "==================================================================="
echo " SUITE A: MANAGEMENT NETWORK (vmbr0) — all-pairs, 10s TCP"
echo " Expectation: ~940 Mbps for every pair (1 GbE line rate)"
echo "==================================================================="
echo ""

# Test each unique pair (i < j to avoid duplicates)
for ((i=0; i<${#HOSTS[@]}; i++)); do
    for ((j=i+1; j<${#HOSTS[@]}; j++)); do
        client="${HOSTS[$i]}"
        server="${HOSTS[$j]}"
        [ -z "${MIP[$client]:-}" ] && continue
        [ -z "${MIP[$server]:-}" ] && continue

        label="mgmt: ${client}→${server}"
        logfile="$LOG_DIR/mgmt-${client}-to-${server}.log"

        start_server "$server" "${MIP[$server]}"
        run_test "$client" "$server" "${MIP[$server]}" "${MIP[$client]}" \
            "$label" "$logfile" "-t 10 -P 1"
    done
done

# ===========================================================================
# SUITE B: Storage network (VLAN1000) — stress test the cross-rack LACP
# ===========================================================================
echo ""
echo "==================================================================="
echo " SUITE B: STORAGE NETWORK (VLAN1000) — stress test cross-rack link"
echo " tsys6 + tsys7 (Rack 3) → tsys4 + tsys5 (Rack 5)"
echo " Expectation: limited by tsys4 USB dongle + tsys5 broken bond"
echo "==================================================================="
echo ""

# --- B.1: Individual tests (one client → one server at a time) ---
echo "--- B.1: Individual tests (sequential) ---"
echo ""

for client in pfv-tsys6 pfv-tsys7; do
    for server in pfv-tsys4 pfv-tsys5; do
        label="stor: ${client}→${server} (8-stream)"
        logfile="$LOG_DIR/stor-indiv-${client}-to-${server}-8stream.log"
        start_server "$server" "${SIP[$server]}"
        run_test "$client" "$server" "${SIP[$server]}" "${SIP[$client]}" \
            "$label" "$logfile" "-P 8 -t 20 -l 128k -O 2"
    done
done

# --- B.2: Reverse direction (tsys4/5 → tsys6/7) ---
echo ""
echo "--- B.2: Reverse direction (tsys4/5 → tsys6/7) ---"
echo ""

for client in pfv-tsys4 pfv-tsys5; do
    for server in pfv-tsys6 pfv-tsys7; do
        label="stor: ${client}→${server} (8-stream rev)"
        logfile="$LOG_DIR/stor-indiv-${client}-to-${server}-8stream.log"
        start_server "$server" "${SIP[$server]}"
        run_test "$client" "$server" "${SIP[$server]}" "${SIP[$client]}" \
            "$label" "$logfile" "-P 8 -t 20 -l 128k -O 2"
    done
done

# --- B.3: Simultaneous stress test (4 flows at once) ---
echo ""
echo "--- B.3: Simultaneous 4-flow stress test ---"
echo "  tsys6→tsys4 + tsys6→tsys5 + tsys7→tsys4 + tsys7→tsys5"
echo "  All running in parallel for 30 seconds"
echo ""

# Start 4 iperf3 servers (one-shot mode won't work for parallel; use persistent)
for server in pfv-tsys4 pfv-tsys5; do
    ssh "${SSH_OPTS[@]}" "root@$server" "pkill -x iperf3 2>/dev/null; nohup iperf3 -s -B ${SIP[$server]} >/dev/null 2>&1 &" 2>/dev/null
    echo "  [server started: $server]"
done
sleep 1

STRESS_LOG="$LOG_DIR/stor-stress-4flow"
mkdir -p "$STRESS_LOG"

# Launch 4 clients in parallel, each writing to its own log
ssh "${SSH_OPTS[@]}" "root@pfv-tsys6" "iperf3 -c ${SIP[pfv-tsys4]} -B ${SIP[pfv-tsys6]} -P 4 -t 30 -l 128k -O 2" > "$STRESS_LOG/tsys6-to-tsys4.log" 2>&1 &
PID1=$!
ssh "${SSH_OPTS[@]}" "root@pfv-tsys6" "iperf3 -c ${SIP[pfv-tsys5]} -B ${SIP[pfv-tsys6]} -P 4 -t 30 -l 128k -O 2" > "$STRESS_LOG/tsys6-to-tsys5.log" 2>&1 &
PID2=$!
ssh "${SSH_OPTS[@]}" "root@pfv-tsys7" "iperf3 -c ${SIP[pfv-tsys4]} -B ${SIP[pfv-tsys7]} -P 4 -t 30 -l 128k -O 2" > "$STRESS_LOG/tsys7-to-tsys4.log" 2>&1 &
PID3=$!
ssh "${SSH_OPTS[@]}" "root@pfv-tsys7" "iperf3 -c ${SIP[pfv-tsys5]} -B ${SIP[pfv-tsys7]} -P 4 -t 30 -l 128k -O 2" > "$STRESS_LOG/tsys7-to-tsys5.log" 2>&1 &
PID4=$!

echo "  [4 clients launched, waiting 40s for completion...]"
wait $PID1 $PID2 $PID3 $PID4 2>/dev/null
echo "  [all 4 flows complete]"

# Kill servers
for server in pfv-tsys4 pfv-tsys5; do
    ssh "${SSH_OPTS[@]}" "root@$server" 'pkill -x iperf3 2>/dev/null; true' 2>/dev/null
done

# ===========================================================================
# SUITE C: All hosts cleanup
# ===========================================================================
echo ""
echo "==================================================================="
echo " Cleanup: killing iperf3 everywhere"
echo "==================================================================="
for host in "${HOSTS[@]}"; do
    ssh "${SSH_OPTS[@]}" "root@$host" 'pkill -x iperf3 2>/dev/null; true' 2>/dev/null
done

# ===========================================================================
# RESULTS SUMMARY
# ===========================================================================
echo ""
echo "==================================================================="
echo " RESULTS SUMMARY"
echo "==================================================================="

echo ""
echo "===== SUITE A: Management network (vmbr0) ====="
echo ""
printf "%-40s %15s %10s\n" "TEST" "THROUGHPUT" "RETRANS"
printf "%-40s %15s %10s\n" "----" "----------" "-------"
for f in "$LOG_DIR"/mgmt-*.log; do
    [ -r "$f" ] || continue
    label=$(head -1 "$f" | sed 's/^=== iperf3: //; s/ ===$//')
    single=$(grep 'sender$' "$f" | tail -1)
    bitrate=$(echo "$single" | awk '{print $7, $8}')
    retrans=$(echo "$single" | awk '{print $9}')
    printf "%-40s %15s %10s\n" "$label" "${bitrate:-?}" "${retrans:--}"
done

echo ""
echo "===== SUITE B.1+B.2: Storage network individual ====="
echo ""
printf "%-45s %15s %10s\n" "TEST" "THROUGHPUT" "RETRANS"
printf "%-45s %15s %10s\n" "----" "----------" "-------"
for f in "$LOG_DIR"/stor-indiv-*.log; do
    [ -r "$f" ] || continue
    label=$(head -1 "$f" | sed 's/^=== iperf3: //; s/ ===$//')
    sum=$(grep '\[SUM\].*sender$' "$f" | tail -1)
    if [ -n "$sum" ]; then
        bitrate=$(echo "$sum" | awk '{print $6, $7}')
        retrans=$(echo "$sum" | awk '{print $8}')
    else
        single=$(grep 'sender$' "$f" | tail -1)
        bitrate=$(echo "$single" | awk '{print $7, $8}')
        retrans=$(echo "$single" | awk '{print $9}')
    fi
    printf "%-45s %15s %10s\n" "$label" "${bitrate:-?}" "${retrans:--}"
done

echo ""
echo "===== SUITE B.3: Simultaneous 4-flow stress test ====="
echo ""
printf "%-30s %15s %10s\n" "FLOW" "THROUGHPUT" "RETRANS"
printf "%-30s %15s %10s\n" "----" "----------" "-------"
total_mbps=0
for f in "$STRESS_LOG"/*.log; do
    [ -r "$f" ] || continue
    flow=$(basename "$f" .log)
    sum=$(grep '\[SUM\].*sender$' "$f" | tail -1)
    if [ -n "$sum" ]; then
        bitrate=$(echo "$sum" | awk '{print $6, $7}')
        retrans=$(echo "$sum" | awk '{print $8}')
        mbps=$(echo "$sum" | awk '{print $6}')
        total_mbps=$(awk "BEGIN{print $total_mbps + $mbps}")
    else
        bitrate="?"
        retrans="-"
    fi
    printf "%-30s %15s %10s\n" "$flow" "$bitrate" "${retrans:--}"
done
printf "%-30s %15s\n" "AGGREGATE (all 4 flows)" "${total_mbps} Mbits/sec"

echo ""
echo "===== CONTEXT ====="
echo "tsys4: USB cdc_ncm dongle (single 1G link, no bond)"
echo "tsys5: bond0 broken (1 active slave, no LACP partner) — cable pending"
echo "tsys6/7: working 2x1G LACP, layer3+4 hash (host side)"
echo "Cross-rack: 4x1G LACP (pfv-r3-tor-stor-01 → pfv-r5-core-01)"
echo ""
echo "All logs in: $LOG_DIR/"
