digdigdig3 0.1.12

//! # Hyperliquid WebSocket Implementation
//!
//! WebSocket connector with auto-reconnection and full event support.
//!
//! ## Features
//!
//! - Auto-reconnect on disconnect
//! - Snapshot + incremental update handling
//! - 19 subscription types supported
//! - Ping/pong heartbeat handling
//! - Broadcast channel for multiple consumers

use std::collections::HashSet;
use std::pin::Pin;
use std::sync::{Arc, Mutex as StdMutex};
use std::time::{Duration, Instant};

use async_trait::async_trait;
use futures_util::{Stream, StreamExt, SinkExt};
use serde::{Deserialize, Serialize};
use serde_json::{json, Value};
use tokio::sync::{mpsc, broadcast, Mutex};
use tokio::time::sleep;
use tokio_tungstenite::{connect_async, tungstenite::Message, WebSocketStream, MaybeTlsStream};

use crate::core::{
    AccountType, ConnectionStatus, StreamEvent, StreamType,
    SubscriptionRequest,
};
use crate::core::types::{WebSocketResult, WebSocketError};
use crate::core::traits::WebSocketConnector;

use super::{HyperliquidUrls, HyperliquidParser};

// ═══════════════════════════════════════════════════════════════════════════════
// WEBSOCKET MESSAGES
// ═══════════════════════════════════════════════════════════════════════════════

/// Outgoing subscription message
#[derive(Debug, Clone, Serialize)]
struct SubscribeMessage {
    method: String,
    subscription: Value,
}

/// Incoming message from Hyperliquid
#[derive(Debug, Clone, Deserialize)]
struct IncomingMessage {
    channel: Option<String>,
    data: Option<Value>,
}

// ═══════════════════════════════════════════════════════════════════════════════
// HYPERLIQUID WEBSOCKET CONNECTOR
// ═══════════════════════════════════════════════════════════════════════════════

type WsStream = WebSocketStream<MaybeTlsStream<tokio::net::TcpStream>>;

/// Hyperliquid WebSocket connector
pub struct HyperliquidWebSocket {
    /// WebSocket URLs
    urls: HyperliquidUrls,
    /// Connection status
    status: Arc<Mutex<ConnectionStatus>>,
    /// Active subscriptions
    subscriptions: Arc<Mutex<HashSet<SubscriptionRequest>>>,
    /// Event sender (internal - for message handler)
    event_tx: Arc<Mutex<Option<mpsc::UnboundedSender<WebSocketResult<StreamEvent>>>>>,
    /// Broadcast sender (for multiple consumers, dropped on disconnect)
    broadcast_tx: Arc<StdMutex<Option<broadcast::Sender<WebSocketResult<StreamEvent>>>>>,
    /// WebSocket stream
    ws_stream: Arc<Mutex<Option<WsStream>>>,
    /// Last ping time
    last_ping: Arc<Mutex<Instant>>,
    /// Most recent ping round-trip time in milliseconds (0 until first pong)
    ws_ping_rtt_ms: Arc<Mutex<u64>>,
}

impl HyperliquidWebSocket {
    /// Create new WebSocket connector
    pub fn new(is_testnet: bool) -> Self {
        let urls = if is_testnet {
            HyperliquidUrls::TESTNET
        } else {
            HyperliquidUrls::MAINNET
        };

        Self {
            urls,
            status: Arc::new(Mutex::new(ConnectionStatus::Disconnected)),
            subscriptions: Arc::new(Mutex::new(HashSet::new())),
            event_tx: Arc::new(Mutex::new(None)),
            broadcast_tx: Arc::new(StdMutex::new(None)),
            ws_stream: Arc::new(Mutex::new(None)),
            last_ping: Arc::new(Mutex::new(Instant::now())),
            ws_ping_rtt_ms: Arc::new(Mutex::new(0)),
        }
    }

    /// Create public WebSocket connector (convenience method)
    pub fn public(is_testnet: bool) -> Self {
        Self::new(is_testnet)
    }

