#!/bin/bash
# PostgreSQL Load Test using pgbench - Foxhunt Trading System
# Tests database performance under increasing concurrent load

set -e

DB_URL="postgresql://foxhunt:foxhunt_dev_password@localhost:5432/foxhunt"
DB_HOST="localhost"
DB_PORT="5432"
DB_NAME="foxhunt"
DB_USER="foxhunt"
DB_PASS="foxhunt_dev_password"

TEST_DURATION=30  # seconds per test
WORKLOAD_FILE="trading_workload.sql"
RESULTS_FILE="DB_LOAD_TEST_RESULTS_$(date +%Y%m%d_%H%M%S).txt"

export PGPASSWORD="$DB_PASS"

echo "========================================" | tee -a "$RESULTS_FILE"
echo "PostgreSQL Load Test - Foxhunt Trading" | tee -a "$RESULTS_FILE"
echo "Using pgbench with custom trading workload" | tee -a "$RESULTS_FILE"
echo "Test Duration: $TEST_DURATION seconds per scenario" | tee -a "$RESULTS_FILE"
echo "Start Time: $(date)" | tee -a "$RESULTS_FILE"
echo "========================================" | tee -a "$RESULTS_FILE"
echo "" | tee -a "$RESULTS_FILE"

# Function to check database health
check_db_health() {
    echo "=== Database Health Check ===" | tee -a "$RESULTS_FILE"

    echo "PostgreSQL Version:" | tee -a "$RESULTS_FILE"
    psql -h "$DB_HOST" -p "$DB_PORT" -U "$DB_USER" -d "$DB_NAME" -t -c "SELECT version();" | tee -a "$RESULTS_FILE"

    echo "" | tee -a "$RESULTS_FILE"
    echo "Configuration:" | tee -a "$RESULTS_FILE"
    psql -h "$DB_HOST" -p "$DB_PORT" -U "$DB_USER" -d "$DB_NAME" -t -c "SHOW max_connections;" | awk '{print "  max_connections: " $1}' | tee -a "$RESULTS_FILE"
    psql -h "$DB_HOST" -p "$DB_PORT" -U "$DB_USER" -d "$DB_NAME" -t -c "SHOW shared_buffers;" | awk '{print "  shared_buffers: " $1}' | tee -a "$RESULTS_FILE"
    psql -h "$DB_HOST" -p "$DB_PORT" -U "$DB_USER" -d "$DB_NAME" -t -c "SHOW synchronous_commit;" | awk '{print "  synchronous_commit: " $1}' | tee -a "$RESULTS_FILE"

    echo "" | tee -a "$RESULTS_FILE"
    echo "Cache Hit Ratio:" | tee -a "$RESULTS_FILE"
    psql -h "$DB_HOST" -p "$DB_PORT" -U "$DB_USER" -d "$DB_NAME" -t -c "SELECT (sum(blks_hit)::float / NULLIF(sum(blks_hit) + sum(blks_read), 0) * 100)::numeric(5,2) || '%' AS cache_hit_ratio FROM pg_stat_database WHERE datname = 'foxhunt';" | awk '{print "  " $1}' | tee -a "$RESULTS_FILE"

    echo "" | tee -a "$RESULTS_FILE"
}

# Function to get baseline table counts
get_baseline_counts() {
    echo "=== Baseline Table Counts ===" | tee -a "$RESULTS_FILE"
    psql -h "$DB_HOST" -p "$DB_PORT" -U "$DB_USER" -d "$DB_NAME" -c "SELECT 'orders' AS table_name, COUNT(*) FROM orders UNION ALL SELECT 'executions', COUNT(*) FROM executions UNION ALL SELECT 'fills', COUNT(*) FROM fills UNION ALL SELECT 'positions', COUNT(*) FROM positions;" | tee -a "$RESULTS_FILE"
    echo "" | tee -a "$RESULTS_FILE"
}

