#!/usr/bin/env bash
# load-publisher.sh — Publish at a sustained target rate for a fixed duration.
#
# Uses `fafb-q bench pub` in 1-second bursts to approximate the target rate.
# Each burst publishes (MSGS_PER_SEC / CONCURRENCY) messages per worker, then
# sleeps the remainder of the second.
#
# Usage:
#   ./scripts/load-publisher.sh
#   MSGS_PER_SEC=5000 DURATION_SECS=300 PAYLOAD_SIZE=512 ./scripts/load-publisher.sh
#
# Requires: fafb-q CLI in PATH (or set FAFB_BIN).

set -euo pipefail

# ── Configurable variables ───────────────────────────────────────────────────
FAFB_BIN="${FAFB_BIN:-fafb-q}"                    # Path to fafb-q binary
SERVER="${SERVER:-localhost:9876}"                 # TCP server address
API_KEY="${API_KEY:-}"                             # API key (required)
NAMESPACE="${NAMESPACE:-default}"                  # Namespace
TOPIC="${TOPIC:-load-test}"                        # Topic name
MSGS_PER_SEC="${MSGS_PER_SEC:-1000}"              # Target messages per second
DURATION_SECS="${DURATION_SECS:-600}"             # Duration in seconds (default 10min)
PAYLOAD_SIZE="${PAYLOAD_SIZE:-256}"                # Payload size in bytes
CONCURRENCY="${CONCURRENCY:-4}"                   # Parallel publisher connections
NO_ACK="${NO_ACK:-false}"                         # Fire-and-forget mode (true/false)
PROFILE="${PROFILE:-}"                            # Config profile (optional)
TLS="${TLS:-false}"                               # Enable TLS (true/false)
TLS_CA="${TLS_CA:-}"                              # CA cert path (optional)
LOG_INTERVAL="${LOG_INTERVAL:-5}"                 # Stats logging interval in seconds
# ─────────────────────────────────────────────────────────────────────────────

if [[ -z "$API_KEY" && -z "$PROFILE" ]]; then
    echo "ERROR: API_KEY or PROFILE must be set" >&2; exit 1
fi
command -v "$FAFB_BIN" &>/dev/null || { echo "ERROR: $FAFB_BIN not found" >&2; exit 1; }

FLAGS=()
if [[ -n "$PROFILE" ]]; then FLAGS+=(--profile "$PROFILE")
else FLAGS+=(--server "$SERVER" --api-key "$API_KEY"); fi
[[ "$TLS" == "true" ]] && FLAGS+=(--tls)
[[ -n "$TLS_CA" ]] && FLAGS+=(--tls-ca "$TLS_CA")

FULL_TOPIC="${NAMESPACE}/${TOPIC}"
TOTAL_MSGS=$((MSGS_PER_SEC * DURATION_SECS))
BURST=$((MSGS_PER_SEC / CONCURRENCY))

echo "═══════════════════════════════════════════════════"
echo " fafb-q Load Publisher"
echo "═══════════════════════════════════════════════════"
echo " Server:        ${SERVER}"
echo " Topic:         ${FULL_TOPIC}"
echo " Rate:          ${MSGS_PER_SEC} msgs/sec (${BURST}/worker x${CONCURRENCY})"
echo " Duration:      ${DURATION_SECS}s"
echo " Total msgs:    ~${TOTAL_MSGS}"
echo " Payload size:  ${PAYLOAD_SIZE} bytes"
echo " No-ack:        ${NO_ACK}"
echo "═══════════════════════════════════════════════════"
echo ""

START_TIME=$(date +%s)
END_TIME=$((START_TIME + DURATION_SECS))
PIDS=()
TMPDIR_WORK=$(mktemp -d)
trap 'kill "${PIDS[@]}" 2>/dev/null; rm -rf "$TMPDIR_WORK"' EXIT INT TERM

# Worker: publish BURST messages via bench pub, sleep remainder of 1s, repeat
run_worker() {
    local worker_id=$1
    local sent=0
    local count_file="${TMPDIR_WORK}/worker_${worker_id}"
    echo 0 > "$count_file"

    local bench_flags=("${FLAGS[@]}")
    [[ "$NO_ACK" == "true" ]] && bench_flags+=(--no-ack)

    while [[ $(date +%s) -lt $END_TIME ]]; do
        local tick_start
        tick_start=$(date +%s%N 2>/dev/null || echo "$(date +%s)000000000")

        "$FAFB_BIN" bench pub "$FULL_TOPIC" \
            "${bench_flags[@]}" \
            --messages "$BURST" \
            --payload-size "$PAYLOAD_SIZE" \
            --concurrency 1 2>/dev/null

        sent=$((sent + BURST))
        echo "$sent" > "$count_file"

        local tick_end
        tick_end=$(date +%s%N 2>/dev/null || echo "$(date +%s)000000000")
        local sleep_ns=$((1000000000 - (tick_end - tick_start)))
        if [[ $sleep_ns -gt 0 ]]; then
            sleep "$(awk "BEGIN {printf \"%.3f\", $sleep_ns/1000000000}")"
        fi
    done
}

# Stats reporter
report_progress() {
    local prev_total=0
    while [[ $(date +%s) -lt $END_TIME ]]; do
        sleep "$LOG_INTERVAL"
        local now elapsed remaining total=0 delta rate
        now=$(date +%s); elapsed=$((now - START_TIME)); remaining=$((END_TIME - now))
        [[ $remaining -lt 0 ]] && remaining=0

        for f in "$TMPDIR_WORK"/worker_*; do
            [[ -f "$f" ]] && total=$((total + $(cat "$f" 2>/dev/null || echo 0)))
        done

        delta=$((total - prev_total))
        if [[ $delta -gt 0 ]]; then
            rate=$((delta / LOG_INTERVAL))
            echo "[${elapsed}s, ~${remaining}s left] ${delta} msgs in ${LOG_INTERVAL}s (${rate} msgs/sec, total: ${total})"
        fi
        prev_total=$total
    done
}

echo "Starting ${CONCURRENCY} workers..."
for i in $(seq 1 "$CONCURRENCY"); do run_worker "$i" & PIDS+=($!); done
echo "Publishing ~${MSGS_PER_SEC} msgs/sec for ${DURATION_SECS}s"
echo ""

report_progress &
REPORTER_PID=$!

FAILED=0
for pid in "${PIDS[@]}"; do wait "$pid" 2>/dev/null || FAILED=$((FAILED + 1)); done
kill "$REPORTER_PID" 2>/dev/null || true; wait "$REPORTER_PID" 2>/dev/null || true

FINAL_TOTAL=0
for f in "$TMPDIR_WORK"/worker_*; do
    [[ -f "$f" ]] && FINAL_TOTAL=$((FINAL_TOTAL + $(cat "$f" 2>/dev/null || echo 0)))
done
WALL_SECS=$(( $(date +%s) - START_TIME ))

echo ""
echo "═══════════════════════════════════════════════════"
echo " Publishing Complete"
echo "═══════════════════════════════════════════════════"
echo " Wall time:     ${WALL_SECS}s"
echo " Total sent:    ${FINAL_TOTAL}"
echo " Failed workers: ${FAILED}/${CONCURRENCY}"
[[ $WALL_SECS -gt 0 ]] && echo " Avg rate:      $((FINAL_TOTAL / WALL_SECS)) msgs/sec"
echo "═══════════════════════════════════════════════════"
