#!/bin/bash
# PostgreSQL Load Test Script for Foxhunt Trading System
# Tests database performance under increasing concurrent load
# Workload: 40% INSERT, 30% SELECT, 20% UPDATE, 10% Complex queries

set -e

DB_URL="postgresql://foxhunt:foxhunt_dev_password@localhost:5432/foxhunt"
TEST_DURATION=60  # seconds per test
RESULTS_FILE="/tmp/db_load_test_results_$(date +%Y%m%d_%H%M%S).txt"

echo "========================================" | tee -a "$RESULTS_FILE"
echo "PostgreSQL Load Test - Foxhunt Trading" | 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"
    psql "$DB_URL" -c "SELECT version();" | tee -a "$RESULTS_FILE"
    psql "$DB_URL" -c "SHOW max_connections;" | tee -a "$RESULTS_FILE"
    psql "$DB_URL" -c "SHOW shared_buffers;" | tee -a "$RESULTS_FILE"
    psql "$DB_URL" -c "SHOW synchronous_commit;" | tee -a "$RESULTS_FILE"

    echo "" | tee -a "$RESULTS_FILE"
    echo "=== Cache Hit Ratio ===" | tee -a "$RESULTS_FILE"
    psql "$DB_URL" -c "SELECT sum(blks_hit)::float / (sum(blks_hit) + sum(blks_read)) * 100 AS cache_hit_ratio FROM pg_stat_database WHERE datname = 'foxhunt';" | tee -a "$RESULTS_FILE"

    echo "" | tee -a "$RESULTS_FILE"
    echo "=== Active Connections ===" | tee -a "$RESULTS_FILE"
    psql "$DB_URL" -c "SELECT count(*) FROM pg_stat_activity WHERE state = 'active';" | 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 "$DB_URL" -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 single-threaded baseline test
run_baseline_test() {
    echo "========================================" | tee -a "$RESULTS_FILE"
    echo "TEST 1: BASELINE (Single-threaded)" | tee -a "$RESULTS_FILE"
    echo "========================================" | tee -a "$RESULTS_FILE"

    local start_time=$(date +%s)
    local end_time=$((start_time + TEST_DURATION))
    local insert_count=0
    local select_count=0
    local update_count=0
    local complex_count=0

    while [ $(date +%s) -lt $end_time ]; do
        # 40% INSERT - orders
        for i in {1..4}; do
            psql "$DB_URL" -c "INSERT INTO orders (order_id, user_id, symbol, side, order_type, quantity, price, status, created_at) VALUES (gen_random_uuid(), gen_random_uuid(), 'BTC/USD', 'buy', 'limit', 1.0, 50000.0, 'pending', NOW());" > /dev/null 2>&1
            ((insert_count++))
        done

        # 30% SELECT - order status
        for i in {1..3}; do
            psql "$DB_URL" -c "SELECT order_id, status, quantity FROM orders WHERE symbol = 'BTC/USD' ORDER BY created_at DESC LIMIT 10;" > /dev/null 2>&1
            ((select_count++))
        done

        # 20% UPDATE - order status
        for i in {1..2}; do
            psql "$DB_URL" -c "UPDATE orders SET status = 'filled', updated_at = NOW() WHERE order_id = (SELECT order_id FROM orders WHERE status = 'pending' LIMIT 1);" > /dev/null 2>&1
            ((update_count++))
        done

        # 10% Complex query - position aggregation
        psql "$DB_URL" -c "SELECT symbol, SUM(quantity) as total_quantity, AVG(entry_price) as avg_price FROM positions WHERE user_id IN (SELECT user_id FROM users LIMIT 5) GROUP BY symbol;" > /dev/null 2>&1
        ((complex_count++))
    done

    local actual_duration=$(($(date +%s) - start_time))
    local total_ops=$((insert_count + select_count + update_count + complex_count))
    local tps=$(echo "scale=2; $total_ops / $actual_duration" | bc)

    echo "Duration: $actual_duration seconds" | tee -a "$RESULTS_FILE"
    echo "Total Operations: $total_ops" | tee -a "$RESULTS_FILE"
    echo "  - INSERTs: $insert_count" | tee -a "$RESULTS_FILE"
    echo "  - SELECTs: $select_count" | tee -a "$RESULTS_FILE"
    echo "  - UPDATEs: $update_count" | tee -a "$RESULTS_FILE"
    echo "  - Complex: $complex_count" | tee -a "$RESULTS_FILE"
    echo "Transactions per Second (TPS): $tps" | tee -a "$RESULTS_FILE"
    echo "" | tee -a "$RESULTS_FILE"
}

