#!/usr/bin/env python3
"""
Huobi (HTX) COB WebSocket Connector

Low-latency order book stream for additional market depth aggregation.
Features:
- GZIP-compressed message handling
- Ping/pong keepalive
- Automatic reconnection with backoff
- Per-symbol tasks; safe for Windows

This module emits standardized COB snapshots via callbacks:
  { 'symbol': 'ETH/USDT', 'bids': [{'price': .., 'size': ..}], 'asks': [...],
    'exchange': 'huobi', 'source': 'huobi_ws', 'timestamp': datetime }
"""

import asyncio
import json
import logging
import time
from datetime import datetime
from typing import Callable, Dict, List, Optional

import gzip

try:
    import websockets
    from websockets.client import connect as websockets_connect
    WEBSOCKETS_AVAILABLE = True
except Exception:
    websockets = None
    websockets_connect = None
    WEBSOCKETS_AVAILABLE = False

logger = logging.getLogger(__name__)


class HuobiCOBWebSocket:
    """Minimal Huobi order book WebSocket for depth updates."""

    def __init__(self, symbols: Optional[List[str]] = None):
        # Expect symbols like 'ETH/USDT', 'BTC/USDT'
        self.symbols = symbols or ['ETH/USDT', 'BTC/USDT']
        self.ws_tasks: Dict[str, asyncio.Task] = {}
        self.callbacks: List[Callable] = []
        self._stopping = False

    def add_cob_callback(self, callback: Callable):
        self.callbacks.append(callback)

    async def start(self):
        if not WEBSOCKETS_AVAILABLE:
            logger.warning("Huobi WS not available (websockets missing)")
            return
        for symbol in self.symbols:
            if symbol not in self.ws_tasks or self.ws_tasks[symbol].done():
                self.ws_tasks[symbol] = asyncio.create_task(self._run_symbol(symbol))

    async def stop(self):
        self._stopping = True
        for _, task in list(self.ws_tasks.items()):
            if task and not task.done():
                task.cancel()
                try:
                    await task
                except asyncio.CancelledError:
                    pass

    def _to_huobi_symbol(self, symbol: str) -> str:
        # 'ETH/USDT' -> 'ethusdt'
        return symbol.replace('/', '').lower()

    async def _run_symbol(self, symbol: str):
        huobi_symbol = self._to_huobi_symbol(symbol)
        sub_msg = {
            "sub": f"market.{huobi_symbol}.depth.step0",
            "id": f"sub_{huobi_symbol}_{int(time.time())}"
        }
        url = "wss://api.huobi.pro/ws"

        backoff = 1
        while not self._stopping:
            try:
                logger.info(f"Huobi: connecting WS for {symbol} -> {url}")
                async with websockets_connect(url, ping_interval=20, ping_timeout=60, close_timeout=10) as ws:
                    # Subscribe
                    await ws.send(json.dumps(sub_msg))
                    logger.info(f"Huobi: subscribed {sub_msg}")
                    backoff = 1

                    while not self._stopping:
                        raw = await ws.recv()
                        # Huobi sends gzip-compressed bytes
                        if isinstance(raw, (bytes, bytearray)):
                            try:
                                message = gzip.decompress(raw).decode('utf-8')
                            except Exception:
                                # Some servers send raw JSON; try decode directly
                                try:
                                    message = raw.decode('utf-8')
                                except Exception:
                                    continue
                        else:
                            message = raw

                        try:
                            data = json.loads(message)
                        except Exception:
                            continue

                        # Ping/pong
                        if 'ping' in data:
                            await ws.send(json.dumps({"pong": data['ping']}))
                            continue

                        # Depth update
                        if 'tick' in data and 'ch' in data:
                            tick = data['tick'] or {}
                            bids = tick.get('bids', []) or []
                            asks = tick.get('asks', []) or []

                            # Convert to standardized structure
                            std_bids = []
                            std_asks = []
                            for b in bids[:1000]:
                                try:
                                    price = float(b[0]); size = float(b[1])
                                    if size > 0:
                                        std_bids.append({'price': price, 'size': size})
                                except Exception:
                                    continue
                            for a in asks[:1000]:
                                try:
                                    price = float(a[0]); size = float(a[1])
                                    if size > 0:
                                        std_asks.append({'price': price, 'size': size})
                                except Exception:
                                    continue

                            if std_bids:
                                std_bids.sort(key=lambda x: x['price'], reverse=True)
                            if std_asks:
                                std_asks.sort(key=lambda x: x['price'])

                            cob = {
                                'symbol': symbol,
                                'timestamp': datetime.now(),
                                'bids': std_bids,
                                'asks': std_asks,
                                'exchange': 'huobi',
                                'source': 'huobi_ws'
                            }

                            # Stats
                            if std_bids and std_asks:
                                best_bid = std_bids[0]['price']
                                best_ask = std_asks[0]['price']
                                mid = (best_bid + best_ask) / 2.0
                                spread = best_ask - best_bid
                                spread_bps = (spread / mid) * 10000 if mid > 0 else 0
                                top_bids = std_bids[:20]; top_asks = std_asks[:20]
                                bid_vol = sum(x['price'] * x['size'] for x in top_bids)
                                ask_vol = sum(x['price'] * x['size'] for x in top_asks)
                                tot = bid_vol + ask_vol
                                cob['stats'] = {
                                    'best_bid': best_bid,
                                    'best_ask': best_ask,
                                    'mid_price': mid,
                                    'spread': spread,
                                    'spread_bps': spread_bps,
                                    'bid_volume': bid_vol,
                                    'ask_volume': ask_vol,
                                    'imbalance': (bid_vol - ask_vol) / tot if tot > 0 else 0.0,
                                }

                            # Notify
                            for cb in self.callbacks:
                                try:
                                    await cb(symbol, cob)
                                except Exception as cb_ex:
                                    logger.debug(f"Huobi callback error: {cb_ex}")

            except asyncio.CancelledError:
                break
            except Exception as e:
                logger.warning(f"Huobi WS error for {symbol}: {e}")
                await asyncio.sleep(backoff)
                backoff = min(backoff * 2, 60)


huobi_cob_websocket: Optional[HuobiCOBWebSocket] = None


async def get_huobi_cob_websocket(symbols: Optional[List[str]] = None) -> HuobiCOBWebSocket:
    global huobi_cob_websocket
    if huobi_cob_websocket is None:
        huobi_cob_websocket = HuobiCOBWebSocket(symbols)
        await huobi_cob_websocket.start()
    return huobi_cob_websocket


async def stop_huobi_cob_websocket():
    global huobi_cob_websocket
    if huobi_cob_websocket:
        await huobi_cob_websocket.stop()
        huobi_cob_websocket = None