#!/bin/bash
# ESP32 Device Testing Script
# Tests all deployed ESP32 devices with iperf

# Configuration
START_IP="${START_IP:-192.168.1.51}"
NUM_DEVICES="${NUM_DEVICES:-10}"
IPERF_TIME="${IPERF_TIME:-5}"
PROTOCOL="${PROTOCOL:-tcp}"  # tcp or udp

# Parse starting IP
IFS='.' read -r -a IP_PARTS <<< "$START_IP"
IP_BASE="${IP_PARTS[0]}.${IP_PARTS[1]}.${IP_PARTS[2]}"
IP_START=${IP_PARTS[3]}

# Colors
GREEN='\033[0;32m'
RED='\033[0;31m'
YELLOW='\033[1;33m'
BLUE='\033[0;34m'
NC='\033[0m'

echo -e "${BLUE}=========================================="
echo "ESP32 Device Testing"
echo -e "==========================================${NC}"
echo "IP Range: ${IP_BASE}.${IP_START} - ${IP_BASE}.$((IP_START + NUM_DEVICES - 1))"
echo "Protocol: $PROTOCOL"
echo "Duration: ${IPERF_TIME}s per device"
echo -e "${BLUE}==========================================${NC}"
echo ""

test_device() {
    local INDEX=$1
    local IP_ADDR="${IP_BASE}.$((IP_START + INDEX))"
    local RESULT_FILE="/tmp/iperf_result_${INDEX}.txt"
    
    # Test ping
    if ! ping -c 1 -W 2 "$IP_ADDR" > /dev/null 2>&1; then
        echo -e "${RED}✗ Device $INDEX ($IP_ADDR): No ping response${NC}"
        echo "FAIL" > "$RESULT_FILE"
        return 1
    fi
    
    # Test iperf
    if [ "$PROTOCOL" = "udp" ]; then
        IPERF_CMD="iperf -c $IP_ADDR -u -b 100M -t $IPERF_TIME"
    else
        IPERF_CMD="iperf -c $IP_ADDR -t $IPERF_TIME"
    fi
    
    if OUTPUT=$($IPERF_CMD 2>&1); then
        # Extract bandwidth
        BANDWIDTH=$(echo "$OUTPUT" | grep -oP '\d+(\.\d+)?\s+(K|M|G)bits/sec' | tail -1)
        echo -e "${GREEN}✓ Device $INDEX ($IP_ADDR): $BANDWIDTH${NC}"
        echo "SUCCESS:$BANDWIDTH" > "$RESULT_FILE"
    else
        echo -e "${RED}✗ Device $INDEX ($IP_ADDR): iperf failed${NC}"
        echo "FAIL" > "$RESULT_FILE"
        return 1
    fi
}

export -f test_device
export IP_BASE IP_START IPERF_TIME PROTOCOL GREEN RED YELLOW BLUE NC

# Clean old results
rm -f /tmp/iperf_result_*.txt

# Quick ping test first
echo -e "${YELLOW}Quick connectivity check...${NC}"
for i in $(seq 0 $((NUM_DEVICES - 1))); do
    IP_ADDR="${IP_BASE}.$((IP_START + i))"
    if ping -c 1 -W 1 "$IP_ADDR" > /dev/null 2>&1; then
        echo -n "."
    else
        echo -n "x"
    fi
done
echo ""
echo ""

# Run iperf tests in parallel (batches of 4 to avoid overwhelming network)
echo -e "${YELLOW}Running iperf tests...${NC}"
BATCH_SIZE=4

for BATCH_START in $(seq 0 $BATCH_SIZE $((NUM_DEVICES - 1))); do
    for i in $(seq $BATCH_START $(( BATCH_START + BATCH_SIZE - 1 ))); do
        if [ $i -lt $NUM_DEVICES ]; then
            test_device "$i" &
        fi
    done
    wait
done

# Summary
echo ""
echo -e "${BLUE}=========================================="
echo "Test Summary"
echo -e "==========================================${NC}"

SUCCESS=0
FAILED=0
TOTAL_BANDWIDTH=0

for i in $(seq 0 $((NUM_DEVICES - 1))); do
    RESULT_FILE="/tmp/iperf_result_${i}.txt"
    if [ -f "$RESULT_FILE" ]; then
        RESULT=$(cat "$RESULT_FILE")
        if [[ "$RESULT" == SUCCESS:* ]]; then
            ((SUCCESS++))
        else
            ((FAILED++))
        fi
    else
        ((FAILED++))
    fi
done

echo "Tested:  $NUM_DEVICES devices"
echo -e "${GREEN}Success: $SUCCESS${NC}"
if [ $FAILED -gt 0 ]; then
    echo -e "${RED}Failed:  $FAILED${NC}"
fi

SUCCESS_RATE=$(awk "BEGIN {printf \"%.1f\", ($SUCCESS / $NUM_DEVICES) * 100}")
echo "Rate:    ${SUCCESS_RATE}%"

echo -e "${BLUE}==========================================${NC}"

# Cleanup
rm -f /tmp/iperf_result_*.txt

if [ $FAILED -gt 0 ]; then
    exit 1
fi