    /// Connect to WebSocket
    async fn connect_ws(&self) -> WebSocketResult<WsStream> {
        let ws_url = self.urls.ws_url();

        let (ws_stream, _) = connect_async(ws_url).await
            .map_err(|e| WebSocketError::ConnectionError(format!("WebSocket connection failed: {}", e)))?;

        Ok(ws_stream)
    }

    /// Start message handling task
    fn start_message_handler(
        ws_stream: Arc<Mutex<Option<WsStream>>>,
        event_tx: mpsc::UnboundedSender<WebSocketResult<StreamEvent>>,
        status: Arc<Mutex<ConnectionStatus>>,
        last_ping: Arc<Mutex<Instant>>,
        ws_ping_rtt_ms: Arc<Mutex<u64>>,
    ) {
        tokio::spawn(async move {
            loop {
                let mut stream_guard = ws_stream.lock().await;
                let stream = match stream_guard.as_mut() {
                    Some(s) => s,
                    None => {
                        drop(stream_guard);
                        sleep(Duration::from_millis(100)).await;
                        continue;
                    }
                };

                match stream.next().await {
                    Some(Ok(Message::Text(text))) => {
                        drop(stream_guard);
                        if let Err(e) = Self::handle_message(&text, &event_tx).await {
                            let _ = event_tx.send(Err(e));
                        }
                    }
                    Some(Ok(Message::Pong(_))) => {
                        // Response to our client-initiated WS Ping frame — measure RTT
                        let rtt = last_ping.lock().await.elapsed().as_millis() as u64;
                        *ws_ping_rtt_ms.lock().await = rtt;
                        drop(stream_guard);
                    }
                    Some(Ok(Message::Ping(data))) => {
                        // Respond to server-initiated ping with pong
                        if let Err(e) = stream.send(Message::Pong(data)).await {
                            drop(stream_guard);
                            let _ = event_tx.send(Err(WebSocketError::ConnectionError(e.to_string())));
                            break;
                        }
                        drop(stream_guard);
                    }
                    Some(Ok(Message::Close(_))) => {
                        drop(stream_guard);
                        *status.lock().await = ConnectionStatus::Disconnected;
                        break;
                    }
                    Some(Err(e)) => {
                        drop(stream_guard);
                        let _ = event_tx.send(Err(WebSocketError::ConnectionError(e.to_string())));
                        break;
                    }
                    None => {
                        drop(stream_guard);
                        *status.lock().await = ConnectionStatus::Disconnected;
                        break;
                    }
                    _ => {
                        drop(stream_guard);
                    }
                }
            }
        });
    }

    /// Handle incoming WebSocket message
    async fn handle_message(
        text: &str,
        event_tx: &mpsc::UnboundedSender<WebSocketResult<StreamEvent>>,
    ) -> WebSocketResult<()> {
        let msg: IncomingMessage = serde_json::from_str(text)
            .map_err(|e| WebSocketError::Parse(format!("Failed to parse message: {}", e)))?;

        // Get channel and data
        let channel = match msg.channel {
            Some(ch) => ch,
            None => return Ok(()), // Ignore messages without channel
        };

        let data = match msg.data {
            Some(d) => d,
            None => return Ok(()), // Ignore messages without data
        };

        // Parse based on channel type
        match channel.as_str() {
            "activeAssetCtx" => {
                if let Some(event) = Self::parse_active_asset_ctx(&data)? {
                    let _ = event_tx.send(Ok(event));
                }
            }
            "allMids" => {
                if let Some(event) = Self::parse_all_mids(&data)? {
                    let _ = event_tx.send(Ok(event));
                }
            }
            "trades" => {
                if let Some(event) = Self::parse_trades(&data)? {
                    let _ = event_tx.send(Ok(event));
                }
            }
            "l2Book" => {
                if let Some(event) = Self::parse_l2_book(&data)? {
                    let _ = event_tx.send(Ok(event));
                }
            }
            "candle" => {
                if let Some(event) = Self::parse_candle(&data)? {
                    let _ = event_tx.send(Ok(event));
                }
            }
            "subscriptionResponse" => {
                // Subscription confirmed - ignore
            }
            "error" => {
                // Error message
                let error_msg = data.get("error")
                    .and_then(|e| e.as_str())
                    .unwrap_or("Unknown error");
                return Err(WebSocketError::ProtocolError(error_msg.to_string()));
            }
            _ => {
                // Unknown channel - ignore for now
            }
        }

        Ok(())
    }

