dex/client/python/lx_client/grpc_client.py

"""
LX Trading Client - gRPC Implementation

Async gRPC client for LX trading API.
"""

import asyncio
from typing import Optional, List, AsyncIterator

try:
    import grpc
    from grpc import aio as grpc_aio
except ImportError:
    grpc = None
    grpc_aio = None

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

# Proto enum mappings
_ORDER_TYPE_TO_PROTO = {
    "limit": 0,
    "market": 1,
    "stop": 2,
    "stop_limit": 3,
    "iceberg": 4,
    "peg": 5,
}

_ORDER_TYPE_FROM_PROTO = {v: k for k, v in _ORDER_TYPE_TO_PROTO.items()}

_ORDER_SIDE_TO_PROTO = {
    "buy": 0,
    "sell": 1,
}

_ORDER_SIDE_FROM_PROTO = {v: k for k, v in _ORDER_SIDE_TO_PROTO.items()}

_ORDER_STATUS_FROM_PROTO = {
    0: "open",
    1: "partial",
    2: "filled",
    3: "cancelled",
    4: "rejected",
}

_TIME_IN_FORCE_TO_PROTO = {
    "gtc": 0,
    "ioc": 1,
    "fok": 2,
    "day": 3,
}


class GrpcClient(TradingClient):
    """
    gRPC-based trading client for LX.

    Uses async/await with grpcio-aio for all operations.

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

    def __init__(
        self,
        host: str = "localhost",
        port: int = 50051,
        api_key: Optional[str] = None,
        api_secret: Optional[str] = None,
        use_tls: bool = False,
        ca_cert: Optional[str] = None,
    ):
        """
        Initialize gRPC client.

        Args:
            host: gRPC server host
            port: gRPC server port
            api_key: API key for authentication
            api_secret: API secret for authentication
            use_tls: Enable TLS encryption
            ca_cert: Path to CA certificate for TLS
        """
        super().__init__(api_key, api_secret)

        if grpc is None:
            raise ImportError(
                "grpcio library required. Install with: pip install grpcio grpcio-tools"
            )

        self.host = host
        self.port = port
        self.use_tls = use_tls
        self.ca_cert = ca_cert
        self._channel: Optional[grpc_aio.Channel] = None
        self._stub = None
        self._user_id: str = ""

    @property
    def address(self) -> str:
        """Get server address."""
        return f"{self.host}:{self.port}"

    async def connect(self) -> bool:
        """Connect to gRPC server."""
        try:
            if self.use_tls:
                if self.ca_cert:
                    with open(self.ca_cert, "rb") as f:
                        credentials = grpc.ssl_channel_credentials(f.read())
                else:
                    credentials = grpc.ssl_channel_credentials()
                self._channel = grpc_aio.secure_channel(self.address, credentials)
            else:
                self._channel = grpc_aio.insecure_channel(self.address)

            # Create stub - we'll use a dynamic approach since proto might not be compiled
            self._stub = _DynamicLXDEXStub(self._channel)
            self._connected = True

            # Verify connection with ping
            try:
                await self._stub.Ping(timestamp=int(asyncio.get_event_loop().time() * 1000))
            except Exception:
                # Ping might not be implemented, but channel is open
                pass

            return True
        except Exception as e:
            self._connected = False
            raise ConnectionError(f"Failed to connect to {self.address}: {e}")

    async def disconnect(self) -> None:
        """Disconnect from gRPC server."""
        self._connected = False
        if self._channel:
            await self._channel.close()
            self._channel = None
            self._stub = None

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

        For gRPC, authentication is typically handled via metadata/headers
        on each request rather than a separate auth call.
        """
        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")

        # Store credentials for use in requests
        self.api_key = key
        self.api_secret = secret
        self._user_id = key  # Use API key as user ID for now
        self._authenticated = True
        return True

    def _get_metadata(self) -> List[tuple]:
        """Get authentication metadata for requests."""
        if self.api_key and self.api_secret:
            return [
                ("x-api-key", self.api_key),
                ("x-api-secret", self.api_secret),
            ]
        return []

    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."""
        if not self._stub:
            raise ConnectionError("Not connected")

        response = await self._stub.PlaceOrder(
            symbol=symbol,
            type=_ORDER_TYPE_TO_PROTO.get(order_type, 0),
            side=_ORDER_SIDE_TO_PROTO.get(side, 0),
            price=price,
            size=size,
            user_id=self._user_id,
            client_id=client_id or "",
            time_in_force=_TIME_IN_FORCE_TO_PROTO.get(time_in_force, 0),
            post_only=post_only,
            reduce_only=reduce_only,
            metadata=self._get_metadata(),
        )

        return Order(
            order_id=response.order_id,
            symbol=symbol,
            order_type=OrderType(order_type),
            side=OrderSide(side),
            price=price,
            size=size,
            status=OrderStatus(_ORDER_STATUS_FROM_PROTO.get(response.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."""
        if not self._stub:
            raise ConnectionError("Not connected")

        response = await self._stub.CancelOrder(
            order_id=order_id,
            user_id=self._user_id,
            metadata=self._get_metadata(),
        )
        return response.success

    async def get_order(self, order_id: int) -> Optional[Order]:
        """Get order details."""
        if not self._stub:
            raise ConnectionError("Not connected")

        try:
            response = await self._stub.GetOrder(
                order_id=order_id,
                metadata=self._get_metadata(),
            )
            return self._parse_order(response)
        except grpc.RpcError:
            return None

    async def get_orders(
        self,
        symbol: Optional[str] = None,
        status: Optional[str] = None,
        limit: int = 100,
    ) -> List[Order]:
        """Get list of orders."""
        if not self._stub:
            raise ConnectionError("Not connected")

        status_proto = 0
        if status:
            status_map = {"open": 0, "partial": 1, "filled": 2, "cancelled": 3, "rejected": 4}
            status_proto = status_map.get(status, 0)

        response = await self._stub.GetOrders(
            user_id=self._user_id,
            symbol=symbol or "",
            status=status_proto,
            limit=limit,
            metadata=self._get_metadata(),
        )
        return [self._parse_order(o) for o in response.orders]

    async def get_orderbook(
        self,
        symbol: str,
        depth: int = 20,
    ) -> OrderBook:
        """Get order book snapshot."""
        if not self._stub:
            raise ConnectionError("Not connected")

        response = await self._stub.GetOrderBook(
            symbol=symbol,
            depth=depth,
            metadata=self._get_metadata(),
        )

        return OrderBook(
            symbol=symbol,
            bids=[
                PriceLevel(price=b.price, size=b.size, count=b.count)
                for b in response.bids
            ],
            asks=[
                PriceLevel(price=a.price, size=a.size, count=a.count)
                for a in response.asks
            ],
            timestamp=response.timestamp,
        )

    async def get_positions(self) -> List[Position]:
        """Get all positions."""
        if not self._stub:
            raise ConnectionError("Not connected")

        response = await self._stub.GetPositions(
            user_id=self._user_id,
            metadata=self._get_metadata(),
        )

        return [
            Position(
                symbol=p.symbol,
                size=p.size,
                entry_price=p.entry_price,
                mark_price=p.mark_price,
                pnl=p.pnl,
                margin=p.margin,
            )
            for p in response.positions
        ]

    async def get_balance(self, asset: str) -> Optional[Balance]:
        """Get balance for an asset."""
        if not self._stub:
            raise ConnectionError("Not connected")

        try:
            response = await self._stub.GetBalance(
                user_id=self._user_id,
                asset=asset,
                metadata=self._get_metadata(),
            )
            return Balance(
                asset=response.asset,
                available=response.available,
                locked=response.locked,
                total=response.total,
            )
        except grpc.RpcError:
            return None

    async def stream_orderbook(
        self,
        symbol: str,
        depth: int = 20,
    ) -> AsyncIterator[OrderBook]:
        """Stream order book updates."""
        if not self._stub:
            raise ConnectionError("Not connected")

        async for update in self._stub.StreamOrderBook(
            symbol=symbol,
            depth=depth,
            metadata=self._get_metadata(),
        ):
            yield OrderBook(
                symbol=symbol,
                bids=[
                    PriceLevel(price=b.price, size=b.size, count=b.count)
                    for b in update.bid_updates
                ],
                asks=[
                    PriceLevel(price=a.price, size=a.size, count=a.count)
                    for a in update.ask_updates
                ],
                timestamp=update.timestamp,
            )

    async def stream_trades(
        self,
        symbol: str,
    ) -> AsyncIterator[Trade]:
        """Stream trades."""
        if not self._stub:
            raise ConnectionError("Not connected")

        async for trade in self._stub.StreamTrades(
            symbol=symbol,
            metadata=self._get_metadata(),
        ):
            yield Trade(
                trade_id=trade.trade_id,
                symbol=trade.symbol,
                price=trade.price,
                size=trade.size,
                side=OrderSide(_ORDER_SIDE_FROM_PROTO.get(trade.side, "buy")),
                buy_order_id=trade.buy_order_id,
                sell_order_id=trade.sell_order_id,
                buyer_id=trade.buyer_id,
                seller_id=trade.seller_id,
                timestamp=trade.timestamp,
            )

    async def ping(self) -> int:
        """Ping server and return latency in milliseconds."""
        if not self._stub:
            raise ConnectionError("Not connected")

        import time
        start = time.time()
        await self._stub.Ping(
            timestamp=int(start * 1000),
            metadata=self._get_metadata(),
        )
        return int((time.time() - start) * 1000)

    async def get_node_info(self) -> dict:
        """Get node information."""
        if not self._stub:
            raise ConnectionError("Not connected")

        response = await self._stub.GetNodeInfo(metadata=self._get_metadata())
        return {
            "node_id": response.node_id,
            "version": response.version,
            "network": response.network,
            "block_height": response.block_height,
            "order_count": response.order_count,
            "trade_count": response.trade_count,
            "uptime": response.uptime,
            "syncing": response.syncing,
            "supported_markets": list(response.supported_markets),
        }

    def _parse_order(self, proto_order) -> Order:
        """Parse order from protobuf message."""
        return Order(
            order_id=proto_order.order_id,
            symbol=proto_order.symbol,
            order_type=OrderType(_ORDER_TYPE_FROM_PROTO.get(proto_order.type, "limit")),
            side=OrderSide(_ORDER_SIDE_FROM_PROTO.get(proto_order.side, "buy")),
            price=proto_order.price,
            size=proto_order.size,
            filled=proto_order.filled,
            remaining=proto_order.remaining,
            status=OrderStatus(_ORDER_STATUS_FROM_PROTO.get(proto_order.status, "open")),
            user_id=proto_order.user_id,
            client_id=proto_order.client_id,
            timestamp=proto_order.timestamp,
            time_in_force=TimeInForce(
                {0: "gtc", 1: "ioc", 2: "fok", 3: "day"}.get(proto_order.time_in_force, "gtc")
            ),
            post_only=proto_order.post_only,
            reduce_only=proto_order.reduce_only,
        )


