rusty_feeder/exchange/bybit/

provider.rs

use rusty_common::collections::FxHashMap;
use smartstring::alias::String;
use std::sync::Arc;
use std::time::{Duration, Instant};

use super::data::{
    orderbook::OrderbookResponse, subscription::SubscriptionRequest, trade::TradeResponse,
};
use super::types::{
    BYBIT_API_URL, BYBIT_RATE_LIMITS, BYBIT_WS_URL_INVERSE, BYBIT_WS_URL_LINEAR,
    BYBIT_WS_URL_OPTION, BYBIT_WS_URL_SPOT,
};
use anyhow::{Context, Result, anyhow};
use async_trait::async_trait;
use parking_lot::RwLock;
use quanta::Clock;
use reqwest::header::{HeaderMap, HeaderValue};
use rusty_common::json::Value;
use rusty_common::websocket::{Message, WebSocketClient, WebSocketConfig, WebSocketError};
use rusty_model::{
    data::{
        book_snapshot::OrderBookSnapshot, market_trade::MarketTrade,
        simd_orderbook::SharedSimdOrderBook,
    },
    instruments::{Instrument, InstrumentId},
    venues::Venue,
};
use simd_json::prelude::{ValueAsArray, ValueAsObject, ValueAsScalar};
use smallvec::SmallVec;
use tokio::sync::{mpsc, watch};
use tokio::task::JoinHandle;

use crate::provider::prelude::*;

/// Bybit V5 product categories
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum BybitCategory {
    /// Spot trading
    Spot,
    /// Linear perpetuals and USDT futures
    Linear,
    /// Inverse contracts
    Inverse,
    /// Options
    Option,
}

impl BybitCategory {
    /// Get the appropriate WebSocket URL for this category
    #[must_use]
    pub const fn websocket_url(&self) -> &'static str {
        match self {
            Self::Spot => BYBIT_WS_URL_SPOT,
            Self::Linear => BYBIT_WS_URL_LINEAR,
            Self::Inverse => BYBIT_WS_URL_INVERSE,
            Self::Option => BYBIT_WS_URL_OPTION,
        }
    }

    /// Get the category string for V5 API calls
    #[must_use]
    pub const fn as_str(&self) -> &'static str {
        match self {
            Self::Spot => "spot",
            Self::Linear => "linear",
            Self::Inverse => "inverse",
            Self::Option => "option",
        }
    }
}

impl Default for BybitCategory {
    fn default() -> Self {
        Self::Spot
    }
}

/// Bybit V5 provider implementation
///
/// High-performance provider for Bybit V5 markets with WebSocket streaming
/// and REST API support. Supports all V5 categories: spot, linear, inverse, option.
/// Optimized for HFT applications with nanosecond precision.
#[derive(Debug)]
pub struct BybitProvider {
    /// Connection configuration
    config: ConnectionConfig,

    /// WebSocket connection status
    connection_status: Arc<RwLock<ConnectionState>>,

    /// Connection statistics
    stats: Arc<RwLock<ConnectionStats>>,

    /// Shared clock for time synchronization
    clock: Clock,

    /// Active subscriptions
    subscriptions: Arc<RwLock<FxHashMap<String, watch::Sender<bool>>>>,

    /// Task handles for WebSocket connections
    ws_handles: Arc<RwLock<FxHashMap<String, JoinHandle<()>>>>,

    /// Last connection attempt timestamp
    last_connection_attempt: Arc<RwLock<Instant>>,

    /// HTTP client for REST API requests
    http_client: reqwest::Client,

    /// V5 category (spot, linear, inverse, option)
    category: BybitCategory,
}

impl Default for BybitProvider {
    fn default() -> Self {
        Self::new()
    }
}

impl BybitProvider {
    /// Get the current category
    #[inline]
    #[must_use]
    pub const fn category(&self) -> BybitCategory {
        self.category
    }

