#!/usr/bin/env bash
# =============================================================================
# benchmark_network.sh — Network Condition Benchmark for fastSyncAI
#
# Compares fastSyncAI and rsync under IDENTICAL network conditions
# by running both over TCP on the loopback interface with tc netem simulation.
#
# 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
#
# Examples:
#   ./benchmark_network.sh --lan
#   ./benchmark_network.sh --wan 200 1
#   ./benchmark_network.sh --wan-loss-1 50
# =============================================================================

set -euo pipefail

# ── Defaults ──────────────────────────────────────────────────────────────
DATA_SIZE_MB=${2:-50}
RUN_COUNT=${3:-1}
SCENARIO="${1:---lan}"

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

case "$SCENARIO" in
    --lan)           LATENCY_MS=10;  PACKET_LOSS="0.1%"; JITTER_MS=2; SCENARIO_NAME="LAN (10ms RTT, 0.1% loss, ±2ms jitter)" ;;
    --wan)           LATENCY_MS=100; PACKET_LOSS="1%";  JITTER_MS=10; SCENARIO_NAME="WAN (100ms RTT, 1% loss, ±10ms jitter)" ;;
    --wan-loss-1)    LATENCY_MS=100; PACKET_LOSS="1%";  SCENARIO_NAME="WAN + 1% loss" ;;
    --wan-loss-5)    LATENCY_MS=100; PACKET_LOSS="5%";  SCENARIO_NAME="WAN + 5% loss" ;;
    --wan-jitter)    LATENCY_MS=100; JITTER_MS=50;   SCENARIO_NAME="WAN + jitter" ;;
    --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"
            exit 1
        fi
        ;;
    *)
        echo "ERROR: Unknown scenario '$SCENARIO'"
        echo "Available: --lan, --wan, --wan-loss-1, --wan-loss-5, --wan-jitter, --custom"
        exit 1
        ;;
esac

# ── Ports ────────────────────────────────────────────────────────────────
FASTSYNC_PORT=19410
RSYNC_PORT=19400
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_BASE="$TEST_DIR/dst"
RSYNC_SHARE="$TEST_DIR/rsync_share"

# ── 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 ────────────────────────────────────────────────────
setup_network() {
    info "Setting up network: ${SCENARIO_NAME}..."
    if ! sudo -n true 2>/dev/null; then
        echo -e "${YELLOW}⚠ sudo required for tc netem${RESET}"
    fi
    sudo tc qdisc del dev lo root 2>/dev/null || true
    local netem_cmd="delay ${LATENCY_MS}ms"
    [ "$JITTER_MS" -gt 0 ] && netem_cmd="$netem_cmd ${JITTER_MS}ms"
    [ "$PACKET_LOSS" != "0%" ] && netem_cmd="$netem_cmd loss ${PACKET_LOSS}"
    sudo tc qdisc add dev lo root netem $netem_cmd || die "tc netem failed"
    ok "Network active: ${LATENCY_MS}ms RTT, loss=${PACKET_LOSS}, jitter=±${JITTER_MS}ms"
}

cleanup_network() {
    sudo tc qdisc del dev lo root 2>/dev/null || true
}

# ── rsync daemon ────────────────────────────────────────────────────────────
RSYNCD_PID=""

start_rsyncd() {
    # Get absolute path for the share
    local abs_share
    abs_share=$(cd "$RSYNC_SHARE" && pwd)
    
    # Write config file with absolute path to share
    cat > "$TEST_DIR/rsyncd.conf" <<EOF
[network_bench]
    path = $abs_share
    read only = false
    use chroot = false
    max connections = 0
EOF
    
    # Start daemon from TEST_DIR using relative config path
    # Note: rsync --daemon forks, so we just check the port
    cd "$TEST_DIR"
    rsync --daemon --port="$RSYNC_PORT" --config=rsyncd.conf 2>/dev/null
    cd - >/dev/null
    sleep 2
    
    # Verify port is listening (rsyncd forks so PID check doesn't work)
    if ! ss -tlnp 2>/dev/null | grep -q ":$RSYNC_PORT "; then
        # Try to get PID
        RSYNCD_PID=$(ss -tlnp 2>/dev/null | grep ":$RSYNC_PORT " | head -1 | grep -oP 'pid=\K[0-9]+' || echo "")
        die "rsyncd not listening on port $RSYNC_PORT (PID=$RSYNCD_PID)"
    fi
    
    # Get actual PID (use head -1 to handle IPv4/IPv6 duplicate)
    RSYNCD_PID=$(ss -tlnp 2>/dev/null | grep ":$RSYNC_PORT " | head -1 | grep -oP 'pid=\K[0-9]+' || echo "")
    
    ok "rsyncd started on port $RSYNC_PORT (PID=$RSYNCD_PID)"
}

