#!/usr/bin/env bash
# =============================================================================
# benchmark_network.sh — Network Condition Benchmark for fastSyncAI
#
# Tests fastSyncAI performance under various realistic network conditions:
#   - LAN (low latency, high bandwidth)
#   - WAN (high latency, moderate bandwidth)
#   - WAN with packet loss
#   - WAN with jitter
#
# Usage:
#   ./benchmark_network.sh [SCENARIO] [DATA_SIZE_MB] [RUN_COUNT]
#
# Scenarios:
#   --lan           : 10ms RTT, no loss (typical LAN)
#   --wan           : 100ms RTT, no loss (typical WAN)
#   --wan-loss-1    : 100ms RTT, 1% packet loss
#   --wan-loss-5    : 100ms RTT, 5% packet loss
#   --wan-jitter    : 100ms RTT, +/-50ms jitter
#   --custom LAT_ms : Custom latency (e.g., --custom 50 for 50ms RTT)
#
# Examples:
#   ./benchmark_network.sh --lan                # LAN test, 100MB, 3 runs
#   ./benchmark_network.sh --wan 200 1         # WAN test, 200MB, 1 run
#   ./benchmark_network.sh --wan-loss-1 50     # WAN with 1% loss, 50MB, 3 runs
#   ./benchmark_network.sh --wan-jitter 100 5   # WAN with jitter, 100MB, 5 runs
#   ./benchmark_network.sh --custom 50 100 3    # Custom 50ms RTT, 100MB, 3 runs
# =============================================================================

set -euo pipefail

# ── Defaults ──────────────────────────────────────────────────────────────
DATA_SIZE_MB=${2:-100}
RUN_COUNT=${3:-3}

# Network scenario configuration
SCENARIO="${1:---lan}"

# ── Parse scenario ────────────────────────────────────────────────────────
LATENCY_MS=0
PACKET_LOSS="0%"
JITTER_MS="0"
BANDWIDTH_MBPS="0"  # 0 = unlimited

case "$SCENARIO" in
    --lan)
        LATENCY_MS=10
        SCENARIO_NAME="LAN (10ms RTT)"
        ;;
    --wan)
        LATENCY_MS=100
        SCENARIO_NAME="WAN (100ms RTT)"
        ;;
    --wan-loss-1)
        LATENCY_MS=100
        PACKET_LOSS="1%"
        SCENARIO_NAME="WAN with 1% loss (100ms RTT)"
        ;;
    --wan-loss-5)
        LATENCY_MS=100
        PACKET_LOSS="5%"
        SCENARIO_NAME="WAN with 5% loss (100ms RTT)"
        ;;
    --wan-loss-10)
        LATENCY_MS=100
        PACKET_LOSS="10%"
        SCENARIO_NAME="WAN with 10% loss (100ms RTT)"
        ;;
    --wan-jitter)
        LATENCY_MS=100
        JITTER_MS=50
        SCENARIO_NAME="WAN with jitter (100ms +/-50ms RTT)"
        ;;
    --custom)
        if [ -n "${2:-}" ] && [[ "${2}" =~ ^[0-9]+$ ]]; then
            LATENCY_MS="$2"
            DATA_SIZE_MB=${3:-100}
            RUN_COUNT=${4:-3}
            SCENARIO_NAME="Custom (${LATENCY_MS}ms RTT)"
        else
            echo "ERROR: --custom requires latency value in ms"
            echo "Usage: ./benchmark_network.sh --custom LATENCY_MS [DATA_SIZE_MB] [RUN_COUNT]"
            exit 1
        fi
        ;;
    *)
        echo "ERROR: Unknown scenario '$SCENARIO'"
        echo "Available scenarios: --lan, --wan, --wan-loss-1, --wan-loss-5, --wan-loss-10, --wan-jitter, --custom"
        exit 1
        ;;
esac

PORT=19399
HOST="127.0.0.1"

SERVER_BIN="./fastsync_server"
CLIENT_BIN="./fastsync_client"

