dex/client/python/lx_client/websocket_client.py

"""
LX Trading Client - WebSocket Implementation

Async WebSocket client for LX trading API.
"""

import asyncio
import json
import uuid
from typing import Optional, List, Dict, Any, AsyncIterator

try:
    import websockets
    from websockets.client import WebSocketClientProtocol
except ImportError:
    websockets = None
    WebSocketClientProtocol = None

from lx_client.base import (
    TradingClient,
    Order,
    OrderType,
    OrderSide,
    OrderStatus,
    TimeInForce,
    OrderBook,
    PriceLevel,
    Position,
    Balance,
    Trade,
)


class WebSocketClient(TradingClient):
    """
    WebSocket-based trading client for LX.

    Uses async/await for all operations with websockets library.

    Example:
        async with WebSocketClient("ws://localhost:8081") as client:
            order = await client.place_order("BTC-USD", "buy", "limit", 50000, 0.1)
            print(f"Order placed: {order.order_id}")
    """

    def __init__(
        self,
        url: str = "ws://localhost:8081",
        api_key: Optional[str] = None,
        api_secret: Optional[str] = None,
        verbose: bool = False,
    ):
        """
        Initialize WebSocket client.

        Args:
            url: WebSocket server URL
            api_key: API key for authentication
            api_secret: API secret for authentication
            verbose: Enable verbose logging
        """
        super().__init__(api_key, api_secret)

        if websockets is None:
            raise ImportError(
                "websockets library required. Install with: pip install websockets"
            )

        self.url = url
        self.verbose = verbose
        self._ws: Optional[WebSocketClientProtocol] = None
        self._responses: Dict[str, asyncio.Future] = {}
        self._recv_task: Optional[asyncio.Task] = None
        self._subscriptions: Dict[str, asyncio.Queue] = {}

    async def connect(self) -> bool:
        """Connect to WebSocket server."""
        try:
            self._ws = await websockets.connect(
                self.url,
                ping_interval=30,
                ping_timeout=10,
            )
            self._recv_task = asyncio.create_task(self._recv_loop())

            # Wait for connected message
            try:
                msg = await asyncio.wait_for(self._ws.recv(), timeout=5.0)
                data = json.loads(msg)
                if self.verbose:
                    print(f"<< {json.dumps(data, indent=2)}")
                if data.get("type") == "connected":
                    self._connected = True
                    return True
            except asyncio.TimeoutError:
                pass

            self._connected = True
            return True
        except Exception as e:
            if self.verbose:
                print(f"Connection failed: {e}")
            return False

    async def disconnect(self) -> None:
        """Disconnect from WebSocket server."""
        self._connected = False
        if self._recv_task:
            self._recv_task.cancel()
            try:
                await self._recv_task
            except asyncio.CancelledError:
                pass
            self._recv_task = None

        if self._ws:
            await self._ws.close()
            self._ws = None

        # Cancel pending responses
        for future in self._responses.values():
            if not future.done():
                future.cancel()
        self._responses.clear()

    async def _recv_loop(self) -> None:
        """Background task to receive messages."""
        while self._connected and self._ws:
            try:
                data = await self._ws.recv()
                msg = json.loads(data)

                if self.verbose:
                    print(f"<< {json.dumps(msg, indent=2)}")

                req_id = msg.get("request_id", "")
                msg_type = msg.get("type", "")

                # Handle auth response
                if msg_type == "auth_success":
                    self._authenticated = True

                # Route response to waiting caller
                if req_id and req_id in self._responses:
                    future = self._responses.pop(req_id)
                    if not future.done():
                        future.set_result(msg)

                # Route subscription updates
                elif msg_type in ("orderbook", "trade", "order_update"):
                    symbol = msg.get("data", {}).get("symbol", "")
                    key = f"{msg_type}:{symbol}"
                    if key in self._subscriptions:
                        await self._subscriptions[key].put(msg)

            except asyncio.CancelledError:
                break
            except Exception as e:
                if self._connected and self.verbose:
                    print(f"Receive error: {e}")
                break

    async def _send(self, msg_type: str, data: Optional[Dict] = None) -> Dict[str, Any]:
        """Send message and wait for response."""
        if not self._ws:
            raise ConnectionError("Not connected")

        req_id = str(uuid.uuid4())[:8]
        msg = {"type": msg_type, "request_id": req_id}
        if data:
            msg.update(data)

        if self.verbose:
            print(f">> {json.dumps(msg, indent=2)}")

        # Create future for response
        future: asyncio.Future = asyncio.get_event_loop().create_future()
        self._responses[req_id] = future

        try:
            await self._ws.send(json.dumps(msg))
            return await asyncio.wait_for(future, timeout=10.0)
        except asyncio.TimeoutError:
            self._responses.pop(req_id, None)
            raise TimeoutError(f"Request {msg_type} timed out")

    async def authenticate(
        self,
        api_key: Optional[str] = None,
        api_secret: Optional[str] = None,
    ) -> bool:
        """Authenticate with API credentials."""
        key = api_key or self.api_key
        secret = api_secret or self.api_secret

        if not key or not secret:
            raise ValueError("API key and secret required")

        resp = await self._send("auth", {"apiKey": key, "apiSecret": secret})
        if resp.get("type") == "auth_success":
            self._authenticated = True
            return True
        return False

    async def place_order(
        self,
        symbol: str,
        side: str,
        order_type: str,
        price: float,
        size: float,
        client_id: Optional[str] = None,
        time_in_force: str = "gtc",
        post_only: bool = False,
        reduce_only: bool = False,
    ) -> Order:
        """Place a new order."""
        order_data = {
            "symbol": symbol,
            "side": side,
            "type": order_type,
            "price": price,
            "size": size,
            "time_in_force": time_in_force,
            "post_only": post_only,
            "reduce_only": reduce_only,
        }
        if client_id:
            order_data["client_id"] = client_id

        resp = await self._send("place_order", {"order": order_data})

        if resp.get("error"):
            raise ValueError(resp["error"])

        data = resp.get("data", resp)
        return Order(
            order_id=data.get("order_id", data.get("orderID", 0)),
            symbol=symbol,
            order_type=OrderType(order_type),
            side=OrderSide(side),
            price=price,
            size=size,
            status=OrderStatus(data.get("status", "open")),
            client_id=client_id or "",
            time_in_force=TimeInForce(time_in_force),
            post_only=post_only,
            reduce_only=reduce_only,
        )

    async def cancel_order(self, order_id: int) -> bool:
        """Cancel an existing order."""
        resp = await self._send("cancel_order", {"orderID": order_id})
        if resp.get("error"):
            return False
        return resp.get("success", True)

    async def get_order(self, order_id: int) -> Optional[Order]:
        """Get order details."""
        resp = await self._send("get_order", {"orderID": order_id})
        if resp.get("error"):
            return None
        return self._parse_order(resp.get("data", resp))

    async def get_orders(
        self,
        symbol: Optional[str] = None,
        status: Optional[str] = None,
        limit: int = 100,
    ) -> List[Order]:
        """Get list of orders."""
        params: Dict[str, Any] = {"limit": limit}
        if symbol:
            params["symbol"] = symbol
        if status:
            params["status"] = status

        resp = await self._send("get_orders", params)
        if resp.get("error"):
            return []

        orders_data = resp.get("data", {}).get("orders", [])
        return [self._parse_order(o) for o in orders_data]

    async def get_orderbook(
        self,
        symbol: str,
        depth: int = 20,
    ) -> OrderBook:
        """Get order book snapshot."""
        resp = await self._send("subscribe", {"symbols": [symbol], "depth": depth})

        # Wait briefly for orderbook data
        await asyncio.sleep(0.5)

        data = resp.get("data", {})
        return OrderBook(
            symbol=symbol,
            bids=[
                PriceLevel(price=b.get("price", 0), size=b.get("size", 0))
                for b in data.get("bids", [])[:depth]
            ],
            asks=[
                PriceLevel(price=a.get("price", 0), size=a.get("size", 0))
                for a in data.get("asks", [])[:depth]
            ],
            timestamp=data.get("timestamp", 0),
        )

    async def get_positions(self) -> List[Position]:
        """Get all positions."""
        resp = await self._send("get_positions", {})
        if resp.get("error"):
            return []

        positions_data = resp.get("data", {}).get("positions", [])
        return [
            Position(
                symbol=p.get("symbol", ""),
                size=p.get("size", 0),
                entry_price=p.get("entry_price", 0),
                mark_price=p.get("mark_price", 0),
                pnl=p.get("pnl", 0),
                margin=p.get("margin", 0),
            )
            for p in positions_data
        ]

    async def get_balance(self, asset: str) -> Optional[Balance]:
        """Get balance for an asset."""
        resp = await self._send("get_balances", {"asset": asset})
        if resp.get("error"):
            return None

        data = resp.get("data", {})
        return Balance(
            asset=asset,
            available=data.get("available", 0),
            locked=data.get("locked", 0),
            total=data.get("total", 0),
        )

    async def stream_orderbook(
        self,
        symbol: str,
        depth: int = 20,
    ) -> AsyncIterator[OrderBook]:
        """Stream order book updates."""
        key = f"orderbook:{symbol}"
        self._subscriptions[key] = asyncio.Queue()

        try:
            await self._send("subscribe", {"symbols": [symbol], "depth": depth})

            while self._connected:
                try:
                    msg = await asyncio.wait_for(
                        self._subscriptions[key].get(),
                        timeout=30.0,
                    )
                    data = msg.get("data", {})
                    yield OrderBook(
                        symbol=symbol,
                        bids=[
                            PriceLevel(price=b.get("price", 0), size=b.get("size", 0))
                            for b in data.get("bids", [])[:depth]
                        ],
                        asks=[
                            PriceLevel(price=a.get("price", 0), size=a.get("size", 0))
                            for a in data.get("asks", [])[:depth]
                        ],
                        timestamp=data.get("timestamp", 0),
                    )
                except asyncio.TimeoutError:
                    continue
        finally:
            self._subscriptions.pop(key, None)
            await self._send("unsubscribe", {"symbols": [symbol]})

    async def stream_trades(
        self,
        symbol: str,
    ) -> AsyncIterator[Trade]:
        """Stream trades."""
        key = f"trade:{symbol}"
        self._subscriptions[key] = asyncio.Queue()

        try:
            await self._send("subscribe_trades", {"symbol": symbol})

            while self._connected:
                try:
                    msg = await asyncio.wait_for(
                        self._subscriptions[key].get(),
                        timeout=30.0,
                    )
                    data = msg.get("data", {})
                    yield Trade(
                        trade_id=data.get("trade_id", 0),
                        symbol=symbol,
                        price=data.get("price", 0),
                        size=data.get("size", 0),
                        side=OrderSide(data.get("side", "buy")),
                        timestamp=data.get("timestamp", 0),
                    )
                except asyncio.TimeoutError:
                    continue
        finally:
            self._subscriptions.pop(key, None)
            await self._send("unsubscribe_trades", {"symbol": symbol})

    def _parse_order(self, data: Dict) -> Order:
        """Parse order from response data."""
        return Order(
            order_id=data.get("order_id", data.get("orderID", 0)),
            symbol=data.get("symbol", ""),
            order_type=OrderType(data.get("type", "limit")),
            side=OrderSide(data.get("side", "buy")),
            price=data.get("price", 0),
            size=data.get("size", 0),
            filled=data.get("filled", 0),
            remaining=data.get("remaining", data.get("size", 0)),
            status=OrderStatus(data.get("status", "open")),
            user_id=data.get("user_id", ""),
            client_id=data.get("client_id", ""),
            timestamp=data.get("timestamp", 0),
            time_in_force=TimeInForce(data.get("time_in_force", "gtc")),
            post_only=data.get("post_only", False),
            reduce_only=data.get("reduce_only", False),
        )