    /// Parse activeAssetCtx message to Ticker event.
    ///
    /// This channel provides per-coin 24h stats including dayNtlVlm, prevDayPx,
    /// markPx, and midPx — far richer than allMids which only has mid-prices.
    ///
    /// Message format:
    /// ```json
    /// {
    ///   "coin": "BTC",
    ///   "ctx": {
    ///     "dayNtlVlm": "1234567890.5",
    ///     "funding": "0.000012345",
    ///     "openInterest": "987654.321",
    ///     "prevDayPx": "49500.0",
    ///     "markPx": "50123.45",
    ///     "midPx": "50123.5",
    ///     "impactPxs": ["50120.0", "50127.0"],
    ///     "premium": "0.5",
    ///     "oraclePx": "50122.95"
    ///   }
    /// }
    /// ```
    fn parse_active_asset_ctx(data: &Value) -> WebSocketResult<Option<StreamEvent>> {
        let coin = data.get("coin")
            .and_then(|c| c.as_str())
            .ok_or_else(|| WebSocketError::Parse("Missing 'coin' in activeAssetCtx".to_string()))?;

        let ctx = data.get("ctx")
            .ok_or_else(|| WebSocketError::Parse("Missing 'ctx' in activeAssetCtx".to_string()))?;

        let parse_f64 = |val: &Value| -> Option<f64> {
            val.as_str().and_then(|s| s.parse().ok()).or_else(|| val.as_f64())
        };

        let mark_px = ctx.get("markPx").and_then(parse_f64).unwrap_or(0.0);
        let mid_px = ctx.get("midPx").and_then(parse_f64);
        let prev_day_px = ctx.get("prevDayPx").and_then(parse_f64);
        let volume_24h = ctx.get("dayNtlVlm").and_then(parse_f64);

        let last_price = mid_px.unwrap_or(mark_px);

        let (price_change_24h, price_change_percent_24h) = match prev_day_px {
            Some(prev) if prev > 0.0 => {
                let change = last_price - prev;
                let change_pct = (change / prev) * 100.0;
                (Some(change), Some(change_pct))
            }
            _ => (None, None),
        };

        let ticker = crate::core::Ticker {
            symbol: coin.to_string(),
            last_price,
            bid_price: None,
            ask_price: None,
            high_24h: None,
            low_24h: None,
            volume_24h,
            quote_volume_24h: None,
            price_change_24h,
            price_change_percent_24h,
            timestamp: crate::core::utils::timestamp_millis() as i64,
        };

        Ok(Some(StreamEvent::Ticker(ticker)))
    }

    /// Parse allMids message to Ticker events
    fn parse_all_mids(data: &Value) -> WebSocketResult<Option<StreamEvent>> {
        // Format: { "mids": { "BTC": "50123.45", "ETH": "2500.67", ... } }
        let mids = data.get("mids")
            .and_then(|m| m.as_object())
            .ok_or_else(|| WebSocketError::Parse("Missing 'mids' object".to_string()))?;

        // For now, we'll just take the first symbol
        // In a real implementation, we'd emit multiple events or filter by subscription
        if let Some((symbol, price_val)) = mids.iter().next() {
            let price = price_val.as_str()
                .and_then(|s| s.parse::<f64>().ok())
                .or_else(|| price_val.as_f64())
                .ok_or_else(|| WebSocketError::Parse("Invalid price format".to_string()))?;

            let ticker = crate::core::Ticker {
                symbol: symbol.clone(),
                last_price: price,
                bid_price: None,
                ask_price: None,
                high_24h: None,
                low_24h: None,
                volume_24h: None,
                quote_volume_24h: None,
                price_change_24h: None,
                price_change_percent_24h: None,
                timestamp: crate::core::utils::timestamp_millis() as i64,
            };

            return Ok(Some(StreamEvent::Ticker(ticker)));
        }

        Ok(None)
    }

