"""WebSocket streaming client for real-time Fireplace data."""

from __future__ import annotations

import asyncio
import json
import time
from typing import Any, AsyncIterator, Dict, List, Optional, Union

try:
    import websockets
    from websockets.asyncio.client import connect as ws_connect
except ImportError:
    websockets = None  # type: ignore[assignment]

_WS_URL = "wss://stream.fireplace.gg"


class Stream:
    """Subscribe to real-time channels over WebSocket.

    All ``subscribe_*`` methods return an async iterator that yields
    parsed JSON messages.  Use them with ``async for``::

        async for msg in fp.stream.trades_by_market(["0x123..."]):
            print(msg)
    """

    def __init__(self, api_key: str, api_secret: str, *, ws_url: str = _WS_URL, ping_interval: float = 30.0) -> None:
        if websockets is None:
            raise ImportError("The 'websockets' package is required for streaming. Install with: pip install fireplace-sdk")
        self._api_key = api_key
        self._api_secret = api_secret
        self._ws_url = ws_url
        self._ping_interval = ping_interval

    def trades_by_market(self, market_ids: Union[str, List[str]]) -> AsyncIterator[Dict[str, Any]]:
        """Stream real-time trades for a list of markets.

        Args:
            market_ids: Single ID or list of market IDs. **Required.**
        """
        return self._subscribe(subscription_type="ext_trades_by_market", market_id=_to_list(market_ids))

    def trades_by_wallet(self, wallets: Union[str, List[str]]) -> AsyncIterator[Dict[str, Any]]:
        """Stream real-time trades involving specific wallet addresses.

        Args:
            wallets: Single address or list. **Required.**
        """
        return self._subscribe(subscription_type="ext_trades_by_wallet", wallets=_to_list(wallets))

    def market_headers(self, market_ids: Union[str, List[str]]) -> AsyncIterator[Dict[str, Any]]:
        """Stream market metadata updates (price, volume, liquidity changes).

        Args:
            market_ids: Single ID or list. **Required.**
        """
        return self._subscribe(subscription_type="ext_market_headers", market_id=_to_list(market_ids))

    def candles(self, market_ids: Union[str, List[str]]) -> AsyncIterator[Dict[str, Any]]:
        """Stream real-time OHLCV candle updates.

        Args:
            market_ids: Single ID or list. **Required.**
        """
        return self._subscribe(subscription_type="ext_candles", market_id=_to_list(market_ids))

    def news(self) -> AsyncIterator[Dict[str, Any]]:
        """Stream market-matched news tweets (global feed, no parameters)."""
        return self._subscribe(subscription_type="ext_news")

    async def _subscribe(self, **sub_msg: Any) -> AsyncIterator[Dict[str, Any]]:
        headers = {"X-API-Key": self._api_key, "X-API-Secret": self._api_secret}
        async for ws in ws_connect(self._ws_url, additional_headers=headers):
            try:
                msg: Dict[str, Any] = {"action": "subscribe", **sub_msg}
                wire = {_to_camel(k): v for k, v in msg.items()}
                await ws.send(json.dumps(wire))
                ping_task = asyncio.create_task(self._ping_loop(ws))
                try:
                    async for raw in ws:
                        parsed = json.loads(raw)
                        msg_type = parsed.get("type")
                        if msg_type == "event":
                            yield parsed
                        elif msg_type == "error":
                            raise StreamError(parsed.get("message", "Unknown WebSocket error"))
                        # ack / pong — skip silently
                finally:
                    ping_task.cancel()
                    try:
                        await ping_task
                    except asyncio.CancelledError:
                        pass
            except websockets.ConnectionClosed:
                continue

    async def _ping_loop(self, ws: Any) -> None:
        while True:
            await asyncio.sleep(self._ping_interval)
            try:
                await ws.send(json.dumps({"action": "ping", "timestamp": int(time.time() * 1000)}))
            except Exception:
                break


class StreamError(Exception):
    """Raised when the WebSocket server sends an error message."""


def _to_list(value: Union[str, List[str]]) -> List[str]:
    return [value] if isinstance(value, str) else list(value)


def _to_camel(snake: str) -> str:
    parts = snake.split("_")
    return parts[0] + "".join(p.capitalize() for p in parts[1:])
