#!/usr/bin/env python3
"""
Stress test for yfinance-plus with multiple tickers
Tests rate limiting, error handling, and performance
"""

import time
import threading
import logging
from concurrent.futures import ThreadPoolExecutor, as_completed
from datetime import datetime
import yfinance_plus as yf

# Setup detailed logging
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
)

# Test tickers - mix of popular and less popular stocks
TEST_TICKERS = [
    'AAPL', 'MSFT', 'GOOGL', 'AMZN', 'TSLA', 'META', 'NVDA', 'NFLX', 'BABA', 'V',
    'JPM', 'JNJ', 'WMT', 'PG', 'UNH', 'DIS', 'HD', 'PYPL', 'BAC', 'MA',
    'ADBE', 'CRM', 'PFE', 'XOM', 'VZ', 'INTC', 'CMCSA', 'T', 'ABT', 'KO'
]

def test_single_ticker(symbol, test_id=""):
    """Test a single ticker with error handling"""
    try:
        start_time = time.time()
        ticker = yf.Ticker(symbol)
        
        # Test multiple operations
        results = {}
        
        # Get basic info
        try:
            info = ticker.info
            results['info'] = info.get('longName', 'N/A') if info else 'No data'
        except Exception as e:
            results['info'] = f"Error: {str(e)[:50]}"
        
        # Get history
        try:
            hist = ticker.history(period='1mo')
            results['history'] = f"{len(hist)} rows" if not hist.empty else "No data"
        except Exception as e:
            results['history'] = f"Error: {str(e)[:50]}"
        
        # Get dividends
        try:
            dividends = ticker.dividends
            results['dividends'] = f"{len(dividends)} records" if not dividends.empty else "No dividends"
        except Exception as e:
            results['dividends'] = f"Error: {str(e)[:50]}"
        
        end_time = time.time()
        duration = end_time - start_time
        
        success = not any('Error:' in str(v) for v in results.values())
        
        print(f"{test_id}✅ {symbol}: {duration:.2f}s - {results}")
        return {
            'symbol': symbol,
            'duration': duration,
            'success': success,
            'results': results
        }
        
    except Exception as e:
        duration = time.time() - start_time
        print(f"{test_id}❌ {symbol}: {duration:.2f}s - FAILED: {str(e)[:100]}")
        return {
            'symbol': symbol,
            'duration': duration,
            'success': False,
            'error': str(e)
        }

def sequential_test():
    """Test tickers sequentially with built-in rate limiting"""
    print("🔄 Starting Sequential Test (30 tickers)")
    print("=" * 50)
    
    # Clear cache for clean test
    yf.clear_cache()
    
    start_time = time.time()
    results = []
    
    for i, symbol in enumerate(TEST_TICKERS, 1):
        test_id = f"[{i:2d}/30] "
        result = test_single_ticker(symbol, test_id)
        results.append(result)
        
        # Small delay to be extra safe
        time.sleep(0.1)
    
    end_time = time.time()
    total_duration = end_time - start_time
    
    # Analyze results
    successful = sum(1 for r in results if r['success'])
    failed = len(results) - successful
    avg_duration = sum(r['duration'] for r in results) / len(results)
    
    print("\n📊 Sequential Test Results:")
    print(f"Total time: {total_duration:.2f}s")
    print(f"Average per ticker: {avg_duration:.2f}s")
    print(f"Successful: {successful}/{len(results)} ({successful/len(results)*100:.1f}%)")
    print(f"Failed: {failed}")
    
    if failed > 0:
        print("\n❌ Failed tickers:")
        for r in results:
            if not r['success']:
                print(f"  {r['symbol']}: {r.get('error', 'Unknown error')}")
    
    return results