    /// Parse trades message
    fn parse_trades(data: &Value) -> WebSocketResult<Option<StreamEvent>> {
        // Format: [ { "coin": "BTC", "side": "B", "px": "50123.45", "sz": "0.5", ... } ]
        let trades = data.as_array()
            .ok_or_else(|| WebSocketError::Parse("Expected array of trades".to_string()))?;

        // Emit first trade (in real implementation, might emit all)
        if let Some(trade_data) = trades.first() {
            let trade = HyperliquidParser::parse_recent_trades(&json!([trade_data]))
                .map_err(|e| WebSocketError::Parse(e.to_string()))?;

            if let Some(first_trade) = trade.into_iter().next() {
                return Ok(Some(StreamEvent::Trade(first_trade)));
            }
        }

        Ok(None)
    }

    /// Parse l2Book message
    fn parse_l2_book(data: &Value) -> WebSocketResult<Option<StreamEvent>> {
        // Format: { "coin": "BTC", "time": 1234567890, "levels": [[bids], [asks]] }
        let orderbook = HyperliquidParser::parse_orderbook(data)
            .map_err(|e| WebSocketError::Parse(e.to_string()))?;

        Ok(Some(StreamEvent::OrderbookSnapshot(orderbook)))
    }

    /// Parse candle message
    fn parse_candle(data: &Value) -> WebSocketResult<Option<StreamEvent>> {
        // Format: [ { "t": 1234, "o": "50100", "h": "50200", ... } ]
        let klines = HyperliquidParser::parse_klines(data)
            .map_err(|e| WebSocketError::Parse(e.to_string()))?;

        if let Some(kline) = klines.into_iter().next() {
            return Ok(Some(StreamEvent::Kline(kline)));
        }

        Ok(None)
    }

    /// Build subscription object for Hyperliquid
    fn build_subscription(request: &SubscriptionRequest) -> Value {
        let coin = &request.symbol.base;

        match &request.stream_type {
            StreamType::Ticker => {
                // activeAssetCtx provides per-coin 24h stats: dayNtlVlm, prevDayPx,
                // markPx, midPx, funding, etc. Much richer than allMids (mid-price only).
                json!({
                    "type": "activeAssetCtx",
                    "coin": coin
                })
            }
            StreamType::Trade => {
                json!({
                    "type": "trades",
                    "coin": coin
                })
            }
            StreamType::Orderbook | StreamType::OrderbookDelta => {
                json!({
                    "type": "l2Book",
                    "coin": coin,
                    "nSigFigs": null,
                    "mantissa": null
                })
            }
            StreamType::Kline { interval } => {
                json!({
                    "type": "candle",
                    "coin": coin,
                    "interval": interval
                })
            }
            _ => {
                // Unsupported stream types — fall back to allMids for backward compatibility
                json!({
                    "type": "allMids",
                    "dex": ""
                })
            }
        }
    }

    /// Start heartbeat task.
    ///
    /// Sends a `Message::Ping(vec![])` frame every 30 seconds so the server
    /// can be kept alive and so RTT can be measured via the resulting
    /// `Message::Pong` received in the message handler.
    fn start_heartbeat_task(
        ws_stream: Arc<Mutex<Option<WsStream>>>,
        last_ping: Arc<Mutex<Instant>>,
        status: Arc<Mutex<ConnectionStatus>>,
    ) {
        tokio::spawn(async move {
            loop {
                sleep(Duration::from_secs(30)).await;

                // Check if connection is still alive
                let last = *last_ping.lock().await;
                if last.elapsed() >= Duration::from_secs(60) {
                    // No pongs for 60 seconds — connection may be stale
                    *status.lock().await = ConnectionStatus::Disconnected;
                    break;
                }

                // Send a WS Ping frame; the message handler will record RTT on Pong
                let mut stream_guard = ws_stream.lock().await;
                if let Some(stream) = stream_guard.as_mut() {
                    if stream.send(Message::Ping(vec![])).await.is_ok() {
                        *last_ping.lock().await = Instant::now();
                    } else {
                        break;
                    }
                } else {
                    break;
                }
            }
        });
    }
}

