#!/bin/bash
set -e

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

NODES=${NODES:-5}
DURATION=${DURATION:-10}
MSG_RATE=${MSG_RATE:-100}
OUTPUT_DIR="${OUTPUT_DIR:-$SCRIPT_DIR/results}"

mkdir -p "$OUTPUT_DIR"
TIMESTAMP=$(date +%Y%m%d_%H%M%S)
RESULT_FILE="$OUTPUT_DIR/throughput_${TIMESTAMP}.json"

echo "=== SWIM Throughput Benchmark Suite ==="
echo "Nodes:    $NODES"
echo "Duration: ${DURATION}s"
echo "Msg Rate: ${MSG_RATE} msg/s per node"
echo "Output:   $RESULT_FILE"
echo ""

cd "$SCRIPT_DIR"

echo "Building Go throughput benchmarks..."
go build -o bin/memberlist_throughput ./cmd/memberlist_throughput 2>/dev/null || {
    echo "Warning: Failed to build memberlist throughput benchmark"
}
go build -o bin/serf_throughput ./cmd/serf_throughput 2>/dev/null || {
    echo "Warning: Failed to build serf throughput benchmark"
}

echo "Building OCaml throughput benchmark..."
cd "$PROJECT_ROOT"
dune build bench/swim_throughput.exe

echo ""
echo "=== Running Throughput Benchmarks ==="
echo ""

RESULTS="[]"

echo "[1/3] Running SWIM OCaml throughput benchmark..."
if SWIM_RESULT=$("$SCRIPT_DIR/swim_throughput_parallel.sh" "$NODES" "$DURATION" "$MSG_RATE" 47946 "-json" 2>/dev/null); then
    RESULTS=$(echo "$RESULTS" | jq --argjson r "$SWIM_RESULT" '. + [$r]')
    echo "  Done: $(echo "$SWIM_RESULT" | jq -r '"Throughput: \(.msgs_per_sec) msg/s"')"
else
    echo "  Failed or skipped"
fi

echo "[2/3] Running Go memberlist throughput benchmark..."
if [ -x "$SCRIPT_DIR/bin/memberlist_throughput" ]; then
    if ML_RESULT=$("$SCRIPT_DIR/bin/memberlist_throughput" -nodes "$NODES" -duration "$DURATION" -rate "$MSG_RATE" -json 2>/dev/null); then
        RESULTS=$(echo "$RESULTS" | jq --argjson r "$ML_RESULT" '. + [$r]')
        echo "  Done: $(echo "$ML_RESULT" | jq -r '"Throughput: \(.msgs_per_sec) msg/s"')"
    else
        echo "  Failed"
    fi
else
    echo "  Skipped (binary not found)"
fi

echo "[3/3] Running Go Serf throughput benchmark..."
if [ -x "$SCRIPT_DIR/bin/serf_throughput" ]; then
    if SERF_RESULT=$("$SCRIPT_DIR/bin/serf_throughput" -nodes "$NODES" -duration "$DURATION" -rate "$MSG_RATE" -json 2>/dev/null); then
        RESULTS=$(echo "$RESULTS" | jq --argjson r "$SERF_RESULT" '. + [$r]')
        echo "  Done: $(echo "$SERF_RESULT" | jq -r '"Throughput: \(.msgs_per_sec) msg/s"')"
    else
        echo "  Failed"
    fi
else
    echo "  Skipped (binary not found)"
fi

FINAL_RESULT=$(cat <<EOF
{
  "timestamp": "$(date -Iseconds)",
  "benchmark_type": "throughput",
  "config": {
    "nodes": $NODES,
    "duration_sec": $DURATION,
    "msg_rate_per_node": $MSG_RATE
  },
  "system": {
    "hostname": "$(hostname)",
    "os": "$(uname -s)",
    "arch": "$(uname -m)",
    "cpu_count": $(nproc 2>/dev/null || sysctl -n hw.ncpu 2>/dev/null || echo 1)
  },
  "results": $RESULTS
}
EOF
)

echo "$FINAL_RESULT" | jq '.' > "$RESULT_FILE"

echo ""
echo "=== Throughput Summary ==="
echo ""

echo "$FINAL_RESULT" | jq -r '
  .results[] |
  "Implementation: \(.implementation)
  Broadcasts Sent:     \(.broadcasts_sent)
  Broadcasts Received: \(.broadcasts_received)
  Throughput:          \(.msgs_per_sec | . * 10 | round / 10) msg/s
"'

echo "Results saved to: $RESULT_FILE"