TCPDashboard/data/base_collector.py

"""
Abstract base class for data collectors.

This module provides a common interface for all data collection implementations,
ensuring consistency across different exchange connectors and data sources.
"""

import asyncio
from abc import ABC, abstractmethod
from datetime import datetime, timezone, timedelta
from decimal import Decimal
from typing import Dict, List, Optional, Any, Callable, Set
from dataclasses import dataclass
from enum import Enum

from utils.logger import get_logger
from .collector.collector_state_telemetry import CollectorStatus, CollectorStateAndTelemetry
from .collector.collector_connection_manager import ConnectionManager
from .collector.collector_callback_dispatcher import CallbackDispatcher
from .common.data_types import DataType, MarketDataPoint
from .common.ohlcv_data import OHLCVData, DataValidationError, validate_ohlcv_data


class DataCollectorError(Exception):
    """Base exception for data collector errors."""
    pass


class ConnectionError(DataCollectorError):
    """Exception raised when connection to data source fails."""
    pass


class BaseDataCollector(ABC):
    """
    Abstract base class for all data collectors.
    
    This class defines the interface that all data collection implementations
    must follow, providing consistency across different exchanges and data sources.
    """
    
    def __init__(self, 
                 exchange_name: str,
                 symbols: List[str],
                 data_types: Optional[List[DataType]] = None,
                 timeframes: Optional[List[str]] = None,
                 component_name: Optional[str] = None,
                 auto_restart: bool = True,
                 health_check_interval: float = 30.0,
                 
                 logger = None,
                 log_errors_only: bool = False):
        """
        Initialize the base data collector.
        
        Args:
            exchange_name: Name of the exchange (e.g., 'okx', 'binance')
            symbols: List of trading symbols to collect data for
            data_types: Types of data to collect (default: [DataType.CANDLE])
            timeframes: List of timeframes to collect (e.g., ['1s', '1m', '5m'])
            component_name: Name for logging (default: based on exchange_name)
            auto_restart: Enable automatic restart on failures (default: True)
            health_check_interval: Seconds between health checks (default: 30.0)
            logger: Logger instance. If None, no logging will be performed.
            log_errors_only: If True and logger is provided, only log error-level messages
        """
        self.exchange_name = exchange_name.lower()
        self.symbols = set(symbols)
        self.data_types = data_types or [DataType.CANDLE]
        self.timeframes = timeframes or ['1m', '5m']  # Default timeframes if none provided
        self.auto_restart = auto_restart
        
        # Initialize logger based on parameters
        if logger is not None:
            self.logger = logger
        else:
            self.logger = get_logger(self.exchange_name) # Ensure a logger is always available
        
        # Initialize state and telemetry manager
        component = component_name or f"{self.exchange_name}_collector"
        self._state_telemetry = CollectorStateAndTelemetry(
            exchange_name=self.exchange_name,
            component_name=component,
            health_check_interval=health_check_interval,
            logger=self.logger, # Pass the actual logger instance
            log_errors_only=log_errors_only
        )
        self.component_name = component # Keep for external access
        
        # Initialize connection manager
        self._connection_manager = ConnectionManager(
            exchange_name=self.exchange_name,
            component_name=component,
            max_reconnect_attempts=5, # Default, can be made configurable later
            reconnect_delay=5.0, # Default, can be made configurable later
            logger=self.logger,
            state_telemetry=self._state_telemetry
        )
        
        # Initialize callback dispatcher
        self._callback_dispatcher = CallbackDispatcher(
            component_name=component,
            logger=self.logger
        )
        
        # Collector state (now managed by _state_telemetry)
        self._tasks: Set[asyncio.Task] = set()
        
        # Log initialization if logger is available
        if self._state_telemetry.logger:
            if not self._state_telemetry.log_errors_only:
                self._state_telemetry._log_info(f"{self.component_name}: Initialized {self.exchange_name} data collector for symbols: {', '.join(symbols)}")
                self._state_telemetry._log_info(f"{self.component_name}: Using timeframes: {', '.join(self.timeframes)}")
    
    @property
    def status(self) -> CollectorStatus:
        return self._state_telemetry.status
        
    def _log_debug(self, message: str) -> None:
        self._state_telemetry._log_debug(message)
    
    def _log_info(self, message: str) -> None:
        self._state_telemetry._log_info(message)
    
    def _log_warning(self, message: str) -> None:
        self._state_telemetry._log_warning(message)
    
    def _log_error(self, message: str, exc_info: bool = False) -> None:
        self._state_telemetry._log_error(message, exc_info=exc_info)
    
    def _log_critical(self, message: str, exc_info: bool = False) -> None:
        self._state_telemetry._log_critical(message, exc_info=exc_info)
    
    @abstractmethod
    async def connect(self) -> bool:
        """
        Establish connection to the data source.
        
        Returns:
            True if connection successful, False otherwise
        """
        pass
    
    @abstractmethod
    async def disconnect(self) -> None:
        """Disconnect from the data source."""
        pass
    
    @abstractmethod
    async def subscribe_to_data(self, symbols: List[str], data_types: List[DataType]) -> bool:
        """
        Subscribe to data streams for specified symbols and data types.
        
        Args:
            symbols: Trading symbols to subscribe to
            data_types: Types of data to subscribe to
            
        Returns:
            True if subscription successful, False otherwise
        """
        pass
    
    @abstractmethod
    async def unsubscribe_from_data(self, symbols: List[str], data_types: List[DataType]) -> bool:
        """
        Unsubscribe from data streams.
        
        Args:
            symbols: Trading symbols to unsubscribe from
            data_types: Types of data to unsubscribe from
            
        Returns:
            True if unsubscription successful, False otherwise
        """
        pass
    
    @abstractmethod
    async def _process_message(self, message: Any) -> Optional[MarketDataPoint]:
        """
        Process incoming message from the data source.
        
        Args:
            message: Raw message from the data source
            
        Returns:
            Processed MarketDataPoint or None if message should be ignored
        """
        pass
    
    async def start(self) -> bool:
        """
        Start the data collector.
        
        Returns:
            True if started successfully, False otherwise
        """
        # Check if already running or starting
        if self._state_telemetry.status in [CollectorStatus.RUNNING, CollectorStatus.STARTING]:
            self._log_warning("Data collector is already running or starting")
            return True
        
        self._log_info(f"Starting {self.exchange_name} data collector")
        self._state_telemetry.update_status(CollectorStatus.STARTING)
        self._state_telemetry.set_should_be_running(True)
        
        try:
            # Connect to data source
            if not await self._connection_manager.connect(self._actual_connect):
                self._log_error("Failed to connect to data source")
                self._state_telemetry.update_status(CollectorStatus.ERROR)
                return False
            
            # Subscribe to data streams
            if not await self.subscribe_to_data(list(self.symbols), self.data_types):
                self._log_error("Failed to subscribe to data streams")
                self._state_telemetry.update_status(CollectorStatus.ERROR)
                await self._connection_manager.disconnect(self._actual_disconnect)
                return False
            
            # Start background tasks
            self._state_telemetry.set_running_state(True)
            self._state_telemetry.update_status(CollectorStatus.RUNNING)
            
            # Start message processing task
            message_task = asyncio.create_task(self._message_loop())
            self._tasks.add(message_task)
            message_task.add_done_callback(self._tasks.discard)
            
            # Start health monitoring task  
            if self._state_telemetry.health_check_interval > 0:
                health_task = asyncio.create_task(self._health_monitor())
                self._tasks.add(health_task)
                health_task.add_done_callback(self._tasks.discard)
            
            self._log_info(f"{self.exchange_name} data collector started successfully")
            return True
            
        except Exception as e:
            self._log_error(f"Failed to start data collector: {e}")
            self._state_telemetry.update_status(CollectorStatus.ERROR)
            self._state_telemetry.set_should_be_running(False)
            return False
    
    async def stop(self, force: bool = False) -> None:
        """
        Stop the data collector and cleanup resources.
        
        Args:
            force: Force stop even if not graceful
        """
        if self._state_telemetry.status == CollectorStatus.STOPPED:
            self._log_warning("Data collector is already stopped")
            return
        
        self._log_info(f"Stopping {self.exchange_name} data collector")
        self._state_telemetry.update_status(CollectorStatus.STOPPING)
        self._state_telemetry.set_should_be_running(False)
        
        try:
            # Stop background tasks
            self._state_telemetry.set_running_state(False)
            
            # Cancel all tasks
            for task in list(self._tasks):
                if not task.done():
                    task.cancel()
                    if not force:
                        try:
                            await task
                        except asyncio.CancelledError:
                            pass
            
            self._tasks.clear()
            
            # Unsubscribe and disconnect
            await self.unsubscribe_from_data(list(self.symbols), self.data_types)
            await self._connection_manager.disconnect(self._actual_disconnect)
            
            self._state_telemetry.update_status(CollectorStatus.STOPPED)
            self._log_info(f"{self.exchange_name} data collector stopped")
            
        except Exception as e:
            self._log_error(f"Error stopping data collector: {e}")
            self._state_telemetry.update_status(CollectorStatus.ERROR)
    
    async def restart(self) -> bool:
        """
        Restart the data collector.
        
        Returns:
            True if restarted successfully, False otherwise
        """
        self._log_info(f"Restarting {self.exchange_name} data collector")
        self._state_telemetry.increment_restarts()
        
        # Stop first
        await self.stop()
        
        # Wait a bit before restarting
        await asyncio.sleep(self._connection_manager._reconnect_delay)
        
        # Start again
        return await self.start()
    
    async def _message_loop(self) -> None:
        """Main message processing loop."""
        try:
            self._log_debug("Starting message processing loop")
            
            while self._state_telemetry._running:
                try:
                    await self._handle_messages()
                except asyncio.CancelledError:
                    break
                except Exception as e:
                    self._state_telemetry.increment_errors(str(e))
                    self._log_error(f"Error processing messages: {e}")
                    
                    # Small delay to prevent tight error loops
                    await asyncio.sleep(0.1)
                    
        except asyncio.CancelledError:
            self._log_debug("Message loop cancelled")
            raise
        except Exception as e:
            self._log_error(f"Error in message loop: {e}")
            self._state_telemetry.update_status(CollectorStatus.ERROR)
    
    async def _health_monitor(self) -> None:
        """Monitor collector health and restart if needed."""
        try:
            self._log_debug("Starting health monitor")
            
            while self._state_telemetry._running:
                try:
                    await asyncio.sleep(self._state_telemetry.health_check_interval)
                    
                    # Check if collector should be running but isn't
                    if self._state_telemetry._should_be_running and self._state_telemetry.status != CollectorStatus.RUNNING:
                        self._log_warning("Collector should be running but isn't - restarting")
                        if self.auto_restart:
                            asyncio.create_task(self.restart())
                        continue
                    
                    # Check heartbeat
                    time_since_heartbeat = datetime.now(timezone.utc) - self._state_telemetry._last_heartbeat
                    if time_since_heartbeat > timedelta(seconds=self._state_telemetry.health_check_interval * 2):
                        self._log_warning(f"No heartbeat for {time_since_heartbeat.total_seconds():.1f}s - restarting")
                        if self.auto_restart:
                            asyncio.create_task(self.restart())
                        continue
                    
                    # Check data reception
                    if self._state_telemetry._last_data_received:
                        time_since_data = datetime.now(timezone.utc) - self._state_telemetry._last_data_received
                        if time_since_data > self._state_telemetry._max_silence_duration:
                            self._log_warning(f"No data received for {time_since_data.total_seconds():.1f}s - restarting")
                            if self.auto_restart:
                                asyncio.create_task(self.restart())
                            continue
                    
                    # Check for error status
                    if self._state_telemetry.status == CollectorStatus.ERROR:
                        self._log_warning(f"Collector in {self._state_telemetry.status.value} status - restarting")
                        if self.auto_restart:
                            asyncio.create_task(self.restart())
                            
                except asyncio.CancelledError:
                    break
                    
        except asyncio.CancelledError:
            self._log_debug("Health monitor cancelled")
            raise
        except Exception as e:
            self._log_error(f"Error in health monitor: {e}")
    
    @abstractmethod
    async def _handle_messages(self) -> None:
        """
        Handle incoming messages from the data source.
        This method should be implemented by subclasses to handle their specific message format.
        """
        pass
    
    async def _handle_connection_error(self) -> bool:
        """
        Handle connection errors and attempt reconnection.
        
        Returns:
            True if reconnection successful, False if max attempts exceeded
        """
        return await self._connection_manager.handle_connection_error(
            connect_logic=self._actual_connect,
            subscribe_logic=self.subscribe_to_data,
            symbols=list(self.symbols),
            data_types=self.data_types
        )
    
    @abstractmethod
    async def _actual_connect(self) -> bool:
        """
        Abstract method for subclasses to implement actual connection logic.
        """
        pass
    
    @abstractmethod
    async def _actual_disconnect(self) -> None:
        """
        Abstract method for subclasses to implement actual disconnection logic.
        """
        pass
    
    def add_data_callback(self, data_type: DataType, callback: Callable[[MarketDataPoint], None]) -> None:
        """
        Add a callback function for specific data type.
        
        Args:
            data_type: Type of data to monitor
            callback: Function to call when data is received
        """
        self._callback_dispatcher.add_data_callback(data_type, callback)
    
    def remove_data_callback(self, data_type: DataType, callback: Callable[[MarketDataPoint], None]) -> None:
        """
        Remove a callback function for specific data type.
        
        Args:
            data_type: Type of data to stop monitoring
            callback: Function to remove
        """
        self._callback_dispatcher.remove_data_callback(data_type, callback)
    
    async def _notify_callbacks(self, data_point: MarketDataPoint) -> None:
        """
        Notify all registered callbacks for a data point.
        
        Args:
            data_point: Market data to distribute
        """
        await self._callback_dispatcher.notify_callbacks(data_point)
        
        # Update statistics
        self._state_telemetry.increment_messages_processed()
        self._state_telemetry._stats['last_message_time'] = data_point.timestamp # Direct update for now, will refactor
        self._state_telemetry.update_data_received_timestamp()
        self._state_telemetry.update_heartbeat()
    
    def get_status(self) -> Dict[str, Any]:
        """
        Get current collector status and statistics.
        
        Returns:
            Dictionary containing status information
        """
        status = self._state_telemetry.get_status()
        
        # Add BaseDataCollector specific information
        status.update({
            'symbols': list(self.symbols),
            'data_types': [dt.value for dt in self.data_types],
            'timeframes': self.timeframes,
            'auto_restart': self.auto_restart
        })
        
        return status
    
    def get_health_status(self) -> Dict[str, Any]:
        """
        Get detailed health status for monitoring.
        
        Returns:
            Dictionary containing health information
        """
        return self._state_telemetry.get_health_status()
    
    def add_symbol(self, symbol: str) -> None:
        """
        Add a new symbol to collect data for.
        
        Args:
            symbol: Trading symbol to add
        """
        if symbol not in self.symbols:
            self.symbols.add(symbol)
            self._log_info(f"Added symbol: {symbol}")
            
            # If collector is running, subscribe to new symbol
            if self._state_telemetry.status == CollectorStatus.RUNNING:
                # Note: This needs to be called from an async context
                # Users should handle this appropriately
                pass
    
    def remove_symbol(self, symbol: str) -> None:
        """
        Remove a symbol from data collection.
        
        Args:
            symbol: Trading symbol to remove
        """
        if symbol in self.symbols:
            self.symbols.remove(symbol)
            self._log_info(f"Removed symbol: {symbol}")
            
            # If collector is running, unsubscribe from symbol
            if self._state_telemetry.status == CollectorStatus.RUNNING:
                # Note: This needs to be called from an async context
                # Users should handle this appropriately
                pass
    
    def validate_ohlcv_data(self, data: Dict[str, Any], symbol: str, timeframe: str) -> OHLCVData:
        """
        Validate and convert raw OHLCV data to standardized format.
        
        Args:
            data: Raw OHLCV data dictionary
            symbol: Trading symbol
            timeframe: Timeframe (e.g., '1m', '5m', '1h')
            
        Returns:
            Validated OHLCVData object
            
        Raises:
            DataValidationError: If data validation fails
        """
        return validate_ohlcv_data(data, symbol, timeframe)
    
    def __repr__(self) -> str:
        """String representation of the collector."""
        return f"<{self.__class__.__name__}({self.exchange_name}, {len(self.symbols)} symbols, {self.status.value})>"