    /// Create a new Bybit provider with default configuration (Spot category)
    ///
    /// This method provides backward compatibility for code expecting a parameterless constructor.
    /// For specific categories, use `new_with_category()` or the category-specific constructors.
    #[inline]
    #[must_use]
    pub fn new() -> Self {
        Self::new_with_category(BybitCategory::Spot)
    }

    /// Create a new Bybit provider with specified category and default configuration
    #[inline]
    #[must_use]
    pub fn new_with_category(category: BybitCategory) -> Self {
        Self::with_config(category, None)
    }

    /// Create a new Bybit spot provider with default configuration
    #[inline]
    #[must_use]
    pub fn new_spot() -> Self {
        Self::new_with_category(BybitCategory::Spot)
    }

    /// Create a new Bybit linear provider with default configuration
    #[inline]
    #[must_use]
    pub fn new_linear() -> Self {
        Self::new_with_category(BybitCategory::Linear)
    }

    /// Create a new Bybit inverse provider with default configuration
    #[inline]
    #[must_use]
    pub fn new_inverse() -> Self {
        Self::new_with_category(BybitCategory::Inverse)
    }

    /// Create a new Bybit option provider with default configuration
    #[inline]
    #[must_use]
    pub fn new_option() -> Self {
        Self::new_with_category(BybitCategory::Option)
    }

    /// Create a new Bybit provider with specified category and custom configuration
    #[inline]
    #[must_use]
    pub fn with_config(category: BybitCategory, config: Option<ConnectionConfig>) -> Self {
        let mut default_config = ConnectionConfig::default();

        // Set default WebSocket configuration based on category
        default_config.websocket_config.base_url = category.websocket_url().into();
        default_config.websocket_config.ping_interval_milliseconds = 20000; // 20 seconds
        default_config.websocket_config.use_compression = false;

        // Set default REST API configuration
        default_config.rest_config.base_url = BYBIT_API_URL.into();
        default_config.rest_config.timeout_milliseconds = 5000; // 5 seconds

        // Use custom config if provided
        let config = config.unwrap_or(default_config);
        let clock = config.clock.clone();

        // Create HTTP client with appropriate settings
        let mut headers = HeaderMap::new();
        headers.insert(
            "User-Agent",
            HeaderValue::from_str(&config.rest_config.user_agent)
                .unwrap_or_else(|_| HeaderValue::from_static("RustyHFT/1.0")),
        );

        let http_client = reqwest::Client::builder()
            .timeout(Duration::from_millis(
                config.rest_config.timeout_milliseconds,
            ))
            .connect_timeout(Duration::from_millis(
                config.rest_config.timeout_milliseconds / 2,
            ))
            .pool_max_idle_per_host(config.rest_config.connection_pool_size)
            .pool_idle_timeout(Duration::from_millis(
                config.rest_config.keep_alive_milliseconds,
            ))
            .default_headers(headers)
            .build()
            .unwrap_or_default();

        Self {
            config,
            connection_status: Arc::new(RwLock::new(ConnectionState::Disconnected)),
            stats: Arc::new(RwLock::new(ConnectionStats::default())),
            clock,
            subscriptions: Arc::new(RwLock::new(FxHashMap::default())),
            ws_handles: Arc::new(RwLock::new(FxHashMap::default())),
            last_connection_attempt: Arc::new(RwLock::new(Instant::now())),
            http_client,
            category,
        }
    }

    /// Helper function to update connection statistics upon message reception
    #[inline]
    fn update_receive_stats(
        stats: Arc<RwLock<ConnectionStats>>,
        message_size: usize,
        local_time: u64,
    ) {
        let mut s = stats.write();
        s.messages_received += 1;
        s.bytes_received += message_size as u64;
        s.last_message_time = local_time;
    }