# ── Directories ────────────────────────────────────────────────────────────
TEST_DIR="network_bench_$$"
SRC_DIR="$TEST_DIR/src"
DST_DIR="$TEST_DIR/dst"

# ── Colours ────────────────────────────────────────────────────────────────
RED='\033[0;31m'; GREEN='\033[0;32m'; YELLOW='\033[1;33m'
CYAN='\033[0;36m'; BOLD='\033[1m'; RESET='\033[0m'

# ── Helpers ────────────────────────────────────────────────────────────────
die()  { echo -e "${RED}ERROR: $*${RESET}" >&2; exit 1; }
info() { echo -e "${CYAN}▶ $*${RESET}"; }
ok()   { echo -e "${GREEN}✓ $*${RESET}"; }
warn() { echo -e "${YELLOW}⚠ $*${RESET}"; }

# ── Network simulation (tc netem) ────────────────────────────────────────
setup_network() {
    info "Setting up network simulation: ${SCENARIO_NAME}..."
    
    if ! sudo -n true 2>/dev/null; then
        echo -e "${YELLOW}⚠ sudo required for tc netem — enter password if prompted${RESET}"
    fi
    
    # Clean existing qdisc
    sudo tc qdisc del dev lo root 2>/dev/null || true
    
    # Build netem command
    local netem_cmd="delay ${LATENCY_MS}ms"
    
    if [ "$JITTER_MS" -gt 0 ]; then
        netem_cmd="$netem_cmd ${JITTER_MS}ms"
    fi
    
    if [ "$PACKET_LOSS" != "0%" ]; then
        netem_cmd="$netem_cmd loss ${PACKET_LOSS}"
    fi
    
    # Apply network conditions
    sudo tc qdisc add dev lo root netem $netem_cmd || \
        die "tc netem failed. Install iproute2 or check parameters."
    
    ok "Network simulation active: latency=${LATENCY_MS}ms, loss=${PACKET_LOSS}, jitter=±${JITTER_MS}ms"
}

cleanup_network() {
    sudo tc qdisc del dev lo root 2>/dev/null || true
    ok "Network simulation removed"
}

# ── Sanity checks ──────────────────────────────────────────────────────────
[ -x "$SERVER_BIN" ] || die "Server binary not found: $SERVER_BIN — run 'make' first."
[ -x "$CLIENT_BIN" ] || die "Client binary not found: $CLIENT_BIN — run 'make' first."

cleanup() {
    cleanup_network
    pkill -f "fastsync_server" 2>/dev/null || true
    pkill -f "fastsync_client" 2>/dev/null || true
    rm -rf "$TEST_DIR"
}
trap cleanup EXIT

# ── Generate test data ──────────────────────────────────────────────────────
generate_test_data() {
    info "Generating ~${DATA_SIZE_MB}MB test data..."
    rm -rf "$TEST_DIR"
    mkdir -p "$SRC_DIR/small" "$SRC_DIR/medium" "$SRC_DIR/large"

    # Small files: many small files to test overhead (~25% of total)
    local small_count=$(( (DATA_SIZE_MB * 1048576 / 4) / 262144 ))
    [ "$small_count" -lt 10 ] && small_count=10
    [ "$small_count" -gt 100 ] && small_count=100
    info "  Creating $small_count small files (256KB each)..."
    for i in $(seq 1 $small_count); do
        dd if=/dev/urandom bs=256K count=1 of="$SRC_DIR/small/file_$i.bin" 2>/dev/null
    done

    # Medium files: 1MB each (~50% of total)
    local medium_count=$(( (DATA_SIZE_MB * 1048576 / 2) / 1048576 ))
    [ "$medium_count" -lt 5 ] && medium_count=5
    [ "$medium_count" -gt 20 ] && medium_count=20
    info "  Creating $medium_count medium files (1MB each)..."
    for i in $(seq 1 $medium_count); do
        dd if=/dev/urandom bs=1M count=1 of="$SRC_DIR/medium/file_$i.bin" 2>/dev/null
    done

    # Large files: remaining ~25%
    local large_size=$(( DATA_SIZE_MB - (small_count * 256 / 1024) - medium_count ))
    [ "$large_size" -lt 5 ] && large_size=5
    local large_count=3
    [ "$large_size" -lt 10 ] && large_count=1
    info "  Creating $large_count large files (~${large_size}MB total)..."
    for i in $(seq 1 $large_count); do
        local this_size=$(( large_size * 1048576 / large_count / 1048576 ))
        [ "$this_size" -lt 1 ] && this_size=1
        dd if=/dev/urandom bs=1M count=$this_size of="$SRC_DIR/large/file_$i.bin" 2>/dev/null
    done

    local total_files total_size
    total_files=$(find "$SRC_DIR" -type f | wc -l)
    total_size=$(du -sh "$SRC_DIR" | cut -f1)
    ok "Generated $total_files files, $total_size total"
}

