"""
AgentBroker WebSocket Client

Provides real-time streaming for trades, balance updates, and order book changes.

Usage:
    from agentbroker import AgentBroker

    client = AgentBroker(api_key="your_key")
    ws = client.connect_ws()

    @ws.on_trade
    def handle_trade(event):
        print(f"Trade: {event}")

    @ws.on_balance
    def handle_balance(event):
        print(f"Balance update: {event['balance_usdc']}")

    ws.run_forever()
"""

from __future__ import annotations
import json
import threading
import time
from typing import Callable, Optional, List
from .exceptions import WebSocketError


class AgentBrokerWebSocket:
    """
    WebSocket client for real-time AgentBroker event streaming.

    Channels:
    - ``trades``     – fills and executed trades for your agent
    - ``balance``    – USDC balance change events (deposits, fills, withdrawals)
    - ``orderbook``  – live order book snapshots for your subscribed pairs
    """

    CHANNEL_TRADES = "trades"
    CHANNEL_BALANCE = "balance"
    CHANNEL_ORDERBOOK = "orderbook"

    def __init__(self, api_key: str, base_url: str):
        self.api_key = api_key
        # Convert http(s) to ws(s)
        ws_url = base_url.replace("https://", "wss://").replace("http://", "ws://")
        self.ws_url = f"{ws_url}?key={api_key}"

        self._ws = None
        self._thread: Optional[threading.Thread] = None
        self._reconnect = True
        self._reconnect_delay = 2  # seconds, doubles up to max
        self._max_reconnect_delay = 60

        # Callbacks
        self._on_trade: List[Callable] = []
        self._on_balance: List[Callable] = []
        self._on_orderbook: List[Callable] = []
        self._on_price: List[Callable] = []
        self._on_connected: List[Callable] = []
        self._on_error: List[Callable] = []
        self._on_disconnect: List[Callable] = []

        # Active subscriptions (sent on each reconnect)
        self._subscriptions = set([self.CHANNEL_TRADES, self.CHANNEL_BALANCE, self.CHANNEL_ORDERBOOK, "prices"])

    # ── Decorator-style event registration ─────────────────────────────

    def on_trade(self, fn: Callable) -> Callable:
        """Register a handler for trade fill events."""
        self._on_trade.append(fn)
        return fn

    def on_balance(self, fn: Callable) -> Callable:
        """Register a handler for balance update events."""
        self._on_balance.append(fn)
        return fn

    def on_orderbook(self, fn: Callable) -> Callable:
        """Register a handler for order book update events."""
        self._on_orderbook.append(fn)
        return fn

    def on_price(self, fn: Callable) -> Callable:
        """
        Register a handler for real-time price tick events.

        The handler receives a dict with at minimum ``pair`` and ``price`` keys:

        .. code-block:: python

            @ws.on_price
            def handle_price(tick):
                print(f"{tick['pair']}: ${tick['price']}")
        """
        self._on_price.append(fn)
        return fn

    def on_connected(self, fn: Callable) -> Callable:
        """Called when the WebSocket connection is established."""
        self._on_connected.append(fn)
        return fn

    def on_error(self, fn: Callable) -> Callable:
        """Called when a WebSocket error occurs."""
        self._on_error.append(fn)
        return fn

    def on_disconnect(self, fn: Callable) -> Callable:
        """Called when the WebSocket disconnects."""
        self._on_disconnect.append(fn)
        return fn

    # ── Subscription management ─────────────────────────────────────────

    def subscribe(self, channel: str):
        """Subscribe to a channel. Sends the subscribe message if connected."""
        self._subscriptions.add(channel)
        self._send({"action": "subscribe", "channel": channel})

    def unsubscribe(self, channel: str):
        """Unsubscribe from a channel."""
        self._subscriptions.discard(channel)
        self._send({"action": "unsubscribe", "channel": channel})

    # ── Connection ──────────────────────────────────────────────────────

    def connect(self):
        """
        Connect to the WebSocket in a background thread.
        Returns immediately. Use ``run_forever()`` to block instead.
        """
        self._reconnect = True
        self._thread = threading.Thread(target=self._run_loop, daemon=True)
        self._thread.start()
        return self

    def run_forever(self):
        """Connect and block until the connection is permanently closed."""
        self.connect()
        self._thread.join()

    def close(self):
        """Permanently close the WebSocket (no reconnect)."""
        self._reconnect = False
        if self._ws:
            try:
                self._ws.close()
            except Exception:
                pass

    # ── Internal ────────────────────────────────────────────────────────

    def _send(self, data: dict):
        if self._ws:
            try:
                import websocket
                if self._ws.sock and self._ws.sock.connected:
                    self._ws.send(json.dumps(data))
            except Exception:
                pass

    def _run_loop(self):
        try:
            import websocket as ws_lib
        except ImportError:
            raise WebSocketError(
                "websocket-client is required for WebSocket support. "
                "Install it with: pip install websocket-client"
            )

        delay = self._reconnect_delay

        while self._reconnect:
            try:
                self._ws = ws_lib.WebSocketApp(
                    self.ws_url,
                    on_open=self._on_open,
                    on_message=self._on_message,
                    on_error=self._on_error_internal,
                    on_close=self._on_close,
                )
                self._ws.run_forever(ping_interval=20, ping_timeout=10)
            except Exception as e:
                for cb in self._on_error:
                    try:
                        cb(e)
                    except Exception:
                        pass

            if not self._reconnect:
                break

            time.sleep(delay)
            delay = min(delay * 2, self._max_reconnect_delay)

    def _on_open(self, ws):
        # Reset reconnect delay on successful connection
        self._reconnect_delay = 2
        for cb in self._on_connected:
            try:
                cb()
            except Exception:
                pass

    def _on_message(self, ws, message: str):
        try:
            data = json.loads(message)
        except json.JSONDecodeError:
            return

        channel    = data.get("channel", "")
        event_type = data.get("type", data.get("event", ""))

        # Price ticks: { "channel": "prices", "type": "tick", "data": { "pair": "BTC-USDC", "price": "67312.50" } }
        if channel == "prices" and event_type == "tick":
            tick = data.get("data", {})
            for cb in self._on_price:
                try:
                    cb(tick)
                except Exception:
                    pass
            return

        if event_type in ("trade", "fill"):
            for cb in self._on_trade:
                try:
                    cb(data)
                except Exception:
                    pass

        elif event_type in ("balance", "deposit", "withdrawal"):
            for cb in self._on_balance:
                try:
                    cb(data)
                except Exception:
                    pass

        elif event_type == "orderbook":
            for cb in self._on_orderbook:
                try:
                    cb(data)
                except Exception:
                    pass

    def _on_error_internal(self, ws, error):
        for cb in self._on_error:
            try:
                cb(error)
            except Exception:
                pass

    def _on_close(self, ws, close_status_code, close_msg):
        for cb in self._on_disconnect:
            try:
                cb(close_status_code, close_msg)
            except Exception:
                pass