    /// Helper function to create WebSocket request with proper headers
    fn create_websocket_config(
        url: String,
        use_compression: bool,
        ping_interval_milliseconds: u64,
        timeout_milliseconds: u64,
    ) -> WebSocketConfig {
        WebSocketConfig::builder(rusty_common::types::Exchange::Bybit, url.to_string())
            .ping_interval(Duration::from_millis(ping_interval_milliseconds))
            .timeout(Duration::from_millis(timeout_milliseconds))
            .compression(if use_compression {
                rusty_common::websocket::CompressionConfig::default()
            } else {
                rusty_common::websocket::CompressionConfig::disabled()
            })
            .build()
    }

    /// Connect to the WebSocket API
    #[inline]
    async fn connect(&self) -> Result<()> {
        if *self.connection_status.read() == ConnectionState::Connected {
            return Ok(());
        }

        // Check if we need to respect reconnection backoff
        {
            let now = Instant::now();
            let last_attempt = *self.last_connection_attempt.read();
            let backoff_ms = self
                .config
                .websocket_config
                .reconnect
                .backoff_initial_milliseconds;

            if now.duration_since(last_attempt) < Duration::from_millis(backoff_ms) {
                // Wait until backoff period is over
                tokio::time::sleep(Duration::from_millis(backoff_ms)).await;
            }

            // Update last attempt time
            *self.last_connection_attempt.write() = Instant::now();
        }

        *self.connection_status.write() = ConnectionState::Connecting;

        // Note: Connection is now handled by individual trade/orderbook connections
        // using WebSocketClient pattern. This method is kept for compatibility.
        *self.connection_status.write() = ConnectionState::Connected;
        self.stats.write().connected_time = self.clock.raw();
        Ok(())
    }
}

#[async_trait]
impl Provider for BybitProvider {
    type TradeMessage = TradeResponse;
    type DepthMessage = OrderbookResponse;
    type InstrumentMessage = Value;

