#!/usr/bin/env bash
# =============================================================================
# compare_rsync.sh — Compare fastSyncAI performance against rsync
#
# Usage:
#   ./compare_rsync.sh [DATA_SIZE_MB] [RUN_COUNT] [LATENCY_MS|--lan|--wan]
#
# Examples:
#   ./compare_rsync.sh                          # 100 MB, 3 runs, no latency
#   ./compare_rsync.sh 500                      # 500 MB, 3 runs, no latency
#   ./compare_rsync.sh 100 5                    # 100 MB, 5 runs, no latency
#   ./compare_rsync.sh 100 1 20                 # 100 MB, 1 run, 20ms RTT (LAN-like)
#   ./compare_rsync.sh 100 1 100                # 100 MB, 1 run, 100ms RTT (WAN-like)
#   ./compare_rsync.sh 100 1 --lan              # 100 MB, 1 run, 10ms RTT (LAN preset)
#   ./compare_rsync.sh 100 1 --wan              # 100 MB, 1 run, 100ms RTT (WAN preset)
#
# Tests:
#   - Small files (many files, 256KB each)
#   - Medium files (fewer files, 1-5MB each)
#   - Large single file
#   - Mixed directory structure
#
# Compares: fastSyncAI (TCP, multi-connection) vs rsync vs rsync with compression
#
# Network Presets:
#   --lan   : 10ms RTT (typical LAN)
#   --wan   : 100ms RTT (typical WAN)
# =============================================================================

set -euo pipefail

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

# ── Parse latency argument ────────────────────────────────────────────────
LATENCY_MS=0
if [ $# -ge 3 ]; then
    case "${3}" in
        --lan)   LATENCY_MS=10  ;;
        --wan)   LATENCY_MS=100 ;;
        *)
            # Check if it's a numeric value
            if [[ "${3}" =~ ^[0-9]+$ ]]; then
                LATENCY_MS="${3}"
            fi
            ;;
    esac
fi

PORT=19199
HOST="127.0.0.1"

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

# ── Directories ────────────────────────────────────────────────────────────
TEST_DIR="compare_test_$$"
SRC_DIR="$TEST_DIR/src"
DST_DIR_BASE="$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}"; }

# ── Latency simulation (tc netem) ────────────────────────────────────────
setup_latency() {
    if [ "$LATENCY_MS" -gt 0 ]; then
        info "Adding ${LATENCY_MS}ms RTT latency to lo via tc netem..."
        if ! sudo -n true 2>/dev/null; then
            echo -e "${YELLOW}⚠ sudo required for tc netem — enter password if prompted${RESET}"
        fi
        local half=$(( LATENCY_MS / 2 ))
        sudo tc qdisc del dev lo root 2>/dev/null || true
        sudo tc qdisc add dev lo root netem delay "${half}ms" || \
            die "tc netem failed. Install iproute2 or run without latency argument."
        ok "Latency applied (${half}ms one-way = ${LATENCY_MS}ms RTT)"
    fi
}

cleanup_latency() {
    if [ "$LATENCY_MS" -gt 0 ]; then
        sudo tc qdisc del dev lo root 2>/dev/null || true
        ok "Latency removed"
    fi
}

# ── 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. Install rsync first."

cleanup() {
    cleanup_latency
    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 size in 256KB files
    local small_count=$(( (DATA_SIZE_MB * 1048576 / 4) / 262144 ))
    if [ "$small_count" -lt 1 ]; then small_count=10; fi
    if [ "$small_count" -gt 100 ]; then small_count=100; fi
    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
    local medium_count=$(( (DATA_SIZE_MB * 1048576 / 2) / 1048576 ))
    if [ "$medium_count" -lt 1 ]; then medium_count=5; fi
    if [ "$medium_count" -gt 20 ]; then medium_count=20; fi
    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 size in fewer large files
    local large_size=$(( DATA_SIZE_MB - (small_count * 256 / 1024) - medium_count ))
    if [ "$large_size" -lt 5 ]; then large_size=5; fi
    local large_count=3
    if [ "$large_size" -lt 10 ]; then large_count=1; fi
    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 ))
        if [ "$this_size" -lt 1 ]; then this_size=1; fi
        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 in $SRC_DIR"
}

# ── Timing helper ───────────────────────────────────────────────────────────
# Returns time in milliseconds
run_with_timer() {
    local start end
    start=$(date +%s%N)
    "$@" > /dev/null 2>&1
    end=$(date +%s%N)
    local elapsed_ms=$(( (end - start) / 1000000 ))
    echo $elapsed_ms
}