def parallel_test(max_workers=5):
    """Test tickers in parallel with thread pool"""
    print(f"\n🚀 Starting Parallel Test ({max_workers} workers, 30 tickers)")
    print("=" * 50)
    
    # Clear cache for clean test
    yf.clear_cache()
    
    start_time = time.time()
    results = []
    
    with ThreadPoolExecutor(max_workers=max_workers) as executor:
        # Submit all tasks
        future_to_symbol = {
            executor.submit(test_single_ticker, symbol, f"[T{i+1:2d}] "): symbol 
            for i, symbol in enumerate(TEST_TICKERS)
        }
        
        # Collect results as they complete
        for future in as_completed(future_to_symbol):
            try:
                result = future.result()
                results.append(result)
            except Exception as e:
                symbol = future_to_symbol[future]
                print(f"❌ {symbol}: Thread exception: {e}")
                results.append({
                    'symbol': symbol,
                    'duration': 0,
                    'success': False,
                    'error': f"Thread exception: {e}"
                })
    
    end_time = time.time()
    total_duration = end_time - start_time
    
    # Analyze results
    successful = sum(1 for r in results if r['success'])
    failed = len(results) - successful
    avg_duration = sum(r['duration'] for r in results) / len(results)
    
    print(f"\n📊 Parallel Test Results ({max_workers} workers):")
    print(f"Total time: {total_duration:.2f}s")
    print(f"Average per ticker: {avg_duration:.2f}s")
    print(f"Successful: {successful}/{len(results)} ({successful/len(results)*100:.1f}%)")
    print(f"Failed: {failed}")
    print(f"Speedup vs sequential: {len(results) * avg_duration / total_duration:.1f}x")
    
    if failed > 0:
        print("\n❌ Failed tickers:")
        for r in results:
            if not r['success']:
                print(f"  {r['symbol']}: {r.get('error', 'Unknown error')}")
    
    return results

def cache_efficiency_test():
    """Test cache efficiency with repeated requests"""
    print(f"\n💾 Cache Efficiency Test")
    print("=" * 50)
    
    # Clear cache
    yf.clear_cache()
    
    # Test with subset of tickers
    test_symbols = TEST_TICKERS[:10]
    
    print("First pass (cold cache):")
    start_time = time.time()
    for symbol in test_symbols:
        ticker = yf.Ticker(symbol)
        hist = ticker.history(period='1mo')
        print(f"  {symbol}: {len(hist)} rows")
    cold_time = time.time() - start_time
    
    print("\nSecond pass (warm cache):")
    start_time = time.time()
    for symbol in test_symbols:
        ticker = yf.Ticker(symbol)
        hist = ticker.history(period='1mo')
        print(f"  {symbol}: {len(hist)} rows")
    warm_time = time.time() - start_time
    
    # Check cache info
    cache_info = yf.get_cache_info()
    
    print(f"\n📊 Cache Performance:")
    print(f"Cold cache time: {cold_time:.2f}s")
    print(f"Warm cache time: {warm_time:.2f}s")
    print(f"Speedup: {cold_time/warm_time:.1f}x")
    print(f"Cache files: {cache_info['file_count']}")
    print(f"Cache size: {cache_info['total_size_mb']:.2f} MB")

def main():
    """Run all stress tests"""
    print("🧪 YFinance Plus Stress Test")
    print("Testing rate limiting, error handling, and performance")
    print("=" * 60)
    
    try:
        # Test 1: Sequential requests
        sequential_results = sequential_test()
        
        # Test 2: Parallel requests (conservative)
        parallel_results_5 = parallel_test(max_workers=5)
        
        # Test 3: Parallel requests (aggressive)
        parallel_results_10 = parallel_test(max_workers=10)
        
        # Test 4: Cache efficiency
        cache_efficiency_test()
        
        print(f"\n🎯 Overall Summary:")
        print("=" * 30)
        seq_success = sum(1 for r in sequential_results if r['success'])
        par5_success = sum(1 for r in parallel_results_5 if r['success'])
        par10_success = sum(1 for r in parallel_results_10 if r['success'])
        
        print(f"Sequential:    {seq_success}/30 ({seq_success/30*100:.1f}%)")
        print(f"Parallel (5):  {par5_success}/30 ({par5_success/30*100:.1f}%)")
        print(f"Parallel (10): {par10_success}/30 ({par10_success/30*100:.1f}%)")
        
        if seq_success >= 28:  # Allow for 2 failures due to network/data issues
            print("✅ Rate limiting appears to be working well!")
        else:
            print("⚠️  High failure rate - may need to adjust rate limiting")
            
    except KeyboardInterrupt:
        print("\n⚠️  Test interrupted by user")
    except Exception as e:
        print(f"\n❌ Test failed with exception: {e}")
        import traceback
        traceback.print_exc()

if __name__ == "__main__":
    main()