#!/usr/bin/env python3
"""
PostgreSQL Load Test for Foxhunt Trading System
Tests database performance under increasing concurrent load
"""

import psycopg2
import psycopg2.pool
import time
import random
import multiprocessing as mp
from datetime import datetime
from typing import Dict, List, Tuple

DB_CONFIG = {
    'host': 'localhost',
    'port': 5432,
    'database': 'foxhunt',
    'user': 'foxhunt',
    'password': 'foxhunt_dev_password'
}

TEST_DURATION = 30  # seconds per scenario
SYMBOLS = ['BTC/USD', 'ETH/USD', 'SOL/USD']

def get_connection():
    """Get database connection"""
    return psycopg2.connect(**DB_CONFIG)

def worker_process(worker_id: int, duration: int, result_queue: mp.Queue):
    """Worker process that executes database operations"""
    conn = get_connection()
    conn.autocommit = True
    cursor = conn.cursor()

    start_time = time.time()
    end_time = start_time + duration

    inserts = 0
    selects = 0
    updates = 0
    complex_queries = 0
    errors = 0

    while time.time() < end_time:
        try:
            operation = random.randint(1, 10)
            symbol = random.choice(SYMBOLS)

            # 40% INSERT
            if operation <= 4:
                cursor.execute("""
                    INSERT INTO orders (account_id, symbol, side, order_type, quantity,
                                        limit_price, venue, created_at, updated_at)
                    VALUES (%s, %s, 'buy', 'limit', 100000000, 5000000000000, 'load_test',
                            EXTRACT(EPOCH FROM NOW())*1000000000,
                            EXTRACT(EPOCH FROM NOW())*1000000000)
                """, (f'worker_{worker_id}', symbol))
                inserts += 1

            # 30% SELECT
            elif operation <= 7:
                cursor.execute("""
                    SELECT id, status, quantity, filled_quantity
                    FROM orders
                    WHERE symbol = %s
                    ORDER BY created_at DESC
                    LIMIT 10
                """, (symbol,))
                cursor.fetchall()
                selects += 1

            # 20% UPDATE
            elif operation <= 9:
                cursor.execute("""
                    UPDATE orders
                    SET status = 'partially_filled',
                        filled_quantity = filled_quantity + 10000000,
                        updated_at = EXTRACT(EPOCH FROM NOW())*1000000000
                    WHERE id = (
                        SELECT id FROM orders
                        WHERE status = 'pending' AND symbol = %s
                        LIMIT 1
                    )
                """, (symbol,))
                updates += 1

            # 10% Complex query
            else:
                cursor.execute("""
                    SELECT symbol,
                           COUNT(*) as order_count,
                           SUM(quantity) as total_quantity,
                           AVG(limit_price) as avg_price
                    FROM orders
                    WHERE created_at > EXTRACT(EPOCH FROM (NOW() - INTERVAL '1 hour'))*1000000000
                    GROUP BY symbol
                    ORDER BY order_count DESC
                """)
                cursor.fetchall()
                complex_queries += 1

        except Exception as e:
            errors += 1
            print(f"Worker {worker_id} error: {e}")

    cursor.close()
    conn.close()

    result_queue.put({
        'worker_id': worker_id,
        'inserts': inserts,
        'selects': selects,
        'updates': updates,
        'complex': complex_queries,
        'errors': errors
    })

def run_load_test(num_workers: int, test_name: str, duration: int) -> Dict:
    """Run load test with specified number of workers"""
    print(f"\n{'='*60}")
    print(f"TEST: {test_name} ({num_workers} concurrent workers)")
    print(f"{'='*60}")

    result_queue = mp.Queue()
    processes = []

    start_time = time.time()

    # Launch workers
    for i in range(num_workers):
        p = mp.Process(target=worker_process, args=(i, duration, result_queue))
        p.start()
        processes.append(p)

    # Wait for all workers
    for p in processes:
        p.join()

    end_time = time.time()
    actual_duration = end_time - start_time

    # Collect results
    total_inserts = 0
    total_selects = 0
    total_updates = 0
    total_complex = 0
    total_errors = 0

    while not result_queue.empty():
        result = result_queue.get()
        total_inserts += result['inserts']
        total_selects += result['selects']
        total_updates += result['updates']
        total_complex += result['complex']
        total_errors += result['errors']

    total_ops = total_inserts + total_selects + total_updates + total_complex
    tps = total_ops / actual_duration if actual_duration > 0 else 0

    results = {
        'test_name': test_name,
        'num_workers': num_workers,
        'duration': actual_duration,
        'total_ops': total_ops,
        'inserts': total_inserts,
        'selects': total_selects,
        'updates': total_updates,
        'complex': total_complex,
        'errors': total_errors,
        'tps': tps
    }

    # Print results
    print(f"Duration: {actual_duration:.2f} seconds")
    print(f"Total Operations: {total_ops}")
    print(f"  - INSERTs: {total_inserts}")
    print(f"  - SELECTs: {total_selects}")
    print(f"  - UPDATEs: {total_updates}")
    print(f"  - Complex: {total_complex}")
    print(f"  - Errors: {total_errors}")
    print(f"Transactions per Second (TPS): {tps:.2f}")

    # Check connection pool and locks
    try:
        conn = get_connection()
        cursor = conn.cursor()

        print("\n=== Connection Pool Status ===")
        cursor.execute("SELECT state, COUNT(*) FROM pg_stat_activity WHERE datname='foxhunt' GROUP BY state")
        for row in cursor.fetchall():
            print(f"  {row[0]}: {row[1]}")

        print("\n=== Lock Contention ===")
        cursor.execute("SELECT COUNT(*) FROM pg_locks WHERE NOT granted")
        locks_waiting = cursor.fetchone()[0]
        print(f"  Locks Waiting: {locks_waiting}")

        print("\n=== Deadlocks ===")
        cursor.execute("SELECT deadlocks FROM pg_stat_database WHERE datname='foxhunt'")
        deadlocks = cursor.fetchone()[0]
        print(f"  Total Deadlocks: {deadlocks}")

        cursor.close()
        conn.close()
    except Exception as e:
        print(f"Error checking pool status: {e}")

    return results