    fn name(&self) -> &'static str {
        "Bybit"
    }

    fn venue(&self) -> Venue {
        Venue::Bybit
    }

    fn config(&self) -> &ConnectionConfig {
        &self.config
    }

    async fn init(&mut self) -> Result<()> {
        self.connect().await
    }

    async fn shutdown(&mut self) -> Result<()> {
        // Unsubscribe from all active subscriptions
        let mut keys_to_remove = Vec::new();

        for (key, _) in self.subscriptions.read().iter() {
            keys_to_remove.push(key.clone());
        }

        // Send stop signal for each subscription
        for key in keys_to_remove {
            if let Some(tx) = self.subscriptions.write().remove(&key) {
                let _ = tx.send(true);
            }

            // Abort task
            if let Some(handle) = self.ws_handles.write().remove(&key) {
                handle.abort();
            }
        }

        // Set connection status to disconnected
        *self.connection_status.write() = ConnectionState::Disconnected;

        Ok(())
    }

    #[inline]
    async fn subscribe_trades(
        &self,
        symbols: SmallVec<[String; 8]>,
    ) -> Result<mpsc::Receiver<Self::TradeMessage>> {
        // Ensure connection is established
        self.connect().await?;

        // Create channel for trade messages
        let (tx, rx) = mpsc::channel(1024);

        // Create subscription message
        let subscription = SubscriptionRequest::trades(symbols.to_vec());
        let connection_id = format!("trade-{}", symbols.join(","));

        // Create WebSocket connection and handle
        let handle = self
            .create_trade_connection(symbols, tx, subscription)
            .await?;

        // Save task handle
        self.ws_handles.write().insert(connection_id.into(), handle);

        Ok(rx)
    }

    #[inline]
    async fn unsubscribe_trades(&self) -> Result<()> {
        // Find all trade subscriptions
        let mut keys_to_remove = Vec::new();

        for (key, _) in self.subscriptions.read().iter() {
            if key.starts_with("trade-") {
                keys_to_remove.push(key.clone());
            }
        }

        // Send stop signal for each subscription
        for key in keys_to_remove {
            if let Some(tx) = self.subscriptions.write().remove(&key) {
                let _ = tx.send(true);
            }

            // Abort task
            if let Some(handle) = self.ws_handles.write().remove(&key) {
                handle.abort();
            }
        }

        Ok(())
    }

    #[inline]
    async fn subscribe_orderbook(
        &self,
        symbols: SmallVec<[String; 8]>,
    ) -> Result<mpsc::Receiver<Self::DepthMessage>> {
        // Ensure connection is established
        self.connect().await?;

        // Create channel for orderbook messages
        let (tx, rx) = mpsc::channel(1024);

        // Create subscription message (depth 50 levels)
        let subscription = SubscriptionRequest::orderbook(symbols.to_vec(), 50);
        let connection_id = format!("orderbook-{}", symbols.join(","));

        // Create WebSocket connection and handle
        let handle = self
            .create_orderbook_connection(symbols, tx, subscription)
            .await?;

        // Save task handle
        self.ws_handles.write().insert(connection_id.into(), handle);

        Ok(rx)
    }

    #[inline]
    async fn unsubscribe_orderbook(&self) -> Result<()> {
        // Find all orderbook subscriptions
        let mut keys_to_remove = Vec::new();

        for (key, _) in self.subscriptions.read().iter() {
            if key.starts_with("orderbook-") {
                keys_to_remove.push(key.clone());
            }
        }

        // Send stop signal for each subscription
        for key in keys_to_remove {
            if let Some(tx) = self.subscriptions.write().remove(&key) {
                let _ = tx.send(true);
            }

            // Abort task
            if let Some(handle) = self.ws_handles.write().remove(&key) {
                handle.abort();
            }
        }

        Ok(())
    }

    #[inline]
    async fn get_instruments(&self) -> Result<Vec<Box<dyn Instrument>>> {
        // Fetch instruments from Bybit V5 API with category parameter
        let url = format!(
            "{}/v5/market/instruments-info?category={}",
            self.config.rest_config.base_url,
            self.category.as_str()
        );

        let response = self
            .http_client
            .get(&url)
            .send()
            .await
            .context("Failed to fetch Bybit instruments info")?;

        // Check if the request was successful
        if !response.status().is_success() {
            return Err(anyhow!(
                "Failed to fetch Bybit instruments info: HTTP {}",
                response.status()
            ));
        }

        // Parse the response using simd_json
        let bytes = response
            .bytes()
            .await
            .context("Failed to get response bytes")?;
        let mut bytes_vec = bytes.to_vec();
        let instruments_info: simd_json::OwnedValue = simd_json::from_slice(&mut bytes_vec)
            .context("Failed to parse Bybit instruments info response")?;

        // Extract the instruments from the response
        let result = instruments_info["result"].as_object().ok_or_else(|| {
            anyhow!("Invalid response format: 'result' field not found or not an object")
        })?;

        let list = result["list"].as_array().ok_or_else(|| {
            anyhow!("Invalid response format: 'list' field not found or not an array")
        })?;

        // Convert to instruments
        let mut instruments = Vec::with_capacity(list.len());

        for instrument_data in list {
            // Skip instruments that are not trading
            let status = instrument_data["status"].as_str().unwrap_or("");
            if status != "Trading" {
                continue;
            }

            let symbol = instrument_data["symbol"].as_str().ok_or_else(|| {
                anyhow!("Invalid instrument data: 'symbol' field not found or not a String")
            })?;

            // Create a basic instrument implementation
            let instrument_id = InstrumentId::new(symbol, Venue::Bybit);
            let instrument = Box::new(BybitInstrument { id: instrument_id }) as Box<dyn Instrument>;
            instruments.push(instrument);
        }

        Ok(instruments)
    }

    #[inline]
    async fn get_historical_trades(
        &self,
        _symbol: &str,
        _limit: Option<u32>,
    ) -> Result<Vec<MarketTrade>> {
        // Fetch historical trades from Bybit API
        // Not fully implemented yet
        Err(anyhow!("Historical trades not implemented yet"))
    }

    #[inline]
    async fn get_orderbook_snapshot(
        &self,
        symbol: &str,
        depth: Option<u32>,
    ) -> Result<OrderBookSnapshot> {
        let limit = depth.unwrap_or(50);
        let url = format!(
            "{}/v5/market/orderbook?category={}&symbol={}&limit={}",
            self.config.rest_config.base_url,
            self.category.as_str(),
            symbol,
            limit
        );

        let response = self.http_client.get(&url).send().await?;

        if !response.status().is_success() {
            let error_text = response.text().await?;
            return Err(anyhow!(
                "Failed to fetch orderbook snapshot: {}",
                error_text
            ));
        }

        // Parse the response using simd_json
        let bytes = response
            .bytes()
            .await
            .map_err(|e| anyhow!("Failed to get response bytes: {}", e))?;
        let mut bytes_vec = bytes.to_vec();
        let snapshot_data: simd_json::OwnedValue = simd_json::from_slice(&mut bytes_vec)
            .map_err(|e| anyhow!("Failed to parse JSON: {}", e))?;

        // Create OrderBookSnapshot with current timestamp (nanoseconds since epoch)
        let timestamp = self.clock.raw();
        let instrument_id = InstrumentId::new(symbol, Venue::Bybit);

        let result = snapshot_data["result"]
            .as_object()
            .ok_or_else(|| anyhow!("Invalid response format: 'result' field not found"))?;

        let update_id = result["u"].as_u64().unwrap_or(0);
        let mut order_book_snapshot =
            OrderBookSnapshot::new_empty(instrument_id, timestamp, update_id);

        // Parse bids
        if let Some(bids) = result["b"].as_array() {
            for bid in bids {
                if let Some(bid_array) = bid.as_array()
                    && bid_array.len() >= 2
                    && let (Some(price_str), Some(quantity_str)) =
                        (bid_array[0].as_str(), bid_array[1].as_str())
                    && let (Ok(price), Ok(quantity)) = (
                        price_str.parse::<rust_decimal::Decimal>(),
                        quantity_str.parse::<rust_decimal::Decimal>(),
                    )
                {
                    order_book_snapshot.add_bid(price, quantity);
                }
            }
        }

        // Parse asks
        if let Some(asks) = result["a"].as_array() {
            for ask in asks {
                if let Some(ask_array) = ask.as_array()
                    && ask_array.len() >= 2
                    && let (Some(price_str), Some(quantity_str)) =
                        (ask_array[0].as_str(), ask_array[1].as_str())
                    && let (Ok(price), Ok(quantity)) = (
                        price_str.parse::<rust_decimal::Decimal>(),
                        quantity_str.parse::<rust_decimal::Decimal>(),
                    )
                {
                    order_book_snapshot.add_ask(price, quantity);
                }
            }
        }

        Ok(order_book_snapshot)
    }

    #[inline]
    async fn get_realtime_orderbook(&self, _symbol: &str) -> Result<SharedSimdOrderBook> {
        // Not yet implemented
        Err(anyhow!("Real-time orderbook not implemented yet"))
    }

    #[inline]
    async fn is_connected(&self) -> bool {
        *self.connection_status.read() == ConnectionState::Connected
    }

    #[inline]
    async fn connection_status(&self) -> ConnectionState {
        *self.connection_status.read()
    }

    #[inline]
    async fn get_stats(&self) -> ConnectionStats {
        self.stats.read().clone()
    }

    #[inline]
    async fn ping(&self) -> Result<u64> {
        // For WebSocketClient pattern, ping/pong is handled automatically
        // Return a simple network latency test via HTTP
        let start = self.clock.raw();

        let response = self
            .http_client
            .get(format!(
                "{}/v5/market/time",
                self.config.rest_config.base_url
            ))
            .send()
            .await?;

        if response.status().is_success() {
            Ok(self.clock.raw().saturating_sub(start))
        } else {
            Err(anyhow!("Ping failed with status: {}", response.status()))
        }
    }

    #[inline]
    async fn reset_connection(&self) -> Result<()> {
        // First, unsubscribe all active subscriptions
        let mut keys_to_remove = Vec::new();

        for (key, _) in self.subscriptions.read().iter() {
            keys_to_remove.push(key.clone());
        }

        // Send stop signal for each subscription
        for key in &keys_to_remove {
            if let Some(tx) = self.subscriptions.write().remove(key) {
                let _ = tx.send(true);
            }

            // Abort task
            if let Some(handle) = self.ws_handles.write().remove(key) {
                handle.abort();
            }
        }

        // Set status to disconnected
        *self.connection_status.write() = ConnectionState::Disconnected;

        // Reset stats
        *self.stats.write() = ConnectionStats::default();

        // Connect again
        self.connect().await
    }

    #[inline]
    fn get_rate_limits(&self) -> Vec<RateLimit> {
        BYBIT_RATE_LIMITS.to_vec()
    }
}

