#!/bin/bash
# PostgreSQL Load Test Script for Foxhunt Trading System (Schema-accurate)
# 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=30  # seconds per test (reduced for faster execution)
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" -t -c "SELECT version();" | tee -a "$RESULTS_FILE"
    psql "$DB_URL" -t -c "SHOW max_connections;" | tee -a "$RESULTS_FILE"
    psql "$DB_URL" -t -c "SHOW shared_buffers;" | tee -a "$RESULTS_FILE"
    psql "$DB_URL" -t -c "SHOW synchronous_commit;" | tee -a "$RESULTS_FILE"

    echo "" | tee -a "$RESULTS_FILE"
    echo "=== Cache Hit Ratio ===" | tee -a "$RESULTS_FILE"
    psql "$DB_URL" -t -c "SELECT ROUND(sum(blks_hit)::float / (sum(blks_hit) + sum(blks_read)) * 100, 2) || '%' 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" -t -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.%N)
    local end_time=$(echo "$start_time + $TEST_DURATION" | bc)
    local insert_count=0
    local select_count=0
    local update_count=0
    local complex_count=0
    local error_count=0

    while (( $(echo "$(date +%s.%N) < $end_time" | bc -l) )); do
        # 40% INSERT - orders with proper schema
        for i in {1..4}; do
            psql "$DB_URL" -c "INSERT INTO orders (account_id, symbol, side, order_type, quantity, limit_price, venue, created_at, updated_at) VALUES ('test_account', 'BTC/USD', 'buy', 'limit', 100000000, 5000000000000, 'test_venue', EXTRACT(EPOCH FROM NOW()) * 1000000000, EXTRACT(EPOCH FROM NOW()) * 1000000000);" > /dev/null 2>&1 && ((insert_count++)) || ((error_count++))
        done

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

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

        # 10% Complex query - order aggregation
        psql "$DB_URL" -c "SELECT symbol, COUNT(*) as order_count, SUM(quantity) as total_qty FROM orders WHERE created_at > EXTRACT(EPOCH FROM (NOW() - INTERVAL '1 hour')) * 1000000000 GROUP BY symbol;" > /dev/null 2>&1 && ((complex_count++)) || ((error_count++))
    done

    local actual_end_time=$(date +%s.%N)
    local actual_duration=$(echo "$actual_end_time - $start_time" | bc)
    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 "  - Errors: $error_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_$$_$num_connections"
    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 worker_errors=0
            local end_time=$(echo "$start_time + $TEST_DURATION" | bc)

            while (( $(echo "$(date +%s.%N) < $end_time" | bc -l) )); do
                local rand=$((RANDOM % 10))

                # 40% INSERT
                if (( rand < 4 )); then
                    psql "$DB_URL" -c "INSERT INTO orders (account_id, symbol, side, order_type, quantity, limit_price, venue, created_at, updated_at) VALUES ('test_account', 'BTC/USD', 'buy', 'limit', 100000000, 5000000000000, 'test_venue', EXTRACT(EPOCH FROM NOW()) * 1000000000, EXTRACT(EPOCH FROM NOW()) * 1000000000);" > /dev/null 2>&1 && ((worker_inserts++)) || ((worker_errors++))
                # 30% SELECT
                elif (( rand < 7 )); then
                    psql "$DB_URL" -c "SELECT id, status, quantity FROM orders WHERE symbol = 'BTC/USD' ORDER BY created_at DESC LIMIT 10;" > /dev/null 2>&1 && ((worker_selects++)) || ((worker_errors++))
                # 20% UPDATE
                elif (( rand < 9 )); then
                    psql "$DB_URL" -c "UPDATE orders SET status = 'partially_filled', updated_at = EXTRACT(EPOCH FROM NOW()) * 1000000000 WHERE id = (SELECT id FROM orders WHERE status = 'pending' LIMIT 1);" > /dev/null 2>&1 && ((worker_updates++)) || ((worker_errors++))
                # 10% Complex
                else
                    psql "$DB_URL" -c "SELECT symbol, COUNT(*) as order_count FROM orders WHERE created_at > EXTRACT(EPOCH FROM (NOW() - INTERVAL '1 hour')) * 1000000000 GROUP BY symbol;" > /dev/null 2>&1 && ((worker_complex++)) || ((worker_errors++))
                fi
            done

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

    # Wait for all workers to complete
    for pid in "${pids[@]}"; do
        wait $pid 2>/dev/null || true
    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
    local total_errors=0

    for ((i=1; i<=num_connections; i++)); do
        if [ -f "$temp_dir/worker_$i.txt" ]; then
            read inserts selects updates complex errors < "$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))
            total_errors=$((total_errors + errors))
        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 "  - Errors: $total_errors" | 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" -t -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 "=== 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 tablename, indexname, idx_scan, idx_tup_read 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" -t -c "SELECT ROUND(sum(blks_hit)::float / (sum(blks_hit) + sum(blks_read)) * 100, 2) || '%' 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"

    echo ""
    echo "Results file: $RESULTS_FILE"
}

# Run main test
main