"""
ML Signal Consumer - Consumes trading signals from ML Engine
"""
from typing import Dict, Callable, Optional, List
import asyncio
import logging
import aiohttp
from datetime import datetime

logger = logging.getLogger(__name__)


class MLSignalConsumer:
    """
    Consumes trading signals from ML Engine via HTTP polling or WebSocket

    This service polls the ML Engine for new signals and dispatches them
    to registered callbacks (trading agents).
    """

    def __init__(
        self,
        ml_engine_url: str,
        api_key: str,
        poll_interval: int = 5
    ):
        self.ml_engine_url = ml_engine_url.rstrip('/')
        self.api_key = api_key
        self.poll_interval = poll_interval
        self.callbacks: Dict[str, List[Callable]] = {}  # symbol -> [callbacks]
        self.is_running = False
        self.session: Optional[aiohttp.ClientSession] = None
        self.logger = logging.getLogger(__name__)

    async def start(self):
        """Start consuming signals"""
        self.is_running = True
        self.session = aiohttp.ClientSession(
            headers={'X-API-Key': self.api_key}
        )
        self.logger.info("ML Signal Consumer started")

        # Start polling task
        asyncio.create_task(self._poll_signals())

    async def stop(self):
        """Stop consuming signals"""
        self.is_running = False
        if self.session:
            await self.session.close()
        self.logger.info("ML Signal Consumer stopped")

    async def subscribe(self, symbol: str, callback: Callable):
        """
        Subscribe to signals for a specific symbol

        Args:
            symbol: Trading symbol (e.g., 'BTCUSDT')
            callback: Async function to call when signal received
        """
        if symbol not in self.callbacks:
            self.callbacks[symbol] = []

        self.callbacks[symbol].append(callback)
        self.logger.info(f"Subscribed to signals for {symbol}")

    async def unsubscribe(self, symbol: str, callback: Callable):
        """
        Unsubscribe from signals for a symbol

        Args:
            symbol: Trading symbol
            callback: Callback function to remove
        """
        if symbol in self.callbacks and callback in self.callbacks[symbol]:
            self.callbacks[symbol].remove(callback)
            self.logger.info(f"Unsubscribed from signals for {symbol}")

    async def get_signal(self, symbol: str) -> Optional[Dict]:
        """
        Get latest signal for a symbol from ML Engine

        Args:
            symbol: Trading symbol

        Returns:
            Signal dict or None
        """
        try:
            url = f"{self.ml_engine_url}/api/v1/signals/latest/{symbol}"

            async with self.session.get(url) as response:
                if response.status == 200:
                    data = await response.json()
                    return data.get('signal')
                elif response.status == 404:
                    # No signal available
                    return None
                else:
                    error_data = await response.text()
                    self.logger.error(f"Error fetching signal for {symbol}: {error_data}")
                    return None

        except Exception as e:
            self.logger.error(f"Failed to get signal for {symbol}: {str(e)}")
            return None

    async def get_signals_batch(self, symbols: List[str]) -> Dict[str, Optional[Dict]]:
        """
        Get signals for multiple symbols in one request

        Args:
            symbols: List of trading symbols

        Returns:
            Dict mapping symbol to signal
        """
        try:
            url = f"{self.ml_engine_url}/api/v1/signals/batch"
            payload = {'symbols': symbols}

            async with self.session.post(url, json=payload) as response:
                if response.status == 200:
                    data = await response.json()
                    return data.get('signals', {})
                else:
                    self.logger.error(f"Batch signals request failed: {response.status}")
                    return {}

        except Exception as e:
            self.logger.error(f"Failed to get batch signals: {str(e)}")
            return {}

    async def _poll_signals(self):
        """
        Poll ML Engine for new signals

        This runs continuously while is_running is True
        """
        while self.is_running:
            try:
                # Get all subscribed symbols
                symbols = list(self.callbacks.keys())

                if not symbols:
                    await asyncio.sleep(self.poll_interval)
                    continue

                # Fetch signals for all symbols
                signals = await self.get_signals_batch(symbols)

                # Dispatch signals to callbacks
                for symbol, signal in signals.items():
                    if signal and symbol in self.callbacks:
                        await self._dispatch_signal(symbol, signal)

                # Wait before next poll
                await asyncio.sleep(self.poll_interval)

            except Exception as e:
                self.logger.error(f"Error in signal polling: {str(e)}")
                await asyncio.sleep(self.poll_interval)

    async def _dispatch_signal(self, symbol: str, signal: Dict):
        """
        Dispatch signal to all registered callbacks

        Args:
            symbol: Trading symbol
            signal: Signal data
        """
        callbacks = self.callbacks.get(symbol, [])

        for callback in callbacks:
            try:
                # Call callback asynchronously
                if asyncio.iscoroutinefunction(callback):
                    await callback(signal)
                else:
                    callback(signal)

            except Exception as e:
                self.logger.error(
                    f"Error dispatching signal to callback for {symbol}: {str(e)}"
                )

    async def get_account_signals(self, account_id: str) -> List[Dict]:
        """
        Get all signals for an account from ML Engine

        Args:
            account_id: Investment account ID

        Returns:
            List of signals
        """
        try:
            url = f"{self.ml_engine_url}/api/v1/signals/account/{account_id}"

            async with self.session.get(url) as response:
                if response.status == 200:
                    data = await response.json()
                    return data.get('signals', [])
                else:
                    self.logger.error(f"Failed to get account signals: {response.status}")
                    return []

        except Exception as e:
            self.logger.error(f"Error getting account signals: {str(e)}")
            return []

    async def health_check(self) -> bool:
        """
        Check if ML Engine is accessible

        Returns:
            True if healthy, False otherwise
        """
        try:
            url = f"{self.ml_engine_url}/health"

            async with self.session.get(url) as response:
                return response.status == 200

        except Exception as e:
            self.logger.error(f"ML Engine health check failed: {str(e)}")
            return False
