#!/usr/bin/env bash
# =============================================================================
# benchmark.sh — fastSyncAI connection-count benchmark
#
# Usage:
#   ./benchmark.sh [latency_ms]
#
# Examples:
#   ./benchmark.sh           # loopback, no extra latency
#   ./benchmark.sh 20        # simulate 20 ms RTT via tc netem (needs sudo)
#   ./benchmark.sh 100       # simulate 100 ms RTT (WAN)
#
# What it does:
#   1. Generates a 50 MB mixed test tree (small / medium / large files)
#   2. Optionally adds RTT latency to lo using tc netem
#   3. Runs fastsync_client with 1, 2, 4, 8, 16 connections
#   4. Prints a formatted throughput comparison table
# =============================================================================

set -euo pipefail

PORT=19099
DATA_DIR="bench_data"
DEST_DIR="bench_dest"
SERVER_BIN="./fastsync_server"
CLIENT_BIN="./fastsync_client"
HOST="127.0.0.1"
LATENCY_MS="${1:-0}"
CONNECTIONS=(1 2 4 8 16)

# ── 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}"; }

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

# ── Generate test data ────────────────────────────────────────────────────────
generate_data() {
    info "Generating test data (~50 MB mixed tree)..."
    rm -rf "$DATA_DIR"
    mkdir -p "$DATA_DIR/small" "$DATA_DIR/medium" "$DATA_DIR/large"

    # 40 × 256 KB = 10 MB (many small files → tests queue/threading overhead)
    for i in $(seq 1 40); do
        dd if=/dev/urandom bs=256K count=1 of="$DATA_DIR/small/file_$i.bin" 2>/dev/null
    done

    # 10 × 1 MB = 10 MB (medium)
    for i in $(seq 1 10); do
        dd if=/dev/urandom bs=1M count=1 of="$DATA_DIR/medium/file_$i.bin" 2>/dev/null
    done

    # 6 × 5 MB = 30 MB (large — stresses TCP buffer sizing)
    for i in $(seq 1 6); do
        dd if=/dev/urandom bs=5M count=1 of="$DATA_DIR/large/file_$i.bin" 2>/dev/null
    done

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

# ── 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
        # Half-latency each way → full RTT = LATENCY_MS
        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
}

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

# ── Run one pass and return "ms|mbs" ──────────────────────────────────────────
run_pass() {
    local nconn=$1
    rm -rf "$DEST_DIR"

    "$SERVER_BIN" -p "$PORT" -d "$DEST_DIR" >/dev/null 2>&1 &
    local srv_pid=$!
    sleep 0.15   # let server bind

    local output
    output=$("$CLIENT_BIN" -h "$HOST" -p "$PORT" -s "$DATA_DIR" -n "$nconn" 2>&1)

    kill "$srv_pid" 2>/dev/null
    wait "$srv_pid" 2>/dev/null

    # Parse the machine-readable BENCH: line
    local ms mbs
    ms=$(echo "$output"  | grep "^BENCH:" | grep -oP 'ms=\K[0-9.]+')
    mbs=$(echo "$output" | grep "^BENCH:" | grep -oP 'throughput_mbs=\K[0-9.]+')

    echo "${ms}|${mbs}"
}

# ── Bar chart helper (ASCII) ──────────────────────────────────────────────────
draw_bar() {
    local val=$1 max=$2 width=30
    local filled
    filled=$(LC_ALL=C awk "BEGIN{printf \"%d\", ($val * $width / ($max == 0 ? 1 : $max))}")
    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 — Multi-Connection Benchmark           ║${RESET}"
echo -e "${BOLD}╚══════════════════════════════════════════════════════════════╝${RESET}"
echo ""

if [ "$LATENCY_MS" -gt 0 ]; then
    echo -e "  Simulated RTT latency : ${YELLOW}${LATENCY_MS} ms${RESET}"
else
    echo -e "  Simulated RTT latency : ${GREEN}none (raw loopback)${RESET}"
fi
echo -e "  Server port           : $PORT"
echo -e "  Connection counts     : ${CONNECTIONS[*]}"
echo ""

generate_data
setup_latency
echo ""

# Collect results
declare -a RES_MS RES_MBS
info "Running benchmark passes..."
echo ""

MAX_MBS="0"
for n in "${CONNECTIONS[@]}"; do
    echo -ne "  Testing ${BOLD}${n}${RESET} connection(s)... "
    result=$(run_pass "$n")
    IFS='|' read -r ms mbs <<< "$result"
    RES_MS+=("$ms")
    RES_MBS+=("$mbs")
    # Track max throughput for bar chart scaling
    if LC_ALL=C awk "BEGIN{exit !(\"$mbs\"+0 > \"$MAX_MBS\"+0)}"; then MAX_MBS="$mbs"; fi
    echo -e "${GREEN}${mbs} MB/s${RESET} in ${ms} ms"
done

teardown_latency

# ── Results table ────────────────────────────────────────────────────────────
echo ""
echo -e "${BOLD}Results${RESET}"
echo ""
echo -e "  ${BOLD}$(printf '%-11s  %-10s  %-10s  %-32s' 'Connections' 'Time (ms)' 'MB/s' 'Throughput')${RESET}"
echo    "  $(printf '%-11s  %-10s  %-10s  %-32s' '-----------' '---------' '----' '---------')"

for i in "${!CONNECTIONS[@]}"; do
    n="${CONNECTIONS[$i]}"
    ms="${RES_MS[$i]}"
    mbs="${RES_MBS[$i]}"
    bar=$(draw_bar "$mbs" "$MAX_MBS" 2>/dev/null || echo "")
    printf "  %-11s  %-10s  %-10s  %s\n" "$n" "$ms" "$mbs" "$bar"
done

echo ""

# ── Speedup analysis ─────────────────────────────────────────────────────────
base_mbs="${RES_MBS[0]}"
echo -e "${BOLD}Speedup vs. 1 connection${RESET}"
echo ""
for i in "${!CONNECTIONS[@]}"; do
    n="${CONNECTIONS[$i]}"
    mbs="${RES_MBS[$i]}"
    speedup=$(LC_ALL=C awk "BEGIN{printf \"%.2f\", $mbs / ($base_mbs == 0 ? 1 : $base_mbs)}")
    if LC_ALL=C awk "BEGIN{exit !($speedup >= 2.0)}"; then
        color="$GREEN"
    elif LC_ALL=C awk "BEGIN{exit !($speedup >= 1.3)}"; then
        color="$YELLOW"
    else
        color="$RESET"
    fi
    LC_ALL=C awk -v n="$n" -v s="$speedup" -v c="$color" -v r="$RESET" \
        'BEGIN{printf "  %2s connection(s): %s%.2fx%s\n", n, c, s+0, r}'
done

echo ""
ok "Benchmark done. Cleaning up..."
rm -rf "$DATA_DIR" "$DEST_DIR"
echo ""