"""
WebSocket-based market data streaming for the Tradier API.

This module provides real-time market data streaming via WebSocket,
supporting quotes, trades, summaries, and timesales for equities and options.

Example::

    from optrabot.broker.tradier import Session
    from optrabot.broker.tradier.streaming import MarketDataStreamer

    session = Session("your-access-token")
    
    async def handle_quote(data: Quote):
        print(f"{data.symbol}: ${data.bid}/{data.ask}")
    
    streamer = await MarketDataStreamer.create(session)
    streamer.add_handler(Quote, handle_quote)
    await streamer.subscribe(["AAPL", "SPY"])
    await streamer.run()
"""

from __future__ import annotations

import asyncio
import json
import ssl
from datetime import datetime
from decimal import Decimal
from enum import Enum
from typing import TYPE_CHECKING, Any, Callable, Coroutine, TypeVar

from loguru import logger
from pydantic import Field

from optrabot.broker.tradier.utils import TradierData

if TYPE_CHECKING:
    from optrabot.broker.tradier.session import Session

# WebSocket URL for market data streaming
WEBSOCKET_URL = 'wss://ws.tradier.com/v1/markets/events'


# ========== Enums ==========


class StreamFilter(str, Enum):
    """Filter types for market data streaming."""

    QUOTE = 'quote'
    TRADE = 'trade'
    SUMMARY = 'summary'
    TIMESALE = 'timesale'
    TRADEX = 'tradex'


# ========== Streaming Data Models ==========


class StreamingSession(TradierData):
    """
    A streaming session returned by the create session endpoint.

    Contains the session ID required to connect to the WebSocket stream.
    """

    url: str
    sessionid: str


class Quote(TradierData):
    """
    Real-time quote data from the market data stream.

    Contains bid/ask prices and sizes with exchange information.

    Attributes:
        symbol: The ticker symbol.
        bid: The current bid price.
        ask: The current ask price.
        bidsz: The bid size (number of shares/contracts).
        asksz: The ask size (number of shares/contracts).
        bidexch: The exchange code for the bid.
        askexch: The exchange code for the ask.
        biddate: Timestamp of the bid.
        askdate: Timestamp of the ask.
    """

    type: str = 'quote'
    symbol: str
    bid: Decimal = Decimal('0')
    ask: Decimal = Decimal('0')
    bidsz: int = Field(default=0, alias='bidsz')
    asksz: int = Field(default=0, alias='asksz')
    bidexch: str = ''
    askexch: str = ''
    biddate: int | None = None
    askdate: int | None = None

    @property
    def bid_datetime(self) -> datetime | None:
        """Convert bid timestamp to datetime."""
        if self.biddate:
            return datetime.fromtimestamp(self.biddate / 1000)
        return None

    @property
    def ask_datetime(self) -> datetime | None:
        """Convert ask timestamp to datetime."""
        if self.askdate:
            return datetime.fromtimestamp(self.askdate / 1000)
        return None

    def __repr__(self) -> str:
        return f'Quote({self.symbol!r}, bid={self.bid}, ask={self.ask})'


class Trade(TradierData):
    """
    Real-time trade data from the market data stream.

    Represents a single trade execution.

    Attributes:
        symbol: The ticker symbol.
        price: The trade price.
        size: The trade size (number of shares/contracts).
        exch: The exchange code where the trade occurred.
        cvol: Cumulative volume for the day.
        last: The last trade price.
        date: Timestamp of the trade.
    """

    type: str = 'trade'
    symbol: str
    price: Decimal = Decimal('0')
    size: int = 0
    exch: str = ''
    cvol: int = Field(default=0, alias='cvol')
    last: Decimal = Decimal('0')
    date: int | None = None

    @property
    def trade_datetime(self) -> datetime | None:
        """Convert trade timestamp to datetime."""
        if self.date:
            return datetime.fromtimestamp(self.date / 1000)
        return None

    def __repr__(self) -> str:
        return f'Trade({self.symbol!r}, price={self.price}, size={self.size})'


