#!/bin/bash

# Configuration for the PostgreSQL instances using absolute paths.
# This script does NOT modify the shell's PATH variable.

# --- Master Config ---
MASTER_NAME="master"
MASTER_PORT=5432
MASTER_BIN_PATH="$HOME/pg-master/bin"
MASTER_DATA="$HOME/pg-master-data"
MASTER_LOG="/tmp/pg-master.log"

# --- Patch v1 Config ---
PATCH_NAME="patch-v1"
PATCH_PORT=5432
PATCH_BIN_PATH="$HOME/pg-patch-v1/bin"
PATCH_DATA="$HOME/pg-patch-v1-data"
PATCH_LOG="/tmp/pg-patch-v1.log"

# Benchmark settings
CHANNEL_NAME="mychannel"
CONNECTIONS=(1 2 4 8 16 32 64 128)
DURATION=10 # Benchmark duration in seconds for each run
MEASUREMENTS=3 # Number of measurements per configuration

# CSV output file
CSV_OUTPUT="benchmark_results.csv"

# Temporary files
PGBENCH_SCRIPT=$(mktemp)

# --- Cleanup Function ---
# Ensures that servers are stopped and temp files are removed on script exit.
cleanup() {
  echo ""
  echo "Cleaning up..."
  # Ensure both servers are stopped, silencing errors if they are not running.
  # Use absolute paths and explicit data directories.
  "$MASTER_BIN_PATH/pg_ctl" -D "$MASTER_DATA" -m fast stop &> /dev/null
  "$PATCH_BIN_PATH/pg_ctl" -D "$PATCH_DATA" -m fast stop &> /dev/null
  rm -f "$PGBENCH_SCRIPT"
  echo "Cleanup complete."
}

# Trap the script's exit (normal or interrupted) to run the cleanup function
trap cleanup EXIT

# Initialize CSV file with headers
echo "version,connections,jobs,tps,run" > "$CSV_OUTPUT"

# --- Benchmark Function ---
# A generic function to run the benchmark for a given configuration.
# It starts, benchmarks, and then stops the specified server instance.
run_benchmark() {
  local name=$1
  local port=$2
  local bin_path=$3
  local data_path=$4
  local log_file=$5

  echo "--- Starting benchmark for: $name ---"

  # Set PGPORT for client tools (pgbench, psql) for this run
  export PGPORT=$port

  # 1. Start the server using absolute path and explicit data directory
  echo "Starting $name server on port $port..."
  "$bin_path/pg_ctl" -D "$data_path" -l "$log_file" -o "-p $port" start
  sleep 2 # Give server a moment to become available

  # Create the pgbench script content
  cat > "$PGBENCH_SCRIPT" << EOF
NOTIFY ${CHANNEL_NAME};
EOF

  # 2. Start the listener in the background for this server
  (echo "LISTEN ${CHANNEL_NAME};"; sleep 300) | "$bin_path/psql" -d postgres &> /dev/null &
  local listener_pid=$!
  
  # 3. Run the benchmark loop
  echo "Running pgbench for connection counts: ${CONNECTIONS[*]}"
  for c in "${CONNECTIONS[@]}"; do
    echo "  Testing with $c connections ($MEASUREMENTS measurements per run)..."
    # Run multiple measurements for each connection count
    for m in $(seq 1 $MEASUREMENTS); do
      # Run pgbench and extract TPS value
      tps=$("$bin_path/pgbench" -d postgres -f "$PGBENCH_SCRIPT" -c "$c" -j "$c" -T "$DURATION" -n \
        | grep -E '^tps' \
        | awk '{printf "%.0f", $3}')
      
      # Write to CSV: version,connections,jobs,tps,run
      echo "$name,$c,$c,$tps,$m" >> "$CSV_OUTPUT"
    done
  done
  
  # 4. Stop the listener and the server
  kill "$listener_pid" &> /dev/null
  echo "Stopping $name server..."
  "$bin_path/pg_ctl" -D "$data_path" -m fast stop &> /dev/null
  echo "--- Benchmark for $name complete ---"
  echo ""
}

# --- Main Execution ---

# 1. Run benchmark for master
run_benchmark "$MASTER_NAME" "$MASTER_PORT" "$MASTER_BIN_PATH" "$MASTER_DATA" "$MASTER_LOG"

# 2. Run benchmark for patch-v1
run_benchmark "$PATCH_NAME" "$PATCH_PORT" "$PATCH_BIN_PATH" "$PATCH_DATA" "$PATCH_LOG"

# 3. Generate report using PostgreSQL
echo "--- Generating Benchmark Report using PostgreSQL ---"

# Start the master server to run the analysis
export PGPORT=$MASTER_PORT
"$MASTER_BIN_PATH/pg_ctl" -D "$MASTER_DATA" -l "$MASTER_LOG" -o "-p $MASTER_PORT" start
sleep 2

# Create analysis database and load data
"$MASTER_BIN_PATH/psql" -d postgres << EOF
-- Create a temporary database for analysis
DROP DATABASE IF EXISTS bench_analysis;
CREATE DATABASE bench_analysis;
\c bench_analysis

-- Create table for benchmark results
CREATE TABLE benchmark_results (
    version TEXT,
    connections INT,
    jobs INT,
    tps NUMERIC,
    run INT
);

-- Load CSV data
\COPY benchmark_results FROM '$CSV_OUTPUT' CSV HEADER

-- Generate comparison report
WITH avg_results AS (
    SELECT 
        version,
        connections,
        AVG(tps) AS avg_tps,
        STDDEV(tps) AS stddev_tps,
        COUNT(*) AS runs
    FROM benchmark_results
    GROUP BY version, connections
),
comparison AS (
    SELECT 
        m.connections,
        m.avg_tps AS master_tps,
        p.avg_tps AS patch_tps,
        CASE 
            WHEN m.avg_tps > 0 THEN ((p.avg_tps - m.avg_tps) / m.avg_tps * 100)
            ELSE 0
        END AS relative_diff_pct,
        m.stddev_tps AS master_stddev,
        p.stddev_tps AS patch_stddev
    FROM avg_results m
    JOIN avg_results p ON m.connections = p.connections
    WHERE m.version = 'master' AND p.version = 'patch-v1'
    ORDER BY m.connections
)
SELECT 
    connections AS "Connections=Jobs",
    ROUND(master_tps) AS "TPS (master)",
    ROUND(patch_tps) AS "TPS (patch)",
    ROUND(relative_diff_pct, 2) || '%' AS "Relative Diff (%)",
    ROUND(master_stddev) AS "StdDev (master)",
    ROUND(patch_stddev) AS "StdDev (patch)"
FROM comparison;

-- Also show raw data summary
\echo ''
\echo 'Raw Data Summary:'
SELECT 
    version AS "Version",
    connections AS "Connections",
    COUNT(*) AS "Runs",
    ROUND(MIN(tps)) AS "Min TPS",
    ROUND(MAX(tps)) AS "Max TPS",
    ROUND(AVG(tps)) AS "Avg TPS"
FROM benchmark_results
GROUP BY version, connections
ORDER BY version, connections;
EOF

# Stop the server
"$MASTER_BIN_PATH/pg_ctl" -D "$MASTER_DATA" -m fast stop &> /dev/null

echo ""
echo "CSV results saved to: $CSV_OUTPUT"