/// Simple instrument implementation for Bybit
#[derive(Debug, Clone)]
pub struct BybitInstrument {
    /// Unique instrument identifier
    pub id: InstrumentId,
}

impl Instrument for BybitInstrument {
    fn id(&self) -> InstrumentId {
        self.id.clone()
    }

    fn symbol(&self) -> String {
        self.id.symbol.clone()
    }

    fn venue(&self) -> Venue {
        self.id.venue
    }

    fn as_any(&self) -> &dyn std::any::Any {
        self
    }

    fn clone_box(&self) -> Box<dyn Instrument> {
        Box::new(self.clone())
    }
}

impl BybitProvider {
    /// Create a WebSocket connection for trade data
    #[inline]
    async fn create_trade_connection(
        &self,
        symbols: SmallVec<[String; 8]>,
        tx: mpsc::Sender<TradeResponse>,
        subscription: SubscriptionRequest,
    ) -> Result<JoinHandle<()>> {
        let connection_id = format!("trade-{}", symbols.join(","));

        // Create stop signal
        let (stop_tx, stop_rx) = watch::channel(false);
        self.subscriptions
            .write()
            .insert(connection_id.clone().into(), stop_tx);

        // Clone parameters for the task
        let url = self.config.websocket_config.base_url.clone();
        let clock = self.clock.clone();
        let timeout_milliseconds = self.config.websocket_config.timeout_milliseconds;
        let ping_interval_milliseconds = self.config.websocket_config.ping_interval_milliseconds;
        let stats = self.stats.clone();
        let connection_status = self.connection_status.clone();
        let use_compression = self.config.websocket_config.use_compression;

        // Spawn a task for the WebSocket connection
        let handle = tokio::spawn(async move {
            loop {
                // Check for stop signal
                if *stop_rx.borrow() {
                    break;
                }

                // Create WebSocket configuration
                let ws_config = Self::create_websocket_config(
                    url.clone(),
                    use_compression,
                    ping_interval_milliseconds,
                    timeout_milliseconds,
                );

                // Create WebSocket client
                let mut client = WebSocketClient::new(ws_config);

                // Update connection status
                *connection_status.write() = ConnectionState::Connecting;

                // Create a simple message handler
                let handler = BybitMessageHandler::new_trade_handler(
                    clock.clone(),
                    stats.clone(),
                    tx.clone(),
                    subscription.clone(),
                );

                // Run the WebSocket client
                if let Err(e) = client.run(handler).await {
                    log::error!("WebSocket client error: {e}");
                    *connection_status.write() = ConnectionState::Error;
                }

                // Check for stop signal before reconnecting
                if *stop_rx.borrow() {
                    break;
                }

                // Reconnect after a delay
                tokio::time::sleep(Duration::from_millis(1000)).await;
            }
        });

        Ok(handle)
    }

