TCPDashboard/examples/collector_demo.py

"""
Demonstration of the enhanced data collector system with health monitoring and auto-restart.

This example shows how to:
1. Create data collectors with health monitoring
2. Use the collector manager for coordinated management
3. Monitor collector health and handle failures
4. Enable/disable collectors dynamically
"""

import asyncio
from datetime import datetime, timezone
from typing import Any, Optional

from data import (
    BaseDataCollector, DataType, CollectorStatus, MarketDataPoint,
    CollectorManager, CollectorConfig
)


class DemoDataCollector(BaseDataCollector):
    """
    Demo implementation of a data collector for demonstration purposes.
    
    This collector simulates receiving market data and can be configured
    to fail periodically to demonstrate auto-restart functionality.
    """
    
    def __init__(self, 
                 exchange_name: str, 
                 symbols: list,
                 fail_every_n_messages: int = 0,
                 connection_delay: float = 0.1):
        """
        Initialize demo collector.
        
        Args:
            exchange_name: Name of the exchange
            symbols: Trading symbols to collect
            fail_every_n_messages: Simulate failure every N messages (0 = no failures)
            connection_delay: Simulated connection delay
        """
        super().__init__(exchange_name, symbols, [DataType.TICKER])
        self.fail_every_n_messages = fail_every_n_messages
        self.connection_delay = connection_delay
        self.message_count = 0
        self.connected = False
        self.subscribed = False
    
    async def connect(self) -> bool:
        """Simulate connection to exchange."""
        print(f"[{self.exchange_name}] Connecting...")
        await asyncio.sleep(self.connection_delay)
        self.connected = True
        print(f"[{self.exchange_name}] Connected successfully")
        return True
    
    async def disconnect(self) -> None:
        """Simulate disconnection from exchange."""
        print(f"[{self.exchange_name}] Disconnecting...")
        await asyncio.sleep(self.connection_delay / 2)
        self.connected = False
        self.subscribed = False
        print(f"[{self.exchange_name}] Disconnected")
    
    async def subscribe_to_data(self, symbols: list, data_types: list) -> bool:
        """Simulate subscription to data streams."""
        if not self.connected:
            return False
        
        print(f"[{self.exchange_name}] Subscribing to {len(symbols)} symbols: {', '.join(symbols)}")
        await asyncio.sleep(0.05)
        self.subscribed = True
        return True
    
    async def unsubscribe_from_data(self, symbols: list, data_types: list) -> bool:
        """Simulate unsubscription from data streams."""
        print(f"[{self.exchange_name}] Unsubscribing from data streams")
        self.subscribed = False
        return True
    
    async def _process_message(self, message: Any) -> Optional[MarketDataPoint]:
        """Process simulated market data message."""
        self.message_count += 1
        
        # Simulate periodic failures if configured
        if (self.fail_every_n_messages > 0 and 
            self.message_count % self.fail_every_n_messages == 0):
            raise Exception(f"Simulated failure after {self.message_count} messages")
        
        # Create mock market data
        data_point = MarketDataPoint(
            exchange=self.exchange_name,
            symbol=message['symbol'],
            timestamp=datetime.now(timezone.utc),
            data_type=DataType.TICKER,
            data={
                'price': message['price'],
                'volume': message.get('volume', 100),
                'timestamp': datetime.now(timezone.utc).isoformat()
            }
        )
        
        return data_point
    
    async def _handle_messages(self) -> None:
        """Simulate receiving and processing messages."""
        if not self.connected or not self.subscribed:
            await asyncio.sleep(0.1)
            return
        
        # Simulate receiving data for each symbol
        for symbol in self.symbols:
            try:
                # Create simulated message
                simulated_message = {
                    'symbol': symbol,
                    'price': 50000 + (self.message_count % 1000),  # Fake price that changes
                    'volume': 1.5
                }
                
                # Process the message
                data_point = await self._process_message(simulated_message)
                if data_point:
                    self._stats['messages_processed'] += 1
                    await self._notify_callbacks(data_point)
                
            except Exception as e:
                # This will trigger reconnection logic
                raise e
        
        # Simulate processing delay
        await asyncio.sleep(1.0)


async def data_callback(data_point: MarketDataPoint):
    """Callback function to handle received data."""
    print(f"📊 Data received: {data_point.exchange} - {data_point.symbol} - "
          f"Price: {data_point.data.get('price')} at {data_point.timestamp.strftime('%H:%M:%S')}")


async def monitor_collectors(manager: CollectorManager, duration: int = 30):
    """Monitor collector status and print updates."""
    print(f"\n🔍 Starting monitoring for {duration} seconds...")
    
    for i in range(duration):
        await asyncio.sleep(1)
        
        status = manager.get_status()
        running = len(manager.get_running_collectors())
        failed = len(manager.get_failed_collectors())
        
        if i % 5 == 0:  # Print status every 5 seconds
            print(f"⏰ Status at {i+1}s: {running} running, {failed} failed, "
                  f"{status['statistics']['restarts_performed']} restarts")
    
    print("🏁 Monitoring complete")