# Function to run concurrent load test
run_concurrent_test() {
    local num_connections=$1
    local test_name=$2

    echo "========================================" | tee -a "$RESULTS_FILE"
    echo "TEST: $test_name ($num_connections concurrent connections)" | tee -a "$RESULTS_FILE"
    echo "========================================" | tee -a "$RESULTS_FILE"

    local pids=()
    local temp_dir="/tmp/db_load_test_$$"
    mkdir -p "$temp_dir"

    # Start time
    local start_time=$(date +%s.%N)

    # Launch concurrent workers
    for ((i=1; i<=num_connections; i++)); do
        (
            local worker_inserts=0
            local worker_selects=0
            local worker_updates=0
            local worker_complex=0
            local end_time=$(echo "$start_time + $TEST_DURATION" | bc)

            while (( $(echo "$(date +%s.%N) < $end_time" | bc -l) )); do
                # 40% INSERT
                if (( RANDOM % 10 < 4 )); then
                    psql "$DB_URL" -c "INSERT INTO orders (order_id, user_id, symbol, side, order_type, quantity, price, status, created_at) VALUES (gen_random_uuid(), gen_random_uuid(), 'BTC/USD', 'buy', 'limit', 1.0, 50000.0, 'pending', NOW());" > /dev/null 2>&1
                    ((worker_inserts++))
                # 30% SELECT
                elif (( RANDOM % 10 < 7 )); then
                    psql "$DB_URL" -c "SELECT order_id, status, quantity FROM orders WHERE symbol = 'BTC/USD' ORDER BY created_at DESC LIMIT 10;" > /dev/null 2>&1
                    ((worker_selects++))
                # 20% UPDATE
                elif (( RANDOM % 10 < 9 )); then
                    psql "$DB_URL" -c "UPDATE orders SET status = 'filled', updated_at = NOW() WHERE order_id = (SELECT order_id FROM orders WHERE status = 'pending' LIMIT 1);" > /dev/null 2>&1
                    ((worker_updates++))
                # 10% Complex
                else
                    psql "$DB_URL" -c "SELECT symbol, COUNT(*) as order_count FROM orders WHERE created_at > NOW() - INTERVAL '1 hour' GROUP BY symbol;" > /dev/null 2>&1
                    ((worker_complex++))
                fi
            done

            # Write worker results
            echo "$worker_inserts $worker_selects $worker_updates $worker_complex" > "$temp_dir/worker_$i.txt"
        ) &
        pids+=($!)
    done

    # Wait for all workers to complete
    for pid in "${pids[@]}"; do
        wait $pid
    done

    local end_time=$(date +%s.%N)
    local actual_duration=$(echo "$end_time - $start_time" | bc)

    # Aggregate results
    local total_inserts=0
    local total_selects=0
    local total_updates=0
    local total_complex=0

    for ((i=1; i<=num_connections; i++)); do
        if [ -f "$temp_dir/worker_$i.txt" ]; then
            read inserts selects updates complex < "$temp_dir/worker_$i.txt"
            total_inserts=$((total_inserts + inserts))
            total_selects=$((total_selects + selects))
            total_updates=$((total_updates + updates))
            total_complex=$((total_complex + complex))
        fi
    done

    local total_ops=$((total_inserts + total_selects + total_updates + total_complex))
    local tps=$(echo "scale=2; $total_ops / $actual_duration" | bc)

    echo "Duration: $actual_duration seconds" | tee -a "$RESULTS_FILE"
    echo "Concurrent Connections: $num_connections" | tee -a "$RESULTS_FILE"
    echo "Total Operations: $total_ops" | tee -a "$RESULTS_FILE"
    echo "  - INSERTs: $total_inserts" | tee -a "$RESULTS_FILE"
    echo "  - SELECTs: $total_selects" | tee -a "$RESULTS_FILE"
    echo "  - UPDATEs: $total_updates" | tee -a "$RESULTS_FILE"
    echo "  - Complex: $total_complex" | tee -a "$RESULTS_FILE"
    echo "Transactions per Second (TPS): $tps" | tee -a "$RESULTS_FILE"

    # Check for connection pool exhaustion
    echo "" | tee -a "$RESULTS_FILE"
    echo "=== Connection Pool Status ===" | tee -a "$RESULTS_FILE"
    psql "$DB_URL" -c "SELECT count(*), state FROM pg_stat_activity GROUP BY state;" | tee -a "$RESULTS_FILE"

    # Check for locks and deadlocks
    echo "" | tee -a "$RESULTS_FILE"
    echo "=== Lock Contention ===" | tee -a "$RESULTS_FILE"
    psql "$DB_URL" -c "SELECT count(*) as lock_count FROM pg_locks WHERE NOT granted;" | tee -a "$RESULTS_FILE"

    echo "" | tee -a "$RESULTS_FILE"
    echo "=== Deadlock Stats ===" | tee -a "$RESULTS_FILE"
    psql "$DB_URL" -c "SELECT datname, deadlocks FROM pg_stat_database WHERE datname = 'foxhunt';" | tee -a "$RESULTS_FILE"

    # Cleanup
    rm -rf "$temp_dir"
    echo "" | tee -a "$RESULTS_FILE"
}

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

    echo "=== Top 10 Slowest Queries ===" | tee -a "$RESULTS_FILE"
    psql "$DB_URL" -c "SELECT query, calls, mean_exec_time, max_exec_time FROM pg_stat_statements ORDER BY mean_exec_time DESC LIMIT 10;" 2>&1 | tee -a "$RESULTS_FILE" || echo "pg_stat_statements extension not available" | tee -a "$RESULTS_FILE"

    echo "" | tee -a "$RESULTS_FILE"
    echo "=== Table Statistics ===" | tee -a "$RESULTS_FILE"
    psql "$DB_URL" -c "SELECT relname, seq_scan, seq_tup_read, idx_scan, idx_tup_fetch, n_tup_ins, n_tup_upd, n_tup_del 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 "$DB_URL" -c "SELECT schemaname, 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 "$DB_URL" -c "SELECT pg_database.datname, pg_size_pretty(pg_database_size(pg_database.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 "$DB_URL" -c "SELECT relname, pg_size_pretty(pg_total_relation_size(relid)) AS total_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 "=== Cache Hit Ratio (Final) ===" | tee -a "$RESULTS_FILE"
    psql "$DB_URL" -c "SELECT sum(blks_hit)::float / (sum(blks_hit) + sum(blks_read)) * 100 AS cache_hit_ratio FROM pg_stat_database WHERE datname = 'foxhunt';" | tee -a "$RESULTS_FILE"

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

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

    # Initial health check
    check_db_health
    get_baseline_counts

    # Run tests with increasing concurrency
    run_baseline_test
    run_concurrent_test 10 "MODERATE LOAD"
    run_concurrent_test 50 "HIGH LOAD"
    run_concurrent_test 100 "STRESS TEST"

    # Post-test analysis
    check_query_performance
    check_resource_utilization

    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"
}

# Run main test
main