def print_database_health():
    """Print database health metrics"""
    conn = get_connection()
    cursor = conn.cursor()

    print("\n=== Database Configuration ===")
    cursor.execute("SHOW max_connections")
    print(f"  Max Connections: {cursor.fetchone()[0]}")

    cursor.execute("SHOW shared_buffers")
    print(f"  Shared Buffers: {cursor.fetchone()[0]}")

    cursor.execute("SHOW synchronous_commit")
    print(f"  Synchronous Commit: {cursor.fetchone()[0]}")

    print("\n=== Cache Hit Ratio ===")
    cursor.execute("""
        SELECT (sum(blks_hit)::float / NULLIF(sum(blks_hit) + sum(blks_read), 0) * 100)::numeric(5,2)
        FROM pg_stat_database WHERE datname='foxhunt'
    """)
    print(f"  {cursor.fetchone()[0]}%")

    print("\n=== Baseline Table Counts ===")
    cursor.execute("SELECT COUNT(*) FROM orders")
    print(f"  Orders: {cursor.fetchone()[0]}")

    cursor.close()
    conn.close()

def print_final_analysis():
    """Print final analysis"""
    conn = get_connection()
    cursor = conn.cursor()

    print("\n" + "="*60)
    print("FINAL ANALYSIS")
    print("="*60)

    print("\n=== Table Statistics ===")
    cursor.execute("""
        SELECT relname, n_tup_ins, n_tup_upd, n_tup_del,
               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 relname IN ('orders', 'executions', 'fills', 'positions')
        ORDER BY n_tup_ins DESC
    """)
    print(f"{'Table':<12} {'Inserts':<10} {'Updates':<10} {'Deletes':<10} {'Seq Scan':<10} {'Idx Scan':<10} {'Idx %':<8}")
    print("-" * 80)
    for row in cursor.fetchall():
        print(f"{row[0]:<12} {row[1]:<10} {row[2]:<10} {row[3]:<10} {row[4]:<10} {row[5]:<10} {row[6]:<8}")

    print("\n=== Index Usage (Top 5) ===")
    cursor.execute("""
        SELECT tablename, indexname, idx_scan, idx_tup_read
        FROM pg_stat_user_indexes
        WHERE tablename IN ('orders', 'executions', 'fills', 'positions')
        AND idx_scan > 0
        ORDER BY idx_scan DESC
        LIMIT 5
    """)
    print(f"{'Table':<12} {'Index':<30} {'Scans':<10} {'Rows Read':<12}")
    print("-" * 70)
    for row in cursor.fetchall():
        print(f"{row[0]:<12} {row[1]:<30} {row[2]:<10} {row[3]:<12}")

    print("\n=== Database Size ===")
    cursor.execute("""
        SELECT pg_size_pretty(pg_database_size('foxhunt'))
    """)
    print(f"  Total Size: {cursor.fetchone()[0]}")

    print("\n=== Table Sizes ===")
    cursor.execute("""
        SELECT relname,
               pg_size_pretty(pg_total_relation_size(relid)) AS total_size
        FROM pg_stat_user_tables
        WHERE relname IN ('orders', 'executions', 'fills', 'positions')
        ORDER BY pg_total_relation_size(relid) DESC
    """)
    for row in cursor.fetchall():
        print(f"  {row[0]}: {row[1]}")

    print("\n=== Final Cache Hit Ratio ===")
    cursor.execute("""
        SELECT (sum(blks_hit)::float / NULLIF(sum(blks_hit) + sum(blks_read), 0) * 100)::numeric(5,2)
        FROM pg_stat_database WHERE datname='foxhunt'
    """)
    print(f"  {cursor.fetchone()[0]}%")

    cursor.close()
    conn.close()

def main():
    """Main test execution"""
    print("="*60)
    print("PostgreSQL Load Test - Foxhunt Trading System")
    print(f"Test Duration: {TEST_DURATION} seconds per scenario")
    print(f"Start Time: {datetime.now()}")
    print("="*60)

    # Initial health check
    print_database_health()

    # Run tests with increasing concurrency
    results = []
    results.append(run_load_test(1, "BASELINE (Single Worker)", TEST_DURATION))
    results.append(run_load_test(10, "MODERATE LOAD", TEST_DURATION))
    results.append(run_load_test(50, "HIGH LOAD", TEST_DURATION))
    results.append(run_load_test(100, "STRESS TEST", TEST_DURATION))

    # Final analysis
    print_final_analysis()

    # Summary
    print("\n" + "="*60)
    print("TEST SUMMARY")
    print("="*60)
    print(f"{'Test':<25} {'Workers':<10} {'TPS':<12} {'Errors':<10}")
    print("-" * 60)
    for r in results:
        print(f"{r['test_name']:<25} {r['num_workers']:<10} {r['tps']:<12.2f} {r['errors']:<10}")

    print("\n" + "="*60)
    print(f"Test Completed: {datetime.now()}")
    print("="*60)

if __name__ == '__main__':
    mp.set_start_method('fork')
    main()