#!/usr/bin/env bash
set -euo pipefail

SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
REPO_ROOT="$(cd "${SCRIPT_DIR}/../../.." && pwd)"
# shellcheck source=../../../baselines/cdc/lib/env.sh
source "${REPO_ROOT}/baselines/cdc/lib/env.sh"

RATE="100"
DURATION_SECS="60"
KEYS="10000"
SEED="106102"
TIMEOUT_AFTER_DONE="120"
DATUM_CDC_SLOT="${DATUM_CDC_SLOT:-datum_cdc_datum_slot}"

usage() {
  cat <<'EOF_USAGE'
Usage: crates/datum-cdc/scripts/run-cdc-benchmark.sh [options]

Options:
  --rate N             committed transactions per second (default: 100)
  --duration-secs N    sustained load duration in seconds (default: 60)
  --keys N             keyspace size (default: 10000)
  --seed N             deterministic workload seed (default: 106102)
  --timeout-secs N     observer wait after loadgen exits (default: 120)
EOF_USAGE
}

while [[ $# -gt 0 ]]; do
  case "$1" in
    --rate)
      RATE="$2"
      shift 2
      ;;
    --duration-secs)
      DURATION_SECS="$2"
      shift 2
      ;;
    --keys)
      KEYS="$2"
      shift 2
      ;;
    --seed)
      SEED="$2"
      shift 2
      ;;
    --timeout-secs)
      TIMEOUT_AFTER_DONE="$2"
      shift 2
      ;;
    -h|--help)
      usage
      exit 0
      ;;
    *)
      usage >&2
      die "unknown argument: $1"
      ;;
  esac
done

require_cmd cargo
require_cmd python3

RUN_ID="datum-cdc-$(date -u +%Y%m%dT%H%M%SZ)-r${RATE}-d${DURATION_SECS}"
RUN_OUT="${RESULTS_DIR}/${RUN_ID}"
EXPECTED="${RUN_OUT}/expected.jsonl"
DONE_FILE="${RUN_OUT}/loadgen.done"
OBSERVER_JSON="${RUN_OUT}/observer.json"
LOADGEN_STDOUT="${RUN_OUT}/loadgen.stdout.json"
LOADGEN_STDERR="${RUN_OUT}/loadgen.stderr.log"

mkdir -p "${RUN_OUT}"

PG_BIN="$(find_pg_bin)" || die "PostgreSQL binary not found; run baselines/cdc/up.sh first"
wait_for "PostgreSQL" 10 "${PG_BIN}/pg_isready" -h "${PG_HOST}" -p "${PG_PORT}" -U "${PG_USER}" -d "${PG_DATABASE}"

"${PG_BIN}/psql" "${PG_URL}" -v ON_ERROR_STOP=1 -v slot="${DATUM_CDC_SLOT}" <<'SQL'
SELECT pg_drop_replication_slot(:'slot')
WHERE EXISTS (
  SELECT 1 FROM pg_replication_slots
  WHERE slot_name = :'slot'
    AND active = false
);
SQL

log "building load generator and datum-cdc observer"
cargo build --release --manifest-path "${LOADGEN_MANIFEST}"
cargo build --release -p datum-cdc --bin cdc_observer

log "starting datum-cdc observer for run_id=${RUN_ID}"
cargo run --release -p datum-cdc --bin cdc_observer -- \
  --database-url "${PG_URL}" \
  --slot "${DATUM_CDC_SLOT}" \
  --publication "${CDC_PUBLICATION}" \
  --run-id "${RUN_ID}" \
  --expected "${EXPECTED}" \
  --done-file "${DONE_FILE}" \
  --out "${OBSERVER_JSON}" \
  --timeout-after-done "${TIMEOUT_AFTER_DONE}" &
OBSERVER_PID="$!"

sleep 2

log "running load generator at ${RATE} tx/s for ${DURATION_SECS}s"
set +e
cargo run --release --manifest-path "${LOADGEN_MANIFEST}" -- \
  --database-url "${PG_URL}" \
  --run-id "${RUN_ID}" \
  --rate "${RATE}" \
  --duration-secs "${DURATION_SECS}" \
  --keys "${KEYS}" \
  --seed "${SEED}" \
  --expected "${EXPECTED}" \
  --reset-table >"${LOADGEN_STDOUT}" 2>"${LOADGEN_STDERR}"
LOADGEN_STATUS="$?"
set -e

touch "${DONE_FILE}"

if [[ "${LOADGEN_STATUS}" -ne 0 ]]; then
  log "loadgen failed; stopping datum-cdc observer"
  kill "${OBSERVER_PID}" >/dev/null 2>&1 || true
  wait "${OBSERVER_PID}" >/dev/null 2>&1 || true
  exit "${LOADGEN_STATUS}"
fi

log "waiting for datum-cdc observer"
wait "${OBSERVER_PID}"

python3 - "${OBSERVER_JSON}" "${LOADGEN_STDOUT}" <<'PY'
import json
import sys

with open(sys.argv[1], "r", encoding="utf-8") as fh:
    observed = json.load(fh)
with open(sys.argv[2], "r", encoding="utf-8") as fh:
    loadgen = json.load(fh)

print(json.dumps({
    "run_id": observed["run_id"],
    "committed_ops": loadgen["committed_ops"],
    "observed_events": observed["observed_events"],
    "throughput_events_per_sec": observed["throughput_events_per_sec"],
    "latency_ms_p50": observed["latency_ms"].get("p50"),
    "latency_ms_p99": observed["latency_ms"].get("p99"),
    "datum_cpu_core_seconds": observed["datum"].get("cpu_core_seconds"),
    "datum_avg_cpu_cores": observed["datum"].get("avg_cpu_cores"),
    "datum_peak_rss_mib": observed["datum"].get("peak_rss_mib"),
    "max_slot_retained_wal_bytes": observed["postgres"].get("max_slot_retained_wal_bytes"),
    "correct": observed["correctness"]["events_observed_eq_committed"]
        and observed["correctness"]["ordered_per_key"],
}, indent=2, sort_keys=True))
PY

cat <<EOF_STATUS

Datum CDC result written to:
  ${OBSERVER_JSON}

Expected operation log:
  ${EXPECTED}
EOF_STATUS