# ── Verify sync ────────────────────────────────────────────────────────────
verify_sync() {
    local src_dir=$1
    local dst_dir=$2
    
    local src_count=$(find "$src_dir" -type f | wc -l)
    local dst_count=$(find "$dst_dir" -type f | wc -l)
    [ "$src_count" -ne "$dst_count" ] && return 1
    
    local src_size=$(find "$src_dir" -type f -exec stat -c %s {} \; | awk '{sum+=$1} END {print sum}')
    local dst_size=$(find "$dst_dir" -type f -exec stat -c %s {} \; | awk '{sum+=$1} END {print sum}')
    [ "$src_size" -ne "$dst_size" ] && return 1
    
    return 0
}

# ── Run single test with connection count ──────────────────────────────────
run_test() {
    local nconn=$1
    local dst_dir="$DST_DIR/fastsync_${nconn}"
    local server_port=$((PORT + nconn))
    
    rm -rf "$dst_dir"
    mkdir -p "$dst_dir"
    
    # Start server
    "$SERVER_BIN" -p "$server_port" -d "$dst_dir" > /dev/null 2>&1 &
    local srv_pid=$!
    sleep 0.2
    
    # Run client
    local client_output
    client_output=$(timeout 300 "$CLIENT_BIN" -h "$HOST" -p "$server_port" -s "$SRC_DIR" -n "$nconn" 2>&1) || true
    
    # Wait for completion
    sleep 3
    kill "$srv_pid" 2>/dev/null || true
    wait "$srv_pid" 2>/dev/null || true
    sync
    
    # Extract results
    local ms mbs
    ms=$(echo "$client_output" | grep "^BENCH:" | grep -oP 'ms=\K[0-9.]+' || echo "0")
    mbs=$(echo "$client_output" | grep "^BENCH:" | grep -oP 'throughput_mbs=\K[0-9.]+' || echo "0")
    
    local verified=1
    if ! verify_sync "$SRC_DIR" "$dst_dir"; then
        verified=0
    fi
    
    echo "${ms}|${mbs}|${verified}"
    
    # Cleanup
    rm -rf "$dst_dir"
}

# ── Bar chart helper ────────────────────────────────────────────────────────
draw_bar() {
    local val=$1
    local max=$2
    local width=30
    
    [ "$max" -eq 0 ] && max=1
    
    local filled
    filled=$(LC_ALL=C awk "BEGIN{printf \"%d\", ($val * $width / $max)}")
    [ "$filled" -gt "$width" ] && filled=$width
    
    local bar=""
    for ((i=0; i<filled; i++)); do bar+="█"; done
    for ((i=filled; i<width; i++)); do bar+="░"; done
    echo "$bar"
}

# ────────────────────────────────────────────────────────────────────────────
# MAIN
# ────────────────────────────────────────────────────────────────────────────

echo ""
echo -e "${BOLD}╔═══════════════════════════════════════════════════════════════════════════╗${RESET}"
echo -e "${BOLD}║     fastSyncAI — Network Condition Benchmark${RESET}"
echo -e "${BOLD}╚═══════════════════════════════════════════════════════════════════════════╝${RESET}"
echo ""

