"""SmaugLabs WebSocket streaming client.

Requires the ``websockets`` package::

    pip install smauglabs[streaming]

Threaded (callback) usage::

    from smauglabs.stream import Stream

    def on_event(event):
        print(f"{event.type}: bid={event.best_bid}, ask={event.best_ask}")

    stream = Stream(on_event=on_event)
    stream.subscribe_all()
    stream.start()       # background thread
    # ... do other work ...
    stream.stop()

Async usage::

    from smauglabs.stream import AsyncStream

    async with AsyncStream() as stream:
        await stream.subscribe_all()
        async for event in stream:
            print(event)
"""

from __future__ import annotations

import json
import threading
import time
from typing import Any, Callable

DEFAULT_WS_URL = "wss://api.smauglabs.com/v1/ws/stream"


class StreamEvent:
    """A parsed event from the SmaugLabs WebSocket stream."""

    __slots__ = ("type", "token_id", "market_id", "data", "raw")

    def __init__(self, raw: dict):
        self.raw = raw
        self.type: str = raw.get("type", "")
        self.token_id: str = raw.get("token_id", "")
        self.market_id: str = raw.get("market_id", "")
        self.data: dict = raw.get("data", {})

    @property
    def best_bid(self) -> float | None:
        v = self.data.get("best_bid")
        return float(v) if v is not None else None

    @property
    def best_ask(self) -> float | None:
        v = self.data.get("best_ask")
        return float(v) if v is not None else None

    @property
    def price(self) -> float | None:
        v = self.data.get("price")
        return float(v) if v is not None else None

    @property
    def size(self) -> float | None:
        v = self.data.get("size")
        return float(v) if v is not None else None

    @property
    def side(self) -> str | None:
        return self.data.get("side")

    def __repr__(self) -> str:
        parts = [f"type={self.type!r}"]
        if self.token_id:
            parts.append(f"token={self.token_id[:12]}...")
        if self.best_bid is not None:
            parts.append(f"bid={self.best_bid}")
        if self.best_ask is not None:
            parts.append(f"ask={self.best_ask}")
        return f"StreamEvent({', '.join(parts)})"


class Stream:
    """Threaded WebSocket stream with auto-reconnect.

    Events are delivered to ``on_event`` in a background daemon thread.
    Call :meth:`start` to connect and :meth:`stop` to disconnect.
    """

    def __init__(
        self,
        on_event: Callable[[StreamEvent], Any] | None = None,
        on_error: Callable[[Exception], Any] | None = None,
        *,
        ws_url: str = DEFAULT_WS_URL,
        auto_reconnect: bool = True,
    ):
        self._ws_url = ws_url
        self._on_event = on_event
        self._on_error = on_error
        self._auto_reconnect = auto_reconnect
        self._running = False
        self._thread: threading.Thread | None = None
        self._ws: Any = None
        self._pending_subs: list[dict] = []

    def subscribe(self, *, tokens: list[str] | None = None, markets: list[str] | None = None) -> None:
        """Subscribe to specific tokens or markets."""
        msg: dict[str, Any] = {"type": "subscribe"}
        if tokens:
            msg["tokens"] = tokens
        if markets:
            msg["markets"] = markets
        self._pending_subs.append(msg)
        self._try_send(msg)

    def subscribe_all(self) -> None:
        """Receive all orderbook updates across every market."""
        msg = {"type": "subscribe_all"}
        self._pending_subs.append(msg)
        self._try_send(msg)

    def unsubscribe(self, *, tokens: list[str] | None = None, markets: list[str] | None = None) -> None:
        """Remove subscriptions."""
        msg: dict[str, Any] = {"type": "unsubscribe"}
        if tokens:
            msg["tokens"] = tokens
        if markets:
            msg["markets"] = markets
        self._try_send(msg)

    def start(self) -> None:
        """Connect and start receiving events in a background thread."""
        if self._running:
            return
        self._running = True
        self._thread = threading.Thread(target=self._loop, daemon=True)
        self._thread.start()

    def stop(self) -> None:
        """Disconnect and stop the background thread."""
        self._running = False
        if self._ws:
            try:
                self._ws.close()
            except Exception:
                pass

    @property
    def connected(self) -> bool:
        return self._ws is not None

    def _try_send(self, msg: dict) -> None:
        if self._ws:
            try:
                self._ws.send(json.dumps(msg))
            except Exception:
                pass

    def _loop(self) -> None:
        from websockets.sync.client import connect

        backoff = 1
        while self._running:
            try:
                with connect(self._ws_url) as ws:
                    self._ws = ws
                    backoff = 1

                    for sub in self._pending_subs:
                        ws.send(json.dumps(sub))

                    while self._running:
                        try:
                            raw = ws.recv(timeout=5)
                        except TimeoutError:
                            continue
                        try:
                            data = json.loads(raw)
                            if self._on_event:
                                self._on_event(StreamEvent(data))
                        except Exception as e:
                            if self._on_error:
                                self._on_error(e)
            except Exception as e:
                if self._on_error:
                    self._on_error(e)

            self._ws = None
            if not self._running or not self._auto_reconnect:
                break

            time.sleep(min(backoff, 30))
            backoff = min(backoff * 2, 30)


class AsyncStream:
    """Async WebSocket stream for use with asyncio.

    Example::

        async with AsyncStream() as stream:
            await stream.subscribe_all()
            async for event in stream:
                print(event.best_bid)
    """

    def __init__(self, *, ws_url: str = DEFAULT_WS_URL):
        self._ws_url = ws_url
        self._ws: Any = None

    async def __aenter__(self) -> AsyncStream:
        import websockets

        self._ws = await websockets.connect(self._ws_url)
        return self

    async def __aexit__(self, *args: Any) -> None:
        if self._ws:
            await self._ws.close()
            self._ws = None

    async def subscribe(self, *, tokens: list[str] | None = None, markets: list[str] | None = None) -> None:
        msg: dict[str, Any] = {"type": "subscribe"}
        if tokens:
            msg["tokens"] = tokens
        if markets:
            msg["markets"] = markets
        await self._ws.send(json.dumps(msg))

    async def subscribe_all(self) -> None:
        await self._ws.send(json.dumps({"type": "subscribe_all"}))

    async def unsubscribe(self, *, tokens: list[str] | None = None, markets: list[str] | None = None) -> None:
        msg: dict[str, Any] = {"type": "unsubscribe"}
        if tokens:
            msg["tokens"] = tokens
        if markets:
            msg["markets"] = markets
        await self._ws.send(json.dumps(msg))

    def __aiter__(self) -> AsyncStream:
        return self

    async def __anext__(self) -> StreamEvent:
        if not self._ws:
            raise StopAsyncIteration
        try:
            raw = await self._ws.recv()
            return StreamEvent(json.loads(raw))
        except Exception:
            raise StopAsyncIteration
