#!/bin/bash

# Comprehensive Stress Testing Script for MEV Bot
# This script runs all stress tests to validate system performance and reliability

set -e  # Exit on any error

# Colors for output
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
BLUE='\033[0;34m'
PURPLE='\033[0;35m'
CYAN='\033[0;36m'
NC='\033[0m' # No Color

# Script information
echo -e "${PURPLE}🧪 MEV Bot Comprehensive Stress Testing Suite${NC}"
echo -e "${PURPLE}==========================================${NC}"
echo ""

# Check if running from project root
if [ ! -f "go.mod" ]; then
    echo -e "${RED}❌ Error: This script must be run from the project root directory${NC}"
    exit 1
fi

# Parse command line arguments
TEST_DURATION=300  # Default 5 minutes
TRANSACTIONS_PER_SECOND=1000
CONCURRENT_WORKERS=10
MEMORY_LIMIT="2G"
CPU_LIMIT="2.0"
VERBOSE=false
RUN_ALL=false
RUN_UNIT=false
RUN_INTEGRATION=false
RUN_PERFORMANCE=false
RUN_LOAD=false
RUN_STRESS=false

while [[ $# -gt 0 ]]; do
    case $1 in
        --duration)
            TEST_DURATION="$2"
            shift 2
            ;;
        --tps)
            TRANSACTIONS_PER_SECOND="$2"
            shift 2
            ;;
        --workers)
            CONCURRENT_WORKERS="$2"
            shift 2
            ;;
        --memory)
            MEMORY_LIMIT="$2"
            shift 2
            ;;
        --cpu)
            CPU_LIMIT="$2"
            shift 2
            ;;
        --verbose|-v)
            VERBOSE=true
            shift
            ;;
        --all)
            RUN_ALL=true
            shift
            ;;
        --unit)
            RUN_UNIT=true
            shift
            ;;
        --integration)
            RUN_INTEGRATION=true
            shift
            ;;
        --performance)
            RUN_PERFORMANCE=true
            shift
            ;;
        --load)
            RUN_LOAD=true
            shift
            ;;
        --stress)
            RUN_STRESS=true
            shift
            ;;
        --help|-h)
            echo "Usage: $0 [OPTIONS]"
            echo ""
            echo "Options:"
            echo "  --duration SECONDS     Test duration in seconds (default: 300)"
            echo "  --tps COUNT           Transactions per second (default: 1000)"
            echo "  --workers COUNT       Concurrent workers (default: 10)"
            echo "  --memory LIMIT        Memory limit (default: 2G)"
            echo "  --cpu LIMIT           CPU limit (default: 2.0)"
            echo "  --verbose, -v         Verbose output"
            echo "  --all                 Run all tests"
            echo "  --unit                Run unit tests"
            echo "  --integration         Run integration tests"
            echo "  --performance         Run performance tests"
            echo "  --load                Run load tests"
            echo "  --stress              Run stress tests"
            echo "  --help, -h            Show this help message"
            echo ""
            exit 0
            ;;
        *)
            echo -e "${RED}❌ Unknown option: $1${NC}"
            exit 1
            ;;
    esac
done

# If no specific tests selected, run all
if [ "$RUN_ALL" = false ] && [ "$RUN_UNIT" = false ] && [ "$RUN_INTEGRATION" = false ] && [ "$RUN_PERFORMANCE" = false ] && [ "$RUN_LOAD" = false ] && [ "$RUN_STRESS" = false ]; then
    RUN_ALL=true
fi

# Set flags based on --all
if [ "$RUN_ALL" = true ]; then
    RUN_UNIT=true
    RUN_INTEGRATION=true
    RUN_PERFORMANCE=true
    RUN_LOAD=true
    RUN_STRESS=true
fi