echo -e "  Scenario       : ${CYAN}${SCENARIO_NAME}${RESET}"
echo -e "  Data size      : ${CYAN}${DATA_SIZE_MB} MB${RESET}"
echo -e "  Runs per test  : ${CYAN}${RUN_COUNT}${RESET}"
echo -e "  Latency        : ${CYAN}${LATENCY_MS}ms RTT${RESET}"
echo -e "  Packet loss    : ${CYAN}${PACKET_LOSS}${RESET}"
echo -e "  Jitter         : ${CYAN}±${JITTER_MS}ms${RESET}"
echo ""

# Setup network conditions
setup_network
echo ""

# Connection counts to test
CONN_COUNTS=(1 2 4 8 16)

# Initialize results
declare -a FS_RESULTS
FS_LABELS=("fastSyncAI (1 conn)" "fastSyncAI (2 conn)" "fastSyncAI (4 conn)" "fastSyncAI (8 conn)" "fastSyncAI (16 conn)")

generate_test_data
echo ""

# Run tests
info "Testing fastSyncAI with different connection counts under ${SCENARIO_NAME}..."
for nconn in "${CONN_COUNTS[@]}"; do
    echo -n "  Testing ${nconn} connections... "
    result=$(run_test "$nconn")
    IFS='|' read -r ms mbs verified <<< "$result"
    FS_RESULTS+=("$ms|$mbs|$verified")
    if [ "$verified" -eq 1 ]; then
        ok "${mbs} MB/s (${ms}ms)"
    else
        warn "${mbs} MB/s (${ms}ms) - VERIFICATION FAILED"
    fi
done

echo ""
echo -e "${BOLD}═══════════════════════════════════════════════════════════════════════════╗${RESET}"
echo -e "${BOLD}                    NETWORK CONDITION RESULTS${RESET}"
echo -e "${BOLD}                    Scenario: ${SCENARIO_NAME}${RESET}"
echo -e "${BOLD}═══════════════════════════════════════════════════════════════════════════╝${RESET}"
echo ""

# Calculate max MB/s for bar scaling
calc_max_mbs() {
    local max=0
    for res in "${FS_RESULTS[@]}"; do
        local mbs=$(echo "$res" | cut -d'|' -f2)
        local int_mbs=$(echo "$mbs" | awk '{printf "%d", $1+0}')
        [ "$int_mbs" -gt "$max" ] && max=$int_mbs
    done
    echo $max
}

MAX_MBS=$(calc_max_mbs)
[ "$MAX_MBS" -eq 0 ] && MAX_MBS=100

printf "  %-25s  %-10s  %-10s  %s  %s\n" "Configuration" "Time (ms)" "MB/s" "Throughput" "Verified"
printf "  %-25s  %-10s  %-10s  %s  %s\n" "-------------------------" "----------" "----------" "-----------" "--------"

for i in "${!FS_RESULTS[@]}"; do
    IFS='|' read -r ms mbs verified <<< "${FS_RESULTS[$i]}"
    bar=$(draw_bar "$mbs" "$MAX_MBS")
    verified_str=$([ "$verified" -eq 1 ] && echo "✓" || echo "✗")
    printf "  %-25s  %-10s  %-10s  %s  %s\n" "${FS_LABELS[$i]}" "$ms" "$mbs" "$bar" "$verified_str"
done

# Find best configuration
best_idx=0
best_mbs=0
for i in "${!FS_RESULTS[@]}"; do
    IFS='|' read -r ms mbs verified <<< "${FS_RESULTS[$i]}"
    mbs_int=$(echo "$mbs" | awk '{printf "%d", $1+0}')
    if [ "$mbs_int" -gt "$best_mbs" ]; then
        best_mbs=$mbs_int
        best_idx=$i
    fi
done

echo ""
echo -e "${BOLD}Best configuration for ${SCENARIO_NAME}:${RESET} ${FS_LABELS[$best_idx]} (${best_mbs} MB/s)"

echo ""
ok "Network benchmark complete!"