# Function to run pgbench test
run_pgbench_test() {
    local clients=$1
    local test_name=$2

    echo "========================================" | tee -a "$RESULTS_FILE"
    echo "TEST: $test_name" | tee -a "$RESULTS_FILE"
    echo "Concurrent Clients: $clients" | tee -a "$RESULTS_FILE"
    echo "========================================" | tee -a "$RESULTS_FILE"

    # Run pgbench
    pgbench -h "$DB_HOST" -p "$DB_PORT" -U "$DB_USER" -d "$DB_NAME" \
        -c "$clients" -j $(( clients > 4 ? 4 : clients )) \
        -T "$TEST_DURATION" \
        -f "$WORKLOAD_FILE" \
        -n 2>&1 | tee -a "$RESULTS_FILE"

    echo "" | tee -a "$RESULTS_FILE"

    # Check connection pool status
    echo "=== Connection Pool Status ===" | tee -a "$RESULTS_FILE"
    psql -h "$DB_HOST" -p "$DB_PORT" -U "$DB_USER" -d "$DB_NAME" -c "SELECT state, COUNT(*) FROM pg_stat_activity WHERE datname = 'foxhunt' GROUP BY state ORDER BY COUNT(*) DESC;" | tee -a "$RESULTS_FILE"

    echo "" | tee -a "$RESULTS_FILE"

    # Check for lock contention
    echo "=== Lock Contention ===" | tee -a "$RESULTS_FILE"
    local locks_waiting=$(psql -h "$DB_HOST" -p "$DB_PORT" -U "$DB_USER" -d "$DB_NAME" -t -c "SELECT COUNT(*) FROM pg_locks WHERE NOT granted;" | tr -d ' ')
    echo "  Locks waiting: $locks_waiting" | tee -a "$RESULTS_FILE"

    echo "" | tee -a "$RESULTS_FILE"

    # Check deadlock count
    echo "=== Deadlock Statistics ===" | tee -a "$RESULTS_FILE"
    psql -h "$DB_HOST" -p "$DB_PORT" -U "$DB_USER" -d "$DB_NAME" -c "SELECT datname, deadlocks, conflicts FROM pg_stat_database WHERE datname = 'foxhunt';" | tee -a "$RESULTS_FILE"

    echo "" | tee -a "$RESULTS_FILE"
}

# Function to analyze query performance
analyze_query_performance() {
    echo "========================================" | tee -a "$RESULTS_FILE"
    echo "QUERY PERFORMANCE ANALYSIS" | tee -a "$RESULTS_FILE"
    echo "========================================" | tee -a "$RESULTS_FILE"

    echo "=== Table Statistics ===" | tee -a "$RESULTS_FILE"
    psql -h "$DB_HOST" -p "$DB_PORT" -U "$DB_USER" -d "$DB_NAME" -c "
    SELECT relname,
           n_tup_ins AS inserts,
           n_tup_upd AS updates,
           n_tup_del AS deletes,
           seq_scan,
           idx_scan,
           CASE WHEN seq_scan + idx_scan > 0
                THEN (idx_scan::float / (seq_scan + idx_scan) * 100)::numeric(5,2)
                ELSE 0 END AS idx_scan_pct
    FROM pg_stat_user_tables
    WHERE schemaname = 'public'
    AND relname IN ('orders', 'executions', 'fills', 'positions')
    ORDER BY n_tup_ins DESC;
    " | tee -a "$RESULTS_FILE"

    echo "" | tee -a "$RESULTS_FILE"

    echo "=== Index Usage ===" | tee -a "$RESULTS_FILE"
    psql -h "$DB_HOST" -p "$DB_PORT" -U "$DB_USER" -d "$DB_NAME" -c "
    SELECT tablename, indexname, idx_scan, idx_tup_read, idx_tup_fetch
    FROM pg_stat_user_indexes
    WHERE schemaname = 'public'
    AND tablename IN ('orders', 'executions', 'fills', 'positions')
    ORDER BY idx_scan DESC
    LIMIT 10;
    " | tee -a "$RESULTS_FILE"

    echo "" | tee -a "$RESULTS_FILE"
}