# Display test configuration
echo -e "${BLUE}📋 Test Configuration:${NC}"
echo -e "  Duration: ${TEST_DURATION}s"
echo -e "  TPS: ${TRANSACTIONS_PER_SECOND}"
echo -e "  Workers: ${CONCURRENT_WORKERS}"
echo -e "  Memory Limit: ${MEMORY_LIMIT}"
echo -e "  CPU Limit: ${CPU_LIMIT}"
echo -e "  Verbose: ${VERBOSE}"
echo ""

# Create results directory
RESULTS_DIR="test_results/stress_tests/$(date +%Y%m%d_%H%M%S)"
mkdir -p "$RESULTS_DIR"

# Function to run a test and capture results
run_test() {
    local test_name="$1"
    local test_cmd="$2"
    local test_file="$3"
    
    echo -e "${BLUE}🚀 Running $test_name...${NC}"
    
    local start_time=$(date +%s)
    local output_file="$RESULTS_DIR/${test_file}.log"
    
    if [ "$VERBOSE" = true ]; then
        eval "$test_cmd" 2>&1 | tee "$output_file"
        local exit_code=${PIPESTATUS[0]}
    else
        eval "$test_cmd" > "$output_file" 2>&1
        local exit_code=$?
    fi
    
    local end_time=$(date +%s)
    local duration=$((end_time - start_time))
    
    if [ $exit_code -eq 0 ]; then
        echo -e "${GREEN}✅ $test_name PASSED in ${duration}s${NC}"
        echo "$test_name,PASSED,$duration" >> "$RESULTS_DIR/test_summary.csv"
    else
        echo -e "${RED}❌ $test_name FAILED in ${duration}s${NC}"
        echo "$test_name,FAILED,$duration" >> "$RESULTS_DIR/test_summary.csv"
        cat "$output_file" | tail -20
    fi
    
    return $exit_code
}

# Initialize test summary
echo "Test Name,Status,Duration (seconds)" > "$RESULTS_DIR/test_summary.csv"

# Run unit tests
if [ "$RUN_UNIT" = true ]; then
    echo -e "${CYAN}🧩 Running Unit Tests${NC}"
    
    # Basic unit tests
    run_test "Basic Unit Tests" "go test -v ./pkg/... -short" "unit_basic"
    
    # Math unit tests
    run_test "Math Unit Tests" "go test -v ./pkg/math/..." "unit_math"
    
    # Scanner unit tests
    run_test "Scanner Unit Tests" "go test -v ./pkg/scanner/..." "unit_scanner"
    
    # Arbitrage unit tests
    run_test "Arbitrage Unit Tests" "go test -v ./pkg/arbitrage/..." "unit_arbitrage"
    
    # Trading unit tests
    run_test "Trading Unit Tests" "go test -v ./pkg/trading/..." "unit_trading"
    
    echo ""
fi

# Run integration tests
if [ "$RUN_INTEGRATION" = true ]; then
    echo -e "${CYAN}🔗 Running Integration Tests${NC}"
    
    # Integration tests with mocked network
    run_test "Integration Tests" "go test -v ./test/integration/..." "integration"
    
    # End-to-end tests
    run_test "End-to-End Tests" "go test -v ./test/e2e/..." "e2e"
    
    echo ""
fi

# Run performance tests
if [ "$RUN_PERFORMANCE" = true ]; then
    echo -e "${CYAN}⚡ Running Performance Tests${NC}"
    
    # Benchmark tests
    run_test "Benchmark Tests" "go test -bench=. -benchmem ./pkg/... -count=3" "benchmarks"
    
    # Performance regression tests
    run_test "Performance Regression Tests" "go test -v ./test/performance_benchmarks_test.go -count=1" "perf_regression"
    
    echo ""
fi