class Tradex(TradierData):
    """
    Extended trade data from the market data stream.

    Similar to Trade but may include additional exchange-specific details.
    """

    type: str = 'tradex'
    symbol: str
    price: Decimal = Decimal('0')
    size: int = 0
    exch: str = ''
    cvol: int = Field(default=0, alias='cvol')
    last: Decimal = Decimal('0')
    date: int | None = None

    @property
    def trade_datetime(self) -> datetime | None:
        """Convert trade timestamp to datetime."""
        if self.date:
            return datetime.fromtimestamp(self.date / 1000)
        return None

    def __repr__(self) -> str:
        return f'Tradex({self.symbol!r}, price={self.price}, size={self.size})'


class Summary(TradierData):
    """
    Daily summary data from the market data stream.

    Contains open, high, low, and previous close prices.

    Attributes:
        symbol: The ticker symbol.
        open: The opening price for the day.
        high: The high price for the day.
        low: The low price for the day.
        prevClose: The previous day's closing price.
    """

    type: str = 'summary'
    symbol: str
    open: Decimal = Field(default=Decimal('0'), alias='open')
    high: Decimal = Decimal('0')
    low: Decimal = Decimal('0')
    prev_close: Decimal = Field(default=Decimal('0'), alias='prevClose')

    def __repr__(self) -> str:
        return f'Summary({self.symbol!r}, open={self.open}, high={self.high}, low={self.low})'


class Timesale(TradierData):
    """
    Time and sales data from the market data stream.

    Provides detailed trade information including bid/ask at time of sale.

    Attributes:
        symbol: The ticker symbol.
        bid: The bid price at time of sale.
        ask: The ask price at time of sale.
        last: The sale price.
        size: The trade size.
        exch: The exchange code.
        date: Timestamp of the timesale.
        seq: Sequence number.
        flag: Trade condition flag.
        cancel: Whether the trade was cancelled.
        correction: Whether the trade was a correction.
        session: Trading session (normal, pre, post).
    """

    type: str = 'timesale'
    symbol: str
    bid: Decimal = Decimal('0')
    ask: Decimal = Decimal('0')
    last: Decimal = Decimal('0')
    size: int = 0
    exch: str = ''
    date: int | None = None
    seq: int = 0
    flag: str = ''
    cancel: bool = False
    correction: bool = False
    session: str = 'normal'

    @property
    def sale_datetime(self) -> datetime | None:
        """Convert timesale timestamp to datetime."""
        if self.date:
            return datetime.fromtimestamp(self.date / 1000)
        return None

    def __repr__(self) -> str:
        return f'Timesale({self.symbol!r}, last={self.last}, size={self.size})'


class StreamError(TradierData):
    """Error message from the streaming API."""

    error: str

    def __repr__(self) -> str:
        return f'StreamError({self.error!r})'


# Type alias for streaming data
StreamData = Quote | Trade | Tradex | Summary | Timesale | StreamError

# Type variable for handlers
T = TypeVar('T', Quote, Trade, Tradex, Summary, Timesale, StreamError)

# Handler type
HandlerType = Callable[[T], Coroutine[Any, Any, None]] | Callable[[T], None]


# ========== Streaming Client ==========


