#!/bin/bash
# Volt Network Benchmark - Latency Tests
# Tests ICMP and TCP latency with percentile analysis

set -e

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

# Parse arguments
SERVER_IP="${1:?Usage: $0 <server-ip> [backend-name] [count]}"
BACKEND="${2:-unknown}"
PING_COUNT="${3:-1000}"
SOCKPERF_DURATION="${4:-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 Latency Benchmark ==="
echo "Server: $SERVER_IP"
echo "Backend: $BACKEND"
echo "Ping count: $PING_COUNT"
echo "Results: $RESULTS_DIR"
echo ""

# Function to calculate percentiles from sorted data
calc_percentiles() {
    local file="$1"
    local count=$(wc -l < "$file")
    
    if [ "$count" -eq 0 ]; then
        echo "N/A N/A N/A N/A N/A"
        return
    fi
    
    # Sort numerically
    sort -n "$file" > "${file}.sorted"
    
    # Calculate indices (1-indexed for sed)
    local p50_idx=$(( (count * 50 + 99) / 100 ))
    local p95_idx=$(( (count * 95 + 99) / 100 ))
    local p99_idx=$(( (count * 99 + 99) / 100 ))
    
    # Ensure indices are at least 1
    [ "$p50_idx" -lt 1 ] && p50_idx=1
    [ "$p95_idx" -lt 1 ] && p95_idx=1
    [ "$p99_idx" -lt 1 ] && p99_idx=1
    
    local min=$(head -1 "${file}.sorted")
    local max=$(tail -1 "${file}.sorted")
    local p50=$(sed -n "${p50_idx}p" "${file}.sorted")
    local p95=$(sed -n "${p95_idx}p" "${file}.sorted")
    local p99=$(sed -n "${p99_idx}p" "${file}.sorted")
    
    # Calculate average
    local avg=$(awk '{sum+=$1} END {printf "%.3f", sum/NR}' "${file}.sorted")
    
    rm -f "${file}.sorted"
    
    echo "$min $avg $p50 $p95 $p99 $max"
}

# ICMP Ping Test
echo "[$(date +%H:%M:%S)] Running ICMP ping test (${PING_COUNT} packets)..."
PING_RAW="${RESULTS_DIR}/ping-raw.txt"
PING_LATENCIES="${RESULTS_DIR}/ping-latencies.txt"

if ping -c "$PING_COUNT" -i 0.01 "$SERVER_IP" > "$PING_RAW" 2>&1; then
    # Extract latency values (time=X.XX ms)
    grep -oP 'time=\K[0-9.]+' "$PING_RAW" > "$PING_LATENCIES"
    
    # Convert to microseconds for consistency
    awk '{print $1 * 1000}' "$PING_LATENCIES" > "${PING_LATENCIES}.us"
    mv "${PING_LATENCIES}.us" "$PING_LATENCIES"
    
    read min avg p50 p95 p99 max <<< $(calc_percentiles "$PING_LATENCIES")
    
    echo "  ICMP Ping Results (µs):"
    printf "    Min: %10.1f\n" "$min"
    printf "    Avg: %10.1f\n" "$avg"
    printf "    P50: %10.1f\n" "$p50"
    printf "    P95: %10.1f\n" "$p95"
    printf "    P99: %10.1f\n" "$p99"
    printf "    Max: %10.1f\n" "$max"
    
    # Save summary
    {
        echo "ICMP_MIN_US=$min"
        echo "ICMP_AVG_US=$avg"
        echo "ICMP_P50_US=$p50"
        echo "ICMP_P95_US=$p95"
        echo "ICMP_P99_US=$p99"
        echo "ICMP_MAX_US=$max"
    } > "${RESULTS_DIR}/ping-summary.env"
else
    echo "  → FAILED (check if ICMP is allowed)"
fi

echo ""

# TCP Latency with sockperf (ping-pong mode)
echo "[$(date +%H:%M:%S)] Running TCP latency test (sockperf pp, ${SOCKPERF_DURATION}s)..."