stop_rsyncd() {
    [ -n "$RSYNCD_PID" ] && kill "$RSYNCD_PID" 2>/dev/null || true
    wait "$RSYNCD_PID" 2>/dev/null || true
    RSYNCD_PID=""
}

# ── 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."
command -v rsync >/dev/null 2>&1 || die "rsync not found."

cleanup() {
    stop_rsyncd
    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

# Clean stale processes
pkill -f "rsync --daemon" 2>/dev/null || true
pkill -f "fastsync_server" 2>/dev/null || true
pkill -f "fastsync_client" 2>/dev/null || true
cleanup_network 2>/dev/null || true
sleep 1

# ── Generate test data ──────────────────────────────────────────────────────
generate_data() {
    info "Generating ~${DATA_SIZE_MB}MB test data..."
    rm -rf "$TEST_DIR"
    mkdir -p "$SRC_DIR/small" "$SRC_DIR/medium" "$SRC_DIR/large" "$RSYNC_SHARE"
    
    # Small files (~25% of total)
    local small_count=$(( (DATA_SIZE_MB * 262144) / 1048576 ))
    [ "$small_count" -lt 10 ] && small_count=10
    [ "$small_count" -gt 100 ] && small_count=100
    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 (~50% of total)
    local medium_count=$(( (DATA_SIZE_MB * 524288) / 1048576 ))
    [ "$medium_count" -lt 5 ] && medium_count=5
    [ "$medium_count" -gt 20 ] && medium_count=20
    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 (~25%)
    local large_mb=$(( DATA_SIZE_MB - (small_count * 256 / 1024) - medium_count ))
    [ "$large_mb" -lt 5 ] && large_mb=5
    local large_count=3
    [ "$large_mb" -lt 10 ] && large_count=1
    for i in $(seq 1 $large_count); do
        local sz=$(( large_mb * 1048576 / large_count / 1048576 ))
        [ "$sz" -lt 1 ] && sz=1
        dd if=/dev/urandom bs=1M count=$sz of="$SRC_DIR/large/file_$i.bin" 2>/dev/null
    done
    
    # Copy to rsync share
    cp -r "$SRC_DIR"/* "$RSYNC_SHARE/"
    
    local total_files=$(find "$SRC_DIR" -type f | wc -l)
    local total_size=$(du -sh "$SRC_DIR" | cut -f1)
    ok "Generated $total_files files, $total_size"
}

# ── Verify ──────────────────────────────────────────────────────────────────
verify() {
    local src=$1 dst=$2
    local sc=$(find "$src" -type f | wc -l)
    local dc=$(find "$dst" -type f | wc -l)
    [ "$sc" -ne "$dc" ] && return 1
    local ss=$(find "$src" -type f -printf '%s\n' | awk '{sum+=$1} END {print sum}')
    local ds=$(find "$dst" -type f -printf '%s\n' | awk '{sum+=$1} END {print sum}')
    [ "$ss" -ne "$ds" ] && return 1
    return 0
}

# ── Run fastSyncAI ────────────────────────────────────────────────────────
run_fastsync() {
    local nconn=$1
    local dst="$DST_DIR_BASE/fastsync_${nconn}"
    local port=$((FASTSYNC_PORT + nconn))
    rm -rf "$dst"
    mkdir -p "$dst"
    
    "$SERVER_BIN" -p "$port" -d "$dst" >/dev/null 2>&1 &
    local spid=$!
    sleep 1
    
    local output
    output=$(timeout 600 "$CLIENT_BIN" -h "$HOST" -p "$port" -s "$SRC_DIR" -n "$nconn" 2>&1) || true
    
    # Extract actual transfer time from client output
    local transfer_ms=$(echo "$output" | grep -oP 'ms=\K[0-9.]+' || echo "0")
    local transfer_ms_int=$(echo "$transfer_ms" | awk '{printf "%d", $1+0}')
    
    # Wait: transfer time + latency-based margin + disk sync margin
    # For 1 connection, add extra margin as it's sequential
    local latency_margin=$(( LATENCY_MS * 2 ))
    local sync_margin=3
    local extra_for_1conn=0
    [ "$nconn" -eq 1 ] && extra_for_1conn=5
    
    local wait_time=$(( transfer_ms_int / 1000 + latency_margin / 10 + sync_margin + extra_for_1conn ))
    [ "$wait_time" -lt 5 ] && wait_time=5
    
    sleep $wait_time
    kill "$spid" 2>/dev/null || true
    wait "$spid" 2>/dev/null || true
    sync
    sleep 1
    
    local ms=$(echo "$output" | grep -oP 'ms=\K[0-9.]+' || echo "0")
    local mbs=$(echo "$output" | grep -oP 'throughput_mbs=\K[0-9.]+' || echo "0")
    local ok=1
    verify "$SRC_DIR" "$dst" || ok=0
    echo "${ms}|${mbs}|${ok}"
    rm -rf "$dst"
}

# ── Run rsync over TCP ───────────────────────────────────────────────────
run_rsync_tcp() {
    local use_compress=${1:-0}
    local dst="$DST_DIR_BASE/rsync_${use_compress}"
    rm -rf "$dst"
    mkdir -p "$dst"
    
    local start=$(date +%s%N)
    local compress_flag=""
    [ "$use_compress" -eq 1 ] && compress_flag="--compress"
    
    rsync -a $compress_flag "rsync://$HOST:$RSYNC_PORT/network_bench/" "$dst/" >/dev/null 2>&1
    
    local end=$(date +%s%N)
    local ms=$(( (end - start) / 1000000 ))
    
    local total_bytes=$(find "$SRC_DIR" -type f -printf '%s\n' | awk '{sum+=$1} END {print sum}')
    local mbs=0
    [ "$ms" -gt 0 ] && mbs=$(awk "BEGIN {printf \"%.2f\", $total_bytes / 1048576.0 / ($ms / 1000.0)}")
    
    local ok=1
    verify "$SRC_DIR" "$dst" || ok=0
    echo "${ms}|${mbs}|${ok}"
    rm -rf "$dst"
}

# ── Bar chart ─────────────────────────────────────────────────────────────
draw_bar() {
    local val=$1 max=$2 width=30
    [ "$max" -eq 0 ] && max=1
    # Convert to integers for arithmetic
    local val_int=$(echo "$val" | awk '{printf "%d", $1+0}')
    local max_int=$(echo "$max" | awk '{printf "%d", $1+0}')
    [ "$max_int" -eq 0 ] && max_int=1
    local filled=$(( val_int * width / max_int ))
    [ "$filled" -gt "$width" ] && filled=$width
    printf -v 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 — Fair Network Comparison Benchmark${RESET}"
echo -e "${BOLD}     (All tools tested with IDENTICAL network conditions)${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 "  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
setup_network

# Generate data and start rsyncd
generate_data
start_rsyncd
echo ""

# Test configurations
CONN_COUNTS=(1 2 4 8 16)
FS_LABELS=("fastSyncAI 1 conn" "fastSyncAI 2 conn" "fastSyncAI 4 conn" "fastSyncAI 8 conn" "fastSyncAI 16 conn")

info "Testing over network (${SCENARIO_NAME})..."
echo ""

# Collect results
declare -a FS_RES RSYNC_RES RSYNCC_RES

# fastSyncAI tests
for i in "${!CONN_COUNTS[@]}"; do
    n=${CONN_COUNTS[$i]}
    echo -n "  ${FS_LABELS[$i]}... "
    result=$(run_fastsync $n)
    IFS='|' read -r ms mbs ok <<< "$result"
    FS_RES+=("$ms|$mbs|$ok")
    [ "$ok" -eq 1 ] && ok "${mbs} MB/s (${ms}ms)" || warn "${mbs} MB/s (${ms}ms) FAILED"
done

# rsync tests (over TCP with network simulation)
echo -n "  rsync (TCP)... "
result=$(run_rsync_tcp 0)
IFS='|' read -r ms mbs ok <<< "$result"
RSYNC_RES=("$ms|$mbs|$ok")
[ "$ok" -eq 1 ] && ok "${mbs} MB/s (${ms}ms)" || warn "${mbs} MB/s (${ms}ms) FAILED"

echo -n "  rsync + compress (TCP)... "
result=$(run_rsync_tcp 1)
IFS='|' read -r ms mbs ok <<< "$result"
RSYNCC_RES=("$ms|$mbs|$ok")
[ "$ok" -eq 1 ] && ok "${mbs} MB/s (${ms}ms)" || warn "${mbs} MB/s (${ms}ms) FAILED"

# Cleanup network
cleanup_network

echo ""
echo -e "${BOLD}═══════════════════════════════════════════════════════════════════════════╗${RESET}"
echo -e "${BOLD}                    FAIR COMPARISON RESULTS${RESET}"
echo -e "${BOLD}                    (All tools over TCP with ${SCENARIO_NAME})${RESET}"
echo -e "${BOLD}═══════════════════════════════════════════════════════════════════════════╝${RESET}"
echo ""

# Calculate max for scaling
calc_max() {
    local max=0
    for res in "${FS_RES[@]}" "${RSYNC_RES[@]}" "${RSYNCC_RES[@]}"; do
        local mbs=$(echo "$res" | cut -d'|' -f2)
        local int=$(echo "$mbs" | awk '{printf "%d", $1+0}')
        [ "$int" -gt "$max" ] 2>/dev/null && max=$int
    done
    echo $max
}
MAX_MBS=$(calc_max)
[ "$MAX_MBS" -eq 0 ] && MAX_MBS=100

printf "  %-22s  %-10s  %-10s  %s\n" "Tool" "Time (ms)" "MB/s" "Throughput"
printf "  %-22s  %-10s  %-10s  %s\n" "----------------------" "----------" "----------" "------------"

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

IFS='|' read -r ms mbs ok <<< "${RSYNC_RES[0]}"
bar=$(draw_bar "$mbs" "$MAX_MBS")
v=$([ "$ok" -eq 1 ] && echo "✓" || echo "✗")
printf "  %-22s  %-10s  %-10s  %s %s\n" "rsync (TCP)" "$ms" "$mbs" "$bar" "$v"

IFS='|' read -r ms mbs ok <<< "${RSYNCC_RES[0]}"
bar=$(draw_bar "$mbs" "$MAX_MBS")
v=$([ "$ok" -eq 1 ] && echo "✓" || echo "✗")
printf "  %-22s  %-10s  %-10s  %s %s\n" "rsync+compress (TCP)" "$ms" "$mbs" "$bar" "$v"

# Analysis
echo ""
echo -e "${BOLD}═══════════════════════════════════════════════════════════════════════════╗${RESET}"
echo -e "${BOLD}                         ANALYSIS${RESET}"
echo -e "${BOLD}═══════════════════════════════════════════════════════════════════════════╝${RESET}"
echo ""

# Find best
best_idx=0
best_mbs=0
for i in "${!FS_RES[@]}"; do
    this_mbs=$(echo "${FS_RES[$i]}" | cut -d'|' -f2 | awk '{printf "%d", $1+0}')
    [ "$this_mbs" -gt "$best_mbs" ] && best_mbs=$this_mbs && best_idx=$i
done
best_label="${FS_LABELS[$best_idx]}"

IFS='|' read -r rs_ms rs_mbs rs_ok <<< "${RSYNC_RES[0]}"
rs_mbs_int=$(echo "$rs_mbs" | awk '{printf "%d", $1+0}')

ok "Best fastSyncAI: $best_label (${best_mbs} MB/s)"

if [ "$rs_mbs_int" -gt 0 ]; then
    echo ""
    echo "Speedup vs rsync (both over ${SCENARIO_NAME}):"
    speedup=$(awk "BEGIN {printf \"%.2f\", $best_mbs / ($rs_mbs_int == 0 ? 0.01 : $rs_mbs_int)}")
    if [ "$best_mbs" -gt "$rs_mbs_int" ]; then
        echo -e "  ${GREEN}fastSyncAI is ${speedup}x FASTER than rsync${RESET}"
    elif [ "$rs_mbs_int" -gt "$best_mbs" ]; then
        speedup=$(awk "BEGIN {printf \"%.2f\", $rs_mbs_int / ($best_mbs == 0 ? 0.01 : $best_mbs)}")
        echo -e "  ${GREEN}rsync is ${speedup}x faster than fastSyncAI${RESET}"
    else
        echo -e "  Both performed similarly"
    fi
fi

echo ""
ok "Benchmark complete!"
