#!/bin/bash
set -e

NUM_NODES=${1:-5}
DURATION=${2:-10}
MSG_RATE=${3:-100}
BASE_PORT=${4:-47946}
JSON_OUTPUT=${5:-false}
USE_DIRECT=${6:-true}

SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
SWIM_THROUGHPUT="${SCRIPT_DIR}/../_build/default/bench/swim_throughput.exe"

if [ ! -f "$SWIM_THROUGHPUT" ]; then
    echo "Error: swim_throughput.exe not found. Run 'dune build bench/swim_throughput.exe'" >&2
    exit 1
fi

PEERS=""
for i in $(seq 0 $((NUM_NODES - 1))); do
    PORT=$((BASE_PORT + i))
    if [ -n "$PEERS" ]; then
        PEERS="${PEERS},"
    fi
    PEERS="${PEERS}${PORT}"
done

DIRECT_FLAG=""
if [ "$USE_DIRECT" = "true" ]; then
    DIRECT_FLAG="-direct"
else
    DIRECT_FLAG="-gossip"
fi

TMPDIR=$(mktemp -d)
trap "rm -rf $TMPDIR" EXIT

PIDS=()
for i in $(seq 0 $((NUM_NODES - 1))); do
    PORT=$((BASE_PORT + i))
    "$SWIM_THROUGHPUT" -port "$PORT" -peers "$PEERS" -duration "$DURATION" -rate "$MSG_RATE" $DIRECT_FLAG > "$TMPDIR/node-$i.json" 2>/dev/null &
    PIDS+=($!)
done

sleep 0.5

FAILED=0
for i in "${!PIDS[@]}"; do
    if ! wait "${PIDS[$i]}"; then
        echo "Warning: Node $i failed" >&2
        FAILED=$((FAILED + 1))
    fi
done

if [ "$FAILED" -eq "$NUM_NODES" ]; then
    echo "Error: All nodes failed" >&2
    exit 1
fi

TOTAL_SENT=0
TOTAL_RECV=0
TOTAL_MPS=0
NODE_COUNT=0

for i in $(seq 0 $((NUM_NODES - 1))); do
    if [ -f "$TMPDIR/node-$i.json" ] && [ -s "$TMPDIR/node-$i.json" ]; then
        JSON=$(cat "$TMPDIR/node-$i.json")
        SENT=$(echo "$JSON" | grep -o '"broadcasts_sent":[0-9]*' | grep -o '[0-9]*')
        RECV=$(echo "$JSON" | grep -o '"broadcasts_received":[0-9]*' | grep -o '[0-9]*')
        MPS=$(echo "$JSON" | grep -o '"msgs_per_sec":[0-9.]*' | grep -o '[0-9.]*')
        
        if [ -n "$SENT" ] && [ -n "$RECV" ]; then
            TOTAL_SENT=$((TOTAL_SENT + SENT))
            TOTAL_RECV=$((TOTAL_RECV + RECV))
            TOTAL_MPS=$(echo "$TOTAL_MPS + $MPS" | bc)
            NODE_COUNT=$((NODE_COUNT + 1))
        fi
    fi
done

if [ "$NODE_COUNT" -eq 0 ]; then
    echo "Error: No valid results collected" >&2
    exit 1
fi

TOTAL_THROUGHPUT=$(echo "scale=1; $TOTAL_RECV / $DURATION" | bc)
DURATION_NS=$(echo "$DURATION * 1000000000" | bc | cut -d. -f1)
CPU_CORES=$(nproc 2>/dev/null || sysctl -n hw.ncpu 2>/dev/null || echo 1)

if [ "$JSON_OUTPUT" = "true" ] || [ "$JSON_OUTPUT" = "-json" ]; then
    cat <<EOF
{
  "implementation": "swim-ocaml",
  "num_nodes": $NUM_NODES,
  "duration_ns": $DURATION_NS,
  "msg_rate": $MSG_RATE,
  "broadcasts_sent": $TOTAL_SENT,
  "broadcasts_received": $TOTAL_RECV,
  "msgs_per_sec": $TOTAL_THROUGHPUT,
  "cpu_cores": $CPU_CORES
}
EOF
else
    echo "=== SWIM OCaml Throughput Results ==="
    echo "Nodes:              $NUM_NODES"
    echo "Duration:           ${DURATION}s"
    echo "Target Rate:        ${MSG_RATE} msg/s per node"
    echo "Broadcasts Sent:    $TOTAL_SENT"
    echo "Broadcasts Recv:    $TOTAL_RECV"
    echo "Throughput:         $TOTAL_THROUGHPUT msg/s"
    echo "CPU Cores:          $CPU_CORES"
fi