# ── Run fastSyncAI ──────────────────────────────────────────────────────────
# Usage: run_fastsync <connection_count>
run_fastsync() {
    local nconn=${1:-4}
    local dst_dir="$DST_DIR_BASE/fastsync_$$"
    mkdir -p "$dst_dir"
    
    # Clean destination
    rm -rf "$dst_dir"/*
    
    # Start server in background
    "$SERVER_BIN" -p "$PORT" -d "$dst_dir" > /dev/null 2>&1 &
    local srv_pid=$!
    sleep 0.2
    
    # Run client and capture output
    local client_output
    client_output=$(timeout 120 "$CLIENT_BIN" -h "$HOST" -p "$PORT" -s "$SRC_DIR" -n "$nconn" 2>&1) || true
    
    # Wait for server to finish writing all files
    # With latency, 1 connection needs more time
    local wait_time=3
    [ "$nconn" -eq 1 ] && [ "$LATENCY_MS" -gt 0 ] && wait_time=$(( 3 + LATENCY_MS / 10 + 2 ))
    sleep $wait_time
    # Clean up server
    kill "$srv_pid" 2>/dev/null || true
    wait "$srv_pid" 2>/dev/null || true
    # Force sync to ensure all disk writes are complete before verification
    sync
    
    # Extract throughput from BENCH line
    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")
    
    # Verify files using fast size-only check (much faster than md5sum)
    local src_count dst_count src_size dst_size
    src_count=$(find "$SRC_DIR" -type f | wc -l)
    dst_count=$(find "$dst_dir" -type f | wc -l)
    src_size=$(find "$SRC_DIR" -type f -printf '%s\n' | awk '{sum+=$1} END {print sum}')
    dst_size=$(find "$dst_dir" -type f -printf '%s\n' | awk '{sum+=$1} END {print sum}')
    
    if [ "$src_count" -eq "$dst_count" ] && [ "$src_size" -eq "$dst_size" ]; then
        echo "${ms}|${mbs}|1"
    else
        echo "${ms}|${mbs}|0"
    fi
    
    # Cleanup
    rm -rf "$dst_dir"
}

# ── Run rsync ───────────────────────────────────────────────────────────────
run_rsync() {
    local use_compress=${1:-0}
    local dst_dir="$DST_DIR_BASE/rsync_$$"
    mkdir -p "$dst_dir"
    
    # Clean destination
    rm -rf "$dst_dir"/*
    
    local start end
    start=$(date +%s%N)
    
    if [ "$use_compress" -eq 1 ]; then
        rsync -a --compress --progress "$SRC_DIR/" "$dst_dir/" > /dev/null 2>&1
    else
        rsync -a --progress "$SRC_DIR/" "$dst_dir/" > /dev/null 2>&1
    fi
    
    end=$(date +%s%N)
    local elapsed_ms=$(( (end - start) / 1000000 ))
    
    # Calculate throughput
    local total_bytes
    total_bytes=$(find "$SRC_DIR" -type f -exec stat -c %s {} \; | awk '{sum+=$1} END {print sum}')
    local mbs=0
    if [ "$elapsed_ms" -gt 0 ]; then
        mbs=$(awk "BEGIN {printf \"%.2f\", $total_bytes / 1048576.0 / ($elapsed_ms / 1000.0)}")
    fi
    
    # Verify files
    local verified=1
    if ! diff <(cd "$SRC_DIR" && find . -type f -exec md5sum {} \; | sort) \
            <(cd "$dst_dir" && find . -type f -exec md5sum {} \; | sort) > /dev/null 2>&1; then
        verified=0
    fi
    
    echo "${elapsed_ms}|${mbs}|${verified}"
    
    # Cleanup
    rm -rf "$dst_dir"
}

# ── Run comparison ──────────────────────────────────────────────────────────
run_comparison() {
    local run=$1
    info "Run $run/$RUN_COUNT..."
    
    # Generate fresh data for each run
    generate_test_data
    
    # Test fastSyncAI
    echo -n "  Testing fastSyncAI... "
    local fs_result
    fs_result=$(run_fastsync)
    IFS='|' read -r fs_ms fs_mbs fs_verified <<< "$fs_result"
    if [ "$fs_verified" -eq 1 ]; then
        ok "${fs_mbs} MB/s (${fs_ms}ms)"
    else
        warn "${fs_mbs} MB/s (${fs_ms}ms) - VERIFICATION FAILED"
    fi
    
    # Test rsync (no compression)
    echo -n "  Testing rsync... "
    local rs_result
    rs_result=$(run_rsync 0)
    IFS='|' read -r rs_ms rs_mbs rs_verified <<< "$rs_result"
    if [ "$rs_verified" -eq 1 ]; then
        ok "${rs_mbs} MB/s (${rs_ms}ms)"
    else
        warn "${rs_mbs} MB/s (${rs_ms}ms) - VERIFICATION FAILED"
    fi
    
    # Test rsync with compression
    echo -n "  Testing rsync + compress... "
    local rsc_result
    rsc_result=$(run_rsync 1)
    IFS='|' read -r rsc_ms rsc_mbs rsc_verified <<< "$rsc_result"
    if [ "$rsc_verified" -eq 1 ]; then
        ok "${rsc_mbs} MB/s (${rsc_ms}ms)"
    else
        warn "${rsc_mbs} MB/s (${rsc_ms}ms) - VERIFICATION FAILED"
    fi
    
    echo ""
    echo "  === Run $run Results ==="
    printf "  %-20s  %-10s  %-10s  %s\n" "Tool" "Time (ms)" "MB/s" "Verified"
    printf "  %-20s  %-10s  %-10s  %s\n" "--------------------" "----------" "----------" "---------"
    printf "  %-20s  %-10s  %-10s  %s\n" "fastSyncAI" "$fs_ms" "$fs_mbs" "($([ "$fs_verified" -eq 1 ] && echo "✓" || echo "✗"))"
    printf "  %-20s  %-10s  %-10s  %s\n" "rsync" "$rs_ms" "$rs_mbs" "($([ "$rs_verified" -eq 1 ] && echo "✓" || echo "✗"))"
    printf "  %-20s  %-10s  %-10s  %s\n" "rsync+compress" "$rsc_ms" "$rsc_mbs" "($([ "$rsc_verified" -eq 1 ] && echo "✓" || echo "✗"))"
    echo ""
    
    # Return formatted results
    echo "RUN:$run|fastSyncAI:${fs_ms}:${fs_mbs}:${fs_verified}|rsync:${rs_ms}:${rs_mbs}:${rs_verified}|rsync_c:${rsc_ms}:${rsc_mbs}:${rsc_verified}"
}

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

echo ""
echo -e "${BOLD}╔══════════════════════════════════════════════════════════════╗${RESET}"
echo -e "${BOLD}║        fastSyncAI vs rsync — Performance Comparison             ║${RESET}"
echo -e "${BOLD}╚══════════════════════════════════════════════════════════════╝${RESET}"
echo ""

# Display network condition
if [ "$LATENCY_MS" -gt 0 ]; then
    echo -e "  Data size     : ${CYAN}${DATA_SIZE_MB} MB${RESET}"
    echo -e "  Runs per test : ${CYAN}${RUN_COUNT}${RESET}"
    echo -e "  Network       : ${CYAN}${LATENCY_MS}ms RTT${RESET}"
    echo -e "  Server port   : ${CYAN}${PORT}${RESET}"
else
    echo -e "  Data size     : ${CYAN}${DATA_SIZE_MB} MB${RESET}"
    echo -e "  Runs per test : ${CYAN}${RUN_COUNT}${RESET}"
    echo -e "  Network       : ${GREEN}no latency (loopback)${RESET}"
    echo -e "  Server port   : ${CYAN}${PORT}${RESET}"
fi
echo ""

# Setup latency if configured
setup_latency
echo ""

# Collect all results
declare -a ALL_RESULTS

for run in $(seq 1 $RUN_COUNT); do
    result=$(run_comparison $run)
    ALL_RESULTS+=("$result")
done

# ── Print summary ────────────────────────────────────────────────────────────
echo ""
echo -e "${BOLD}════════════════════════════════════════════════════════════════════════╗${RESET}"
echo -e "${BOLD}                         SUMMARY${RESET}"
echo -e "${BOLD}════════════════════════════════════════════════════════════════════════╝${RESET}"
echo ""

# Parse all results
declare -A FS_MS FS_MBS FS_VERIFIED
declare -A RS_MS RS_MBS RS_VERIFIED
declare -A RSC_MS RSC_MBS RSC_VERIFIED

for result in "${ALL_RESULTS[@]}"; do
    # Extract run number
    run_num=$(echo "$result" | grep -oP 'RUN:\K[0-9]+')
    
    # Extract fastSyncAI
    fs_data=$(echo "$result" | grep -oP 'fastSyncAI:\K[^|]+')
    IFS=':' read -r fs_ms fs_mbs fs_v <<< "$fs_data"
    FS_MS[$run_num]="$fs_ms"
    FS_MBS[$run_num]="$fs_mbs"
    FS_VERIFIED[$run_num]="$fs_v"
    
    # Extract rsync
    rs_data=$(echo "$result" | grep -oP 'rsync:\K[^|]+')
    IFS=':' read -r rs_ms rs_mbs rs_v <<< "$rs_data"
    RS_MS[$run_num]="$rs_ms"
    RS_MBS[$run_num]="$rs_mbs"
    RS_VERIFIED[$run_num]="$rs_v"
    
    # Extract rsync+compress
    rsc_data=$(echo "$result" | grep -oP 'rsync_c:\K[^|]+')
    IFS=':' read -r rsc_ms rsc_mbs rsc_v <<< "$rsc_data"
    RSC_MS[$run_num]="$rsc_ms"
    RSC_MBS[$run_num]="$rsc_mbs"
    RSC_VERIFIED[$run_num]="$rsc_v"
done

# Calculate averages
calc_avg() {
    local -n arr=$1
    local sum=0 count=0
    for val in "${arr[@]}"; do
        sum=$(awk "BEGIN {print $sum + $val}")
        count=$((count + 1))
    done
    if [ "$count" -gt 0 ]; then
        awk "BEGIN {printf \"%.1f\", $sum / $count}"
    else
        echo "0"
    fi
}

fs_avg_ms=$(calc_avg FS_MS)
fs_avg_mbs=$(calc_avg FS_MBS)
rs_avg_ms=$(calc_avg RS_MS)
rs_avg_mbs=$(calc_avg RS_MBS)
rsc_avg_ms=$(calc_avg RSC_MS)
rsc_avg_mbs=$(calc_avg RSC_MBS)

# Check all verified
all_fs_ok=1; for v in "${FS_VERIFIED[@]}"; do [ "$v" -ne 1 ] && all_fs_ok=0; done
all_rs_ok=1; for v in "${RS_VERIFIED[@]}"; do [ "$v" -ne 1 ] && all_rs_ok=0; done
all_rsc_ok=1; for v in "${RSC_VERIFIED[@]}"; do [ "$v" -ne 1 ] && all_rsc_ok=0; done

# Print averages
printf "  %-20s  %-10s  %-10s  %s\n" "Tool" "Avg Time (ms)" "Avg MB/s" "All OK"
printf "  %-20s  %-10s  %-10s  %s\n" "--------------------" "------------" "----------" "-------"
printf "  %-20s  %-10s  %-10s  %s\n" "fastSyncAI" "$fs_avg_ms" "$fs_avg_mbs" "($([ "$all_fs_ok" -eq 1 ] && echo "✓" || echo "✗"))"
printf "  %-20s  %-10s  %-10s  %s\n" "rsync" "$rs_avg_ms" "$rs_avg_mbs" "($([ "$all_rs_ok" -eq 1 ] && echo "✓" || echo "✗"))"
printf "  %-20s  %-10s  %-10s  %s\n" "rsync+compress" "$rsc_avg_ms" "$rsc_avg_mbs" "($([ "$all_rsc_ok" -eq 1 ] && echo "✓" || echo "✗"))"

# Speedup calculation
echo ""
if [ "$(awk "BEGIN {print ($rs_avg_mbs > 0) ? 1 : 0}")" = "1" ]; then
    fs_speedup=$(awk "BEGIN {printf \"%.2f\", $fs_avg_mbs / ($rs_avg_mbs == 0 ? 0.01 : $rs_avg_mbs)}")
    
    # Compare as strings to avoid floating point issues
    if [ "$(awk "BEGIN {print ($fs_avg_mbs > $rs_avg_mbs) ? 1 : 0}")" = "1" ]; then
        echo -e "  fastSyncAI is ${GREEN}${fs_speedup}x faster than rsync${RESET}"
    elif [ "$(awk "BEGIN {print ($rs_avg_mbs > $fs_avg_mbs) ? 1 : 0}")" = "1" ]; then
        speedup_inv=$(awk "BEGIN {printf \"%.2f\", $rs_avg_mbs / ($fs_avg_mbs == 0 ? 0.01 : $fs_avg_mbs)}")
        echo -e "  rsync is ${GREEN}${speedup_inv}x faster than fastSyncAI${RESET}"
    else
        echo -e "  Both tools performed similarly"
    fi
fi

echo ""
ok "Comparison complete!"