    /// Create a WebSocket connection for orderbook data
    #[inline]
    async fn create_orderbook_connection(
        &self,
        symbols: SmallVec<[String; 8]>,
        tx: mpsc::Sender<OrderbookResponse>,
        subscription: SubscriptionRequest,
    ) -> Result<JoinHandle<()>> {
        let connection_id = format!("orderbook-{}", symbols.join(","));

        // Create stop signal
        let (stop_tx, stop_rx) = watch::channel(false);
        self.subscriptions
            .write()
            .insert(connection_id.clone().into(), stop_tx);

        // Clone parameters for the task
        let url = self.config.websocket_config.base_url.clone();
        let clock = self.clock.clone();
        let timeout_milliseconds = self.config.websocket_config.timeout_milliseconds;
        let ping_interval_milliseconds = self.config.websocket_config.ping_interval_milliseconds;
        let stats = self.stats.clone();
        let connection_status = self.connection_status.clone();
        let use_compression = self.config.websocket_config.use_compression;

        // Spawn a task for the WebSocket connection
        let handle = tokio::spawn(async move {
            loop {
                // Check for stop signal
                if *stop_rx.borrow() {
                    break;
                }

                // Create WebSocket configuration
                let ws_config = Self::create_websocket_config(
                    url.clone(),
                    use_compression,
                    ping_interval_milliseconds,
                    timeout_milliseconds,
                );

                // Create WebSocket client
                let mut client = WebSocketClient::new(ws_config);

                // Update connection status
                *connection_status.write() = ConnectionState::Connecting;

                // Create a simple message handler
                let handler = BybitMessageHandler::new_orderbook_handler(
                    clock.clone(),
                    stats.clone(),
                    tx.clone(),
                    subscription.clone(),
                );

                // Run the WebSocket client
                if let Err(e) = client.run(handler).await {
                    log::error!("WebSocket client error: {e}");
                    *connection_status.write() = ConnectionState::Error;
                }

                // Check for stop signal before reconnecting
                if *stop_rx.borrow() {
                    break;
                }

                // Reconnect after a delay
                tokio::time::sleep(Duration::from_millis(1000)).await;
            }
        });

        Ok(handle)
    }
}

