#!/bin/bash
set -e

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

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

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

echo "=== SWIM Benchmark Suite ==="
echo "Nodes: $NODES"
echo "Duration: ${DURATION}s"
echo "Output: $RESULT_FILE"
echo ""

cd "$SCRIPT_DIR"

echo "Building Go benchmarks..."
go mod tidy 2>/dev/null || true
go build -o bin/memberlist_bench ./cmd/memberlist 2>/dev/null || {
    echo "Warning: Failed to build memberlist benchmark"
}
go build -o bin/serf_bench ./cmd/serf 2>/dev/null || {
    echo "Warning: Failed to build serf benchmark"
}

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

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

RESULTS="[]"

echo "[1/3] Running SWIM OCaml benchmark..."
if SWIM_RESULT=$("$SCRIPT_DIR/swim_parallel.sh" "$NODES" "$DURATION" 37946 "-json" 2>/dev/null); then
    RESULTS=$(echo "$RESULTS" | jq --argjson r "$SWIM_RESULT" '. + [$r]')
    echo "  Done: $(echo "$SWIM_RESULT" | jq -r '.convergence_time_ns / 1e9 | "Convergence: \(.)s"')"
else
    echo "  Failed or skipped"
fi

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

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

FINAL_RESULT=$(cat <<EOF
{
  "timestamp": "$(date -Iseconds)",
  "config": {
    "nodes": $NODES,
    "duration_sec": $DURATION
  },
  "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 "=== Summary ==="
echo ""

echo "$FINAL_RESULT" | jq -r '
  .results[] |
  "Implementation: \(.implementation)
  Convergence:  \(.convergence_time_ns / 1e9 | . * 1000 | round / 1000)s
  Memory:       \(.memory_used_bytes / 1048576 | . * 100 | round / 100) MB
  Messages:     sent=\(.messages_sent // .events_received // "N/A") recv=\(.messages_received // .queries_processed // "N/A")
"'

echo "Results saved to: $RESULT_FILE"