"""
WebSocket client for exchange-data-manager real-time candle streaming.
"""

import asyncio
import json
import logging
import uuid
from typing import Optional, Dict, List, Callable, Any

import websockets
from websockets.exceptions import ConnectionClosed

from .models import Candle, Subscription, EdmConfig
from .exceptions import EdmConnectionError

logger = logging.getLogger(__name__)


class EdmWebSocketClient:
    """
    WebSocket client for real-time candle streaming from EDM.

    Connects to EDM's WebSocket server and subscribes to candle streams.
    Handles automatic reconnection and subscription management.
    """

    def __init__(
        self,
        ws_url: str = "ws://localhost:8765",
        reconnect_interval: float = 5.0,
        max_reconnect_attempts: int = 10,
        config: Optional[EdmConfig] = None
    ):
        """
        Initialize WebSocket client.

        Args:
            ws_url: WebSocket URL for EDM
            reconnect_interval: Seconds between reconnection attempts
            max_reconnect_attempts: Maximum reconnection attempts (0 = infinite)
            config: Optional EdmConfig object
        """
        if config:
            self.ws_url = config.ws_url
            self.reconnect_interval = config.reconnect_interval
            self.max_reconnect_attempts = config.max_reconnect_attempts
        else:
            self.ws_url = ws_url
            self.reconnect_interval = reconnect_interval
            self.max_reconnect_attempts = max_reconnect_attempts

        self._ws: Optional[websockets.WebSocketClientProtocol] = None
        self._subscriptions: Dict[str, Subscription] = {}
        self._callbacks: Dict[str, List[Callable[[dict], None]]] = {}
        self._running = False
        self._reconnect_count = 0
        self._task: Optional[asyncio.Task] = None

    @property
    def connected(self) -> bool:
        """Check if WebSocket is connected."""
        return self._ws is not None and self._ws.open

    async def connect(self) -> None:
        """
        Establish WebSocket connection to EDM.

        Raises:
            EdmConnectionError: If connection fails
        """
        try:
            logger.info(f"Connecting to EDM WebSocket at {self.ws_url}")
            self._ws = await websockets.connect(self.ws_url)
            self._reconnect_count = 0
            logger.info("Connected to EDM WebSocket")

            # Resubscribe to any existing subscriptions after reconnect
            await self._resubscribe_all()

        except Exception as e:
            logger.error(f"Failed to connect to EDM WebSocket: {e}")
            raise EdmConnectionError(f"WebSocket connection failed: {e}")

    async def _resubscribe_all(self) -> None:
        """Resubscribe to all active subscriptions after reconnect."""
        for sub in self._subscriptions.values():
            if sub.active:
                await self._send_subscribe(sub.exchange, sub.symbol, sub.timeframe)
                logger.info(f"Resubscribed to {sub.to_key()}")

    async def _send_subscribe(self, exchange: str, symbol: str, timeframe: str) -> None:
        """Send subscribe message to EDM."""
        if not self._ws or not self._ws.open:
            raise EdmConnectionError("WebSocket not connected")

        message = {
            'action': 'subscribe',
            'exchange': exchange,
            'symbol': symbol,
            'timeframe': timeframe
        }
        await self._ws.send(json.dumps(message))
        logger.debug(f"Sent subscribe: {message}")

    async def subscribe(
        self,
        exchange: str,
        symbol: str,
        timeframe: str,
        callback: Callable[[dict], None]
    ) -> str:
        """
        Subscribe to a candle stream.

        Args:
            exchange: Exchange name (e.g., 'binance')
            symbol: Trading pair (e.g., 'BTC/USDT')
            timeframe: Candle timeframe (e.g., '1m')
            callback: Function to call with candle updates

        Returns:
            Subscription ID for unsubscribe

        Raises:
            EdmConnectionError: If not connected
        """
        sub_id = str(uuid.uuid4())
        sub = Subscription(
            id=sub_id,
            exchange=exchange,
            symbol=symbol,
            timeframe=timeframe,
            active=True
        )

        key = sub.to_key()
        self._subscriptions[sub_id] = sub

        # Register callback
        if key not in self._callbacks:
            self._callbacks[key] = []
        self._callbacks[key].append(callback)

        # Send subscribe if connected
        if self.connected:
            await self._send_subscribe(exchange, symbol, timeframe)

        logger.info(f"Subscribed to {key} with ID {sub_id}")
        return sub_id

    async def unsubscribe(self, subscription_id: str) -> None:
        """
        Unsubscribe from a candle stream.

        Args:
            subscription_id: ID returned from subscribe()
        """
        if subscription_id not in self._subscriptions:
            logger.warning(f"Subscription {subscription_id} not found")
            return

        sub = self._subscriptions[subscription_id]
        sub.active = False
        del self._subscriptions[subscription_id]

        # Send unsubscribe message
        if self.connected:
            message = {
                'action': 'unsubscribe',
                'exchange': sub.exchange,
                'symbol': sub.symbol,
                'timeframe': sub.timeframe
            }
            await self._ws.send(json.dumps(message))

        logger.info(f"Unsubscribed from {sub.to_key()}")

    async def unsubscribe_all(self) -> None:
        """Unsubscribe from all streams."""
        for sub_id in list(self._subscriptions.keys()):
            await self.unsubscribe(sub_id)
        self._callbacks.clear()

    def _dispatch_candle(self, data: dict) -> None:
        """Dispatch candle update to registered callbacks."""
        exchange = data.get('exchange', '')
        symbol = data.get('symbol', '')
        timeframe = data.get('timeframe', '')
        key = f"{exchange}:{symbol}:{timeframe}"

        callbacks = self._callbacks.get(key, [])
        for callback in callbacks:
            try:
                callback(data)
            except Exception as e:
                logger.error(f"Error in candle callback: {e}")

    async def _handle_message(self, message: str) -> None:
        """Handle incoming WebSocket message."""
        try:
            data = json.loads(message)
            msg_type = data.get('type', '')

            if msg_type == 'candle':
                candle_data = data.get('data', data)
                self._dispatch_candle(candle_data)
            elif msg_type == 'subscribed':
                logger.debug(f"Subscription confirmed: {data}")
            elif msg_type == 'unsubscribed':
                logger.debug(f"Unsubscription confirmed: {data}")
            elif msg_type == 'error':
                logger.error(f"EDM WebSocket error: {data.get('message', data)}")
            elif msg_type == 'pong':
                logger.debug("Received pong")
            else:
                logger.debug(f"Unknown message type: {msg_type}")

        except json.JSONDecodeError as e:
            logger.error(f"Failed to parse WebSocket message: {e}")

    async def run(self) -> None:
        """
        Main event loop - receive and dispatch messages.

        This method runs indefinitely, handling reconnection on disconnect.
        Use start() and stop() for managed lifecycle.
        """
        self._running = True

        while self._running:
            try:
                if not self.connected:
                    await self.connect()

                async for message in self._ws:
                    if not self._running:
                        break
                    await self._handle_message(message)

            except ConnectionClosed as e:
                logger.warning(f"WebSocket connection closed: {e}")
                self._ws = None

                if not self._running:
                    break

                # Attempt reconnection
                self._reconnect_count += 1
                if self.max_reconnect_attempts > 0 and self._reconnect_count > self.max_reconnect_attempts:
                    logger.error(f"Max reconnection attempts ({self.max_reconnect_attempts}) exceeded")
                    break

                wait_time = min(self.reconnect_interval * (2 ** (self._reconnect_count - 1)), 60)
                logger.info(f"Reconnecting in {wait_time}s (attempt {self._reconnect_count})")
                await asyncio.sleep(wait_time)

            except Exception as e:
                logger.error(f"WebSocket error: {e}")
                self._ws = None
                await asyncio.sleep(self.reconnect_interval)

    def start(self) -> asyncio.Task:
        """
        Start the WebSocket client in the background.

        Returns:
            asyncio.Task that can be awaited or cancelled
        """
        if self._task and not self._task.done():
            logger.warning("WebSocket client already running")
            return self._task

        self._task = asyncio.create_task(self.run())
        return self._task

    async def stop(self) -> None:
        """Stop the WebSocket client gracefully."""
        self._running = False

        if self._task:
            self._task.cancel()
            try:
                await self._task
            except asyncio.CancelledError:
                pass
            self._task = None

        await self.close()

    async def close(self) -> None:
        """Close WebSocket connection gracefully."""
        if self._ws and self._ws.open:
            await self._ws.close()
            logger.info("WebSocket connection closed")
        self._ws = None

    async def ping(self) -> bool:
        """
        Send ping to EDM WebSocket.

        Returns:
            True if pong received, False otherwise
        """
        if not self.connected:
            return False

        try:
            await self._ws.send(json.dumps({'action': 'ping'}))
            return True
        except Exception as e:
            logger.error(f"Ping failed: {e}")
            return False