/// Simple message handler for Bybit WebSocket messages
#[derive(Debug)]
struct BybitMessageHandler {
    handler_type: HandlerType,
    clock: Clock,
    stats: Arc<RwLock<ConnectionStats>>,
    trade_tx: Option<mpsc::Sender<TradeResponse>>,
    orderbook_tx: Option<mpsc::Sender<OrderbookResponse>>,
    subscription: SubscriptionRequest,
    sender: Option<mpsc::UnboundedSender<Message>>,
}

#[derive(Debug)]
enum HandlerType {
    Trade,
    Orderbook,
}

impl BybitMessageHandler {
    const fn new_trade_handler(
        clock: Clock,
        stats: Arc<RwLock<ConnectionStats>>,
        tx: mpsc::Sender<TradeResponse>,
        subscription: SubscriptionRequest,
    ) -> Self {
        Self {
            handler_type: HandlerType::Trade,
            clock,
            stats,
            trade_tx: Some(tx),
            orderbook_tx: None,
            subscription,
            sender: None,
        }
    }

    const fn new_orderbook_handler(
        clock: Clock,
        stats: Arc<RwLock<ConnectionStats>>,
        tx: mpsc::Sender<OrderbookResponse>,
        subscription: SubscriptionRequest,
    ) -> Self {
        Self {
            handler_type: HandlerType::Orderbook,
            clock,
            stats,
            trade_tx: None,
            orderbook_tx: Some(tx),
            subscription,
            sender: None,
        }
    }
}

