gogo2/COBY/connectors/kucoin_connector.py

"""
KuCoin exchange connector implementation.
Supports WebSocket connections to KuCoin with proper token-based authentication.
"""

import json
import hmac
import hashlib
import base64
import time
from typing import Dict, List, Optional, Any
from datetime import datetime, timezone

from ..models.core import OrderBookSnapshot, TradeEvent, PriceLevel
from ..utils.logging import get_logger, set_correlation_id
from ..utils.exceptions import ValidationError, ConnectionError
from ..utils.validation import validate_symbol, validate_price, validate_volume
from .base_connector import BaseExchangeConnector

logger = get_logger(__name__)


class KuCoinConnector(BaseExchangeConnector):
    """
    KuCoin WebSocket connector implementation.
    
    Supports:
    - Token-based authentication
    - Order book streams
    - Trade streams
    - Symbol normalization
    - Bullet connection protocol
    """
    
    # KuCoin API URLs
    API_URL = "https://api.kucoin.com"
    SANDBOX_API_URL = "https://openapi-sandbox.kucoin.com"
    
    def __init__(self, use_sandbox: bool = False, api_key: str = None, 
                 api_secret: str = None, passphrase: str = None):
        """
        Initialize KuCoin connector.
        
        Args:
            use_sandbox: Whether to use sandbox environment
            api_key: API key for authentication (optional)
            api_secret: API secret for authentication (optional)
            passphrase: API passphrase for authentication (optional)
        """
        # KuCoin requires getting WebSocket URL from REST API
        super().__init__("kucoin", "")  # URL will be set after token retrieval
        
        # Authentication credentials (optional)
        self.api_key = api_key
        self.api_secret = api_secret
        self.passphrase = passphrase
        self.use_sandbox = use_sandbox
        
        # KuCoin-specific attributes
        self.token = None
        self.connect_id = None
        self.ping_interval = 18000  # 18 seconds (KuCoin requirement)
        self.ping_timeout = 10000   # 10 seconds
        
        # KuCoin-specific message handlers
        self.message_handlers.update({
            'message': self._handle_data_message,
            'welcome': self._handle_welcome_message,
            'ack': self._handle_ack_message,
            'error': self._handle_error_message,
            'pong': self._handle_pong_message
        })
        
        # Subscription tracking
        self.subscribed_topics = set()
        self.subscription_id = 1
        
        logger.info(f"KuCoin connector initialized ({'sandbox' if use_sandbox else 'live'})")
    
    def _get_message_type(self, data: Dict) -> str:
        """
        Determine message type from KuCoin message data.
        
        Args:
            data: Parsed message data
            
        Returns:
            str: Message type identifier
        """
        # KuCoin message format
        if 'type' in data:
            return data['type']  # 'message', 'welcome', 'ack', 'error', 'pong'
        elif 'subject' in data:
            # Data message with subject
            return 'message'
        
        return 'unknown'
    
    def normalize_symbol(self, symbol: str) -> str:
        """
        Normalize symbol to KuCoin format.
        
        Args:
            symbol: Standard symbol format (e.g., 'BTCUSDT')
            
        Returns:
            str: KuCoin symbol format (e.g., 'BTC-USDT')
        """
        # KuCoin uses dash-separated format
        if symbol.upper() == 'BTCUSDT':
            return 'BTC-USDT'
        elif symbol.upper() == 'ETHUSDT':
            return 'ETH-USDT'
        elif symbol.upper().endswith('USDT'):
            base = symbol[:-4].upper()
            return f"{base}-USDT"
        elif symbol.upper().endswith('USD'):
            base = symbol[:-3].upper()
            return f"{base}-USD"
        else:
            # Assume it's already in correct format or add dash
            if '-' not in symbol:
                # Try to split common patterns
                if len(symbol) >= 6:
                    # Assume last 4 chars are quote currency
                    base = symbol[:-4].upper()
                    quote = symbol[-4:].upper()
                    return f"{base}-{quote}"
                else:
                    return symbol.upper()
            else:
                return symbol.upper()
    
    def _denormalize_symbol(self, kucoin_symbol: str) -> str:
        """
        Convert KuCoin symbol back to standard format.
        
        Args:
            kucoin_symbol: KuCoin symbol format (e.g., 'BTC-USDT')
            
        Returns:
            str: Standard symbol format (e.g., 'BTCUSDT')
        """
        if '-' in kucoin_symbol:
            return kucoin_symbol.replace('-', '')
        return kucoin_symbol
    
    async def _get_websocket_token(self) -> Optional[Dict[str, Any]]:
        """
        Get WebSocket connection token from KuCoin REST API.
        
        Returns:
            Dict: Token information including WebSocket URL
        """
        try:
            import aiohttp
            
            api_url = self.SANDBOX_API_URL if self.use_sandbox else self.API_URL
            endpoint = "/api/v1/bullet-public"
            
            # Use private endpoint if authenticated
            if self.api_key and self.api_secret and self.passphrase:
                endpoint = "/api/v1/bullet-private"
                headers = self._get_auth_headers("POST", endpoint, "")
            else:
                headers = {}
            
            async with aiohttp.ClientSession() as session:
                async with session.post(f"{api_url}{endpoint}", headers=headers) as response:
                    if response.status == 200:
                        data = await response.json()
                        
                        if data.get('code') != '200000':
                            logger.error(f"KuCoin token error: {data.get('msg')}")
                            return None
                        
                        return data.get('data')
                    else:
                        logger.error(f"Failed to get KuCoin token: HTTP {response.status}")
                        return None
                        
        except Exception as e:
            logger.error(f"Error getting KuCoin WebSocket token: {e}")
            return None
    
    async def connect(self) -> bool:
        """Override connect to get token first."""
        try:
            # Get WebSocket token and URL
            token_data = await self._get_websocket_token()
            if not token_data:
                logger.error("Failed to get KuCoin WebSocket token")
                return False
            
            self.token = token_data.get('token')
            servers = token_data.get('instanceServers', [])
            
            if not servers:
                logger.error("No KuCoin WebSocket servers available")
                return False
            
            # Use first available server
            server = servers[0]
            self.websocket_url = f"{server['endpoint']}?token={self.token}&connectId={int(time.time() * 1000)}"
            self.ping_interval = server.get('pingInterval', 18000)
            self.ping_timeout = server.get('pingTimeout', 10000)
            
            logger.info(f"KuCoin WebSocket URL: {server['endpoint']}")
            
            # Now connect using the base connector method
            return await super().connect()
            
        except Exception as e:
            logger.error(f"Error connecting to KuCoin: {e}")
            return False
    
    async def subscribe_orderbook(self, symbol: str) -> None:
        """
        Subscribe to order book updates for a symbol.
        
        Args:
            symbol: Trading symbol (e.g., 'BTCUSDT')
        """
        try:
            set_correlation_id()
            kucoin_symbol = self.normalize_symbol(symbol)
            topic = f"/market/level2:{kucoin_symbol}"
            
            # Create subscription message
            subscription_msg = {
                "id": str(self.subscription_id),
                "type": "subscribe",
                "topic": topic,
                "privateChannel": False,
                "response": True
            }
            self.subscription_id += 1
            
            # Send subscription
            success = await self._send_message(subscription_msg)
            if success:
                # Track subscription
                if symbol not in self.subscriptions:
                    self.subscriptions[symbol] = []
                if 'orderbook' not in self.subscriptions[symbol]:
                    self.subscriptions[symbol].append('orderbook')
                
                self.subscribed_topics.add(topic)
                
                logger.info(f"Subscribed to order book for {symbol} ({kucoin_symbol}) on KuCoin")
            else:
                logger.error(f"Failed to subscribe to order book for {symbol} on KuCoin")
                
        except Exception as e:
            logger.error(f"Error subscribing to order book for {symbol}: {e}")
            raise
    
    async def subscribe_trades(self, symbol: str) -> None:
        """
        Subscribe to trade updates for a symbol.
        
        Args:
            symbol: Trading symbol (e.g., 'BTCUSDT')
        """
        try:
            set_correlation_id()
            kucoin_symbol = self.normalize_symbol(symbol)
            topic = f"/market/match:{kucoin_symbol}"
            
            # Create subscription message
            subscription_msg = {
                "id": str(self.subscription_id),
                "type": "subscribe",
                "topic": topic,
                "privateChannel": False,
                "response": True
            }
            self.subscription_id += 1
            
            # Send subscription
            success = await self._send_message(subscription_msg)
            if success:
                # Track subscription
                if symbol not in self.subscriptions:
                    self.subscriptions[symbol] = []
                if 'trades' not in self.subscriptions[symbol]:
                    self.subscriptions[symbol].append('trades')
                
                self.subscribed_topics.add(topic)
                
                logger.info(f"Subscribed to trades for {symbol} ({kucoin_symbol}) on KuCoin")
            else:
                logger.error(f"Failed to subscribe to trades for {symbol} on KuCoin")
                
        except Exception as e:
            logger.error(f"Error subscribing to trades for {symbol}: {e}")
            raise
    
    async def unsubscribe_orderbook(self, symbol: str) -> None:
        """
        Unsubscribe from order book updates for a symbol.
        
        Args:
            symbol: Trading symbol (e.g., 'BTCUSDT')
        """
        try:
            kucoin_symbol = self.normalize_symbol(symbol)
            topic = f"/market/level2:{kucoin_symbol}"
            
            # Create unsubscription message
            unsubscription_msg = {
                "id": str(self.subscription_id),
                "type": "unsubscribe",
                "topic": topic,
                "privateChannel": False,
                "response": True
            }
            self.subscription_id += 1
            
            # Send unsubscription
            success = await self._send_message(unsubscription_msg)
            if success:
                # Remove from tracking
                if symbol in self.subscriptions and 'orderbook' in self.subscriptions[symbol]:
                    self.subscriptions[symbol].remove('orderbook')
                    if not self.subscriptions[symbol]:
                        del self.subscriptions[symbol]
                
                self.subscribed_topics.discard(topic)
                
                logger.info(f"Unsubscribed from order book for {symbol} ({kucoin_symbol}) on KuCoin")
            else:
                logger.error(f"Failed to unsubscribe from order book for {symbol} on KuCoin")
                
        except Exception as e:
            logger.error(f"Error unsubscribing from order book for {symbol}: {e}")
            raise
    
    async def unsubscribe_trades(self, symbol: str) -> None:
        """
        Unsubscribe from trade updates for a symbol.
        
        Args:
            symbol: Trading symbol (e.g., 'BTCUSDT')
        """
        try:
            kucoin_symbol = self.normalize_symbol(symbol)
            topic = f"/market/match:{kucoin_symbol}"
            
            # Create unsubscription message
            unsubscription_msg = {
                "id": str(self.subscription_id),
                "type": "unsubscribe",
                "topic": topic,
                "privateChannel": False,
                "response": True
            }
            self.subscription_id += 1
            
            # Send unsubscription
            success = await self._send_message(unsubscription_msg)
            if success:
                # Remove from tracking
                if symbol in self.subscriptions and 'trades' in self.subscriptions[symbol]:
                    self.subscriptions[symbol].remove('trades')
                    if not self.subscriptions[symbol]:
                        del self.subscriptions[symbol]
                
                self.subscribed_topics.discard(topic)
                
                logger.info(f"Unsubscribed from trades for {symbol} ({kucoin_symbol}) on KuCoin")
            else:
                logger.error(f"Failed to unsubscribe from trades for {symbol} on KuCoin")
                
        except Exception as e:
            logger.error(f"Error unsubscribing from trades for {symbol}: {e}")
            raise
    
    async def get_symbols(self) -> List[str]:
        """
        Get list of available trading symbols from KuCoin.
        
        Returns:
            List[str]: List of available symbols in standard format
        """
        try:
            import aiohttp
            
            api_url = self.SANDBOX_API_URL if self.use_sandbox else self.API_URL
            
            async with aiohttp.ClientSession() as session:
                async with session.get(f"{api_url}/api/v1/symbols") as response:
                    if response.status == 200:
                        data = await response.json()
                        
                        if data.get('code') != '200000':
                            logger.error(f"KuCoin API error: {data.get('msg')}")
                            return []
                        
                        symbols = []
                        symbol_data = data.get('data', [])
                        
                        for symbol_info in symbol_data:
                            if symbol_info.get('enableTrading'):
                                symbol = symbol_info.get('symbol', '')
                                # Convert to standard format
                                standard_symbol = self._denormalize_symbol(symbol)
                                symbols.append(standard_symbol)
                        
                        logger.info(f"Retrieved {len(symbols)} symbols from KuCoin")
                        return symbols
                    else:
                        logger.error(f"Failed to get symbols from KuCoin: HTTP {response.status}")
                        return []
                        
        except Exception as e:
            logger.error(f"Error getting symbols from KuCoin: {e}")
            return []
    
    async def get_orderbook_snapshot(self, symbol: str, depth: int = 20) -> Optional[OrderBookSnapshot]:
        """
        Get current order book snapshot from KuCoin REST API.
        
        Args:
            symbol: Trading symbol
            depth: Number of price levels to retrieve
            
        Returns:
            OrderBookSnapshot: Current order book or None if unavailable
        """
        try:
            import aiohttp
            
            kucoin_symbol = self.normalize_symbol(symbol)
            api_url = self.SANDBOX_API_URL if self.use_sandbox else self.API_URL
            
            url = f"{api_url}/api/v1/market/orderbook/level2_20"
            params = {'symbol': kucoin_symbol}
            
            async with aiohttp.ClientSession() as session:
                async with session.get(url, params=params) as response:
                    if response.status == 200:
                        data = await response.json()
                        
                        if data.get('code') != '200000':
                            logger.error(f"KuCoin API error: {data.get('msg')}")
                            return None
                        
                        result = data.get('data', {})
                        return self._parse_orderbook_snapshot(result, symbol)
                    else:
                        logger.error(f"Failed to get order book for {symbol}: HTTP {response.status}")
                        return None
                        
        except Exception as e:
            logger.error(f"Error getting order book snapshot for {symbol}: {e}")
            return None   
 
    def _parse_orderbook_snapshot(self, data: Dict, symbol: str) -> OrderBookSnapshot:
        """
        Parse KuCoin order book data into OrderBookSnapshot.
        
        Args:
            data: Raw KuCoin order book data
            symbol: Trading symbol
            
        Returns:
            OrderBookSnapshot: Parsed order book
        """
        try:
            # Parse bids and asks
            bids = []
            for bid_data in data.get('bids', []):
                price = float(bid_data[0])
                size = float(bid_data[1])
                
                if validate_price(price) and validate_volume(size):
                    bids.append(PriceLevel(price=price, size=size))
            
            asks = []
            for ask_data in data.get('asks', []):
                price = float(ask_data[0])
                size = float(ask_data[1])
                
                if validate_price(price) and validate_volume(size):
                    asks.append(PriceLevel(price=price, size=size))
            
            # Create order book snapshot
            orderbook = OrderBookSnapshot(
                symbol=symbol,
                exchange=self.exchange_name,
                timestamp=datetime.fromtimestamp(int(data.get('time', 0)) / 1000, tz=timezone.utc),
                bids=bids,
                asks=asks,
                sequence_id=int(data.get('sequence', 0))
            )
            
            return orderbook
            
        except Exception as e:
            logger.error(f"Error parsing order book snapshot: {e}")
            raise ValidationError(f"Invalid order book data: {e}", "PARSE_ERROR")
    
    async def _handle_data_message(self, data: Dict) -> None:
        """
        Handle data message from KuCoin.
        
        Args:
            data: Data message
        """
        try:
            set_correlation_id()
            
            subject = data.get('subject', '')
            topic = data.get('topic', '')
            message_data = data.get('data', {})
            
            if 'level2' in subject:
                await self._handle_orderbook_update(data)
            elif 'match' in subject:
                await self._handle_trade_update(data)
            else:
                logger.debug(f"Unhandled KuCoin subject: {subject}")
                
        except Exception as e:
            logger.error(f"Error handling data message: {e}")
    
    async def _handle_orderbook_update(self, data: Dict) -> None:
        """
        Handle order book update from KuCoin.
        
        Args:
            data: Order book update data
        """
        try:
            topic = data.get('topic', '')
            if not topic:
                logger.warning("Order book update missing topic")
                return
            
            # Extract symbol from topic: /market/level2:BTC-USDT
            parts = topic.split(':')
            if len(parts) < 2:
                logger.warning("Invalid order book topic format")
                return
            
            kucoin_symbol = parts[1]
            symbol = self._denormalize_symbol(kucoin_symbol)
            
            message_data = data.get('data', {})
            changes = message_data.get('changes', {})
            
            # Parse bids and asks changes
            bids = []
            for bid_data in changes.get('bids', []):
                price = float(bid_data[0])
                size = float(bid_data[1])
                
                if validate_price(price) and validate_volume(size):
                    bids.append(PriceLevel(price=price, size=size))
            
            asks = []
            for ask_data in changes.get('asks', []):
                price = float(ask_data[0])
                size = float(ask_data[1])
                
                if validate_price(price) and validate_volume(size):
                    asks.append(PriceLevel(price=price, size=size))
            
            # Create order book snapshot
            orderbook = OrderBookSnapshot(
                symbol=symbol,
                exchange=self.exchange_name,
                timestamp=datetime.fromtimestamp(int(message_data.get('time', 0)) / 1000, tz=timezone.utc),
                bids=bids,
                asks=asks,
                sequence_id=int(message_data.get('sequenceEnd', 0))
            )
            
            # Notify callbacks
            self._notify_data_callbacks(orderbook)
            
            logger.debug(f"Processed order book update for {symbol}")
            
        except Exception as e:
            logger.error(f"Error handling order book update: {e}")
    
    async def _handle_trade_update(self, data: Dict) -> None:
        """
        Handle trade update from KuCoin.
        
        Args:
            data: Trade update data
        """
        try:
            topic = data.get('topic', '')
            if not topic:
                logger.warning("Trade update missing topic")
                return
            
            # Extract symbol from topic: /market/match:BTC-USDT
            parts = topic.split(':')
            if len(parts) < 2:
                logger.warning("Invalid trade topic format")
                return
            
            kucoin_symbol = parts[1]
            symbol = self._denormalize_symbol(kucoin_symbol)
            
            message_data = data.get('data', {})
            
            price = float(message_data.get('price', 0))
            size = float(message_data.get('size', 0))
            
            # Validate data
            if not validate_price(price) or not validate_volume(size):
                logger.warning(f"Invalid trade data: price={price}, size={size}")
                return
            
            # Determine side (KuCoin uses 'side' field)
            side = message_data.get('side', 'unknown').lower()
            
            # Create trade event
            trade = TradeEvent(
                symbol=symbol,
                exchange=self.exchange_name,
                timestamp=datetime.fromtimestamp(int(message_data.get('time', 0)) / 1000, tz=timezone.utc),
                price=price,
                size=size,
                side=side,
                trade_id=str(message_data.get('tradeId', ''))
            )
            
            # Notify callbacks
            self._notify_data_callbacks(trade)
            
            logger.debug(f"Processed trade for {symbol}: {side} {size} @ {price}")
            
        except Exception as e:
            logger.error(f"Error handling trade update: {e}")
    
    async def _handle_welcome_message(self, data: Dict) -> None:
        """
        Handle welcome message from KuCoin.
        
        Args:
            data: Welcome message data
        """
        try:
            connect_id = data.get('id')
            if connect_id:
                self.connect_id = connect_id
                logger.info(f"KuCoin connection established with ID: {connect_id}")
            
        except Exception as e:
            logger.error(f"Error handling welcome message: {e}")
    
    async def _handle_ack_message(self, data: Dict) -> None:
        """
        Handle acknowledgment message from KuCoin.
        
        Args:
            data: Ack message data
        """
        try:
            msg_id = data.get('id', '')
            logger.debug(f"KuCoin ACK received for message ID: {msg_id}")
            
        except Exception as e:
            logger.error(f"Error handling ack message: {e}")
    
    async def _handle_error_message(self, data: Dict) -> None:
        """
        Handle error message from KuCoin.
        
        Args:
            data: Error message data
        """
        try:
            code = data.get('code', 'unknown')
            message = data.get('data', 'Unknown error')
            
            logger.error(f"KuCoin error {code}: {message}")
            
        except Exception as e:
            logger.error(f"Error handling error message: {e}")
    
    async def _handle_pong_message(self, data: Dict) -> None:
        """
        Handle pong message from KuCoin.
        
        Args:
            data: Pong message data
        """
        logger.debug("Received KuCoin pong")
    
    def _get_auth_headers(self, method: str, endpoint: str, body: str) -> Dict[str, str]:
        """
        Generate authentication headers for KuCoin API.
        
        Args:
            method: HTTP method
            endpoint: API endpoint
            body: Request body
            
        Returns:
            Dict: Authentication headers
        """
        if not all([self.api_key, self.api_secret, self.passphrase]):
            return {}
        
        try:
            timestamp = str(int(time.time() * 1000))
            
            # Create signature string
            str_to_sign = timestamp + method + endpoint + body
            signature = base64.b64encode(
                hmac.new(
                    self.api_secret.encode('utf-8'),
                    str_to_sign.encode('utf-8'),
                    hashlib.sha256
                ).digest()
            ).decode('utf-8')
            
            # Create passphrase signature
            passphrase_signature = base64.b64encode(
                hmac.new(
                    self.api_secret.encode('utf-8'),
                    self.passphrase.encode('utf-8'),
                    hashlib.sha256
                ).digest()
            ).decode('utf-8')
            
            return {
                'KC-API-SIGN': signature,
                'KC-API-TIMESTAMP': timestamp,
                'KC-API-KEY': self.api_key,
                'KC-API-PASSPHRASE': passphrase_signature,
                'KC-API-KEY-VERSION': '2',
                'Content-Type': 'application/json'
            }
            
        except Exception as e:
            logger.error(f"Error generating auth headers: {e}")
            return {}
    
    async def _send_ping(self) -> None:
        """Send ping to keep connection alive."""
        try:
            ping_msg = {
                "id": str(self.subscription_id),
                "type": "ping"
            }
            self.subscription_id += 1
            
            await self._send_message(ping_msg)
            logger.debug("Sent ping to KuCoin")
            
        except Exception as e:
            logger.error(f"Error sending ping: {e}")
    
    def get_kucoin_stats(self) -> Dict[str, Any]:
        """Get KuCoin-specific statistics."""
        base_stats = self.get_stats()
        
        kucoin_stats = {
            'subscribed_topics': list(self.subscribed_topics),
            'use_sandbox': self.use_sandbox,
            'authenticated': bool(self.api_key and self.api_secret and self.passphrase),
            'connect_id': self.connect_id,
            'token_available': bool(self.token),
            'next_subscription_id': self.subscription_id
        }
        
        base_stats.update(kucoin_stats)
        return base_stats