# Run load tests
if [ "$RUN_LOAD" = true ]; then
    echo -e "${CYAN}🏋️ Running Load Tests${NC}"
    
    # High TPS load test
    run_test "High TPS Load Test" "timeout ${TEST_DURATION}s go run test/load/high_tps_test.go --tps ${TRANSACTIONS_PER_SECOND} --duration ${TEST_DURATION}" "load_high_tps"
    
    # Concurrent workers load test
    run_test "Concurrent Workers Load Test" "timeout ${TEST_DURATION}s go run test/load/concurrent_workers_test.go --workers ${CONCURRENT_WORKERS} --duration ${TEST_DURATION}" "load_concurrent"
    
    # Memory intensive load test
    run_test "Memory Intensive Load Test" "timeout ${TEST_DURATION}s go run test/load/memory_intensive_test.go --memory ${MEMORY_LIMIT} --duration ${TEST_DURATION}" "load_memory"
    
    echo ""
fi

# Run stress tests
if [ "$RUN_STRESS" = true ]; then
    echo -e "${CYAN}🔥 Running Stress Tests${NC}"
    
    # Stress test runner
    run_test "Stress Test Suite" "timeout ${TEST_DURATION}s go run test/stress/stress_test_runner.go --full-suite --duration ${TEST_DURATION} --tps ${TRANSACTIONS_PER_SECOND}" "stress_suite"
    
    # Market scanner stress test
    run_test "Market Scanner Stress Test" "timeout ${TEST_DURATION}s go run test/stress/market_scanner_test.go --duration ${TEST_DURATION}" "stress_market_scanner"
    
    # Swap analyzer stress test
    run_test "Swap Analyzer Stress Test" "timeout ${TEST_DURATION}s go run test/stress/swap_analyzer_test.go --duration ${TEST_DURATION}" "stress_swap_analyzer"
    
    # Pool discovery stress test
    run_test "Pool Discovery Stress Test" "timeout ${TEST_DURATION}s go run test/stress/pool_discovery_test.go --duration ${TEST_DURATION}" "stress_pool_discovery"
    
    # Arbitrage engine stress test
    run_test "Arbitrage Engine Stress Test" "timeout ${TEST_DURATION}s go run test/stress/arbitrage_engine_test.go --duration ${TEST_DURATION}" "stress_arbitrage_engine"
    
    echo ""
fi

# Generate test summary report
echo -e "${PURPLE}📊 Test Summary Report${NC}"
echo -e "${PURPLE}====================${NC}"

PASSED_COUNT=0
FAILED_COUNT=0

while IFS=, read -r name status duration; do
    if [ "$name" != "Test Name" ]; then
        if [ "$status" = "PASSED" ]; then
            PASSED_COUNT=$((PASSED_COUNT + 1))
            echo -e "  ✅ $name - ${GREEN}PASSED${NC} (${duration}s)"
        else
            FAILED_COUNT=$((FAILED_COUNT + 1))
            echo -e "  ❌ $name - ${RED}FAILED${NC} (${duration}s)"
        fi
    fi
done < "$RESULTS_DIR/test_summary.csv"

TOTAL_TESTS=$((PASSED_COUNT + FAILED_COUNT))
SUCCESS_RATE=0
if [ $TOTAL_TESTS -gt 0 ]; then
    SUCCESS_RATE=$((PASSED_COUNT * 100 / TOTAL_TESTS))
fi

echo ""
echo -e "${BLUE}📈 Test Results:${NC}"
echo -e "  Total Tests: ${TOTAL_TESTS}"
echo -e "  Passed: ${PASSED_COUNT}"
echo -e "  Failed: ${FAILED_COUNT}"
echo -e "  Success Rate: ${SUCCESS_RATE}%"
echo ""

# Final status
if [ $FAILED_COUNT -eq 0 ]; then
    echo -e "${GREEN}🎉 All stress tests passed! System is ready for production deployment.${NC}"
    echo -e "${GREEN}   Results saved to: $RESULTS_DIR${NC}"
    exit 0
else
    echo -e "${RED}⚠️  $FAILED_COUNT stress tests failed. Please review results and fix issues before production deployment.${NC}"
    echo -e "${YELLOW}   Results saved to: $RESULTS_DIR${NC}"
    exit 1
fi