#[async_trait]
impl rusty_common::websocket::MessageHandler for BybitMessageHandler {
    async fn on_message(&mut self, message: Message) -> std::result::Result<(), WebSocketError> {
        let local_time = self.clock.raw();

        // Only process text messages
        if let Some(text) = message.as_text() {
            BybitProvider::update_receive_stats(self.stats.clone(), text.len(), local_time);

            // PERFORMANCE: Avoid heap allocation in hot path
            // First, do a quick check for the topic without full parsing
            // This avoids parsing messages we don't care about
            let should_process = match self.handler_type {
                HandlerType::Trade => text.contains("\"topic\":\"publicTrade"),
                HandlerType::Orderbook => text.contains("\"topic\":\"orderbook"),
            };

            if should_process {
                // Use a stack-allocated buffer for small messages
                // Falls back to heap allocation for larger messages
                const STACK_BUFFER_SIZE: usize = 8192;

                if text.len() <= STACK_BUFFER_SIZE {
                    // Stack allocation path for small messages (most common case)
                    let mut stack_buffer = [0u8; STACK_BUFFER_SIZE];
                    let bytes = text.as_bytes();
                    stack_buffer[..bytes.len()].copy_from_slice(bytes);
                    let buffer_slice = &mut stack_buffer[..bytes.len()];

                    match self.handler_type {
                        HandlerType::Trade => {
                            if let Ok(trade_response) =
                                simd_json::from_slice::<TradeResponse>(buffer_slice)
                                && let Some(ref tx) = self.trade_tx
                                && let Err(e) = tx.send(trade_response).await
                            {
                                log::warn!("Failed to send trade message: {e}");
                            }
                        }
                        HandlerType::Orderbook => {
                            if let Ok(orderbook_response) =
                                simd_json::from_slice::<OrderbookResponse>(buffer_slice)
                                && let Some(ref tx) = self.orderbook_tx
                                && let Err(e) = tx.send(orderbook_response).await
                            {
                                log::warn!("Failed to send orderbook message: {e}");
                            }
                        }
                    }
                } else {
                    // Heap allocation path for large messages (rare case)
                    let mut message_bytes = text.as_bytes().to_vec();

                    match self.handler_type {
                        HandlerType::Trade => {
                            if let Ok(trade_response) =
                                simd_json::from_slice::<TradeResponse>(&mut message_bytes)
                                && let Some(ref tx) = self.trade_tx
                                && let Err(e) = tx.send(trade_response).await
                            {
                                log::warn!("Failed to send trade message: {e}");
                            }
                        }
                        HandlerType::Orderbook => {
                            if let Ok(orderbook_response) =
                                simd_json::from_slice::<OrderbookResponse>(&mut message_bytes)
                                && let Some(ref tx) = self.orderbook_tx
                                && let Err(e) = tx.send(orderbook_response).await
                            {
                                log::warn!("Failed to send orderbook message: {e}");
                            }
                        }
                    }
                }
            }
        }

        Ok(())
    }

    async fn on_connected(&mut self) -> std::result::Result<(), WebSocketError> {
        // Send subscription message after connection
        let subscription_json = simd_json::to_string(&self.subscription).map_err(|e| {
            WebSocketError::MessageProcessingError(format!("Failed to serialize subscription: {e}"))
        })?;

        log::info!("Sending subscription: {subscription_json}");

        // Actually send the subscription message
        let message = Message::Text(subscription_json.into());
        self.send_message(message)?;

        Ok(())
    }

    async fn on_disconnected(&mut self) -> std::result::Result<(), WebSocketError> {
        log::info!("Disconnected from Bybit WebSocket");
        Ok(())
    }

    async fn on_error(&mut self, error: WebSocketError) -> std::result::Result<(), WebSocketError> {
        log::error!("Bybit WebSocket error: {error}");
        Ok(())
    }

    fn set_sender(&mut self, sender: mpsc::UnboundedSender<Message>) {
        self.sender = Some(sender);
    }

    fn send_message(&self, message: Message) -> std::result::Result<(), WebSocketError> {
        match &self.sender {
            Some(sender) => {
                sender
                    .send(message)
                    .map_err(|_| WebSocketError::NotConnected)?;
                Ok(())
            }
            None => Err(WebSocketError::NotConnected),
        }
    }
}