async def demo_basic_usage():
    """Demonstrate basic collector usage."""
    print("=" * 60)
    print("🚀 Demo 1: Basic Data Collector Usage")
    print("=" * 60)
    
    # Create a stable collector
    collector = DemoDataCollector("demo_exchange", ["BTC-USDT", "ETH-USDT"])
    
    # Add data callback
    collector.add_data_callback(DataType.TICKER, data_callback)
    
    # Start the collector
    print("Starting collector...")
    success = await collector.start()
    if success:
        print("✅ Collector started successfully")
        
        # Let it run for a few seconds
        await asyncio.sleep(5)
        
        # Show status
        status = collector.get_status()
        print(f"📈 Messages processed: {status['statistics']['messages_processed']}")
        print(f"⏱️  Uptime: {status['statistics']['uptime_seconds']:.1f}s")
        
        # Stop the collector
        await collector.stop()
        print("✅ Collector stopped")
    else:
        print("❌ Failed to start collector")


async def demo_manager_usage():
    """Demonstrate collector manager usage."""
    print("\n" + "=" * 60)
    print("🎛️  Demo 2: Collector Manager Usage")
    print("=" * 60)
    
    # Create manager
    manager = CollectorManager("demo_manager", global_health_check_interval=3.0)
    
    # Create multiple collectors
    stable_collector = DemoDataCollector("stable_exchange", ["BTC-USDT"])
    failing_collector = DemoDataCollector("failing_exchange", ["ETH-USDT"], 
                                        fail_every_n_messages=5)  # Fails every 5 messages
    
    # Add data callbacks
    stable_collector.add_data_callback(DataType.TICKER, data_callback)
    failing_collector.add_data_callback(DataType.TICKER, data_callback)
    
    # Add collectors to manager
    manager.add_collector(stable_collector)
    manager.add_collector(failing_collector)
    
    print(f"📝 Added {len(manager.list_collectors())} collectors to manager")
    
    # Start manager
    success = await manager.start()
    if success:
        print("✅ Manager started successfully")
        
        # Monitor for a while
        await monitor_collectors(manager, duration=15)
        
        # Show final status
        status = manager.get_status()
        print(f"\n📊 Final Statistics:")
        print(f"   - Total restarts: {status['statistics']['restarts_performed']}")
        print(f"   - Running collectors: {len(manager.get_running_collectors())}")
        print(f"   - Failed collectors: {len(manager.get_failed_collectors())}")
        
        # Stop manager
        await manager.stop()
        print("✅ Manager stopped")
    else:
        print("❌ Failed to start manager")


async def demo_dynamic_management():
    """Demonstrate dynamic collector management."""
    print("\n" + "=" * 60)
    print("🔄 Demo 3: Dynamic Collector Management")
    print("=" * 60)
    
    # Create manager
    manager = CollectorManager("dynamic_manager", global_health_check_interval=2.0)
    
    # Start with one collector
    collector1 = DemoDataCollector("exchange_1", ["BTC-USDT"])
    collector1.add_data_callback(DataType.TICKER, data_callback)
    
    manager.add_collector(collector1)
    await manager.start()
    
    print("✅ Started with 1 collector")
    await asyncio.sleep(3)
    
    # Add second collector
    collector2 = DemoDataCollector("exchange_2", ["ETH-USDT"])
    collector2.add_data_callback(DataType.TICKER, data_callback)
    manager.add_collector(collector2)
    
    print("➕ Added second collector")
    await asyncio.sleep(3)
    
    # Disable first collector
    collector_names = manager.list_collectors()
    manager.disable_collector(collector_names[0])
    
    print("⏸️  Disabled first collector")
    await asyncio.sleep(3)
    
    # Re-enable first collector
    manager.enable_collector(collector_names[0])
    
    print("▶️  Re-enabled first collector")
    await asyncio.sleep(3)
    
    # Show final status
    status = manager.get_status()
    print(f"📊 Final state: {len(manager.get_running_collectors())} running collectors")
    
    await manager.stop()
    print("✅ Dynamic demo complete")


async def main():
    """Run all demonstrations."""
    print("🎯 Data Collector System Demonstration")
    print("This demo shows health monitoring and auto-restart capabilities\n")
    
    try:
        # Run demonstrations
        await demo_basic_usage()
        await demo_manager_usage()
        await demo_dynamic_management()
        
        print("\n" + "=" * 60)
        print("🎉 All demonstrations completed successfully!")
        print("=" * 60)
        
    except Exception as e:
        print(f"❌ Demo failed with error: {e}")
        import traceback
        traceback.print_exc()


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