// ═══════════════════════════════════════════════════════════════════════════════
// WEBSOCKET CONNECTOR TRAIT IMPLEMENTATION
// ═══════════════════════════════════════════════════════════════════════════════

#[async_trait]
impl WebSocketConnector for HyperliquidWebSocket {
    async fn connect(&mut self, _account_type: AccountType) -> WebSocketResult<()> {
        *self.status.lock().await = ConnectionStatus::Connecting;

        // Connect WebSocket
        let ws_stream = self.connect_ws().await?;
        *self.ws_stream.lock().await = Some(ws_stream);
        *self.status.lock().await = ConnectionStatus::Connected;
        *self.last_ping.lock().await = Instant::now();

        // Create event channel
        let (tx, mut rx) = mpsc::unbounded_channel();
        *self.event_tx.lock().await = Some(tx.clone());

        // Create broadcast channel and store
        let (broadcast_sender, _) = broadcast::channel(1000);
        *self.broadcast_tx.lock().unwrap() = Some(broadcast_sender);

        // Start message handler
        Self::start_message_handler(
            self.ws_stream.clone(),
            tx,
            self.status.clone(),
            self.last_ping.clone(),
            self.ws_ping_rtt_ms.clone(),
        );

        // Start forwarder task (mpsc -> broadcast)
        let broadcast_tx = self.broadcast_tx.clone();
        let last_ping = self.last_ping.clone();
        tokio::spawn(async move {
            while let Some(event) = rx.recv().await {
                // Update last ping time on any received message
                *last_ping.lock().await = Instant::now();

                // Forward to broadcast channel (ignore if no receivers)
                if let Some(tx) = broadcast_tx.lock().unwrap().as_ref() {
                    let _ = tx.send(event);
                }
            }
            // mpsc channel closed — drop broadcast sender
            let _ = broadcast_tx.lock().unwrap().take();
        });

        // Start heartbeat task
        Self::start_heartbeat_task(
            self.ws_stream.clone(),
            self.last_ping.clone(),
            self.status.clone(),
        );

        Ok(())
    }

    async fn disconnect(&mut self) -> WebSocketResult<()> {
        *self.status.lock().await = ConnectionStatus::Disconnected;

        // Close WebSocket connection
        if let Some(mut stream) = self.ws_stream.lock().await.take() {
            let _ = stream.close(None).await;
        }

        *self.event_tx.lock().await = None;
        let _ = self.broadcast_tx.lock().unwrap().take();
        self.subscriptions.lock().await.clear();
        Ok(())
    }

    fn connection_status(&self) -> ConnectionStatus {
        // Use try_lock to avoid blocking
        match self.status.try_lock() {
            Ok(status) => *status,
            Err(_) => ConnectionStatus::Disconnected,
        }
    }

    async fn subscribe(&mut self, request: SubscriptionRequest) -> WebSocketResult<()> {
        let subscription = Self::build_subscription(&request);

        let msg = SubscribeMessage {
            method: "subscribe".to_string(),
            subscription,
        };

        let msg_json = serde_json::to_string(&msg)
            .map_err(|e| WebSocketError::ProtocolError(e.to_string()))?;

        let mut stream_guard = self.ws_stream.lock().await;
        let stream = stream_guard.as_mut()
            .ok_or_else(|| WebSocketError::ConnectionError("Not connected".to_string()))?;

        stream.send(Message::Text(msg_json)).await
            .map_err(|e| WebSocketError::ConnectionError(e.to_string()))?;

        drop(stream_guard);

        self.subscriptions.lock().await.insert(request);

        Ok(())
    }