class MarketDataStreamer:
    """
    WebSocket-based market data streamer.

    Provides real-time streaming of market data including quotes, trades,
    summaries, and timesales for equities and options.

    The streamer uses an event-driven pattern with handlers for each data type.
    You can register handlers for specific data types (Quote, Trade, etc.) and
    they will be called whenever matching data is received.

    Example::

        session = Session("your-token")
        streamer = await MarketDataStreamer.create(session)
        
        # Register handlers
        streamer.add_handler(Quote, lambda q: print(f"{q.symbol}: {q.bid}/{q.ask}"))
        streamer.add_handler(Trade, lambda t: print(f"Trade: {t.price}"))
        
        # Subscribe and run
        await streamer.subscribe(["AAPL", "SPY"])
        await streamer.run()

    Note:
        You must first create a streaming session before connecting.
        Session IDs expire after 5 minutes if not connected.
    """

    def __init__(
        self,
        session_id: str,
        websocket_url: str = WEBSOCKET_URL,
        *,
        ssl_context: ssl.SSLContext | bool | None = None,
    ) -> None:
        """
        Initialize the market data streamer.

        Prefer using the `create()` class method which handles session creation.

        :param session_id: The streaming session ID from create_market_session.
        :param websocket_url: The WebSocket URL to connect to.
        :param ssl_context: Optional SSL context for the WebSocket connection.
                           Pass an `ssl.SSLContext` for custom TLS configuration,
                           `True` for default TLS (same as None), or `False`
                           to disable TLS verification (not recommended).
        """
        self._session_id = session_id
        self._websocket_url = websocket_url
        self._websocket: Any = None
        self._running = False
        self._symbols: list[str] = []
        self._filters: list[StreamFilter] = []
        self._linebreak = True
        self._valid_only = True
        self._advanced_details = False
        # Optional SSL context to pass to websockets.connect. If True/None,
        # the default TLS handling will be used by the websockets library.
        self._ssl: ssl.SSLContext | bool | None = ssl_context

        # Handlers for different data types
        self._handlers: dict[type, list[HandlerType[Any]]] = {
            Quote: [],
            Trade: [],
            Tradex: [],
            Summary: [],
            Timesale: [],
            StreamError: [],
        }

    @classmethod
    async def create(
        cls,
        session: Session,
        websocket_url: str = WEBSOCKET_URL,
        *,
        ssl_context: ssl.SSLContext | bool | None = None,
    ) -> MarketDataStreamer:
        """
        Create a new MarketDataStreamer with an authenticated session.

        This method handles the creation of the streaming session automatically.

        :param session: An authenticated Tradier Session.
        :param websocket_url: The WebSocket URL to connect to.
        :param ssl_context: Optional SSL context for the WebSocket connection.
                           Pass an `ssl.SSLContext` for custom TLS configuration,
                           `True` for default TLS (same as None), or `False`
                           to disable TLS verification (not recommended).
        :returns: A configured MarketDataStreamer ready to subscribe and run.

        Example::

            session = Session("your-token")
            streamer = await MarketDataStreamer.create(session)
            await streamer.subscribe(["AAPL"])
            await streamer.run()
        """
        streaming_session = await create_market_session(session)
        logger.debug(
            'Created streaming session: {}...', streaming_session.sessionid[:8]
        )
        return cls(
            session_id=streaming_session.sessionid,
            websocket_url=websocket_url,
            ssl_context=ssl_context,
        )

    @classmethod
    def create_sync(
        cls,
        session: Session,
        websocket_url: str = WEBSOCKET_URL,
        *,
        ssl_context: ssl.SSLContext | bool | None = None,
    ) -> MarketDataStreamer:
        """
        Create a new MarketDataStreamer with an authenticated session (sync).

        :param session: An authenticated Tradier Session.
        :param websocket_url: The WebSocket URL to connect to.
        :param ssl_context: Optional SSL context for the WebSocket connection.
                           Pass an `ssl.SSLContext` for custom TLS configuration,
                           `True` for default TLS (same as None), or `False`
                           to disable TLS verification (not recommended).
        :returns: A configured MarketDataStreamer ready to subscribe and run.
        """
        streaming_session = create_market_session_sync(session)
        logger.debug(
            'Created streaming session: {}...', streaming_session.sessionid[:8]
        )
        return cls(
            session_id=streaming_session.sessionid,
            websocket_url=websocket_url,
            ssl_context=ssl_context,
        )

    def add_handler(
        self,
        data_type: type[T],
        handler: HandlerType[T],
    ) -> None:
        """
        Register a handler for a specific data type.

        The handler will be called whenever data of the specified type is received.
        Handlers can be sync or async functions.

        :param data_type: The type of data to handle (Quote, Trade, etc.).
        :param handler: A callable that receives the data object.

        Example::

            def on_quote(quote: Quote):
                print(f"{quote.symbol}: {quote.bid}/{quote.ask}")
            
            streamer.add_handler(Quote, on_quote)
        """
        if data_type in self._handlers:
            self._handlers[data_type].append(handler)
            logger.debug('Added handler for {}', data_type.__name__)
        else:
            raise ValueError(f'Unknown data type: {data_type}')

    def remove_handler(
        self,
        data_type: type[T],
        handler: HandlerType[T],
    ) -> None:
        """
        Remove a previously registered handler.

        :param data_type: The type of data the handler was registered for.
        :param handler: The handler to remove.
        """
        if data_type in self._handlers and handler in self._handlers[data_type]:
            self._handlers[data_type].remove(handler)
            logger.debug('Removed handler for {}', data_type.__name__)

    async def subscribe(
        self,
        symbols: list[str],
        filters: list[StreamFilter] | None = None,
        *,
        valid_only: bool = True,
        advanced_details: bool = False,
    ) -> None:
        """
        Subscribe to market data for the specified symbols.

        Can be called multiple times to update the subscription.
        If already connected, the subscription will be updated on the fly.

        :param symbols: List of symbols to subscribe to (equities or options).
        :param filters: Optional list of data types to receive. If None, receives all.
        :param valid_only: Only include ticks considered valid by exchanges.
        :param advanced_details: Include advanced details in timesale payloads.

        Example::

            # Subscribe to quotes and trades only
            await streamer.subscribe(
                ["AAPL", "SPY"],
                filters=[StreamFilter.QUOTE, StreamFilter.TRADE]
            )
            
            # Subscribe to an option
            await streamer.subscribe(["SPY250117C00500000"])
        """
        self._symbols = symbols
        self._filters = filters or []
        self._valid_only = valid_only
        self._advanced_details = advanced_details

        logger.trace('Subscribed to symbols: {}', symbols)

        # If already connected, update subscription
        if self._websocket is not None and self._running:
            await self._send_subscription()

    async def unsubscribe(self, symbols: list[str] | None = None) -> None:
        """
        Unsubscribe from market data for the specified symbols.

        If no symbols are provided, unsubscribes from all symbols.

        :param symbols: List of symbols to unsubscribe from. If None or empty,
                       unsubscribes from all current subscriptions.

        Example::

            # Unsubscribe from specific symbols
            await streamer.unsubscribe(["AAPL"])
            
            # Unsubscribe from all symbols
            await streamer.unsubscribe()
        """
        if symbols is None or len(symbols) == 0:
            # Unsubscribe from all
            removed = self._symbols.copy()
            self._symbols = []
            logger.trace('Unsubscribed from all symbols: {}', removed)
        else:
            # Remove specific symbols
            for symbol in symbols:
                if symbol in self._symbols:
                    self._symbols.remove(symbol)
            logger.trace('Unsubscribed from symbols: {}', symbols)

        # If connected, send updated subscription (or empty list)
        if self._websocket is not None and self._running:
            await self._send_subscription()

    @property
    def symbols(self) -> list[str]:
        """
        Get the list of currently subscribed symbols.

        :returns: A copy of the current symbol subscription list.
        """
        return self._symbols.copy()

    def _build_payload(self) -> str:
        """Build the subscription payload JSON."""
        payload: dict[str, Any] = {
            'symbols': self._symbols,
            'sessionid': self._session_id,
            'linebreak': self._linebreak,
            'validOnly': self._valid_only,
            'advancedDetails': self._advanced_details,
        }

        if self._filters:
            payload['filter'] = [f.value for f in self._filters]

        return json.dumps(payload)

    async def _send_subscription(self) -> None:
        """Send the subscription payload to the WebSocket."""
        if self._websocket is not None:
            payload = self._build_payload()
            await self._websocket.send(payload)
            logger.trace('Sent subscription: {}', payload)

    def _parse_message(self, message: str) -> StreamData | None:
        """Parse a streaming message into the appropriate data type."""
        try:
            data = json.loads(message)
        except json.JSONDecodeError:
            logger.warning('Failed to parse message: {}', message[:100])
            return None

        msg_type = data.get('type')

        if msg_type == 'quote':
            return Quote.model_validate(data)
        elif msg_type == 'trade':
            return Trade.model_validate(data)
        elif msg_type == 'tradex':
            return Tradex.model_validate(data)
        elif msg_type == 'summary':
            return Summary.model_validate(data)
        elif msg_type == 'timesale':
            return Timesale.model_validate(data)
        elif 'error' in data:
            return StreamError.model_validate(data)
        else:
            logger.warning('Unknown message type: {}', msg_type)
            return None

    async def _dispatch(self, data: StreamData) -> None:
        """Dispatch data to registered handlers."""
        handlers = self._handlers.get(type(data), [])

        for handler in handlers:
            try:
                result = handler(data)
                # Handle async handlers
                if asyncio.iscoroutine(result):
                    await result
            except Exception:
                logger.exception('Handler error for {}', type(data).__name__)

    async def run(self) -> None:
        """
        Start the streaming connection and process messages.

        This method runs indefinitely until `stop()` is called or
        the connection is closed.

        Raises:
            RuntimeError: If no symbols have been subscribed.
            websockets.exceptions.ConnectionClosed: If connection is lost.

        Example::

            async def main():
                streamer = await MarketDataStreamer.create(session)
                streamer.add_handler(Quote, handle_quote)
                await streamer.subscribe(["AAPL"])
                
                try:
                    await streamer.run()
                except KeyboardInterrupt:
                    await streamer.stop()
        """
        if not self._symbols:
            raise RuntimeError('No symbols subscribed. Call subscribe() first.')

        try:
            import websockets
        except ImportError:
            raise ImportError(
                'websockets package is required for streaming. '
                'Install it with: pip install websockets'
            ) from None

        self._running = True
        logger.info('Connecting to {}', self._websocket_url)

        async with websockets.connect(
            self._websocket_url,
            ssl=self._ssl,
            compression=None,
        ) as websocket:
            self._websocket = websocket
            await self._send_subscription()

            logger.info('Streaming started for {} symbols', len(self._symbols))

            async for message in websocket:
                if not self._running:
                    break

                data = self._parse_message(str(message))
                if data:
                    await self._dispatch(data)

        self._websocket = None
        self._running = False
        logger.info('Streaming stopped')

    async def stop(self) -> None:
        """
        Stop the streaming connection gracefully.

        Example::

            async def handle_signal():
                await streamer.stop()
        """
        self._running = False
        if self._websocket is not None:
            await self._websocket.close()
            logger.debug('WebSocket connection closed')

    @property
    def is_running(self) -> bool:
        """
        Check if the streamer is currently running.

        :returns: True if the streamer is connected and processing messages.
        """
        return self._running


# ========== Session Creation ==========


async def create_market_session(session: Session) -> StreamingSession:
    """
    Create a streaming session for market data.

    This must be called before connecting to the WebSocket stream.
    Session IDs expire after 5 minutes if not connected.

    :param session: An authenticated Tradier Session.
    :returns: A StreamingSession with the session ID and URL.

    Example::

        session = Session("your-token")
        streaming = await create_market_session(session)
        print(f"Session ID: {streaming.sessionid}")
    """
    data = await session._a_post('/markets/events/session')
    stream_data = data.get('stream', {})
    return StreamingSession.model_validate(stream_data)


def create_market_session_sync(session: Session) -> StreamingSession:
    """
    Create a streaming session for market data (synchronous).

    :param session: An authenticated Tradier Session.
    :returns: A StreamingSession with the session ID and URL.
    """
    data = session._post('/markets/events/session')
    stream_data = data.get('stream', {})
    return StreamingSession.model_validate(stream_data)