# Function to check resource utilization
check_resource_utilization() {
    echo "========================================" | tee -a "$RESULTS_FILE"
    echo "RESOURCE UTILIZATION" | tee -a "$RESULTS_FILE"
    echo "========================================" | tee -a "$RESULTS_FILE"

    echo "=== Database Size ===" | tee -a "$RESULTS_FILE"
    psql -h "$DB_HOST" -p "$DB_PORT" -U "$DB_USER" -d "$DB_NAME" -c "
    SELECT datname, pg_size_pretty(pg_database_size(datname)) AS size
    FROM pg_database
    WHERE datname = 'foxhunt';
    " | tee -a "$RESULTS_FILE"

    echo "" | tee -a "$RESULTS_FILE"

    echo "=== Table Sizes ===" | tee -a "$RESULTS_FILE"
    psql -h "$DB_HOST" -p "$DB_PORT" -U "$DB_USER" -d "$DB_NAME" -c "
    SELECT relname,
           pg_size_pretty(pg_total_relation_size(relid)) AS total_size,
           pg_size_pretty(pg_relation_size(relid)) AS table_size,
           pg_size_pretty(pg_total_relation_size(relid) - pg_relation_size(relid)) AS index_size
    FROM pg_stat_user_tables
    WHERE schemaname = 'public'
    AND relname IN ('orders', 'executions', 'fills', 'positions')
    ORDER BY pg_total_relation_size(relid) DESC;
    " | tee -a "$RESULTS_FILE"

    echo "" | tee -a "$RESULTS_FILE"

    echo "=== Final Cache Hit Ratio ===" | tee -a "$RESULTS_FILE"
    psql -h "$DB_HOST" -p "$DB_PORT" -U "$DB_USER" -d "$DB_NAME" -t -c "
    SELECT (sum(blks_hit)::float / NULLIF(sum(blks_hit) + sum(blks_read), 0) * 100)::numeric(5,2) || '%' AS cache_hit_ratio
    FROM pg_stat_database
    WHERE datname = 'foxhunt';
    " | awk '{print "  " $1}' | tee -a "$RESULTS_FILE"

    echo "" | tee -a "$RESULTS_FILE"
}

# Function to get final table counts
get_final_counts() {
    echo "=== Final Table Counts ===" | tee -a "$RESULTS_FILE"
    psql -h "$DB_HOST" -p "$DB_PORT" -U "$DB_USER" -d "$DB_NAME" -c "
    SELECT 'orders' AS table_name, COUNT(*) FROM orders
    UNION ALL SELECT 'executions', COUNT(*) FROM executions
    UNION ALL SELECT 'fills', COUNT(*) FROM fills
    UNION ALL SELECT 'positions', COUNT(*) FROM positions;
    " | tee -a "$RESULTS_FILE"
    echo "" | tee -a "$RESULTS_FILE"
}

# Main test execution
main() {
    echo "Starting PostgreSQL Load Test with pgbench..."

    # Check if workload file exists
    if [ ! -f "$WORKLOAD_FILE" ]; then
        echo "ERROR: Workload file $WORKLOAD_FILE not found!" | tee -a "$RESULTS_FILE"
        exit 1
    fi

    # Initial health check
    check_db_health
    get_baseline_counts

    # Run tests with increasing concurrency
    run_pgbench_test 1 "BASELINE (Single Client)"
    run_pgbench_test 10 "MODERATE LOAD (10 Clients)"
    run_pgbench_test 50 "HIGH LOAD (50 Clients)"
    run_pgbench_test 100 "STRESS TEST (100 Clients)"

    # Post-test analysis
    analyze_query_performance
    check_resource_utilization
    get_final_counts

    echo "========================================" | tee -a "$RESULTS_FILE"
    echo "Test Completed: $(date)" | tee -a "$RESULTS_FILE"
    echo "Results saved to: $RESULTS_FILE" | tee -a "$RESULTS_FILE"
    echo "========================================" | tee -a "$RESULTS_FILE"

    echo ""
    echo "Test complete! Results saved to: $RESULTS_FILE"
}

# Run main test
main