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

# -----------------------------------
# Caller should adjust these values
# -----------------------------------
# Scripts path
BASE_DIR="$HOME/scripts/test_autovacuum_prioritization"
# Workload parameters
OLTP_TABLES=100
OLTP_ROWS_PER_TABLE=100000
BATCH_TABLES=5
BATCH_SIZE=1000
BATCH_CONNECTIONS=1
OLTP_CONNECTIONS=3
TIMEOUT=600

# ----------------------------
# DO NOT TOUCH!
# ----------------------------
TMP_DIR="$BASE_DIR/tmp"
TMP_OLTP_FILE="$BASE_DIR/oltp_workload.sql"
TMP_BATCH_FILES="$BASE_DIR/batch_workload"

OLTP_SCRIPT="$BASE_DIR/oltp.sh"
BATCH_SCRIPT="$BASE_DIR/batch.sh"
SUMMARY_SCRIPT="$BASE_DIR/summary.sh"

OLTP_LOG="$TMP_DIR/oltp.out"
BATCH_LOG="$TMP_DIR/batch.out"
READY_FILE="$TMP_DIR/oltp_ready"

# ----------------------------
# Ensure temp directory exists
# ----------------------------
if [ -n "$TMP_DIR" ] && [ "$TMP_DIR" != "/" ]; then
    if [ -d "$TMP_DIR" ]; then
        echo "Removing existing temp directory: $TMP_DIR"
        rm -rf "$TMP_DIR"
    fi
    mkdir -p "$TMP_DIR" || { echo "Error: failed to create $TMP_DIR" >&2; exit 1; }
else
    echo "Error: TMP_DIR is invalid!" >&2
    exit 1
fi

# ----------------------------
# Start OLTP workload in background
# ----------------------------
echo "Starting OLTP workload..."
bash "$OLTP_SCRIPT" "$TMP_OLTP_FILE" "$OLTP_TABLES" "$OLTP_ROWS_PER_TABLE" "$READY_FILE" "$OLTP_CONNECTIONS" "$TIMEOUT" | tee "$OLTP_LOG" 2>&1 &
OLTP_PID=$!

# ----------------------------
# Wait for OLTP ready file
# ----------------------------
echo "Waiting for OLTP ready file..."
while [ ! -f "$READY_FILE" ]; do
    sleep 1
done
echo "OLTP ready file detected. Starting batch workload..."

# ----------------------------
# Start batch workload in background
# ----------------------------
bash "$BATCH_SCRIPT" "$BATCH_SIZE" "$TMP_BATCH_FILES" "$BATCH_TABLES" "$BATCH_CONNECTIONS" "$TIMEOUT" | tee "$BATCH_LOG" 2>&1 &
BATCH_PID=$!

echo "OLTP PID: $OLTP_PID | Batch PID: $BATCH_PID"

# ----------------------------
# Monitor both workloads
# ----------------------------
START_TIME=$(date +%s)

while true; do
    CURRENT_TIME=$(date +%s)
    ELAPSED=$((CURRENT_TIME - START_TIME))

    # If OLTP finished, terminate batch
    if ! kill -0 "$OLTP_PID" 2>/dev/null; then
        echo "OLTP workload completed, terminating batch..."
        kill -TERM "$BATCH_PID" 2>/dev/null || true
        wait "$BATCH_PID" 2>/dev/null || true
        break
    fi

    # If batch finished, terminate OLTP
    if ! kill -0 "$BATCH_PID" 2>/dev/null; then
        echo "Batch workload completed, terminating OLTP..."
        kill -TERM "$OLTP_PID" 2>/dev/null || true
        wait "$OLTP_PID" 2>/dev/null || true
        break
    fi

    # Check timeout
    if [ "$ELAPSED" -ge "$TIMEOUT" ]; then
        echo "⏱ Timeout reached (${ELAPSED}s), stopping both workloads..."
        kill -TERM "$OLTP_PID" "$BATCH_PID" 2>/dev/null || true
        wait "$OLTP_PID" 2>/dev/null || true
        wait "$BATCH_PID" 2>/dev/null || true
        break
    fi

    sleep 5
done

echo "Both workloads stopped."

# ----------------------------
# Restart PostgreSQL
# ----------------------------
echo "Restarting Postgres..."
pg_ctl restart
echo "Postgres restarted."

# ----------------------------
# Run summary script
# ----------------------------
"$SUMMARY_SCRIPT"