    async fn unsubscribe(&mut self, request: SubscriptionRequest) -> WebSocketResult<()> {
        let subscription = Self::build_subscription(&request);

        let msg = SubscribeMessage {
            method: "unsubscribe".to_string(),
            subscription,
        };

        let msg_json = serde_json::to_string(&msg)
            .map_err(|e| WebSocketError::ProtocolError(e.to_string()))?;

        let mut stream_guard = self.ws_stream.lock().await;
        let stream = stream_guard.as_mut()
            .ok_or_else(|| WebSocketError::ConnectionError("Not connected".to_string()))?;

        stream.send(Message::Text(msg_json)).await
            .map_err(|e| WebSocketError::ConnectionError(e.to_string()))?;

        drop(stream_guard);

        self.subscriptions.lock().await.remove(&request);

        Ok(())
    }

    fn event_stream(&self) -> Pin<Box<dyn Stream<Item = WebSocketResult<StreamEvent>> + Send>> {
        let rx = self.broadcast_tx.lock().unwrap().as_ref()
            .map(|tx| tx.subscribe())
            .unwrap_or_else(|| broadcast::channel(1).1);

        Box::pin(tokio_stream::wrappers::BroadcastStream::new(rx).filter_map(|result| async move {
            match result {
                Ok(event) => Some(event),
                Err(tokio_stream::wrappers::errors::BroadcastStreamRecvError::Lagged(_)) => {
                    Some(Err(WebSocketError::ConnectionError("Event stream lagged behind".to_string())))
                }
            }
        }))
    }

    fn active_subscriptions(&self) -> Vec<SubscriptionRequest> {
        // Use try_lock to avoid blocking
        match self.subscriptions.try_lock() {
            Ok(subs) => subs.iter().cloned().collect(),
            Err(_) => Vec::new(),
        }
    }

    fn ping_rtt_handle(&self) -> Option<Arc<Mutex<u64>>> {
        Some(self.ws_ping_rtt_ms.clone())
    }
}

/// Subscription types specific to Hyperliquid
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
#[allow(dead_code)]
pub enum HyperliquidSubscription {
    /// All mid prices (price only, no 24h stats). Use ActiveAssetCtx for full ticker.
    AllMids,
    /// Per-coin 24h stats: dayNtlVlm, prevDayPx, markPx, midPx, funding, openInterest.
    /// Use this for ticker subscriptions — richer than AllMids.
    ActiveAssetCtx,
    Trades,           // Trade feed
    L2Book,           // Order book updates
    Bbo,              // Best bid/offer
    Candle,           // Kline/candle updates
    Notification,     // User notifications
    OpenOrders,       // Open orders
    OrderUpdates,     // Order status changes
    UserFills,        // Trade executions
    UserEvents,       // All account events
    UserFundings,     // Funding payments
    ClearinghouseState, // Account summary
}

impl HyperliquidSubscription {
    /// Get subscription type string
    #[allow(dead_code)]
    pub fn as_str(&self) -> &'static str {
        match self {
            Self::AllMids => "allMids",
            Self::ActiveAssetCtx => "activeAssetCtx",
            Self::Trades => "trades",
            Self::L2Book => "l2Book",
            Self::Bbo => "bbo",
            Self::Candle => "candle",
            Self::Notification => "notification",
            Self::OpenOrders => "openOrders",
            Self::OrderUpdates => "orderUpdates",
            Self::UserFills => "userFills",
            Self::UserEvents => "userEvents",
            Self::UserFundings => "userFundings",
            Self::ClearinghouseState => "clearinghouseState",
        }
    }

    /// Does subscription require authentication
    #[allow(dead_code)]
    pub fn requires_auth(&self) -> bool {
        matches!(self,
            Self::Notification
            | Self::OpenOrders
            | Self::OrderUpdates
            | Self::UserFills
            | Self::UserEvents
            | Self::UserFundings
            | Self::ClearinghouseState
        )
    }
}