#!/usr/bin/env python3
"""
Comprehensive Load Test for Trading Service

Tests:
1. Baseline latency (single client)
2. Concurrent connections (100 clients)
3. Sustained load (5+ minutes)
4. Database performance
5. Resource monitoring
6. Production readiness assessment
"""

import asyncio
import time
import statistics
import sys
import uuid
from concurrent.futures import ThreadPoolExecutor
from typing import List, Tuple
import requests

# Metrics storage
class PerformanceMetrics:
    def __init__(self):
        self.latencies_ms = []
        self.successful = 0
        self.failed = 0
        self.start_time = None
        self.end_time = None

    def record_success(self, latency_ms: float):
        self.latencies_ms.append(latency_ms)
        self.successful += 1

    def record_failure(self):
        self.failed += 1

    def get_percentiles(self) -> Tuple[float, float, float, float, float]:
        if not self.latencies_ms:
            return (0, 0, 0, 0, 0)

        sorted_lat = sorted(self.latencies_ms)
        n = len(sorted_lat)

        return (
            sorted_lat[0],  # min
            sorted_lat[n // 2],  # p50
            sorted_lat[int(n * 0.95)],  # p95
            sorted_lat[int(n * 0.99)],  # p99
            sorted_lat[-1]  # max
        )

    def print_summary(self):
        duration = (self.end_time - self.start_time) if self.end_time and self.start_time else 0
        total = self.successful + self.failed
        success_rate = (self.successful / total * 100) if total > 0 else 0
        throughput = self.successful / duration if duration > 0 else 0

        min_lat, p50, p95, p99, max_lat = self.get_percentiles()

        print("\n" + "="*70)
        print("              TRADING SERVICE LOAD TEST RESULTS")
        print("="*70)
        print(f"Test Duration:       {duration:.2f}s")
        print(f"Total Orders:        {total}")
        print(f"Successful Orders:   {self.successful} ({success_rate:.2f}%)")
        print(f"Failed Orders:       {self.failed}")
        print(f"Throughput:          {throughput:.0f} orders/sec")
        print("-"*70)
        print("                    LATENCY METRICS")
        print("-"*70)
        print(f"Min Latency:         {min_lat:.2f}ms")
        print(f"P50 Latency:         {p50:.2f}ms")
        print(f"P95 Latency:         {p95:.2f}ms")
        print(f"P99 Latency:         {p99:.2f}ms")
        print(f"Max Latency:         {max_lat:.2f}ms")
        print("="*70)

        print("\n📊 PERFORMANCE ASSESSMENT:")

        if throughput >= 10000:
            print(f"✅ Throughput target ACHIEVED: {throughput:.0f} orders/sec (target: 10K orders/sec)")
        elif throughput >= 5000:
            print(f"⚠️  Throughput ACCEPTABLE: {throughput:.0f} orders/sec (target: 10K orders/sec)")
        else:
            print(f"❌ Throughput BELOW target: {throughput:.0f} orders/sec (target: 10K orders/sec)")

        if p99 < 100:
            print(f"✅ P99 latency EXCELLENT: {p99:.2f}ms (< 100ms)")
        elif p99 < 500:
            print(f"⚠️  P99 latency ACCEPTABLE: {p99:.2f}ms (< 500ms)")
        else:
            print(f"❌ P99 latency HIGH: {p99:.2f}ms (> 500ms)")

        if success_rate >= 99.0:
            print(f"✅ Success rate EXCELLENT: {success_rate:.2f}%")
        elif success_rate >= 95.0:
            print(f"⚠️  Success rate ACCEPTABLE: {success_rate:.2f}%")
        else:
            print(f"❌ Success rate POOR: {success_rate:.2f}%")


def submit_order_http(order_id: str, symbol: str) -> Tuple[bool, float]:
    """Submit order via HTTP (fallback if gRPC unavailable)"""
    url = "http://localhost:8081/api/v1/orders"

    payload = {
        "order_id": order_id,
        "symbol": symbol,
        "side": "buy",
        "order_type": "limit",
        "quantity": "1.0",
        "price": "50000.0",
        "time_in_force": "gtc"
    }

    start = time.time()
    try:
        response = requests.post(url, json=payload, timeout=5)
        latency_ms = (time.time() - start) * 1000
        return response.status_code == 200, latency_ms
    except Exception as e:
        latency_ms = (time.time() - start) * 1000
        return False, latency_ms


async def test_1_baseline_latency():
    """Test 1: Baseline latency with single client"""
    print("\n" + "="*70)
    print("  TEST 1: BASELINE LATENCY (Single Client)")
    print("="*70)

    metrics = PerformanceMetrics()
    num_requests = 1000
    symbols = ["BTC/USD", "ETH/USD", "SOL/USD", "AVAX/USD", "MATIC/USD"]

    print(f"📊 Sending {num_requests} orders sequentially...")

    metrics.start_time = time.time()

    for i in range(num_requests):
        order_id = str(uuid.uuid4())
        symbol = symbols[i % len(symbols)]

        success, latency_ms = submit_order_http(order_id, symbol)

        if success:
            metrics.record_success(latency_ms)
        else:
            metrics.record_failure()
            if metrics.failed <= 5:
                print(f"❌ Order {i} failed")

    metrics.end_time = time.time()
    metrics.print_summary()


async def test_2_concurrent_connections():
    """Test 2: Concurrent connections (100 clients)"""
    print("\n" + "="*70)
    print("  TEST 2: CONCURRENT CONNECTIONS (100 Clients)")
    print("="*70)

    num_clients = 100
    orders_per_client = 100
    metrics = PerformanceMetrics()
    symbols = ["BTC/USD", "ETH/USD", "SOL/USD", "AVAX/USD", "MATIC/USD"]

    print(f"🚀 Spawning {num_clients} concurrent clients ({orders_per_client} orders each)...")

    metrics.start_time = time.time()

    def client_worker(client_id: int):
        for order_idx in range(orders_per_client):
            order_id = str(uuid.uuid4())
            symbol = symbols[(client_id * orders_per_client + order_idx) % len(symbols)]

            success, latency_ms = submit_order_http(order_id, symbol)

            if success:
                metrics.record_success(latency_ms)
            else:
                metrics.record_failure()

    with ThreadPoolExecutor(max_workers=num_clients) as executor:
        futures = [executor.submit(client_worker, i) for i in range(num_clients)]
        for future in futures:
            future.result()

    metrics.end_time = time.time()
    metrics.print_summary()


async def test_3_sustained_load():
    """Test 3: Sustained load (5 minutes)"""
    print("\n" + "="*70)
    print("  TEST 3: SUSTAINED LOAD (5 Minutes)")
    print("="*70)

    test_duration_secs = 300
    num_clients = 50
    target_rate_per_client = 200  # 10K total / 50 clients = 200 per client

    metrics = PerformanceMetrics()
    symbols = ["BTC/USD", "ETH/USD", "SOL/USD", "AVAX/USD", "MATIC/USD"]

    print(f"🚀 Starting {num_clients} clients for {test_duration_secs} seconds...")
    print(f"🎯 Target: {num_clients * target_rate_per_client} orders/sec total")

    shutdown_flag = [False]
    metrics.start_time = time.time()

    def client_worker(client_id: int):
        order_count = 0
        delay_secs = 1.0 / target_rate_per_client

        while not shutdown_flag[0]:
            order_id = str(uuid.uuid4())
            symbol = symbols[order_count % len(symbols)]

            success, latency_ms = submit_order_http(order_id, symbol)

            if success:
                metrics.record_success(latency_ms)
            else:
                metrics.record_failure()

            order_count += 1
            time.sleep(delay_secs)

    with ThreadPoolExecutor(max_workers=num_clients) as executor:
        futures = [executor.submit(client_worker, i) for i in range(num_clients)]

        # Run for specified duration
        time.sleep(test_duration_secs)
        shutdown_flag[0] = True

        # Wait for all clients to finish
        for future in futures:
            try:
                future.result(timeout=5)
            except:
                pass

    metrics.end_time = time.time()
    metrics.print_summary()


async def test_4_database_performance():
    """Test 4: Database performance"""
    print("\n" + "="*70)
    print("  TEST 4: DATABASE PERFORMANCE")
    print("="*70)

    num_orders = 5000
    symbols = ["BTC/USD", "ETH/USD", "SOL/USD"]

    print(f"📊 Submitting {num_orders} orders to measure database performance...")

    start_time = time.time()
    success_count = 0
    failure_count = 0

    for i in range(num_orders):
        order_id = str(uuid.uuid4())
        symbol = symbols[i % len(symbols)]

        success, _ = submit_order_http(order_id, symbol)

        if success:
            success_count += 1
        else:
            failure_count += 1
            if failure_count <= 5:
                print(f"❌ Order {i} failed")

    duration = time.time() - start_time
    throughput = success_count / duration

    print("\n📈 DATABASE PERFORMANCE:")
    print(f"  Duration:         {duration:.2f}s")
    print(f"  Successful:       {success_count}")
    print(f"  Failed:           {failure_count}")
    print(f"  DB Writes/sec:    {throughput:.0f}")

    if throughput >= 2000:
        print(f"✅ Database performance EXCELLENT: {throughput:.0f} writes/sec")
    elif throughput >= 1000:
        print(f"⚠️  Database performance ACCEPTABLE: {throughput:.0f} writes/sec")
    else:
        print(f"❌ Database performance LOW: {throughput:.0f} writes/sec (expected >2000)")


async def test_5_resource_monitoring():
    """Test 5: Resource monitoring"""
    print("\n" + "="*70)
    print("  TEST 5: RESOURCE MONITORING")
    print("="*70)

    # Check service health
    health_url = "http://localhost:8081/health"
    print(f"🏥 Checking service health at {health_url}...")

    try:
        response = requests.get(health_url, timeout=5)
        print(f"✅ Health check response: {response.status_code}")
        print(f"   Body: {response.text[:200]}")
    except Exception as e:
        print(f"⚠️  Health check failed: {e}")

    # Check Prometheus metrics
    metrics_url = "http://localhost:9092/metrics"
    print(f"\n📊 Checking Prometheus metrics at {metrics_url}...")

    try:
        response = requests.get(metrics_url, timeout=5)
        lines = [l for l in response.text.split('\n') if l and not l.startswith('#')]
        print(f"✅ Found {len(lines)} metric entries")

        # Show key metrics
        for line in lines[:20]:
            if any(keyword in line for keyword in ['orders', 'latency', 'cpu', 'memory']):
                print(f"   {line[:100]}")
    except Exception as e:
        print(f"⚠️  Metrics check failed: {e}")


async def test_6_production_readiness():
    """Test 6: Production readiness assessment"""
    print("\n" + "="*70)
    print("  TEST 6: PRODUCTION READINESS ASSESSMENT")
    print("="*70)

    num_clients = 50
    orders_per_client = 200
    metrics = PerformanceMetrics()
    symbols = ["BTC/USD", "ETH/USD", "SOL/USD", "AVAX/USD", "MATIC/USD"]

    print(f"🎯 Production simulation: {num_clients} clients, {orders_per_client} orders each")

    metrics.start_time = time.time()

    def client_worker(client_id: int):
        for order_idx in range(orders_per_client):
            order_id = str(uuid.uuid4())
            symbol = symbols[(client_id * orders_per_client + order_idx) % len(symbols)]

            success, latency_ms = submit_order_http(order_id, symbol)

            if success:
                metrics.record_success(latency_ms)
            else:
                metrics.record_failure()

    with ThreadPoolExecutor(max_workers=num_clients) as executor:
        futures = [executor.submit(client_worker, i) for i in range(num_clients)]
        for future in futures:
            future.result()

    metrics.end_time = time.time()
    metrics.print_summary()

    # Production readiness criteria
    total = metrics.successful + metrics.failed
    success_rate = (metrics.successful / total * 100) if total > 0 else 0
    duration = metrics.end_time - metrics.start_time
    throughput = metrics.successful / duration if duration > 0 else 0
    _, _, _, p99, _ = metrics.get_percentiles()

    print("\n🎯 PRODUCTION READINESS:")

    passed = 0
    total_checks = 3

    # Check 1: Success rate
    if success_rate >= 99.0:
        print(f"✅ Success rate: {success_rate:.2f}% (>= 99%)")
        passed += 1
    else:
        print(f"❌ Success rate: {success_rate:.2f}% (< 99%)")

    # Check 2: Throughput
    if throughput >= 5000.0:
        print(f"✅ Throughput: {throughput:.0f} orders/sec (>= 5000)")
        passed += 1
    elif throughput >= 3000.0:
        print(f"⚠️  Throughput: {throughput:.0f} orders/sec (>= 3000)")
        passed += 0.5
    else:
        print(f"❌ Throughput: {throughput:.0f} orders/sec (< 3000)")

    # Check 3: P99 latency
    if p99 < 100.0:
        print(f"✅ P99 latency: {p99:.2f}ms (< 100ms)")
        passed += 1
    elif p99 < 500.0:
        print(f"⚠️  P99 latency: {p99:.2f}ms (< 500ms)")
        passed += 0.5
    else:
        print(f"❌ P99 latency: {p99:.2f}ms (>= 500ms)")

    print(f"\n📊 OVERALL: {passed}/{total_checks} checks passed")

    if passed >= 2.5:
        print("🎉 PRODUCTION READY!")
    elif passed >= 2.0:
        print("⚠️  Acceptable for production with monitoring")
    else:
        print("❌ Not ready for production deployment")


async def main():
    """Run all load tests"""
    print("\n" + "="*70)
    print("     FOXHUNT TRADING SERVICE - COMPREHENSIVE LOAD TEST")
    print("="*70)
    print("\nTarget: Trading Service at http://localhost:8081")
    print("gRPC Port: 50052 | Health Port: 8081 | Metrics Port: 9092")

    tests = [
        ("Test 1: Baseline Latency", test_1_baseline_latency),
        ("Test 2: Concurrent Connections", test_2_concurrent_connections),
        ("Test 4: Database Performance", test_4_database_performance),
        ("Test 5: Resource Monitoring", test_5_resource_monitoring),
        ("Test 6: Production Readiness", test_6_production_readiness),
        # Test 3 (sustained load) commented out for quick runs
        # ("Test 3: Sustained Load", test_3_sustained_load),
    ]

    for test_name, test_func in tests:
        try:
            print(f"\n\n{'='*70}")
            print(f"Starting: {test_name}")
            print(f"{'='*70}")
            await test_func()
        except KeyboardInterrupt:
            print("\n\n⚠️  Test interrupted by user")
            break
        except Exception as e:
            print(f"\n❌ Test failed with error: {e}")
            import traceback
            traceback.print_exc()

    print("\n" + "="*70)
    print("     LOAD TEST SUITE COMPLETED")
    print("="*70)


if __name__ == "__main__":
    asyncio.run(main())