# Check if sockperf server is reachable
if timeout 5 bash -c "echo > /dev/tcp/$SERVER_IP/11111" 2>/dev/null; then
    SOCKPERF_RAW="${RESULTS_DIR}/sockperf-raw.txt"
    SOCKPERF_LATENCIES="${RESULTS_DIR}/sockperf-latencies.txt"
    
    # Run sockperf in ping-pong mode
    if sockperf pp -i "$SERVER_IP" -t "$SOCKPERF_DURATION" --full-log "$SOCKPERF_RAW" > "${RESULTS_DIR}/sockperf-output.txt" 2>&1; then
        
        # Extract latency values from full log (if available)
        if [ -f "$SOCKPERF_RAW" ]; then
            # sockperf full-log format: txTime, rxTime, latency (nsec)
            awk '{print $3/1000}' "$SOCKPERF_RAW" > "$SOCKPERF_LATENCIES"
        else
            # Parse from summary output
            grep -oP 'latency=\K[0-9.]+' "${RESULTS_DIR}/sockperf-output.txt" > "$SOCKPERF_LATENCIES" 2>/dev/null || true
        fi
        
        if [ -s "$SOCKPERF_LATENCIES" ]; then
            read min avg p50 p95 p99 max <<< $(calc_percentiles "$SOCKPERF_LATENCIES")
            
            echo "  TCP Latency Results (µs):"
            printf "    Min: %10.1f\n" "$min"
            printf "    Avg: %10.1f\n" "$avg"
            printf "    P50: %10.1f\n" "$p50"
            printf "    P95: %10.1f\n" "$p95"
            printf "    P99: %10.1f\n" "$p99"
            printf "    Max: %10.1f\n" "$max"
            
            {
                echo "TCP_MIN_US=$min"
                echo "TCP_AVG_US=$avg"
                echo "TCP_P50_US=$p50"
                echo "TCP_P95_US=$p95"
                echo "TCP_P99_US=$p99"
                echo "TCP_MAX_US=$max"
            } > "${RESULTS_DIR}/sockperf-summary.env"
        else
            # Parse summary from sockperf output
            echo "  → Parsing summary output..."
            grep -E "(avg|percentile|latency)" "${RESULTS_DIR}/sockperf-output.txt" || true
        fi
    else
        echo "  → FAILED"
    fi
else
    echo "  → SKIPPED (sockperf server not running on $SERVER_IP:11111)"
    echo "  → Run 'sockperf sr' on the server"
fi

echo ""

# UDP Latency with sockperf
echo "[$(date +%H:%M:%S)] Running UDP latency test (sockperf under-load, ${SOCKPERF_DURATION}s)..."

if timeout 5 bash -c "echo > /dev/udp/$SERVER_IP/11111" 2>/dev/null || true; then
    SOCKPERF_UDP_RAW="${RESULTS_DIR}/sockperf-udp-raw.txt"
    
    if sockperf under-load -i "$SERVER_IP" -t "$SOCKPERF_DURATION" --full-log "$SOCKPERF_UDP_RAW" > "${RESULTS_DIR}/sockperf-udp-output.txt" 2>&1; then
        echo "  → Complete"
        # Parse percentiles from sockperf output
        grep -E "(percentile|avg-latency)" "${RESULTS_DIR}/sockperf-udp-output.txt" | head -10
    else
        echo "  → FAILED or server not running"
    fi
fi

# Generate overall summary
echo ""
echo "=== Latency Summary ==="
SUMMARY_FILE="${RESULTS_DIR}/latency-summary.txt"
{
    echo "Volt Latency Benchmark Results"
    echo "===================================="
    echo "Backend: $BACKEND"
    echo "Server: $SERVER_IP"
    echo "Date: $(date)"
    echo ""
    
    if [ -f "${RESULTS_DIR}/ping-summary.env" ]; then
        echo "ICMP Ping Latency (µs):"
        source "${RESULTS_DIR}/ping-summary.env"
        printf "  %-8s %10.1f\n" "Min:" "$ICMP_MIN_US"
        printf "  %-8s %10.1f\n" "Avg:" "$ICMP_AVG_US"
        printf "  %-8s %10.1f\n" "P50:" "$ICMP_P50_US"
        printf "  %-8s %10.1f\n" "P95:" "$ICMP_P95_US"
        printf "  %-8s %10.1f\n" "P99:" "$ICMP_P99_US"
        printf "  %-8s %10.1f\n" "Max:" "$ICMP_MAX_US"
        echo ""
    fi
    
    if [ -f "${RESULTS_DIR}/sockperf-summary.env" ]; then
        echo "TCP Latency (µs):"
        source "${RESULTS_DIR}/sockperf-summary.env"
        printf "  %-8s %10.1f\n" "Min:" "$TCP_MIN_US"
        printf "  %-8s %10.1f\n" "Avg:" "$TCP_AVG_US"
        printf "  %-8s %10.1f\n" "P50:" "$TCP_P50_US"
        printf "  %-8s %10.1f\n" "P95:" "$TCP_P95_US"
        printf "  %-8s %10.1f\n" "P99:" "$TCP_P99_US"
        printf "  %-8s %10.1f\n" "Max:" "$TCP_MAX_US"
    fi
} | tee "$SUMMARY_FILE"

echo ""
echo "Full results saved to: $RESULTS_DIR"