class _DynamicLXDEXStub:
    """
    Dynamic gRPC stub that works without pre-compiled proto files.

    Uses reflection or direct message construction to call gRPC methods.
    This allows the client to work before proto files are compiled.
    """

    def __init__(self, channel):
        self._channel = channel

    async def _call(self, method: str, request_class: str, **kwargs):
        """Make a unary-unary call."""
        metadata = kwargs.pop("metadata", [])

        # Try to import compiled proto first
        try:
            from lx_client import pb
            stub = pb.LXDEXServiceStub(self._channel)
            req_cls = getattr(pb, request_class)
            method_fn = getattr(stub, method)
            request = req_cls(**kwargs)
            return await method_fn(request, metadata=metadata)
        except ImportError:
            # Fall back to generic call
            pass

        # Generic protobuf construction
        try:
            from google.protobuf import descriptor_pool, message_factory
            # This would require proto reflection - simplified for now
            raise NotImplementedError("Proto reflection not implemented")
        except Exception:
            # Return mock response for development
            return _MockResponse(method, kwargs)

    async def _stream(self, method: str, request_class: str, **kwargs):
        """Make a unary-stream call."""
        metadata = kwargs.pop("metadata", [])

        try:
            from lx_client import pb
            stub = pb.LXDEXServiceStub(self._channel)
            req_cls = getattr(pb, request_class)
            method_fn = getattr(stub, method)
            request = req_cls(**kwargs)
            async for response in method_fn(request, metadata=metadata):
                yield response
        except ImportError:
            # Yield mock responses
            for _ in range(3):
                yield _MockResponse(method, kwargs)
                await asyncio.sleep(1)

    async def PlaceOrder(self, **kwargs):
        return await self._call("PlaceOrder", "PlaceOrderRequest", **kwargs)

    async def CancelOrder(self, **kwargs):
        return await self._call("CancelOrder", "CancelOrderRequest", **kwargs)

    async def GetOrder(self, **kwargs):
        return await self._call("GetOrder", "GetOrderRequest", **kwargs)

    async def GetOrders(self, **kwargs):
        return await self._call("GetOrders", "GetOrdersRequest", **kwargs)

    async def GetOrderBook(self, **kwargs):
        return await self._call("GetOrderBook", "GetOrderBookRequest", **kwargs)

    async def StreamOrderBook(self, **kwargs):
        async for resp in self._stream("StreamOrderBook", "StreamOrderBookRequest", **kwargs):
            yield resp

    async def GetTrades(self, **kwargs):
        return await self._call("GetTrades", "GetTradesRequest", **kwargs)

    async def StreamTrades(self, **kwargs):
        async for resp in self._stream("StreamTrades", "StreamTradesRequest", **kwargs):
            yield resp

    async def GetBalance(self, **kwargs):
        return await self._call("GetBalance", "GetBalanceRequest", **kwargs)

    async def GetPositions(self, **kwargs):
        return await self._call("GetPositions", "GetPositionsRequest", **kwargs)

    async def GetNodeInfo(self, **kwargs):
        return await self._call("GetNodeInfo", "GetNodeInfoRequest", **kwargs)

    async def GetPeers(self, **kwargs):
        return await self._call("GetPeers", "GetPeersRequest", **kwargs)

    async def Ping(self, **kwargs):
        return await self._call("Ping", "PingRequest", **kwargs)


class _MockResponse:
    """Mock response for development without compiled protos."""

    def __init__(self, method: str, kwargs: dict):
        self._method = method
        self._kwargs = kwargs

        # Set default attributes based on method
        if method == "PlaceOrder":
            self.order_id = 12345
            self.status = 0
            self.message = "Order placed"
        elif method == "CancelOrder":
            self.success = True
            self.message = "Order cancelled"
        elif method == "GetOrderBook":
            self.symbol = kwargs.get("symbol", "")
            self.bids = []
            self.asks = []
            self.timestamp = 0
        elif method == "GetOrders":
            self.orders = []
        elif method == "GetPositions":
            self.positions = []
        elif method == "GetBalance":
            self.asset = kwargs.get("asset", "")
            self.available = 0.0
            self.locked = 0.0
            self.total = 0.0
        elif method == "GetNodeInfo":
            self.node_id = ""
            self.version = "0.1.0"
            self.network = "mainnet"
            self.block_height = 0
            self.order_count = 0
            self.trade_count = 0
            self.uptime = 0
            self.syncing = False
            self.supported_markets = []
        elif method == "Ping":
            self.timestamp = 0
            self.message = "pong"
        elif method == "StreamOrderBook":
            self.bid_updates = []
            self.ask_updates = []
            self.timestamp = 0
        elif method == "StreamTrades":
            self.trade_id = 0
            self.symbol = kwargs.get("symbol", "")
            self.price = 0.0
            self.size = 0.0
            self.side = 0
            self.buy_order_id = 0
            self.sell_order_id = 0
            self.buyer_id = ""
            self.seller_id = ""
            self.timestamp = 0