rusty_feeder/exchange/upbit/

provider.rs

//! Upbit exchange provider implementation
//!
//! High-performance provider for Upbit markets with WebSocket streaming
//! and REST API support. Optimized for HFT applications with nanosecond precision.

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

use super::auth::{UpbitAuth, UpbitAuthConfig};
use super::data::{orderbook::OrderbookMessage, prelude::*, trade::TradeMessage};
use super::types::{UPBIT_API_URL, UPBIT_RATE_LIMITS, UPBIT_WS_URL};
use anyhow::{Result, anyhow};
use async_trait::async_trait;
use parking_lot::RwLock;
use quanta::Clock;
use reqwest::header::{HeaderMap, HeaderValue};
use rust_decimal::prelude::FromPrimitive;
use rusty_common::json::Value;
use rusty_common::websocket::{Message, WebSocketClient, WebSocketConfig, WebSocketError};
use rusty_model::{
    data::{
        book_snapshot::OrderBookSnapshot,
        market_trade::MarketTrade,
        orderbook::OrderBook,
        simd_orderbook::{SharedSimdOrderBook, SimdPriceLevels},
    },
    enums::OrderSide,
    instruments::Instrument,
};
use simd_json::prelude::{ValueAsArray, ValueAsScalar, ValueObjectAccess};
use smallvec::SmallVec;
use std::str::FromStr;
use tokio::sync::{mpsc, watch};
use tokio::task::JoinHandle;

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

/// Upbit exchange provider implementation
///
/// High-performance provider for Upbit markets with WebSocket streaming
/// and REST API support. Optimized for HFT applications with nanosecond precision.
#[derive(Debug)]
pub struct UpbitProvider {
    /// 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>>>>,

    /// Cached instruments
    instruments: Arc<RwLock<FxHashMap<String, Box<dyn Instrument>>>>,

    /// 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,

    /// Authentication handler (optional, for private endpoints)
    auth: Option<UpbitAuth>,
}

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

impl UpbitProvider {
    /// Create a new Upbit provider with default configuration
    #[inline]
    #[must_use]
    pub fn new() -> Self {
        Self::with_config(None)
    }

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

        // Set default WebSocket configuration
        default_config.websocket_config.base_url = UPBIT_WS_URL.into();
        default_config.websocket_config.ping_interval_milliseconds = 30000; // 30 seconds
        default_config.websocket_config.use_compression = true;

        // Set default REST API configuration
        default_config.rest_config.base_url = UPBIT_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();

        // Extract auth config if present
        let auth = if let (Some(api_key), Some(secret_key)) =
            (&config.auth_config.api_key, &config.auth_config.secret_key)
        {
            if !api_key.is_empty() && !secret_key.is_empty() {
                Some(UpbitAuth::new(UpbitAuthConfig::new(
                    api_key.clone(),
                    secret_key.clone(),
                )))
            } else {
                None
            }
        } else {
            None
        };

        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())),
            instruments: 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,
            auth,
        }
    }

    /// 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::Upbit, 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(())
    }

    /// Create a WebSocket connection for trade data
    #[inline]
    async fn create_trade_connection(
        &self,
        symbols: SmallVec<[String; 8]>,
        tx: mpsc::Sender<TradeMessage>,
        subscription: Vec<simd_json::OwnedValue>,
    ) -> 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 = UpbitMessageHandler::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<OrderbookMessage>,
        subscription: Vec<simd_json::OwnedValue>,
    ) -> 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 = UpbitMessageHandler::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)
    }
}

impl HttpClientProvider for UpbitProvider {
    fn http_client(&self) -> &reqwest::Client {
        &self.http_client
    }
}

#[async_trait]
impl Provider for UpbitProvider {
    type TradeMessage = TradeMessage;
    type DepthMessage = OrderbookMessage;
    type InstrumentMessage = Value;

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

    fn venue(&self) -> rusty_model::venues::Venue {
        rusty_model::venues::Venue::Upbit
    }

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

    async fn init(&mut self) -> Result<()> {
        // Initialize the provider
        match self.get_instruments().await {
            Ok(_) => Ok(()),
            Err(e) => Err(anyhow!("Failed to initialize Upbit provider: {}", e)),
        }
    }

    #[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 connection_id = format!("trade-{}", symbols.join(","));
        let subscription =
            create_trade_subscription(&connection_id, symbols.clone(), None, None, false);

        // 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 for level2 updates
        let connection_id = format!("orderbook-{}", symbols.join(","));
        let subscription =
            create_orderbook_subscription(&connection_id, symbols.clone(), None, None, None, false);

        // 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>>> {
        // Check if instruments are already cached
        if !self.instruments.read().is_empty() {
            // Use clone_box to create clones of boxed trait objects
            let mut instruments = Vec::new();
            for instrument in self.instruments.read().values() {
                instruments.push(instrument.clone_box());
            }
            return Ok(instruments);
        }

        // Fetch instruments from Upbit API
        let url = format!(
            "{}/market/all?isDetails=true",
            self.config.rest_config.base_url
        );

        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 instruments: {}", error_text));
        }

        // Parse the response as JSON array 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 markets: Vec<simd_json::OwnedValue> = simd_json::from_slice(&mut bytes_vec)
            .map_err(|e| anyhow!("Failed to parse JSON: {}", e))?;

        // Convert to instruments and cache them
        let mut instruments = Vec::with_capacity(markets.len());
        let mut instruments_map = self.instruments.write();

        for market in markets {
            // Extract necessary fields
            let market_id = market["market"]
                .as_str()
                .ok_or_else(|| anyhow!("Missing market field"))?;
            let base_currency = market["base_asset"].as_str().unwrap_or_else(|| {
                market["baseCurrency"]
                    .as_str()
                    .unwrap_or_else(|| market_id.split('-').nth(1).unwrap_or(""))
            });
            let quote_currency = market["quote_asset"].as_str().unwrap_or_else(|| {
                market["quoteCurrency"]
                    .as_str()
                    .unwrap_or_else(|| market_id.split('-').nth(0).unwrap_or(""))
            });

            // Create instrument
            let instrument = Box::new(rusty_model::instruments::SpotInstrument::new(
                market_id,
                base_currency,
                quote_currency,
                rusty_model::venues::Venue::Upbit,
            ));

            // Add to result and cache
            instruments.push(instrument.clone_box());
            instruments_map.insert(market_id.into(), instrument);
        }

        Ok(instruments)
    }

    #[inline]
    async fn get_historical_trades(
        &self,
        symbol: &str,
        limit: Option<u32>,
    ) -> Result<Vec<MarketTrade>> {
        // Construct the URL
        let limit = limit.unwrap_or(100).min(500); // Upbit typically limits to 500
        let url = format!(
            "{}/trades/ticks?market={}&count={}",
            self.config.rest_config.base_url, symbol, limit
        );

        // Make the request
        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 historical trades: {}", 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 trades_data: Vec<simd_json::OwnedValue> = simd_json::from_slice(&mut bytes_vec)
            .map_err(|e| anyhow!("Failed to parse JSON: {}", e))?;

        // Convert to Trade objects
        let mut trades = Vec::with_capacity(trades_data.len());
        let instrument_id =
            rusty_model::instruments::InstrumentId::new(symbol, rusty_model::venues::Venue::Upbit);

        for trade_data in trades_data {
            // Extract trade info
            let timestamp_ms = match trade_data["timestamp"].as_u64() {
                Some(ts) => ts,
                None => {
                    // Try parsing from sequential id or trade time
                    if let Some(trade_time_str) = trade_data["trade_time_utc"].as_str() {
                        // Parse ISO8601 timestamp
                        if let Some(ts) = timestamp::iso8601_to_nanos(trade_time_str) {
                            ts / 1_000_000 // Convert to milliseconds
                        } else {
                            // Use current time as fallback
                            self.clock.raw() / 1_000_000 // Convert ns to ms
                        }
                    } else {
                        // Use current time as fallback
                        self.clock.raw() / 1_000_000 // Convert ns to ms
                    }
                }
            };

            // Convert to nanoseconds using optimized method
            let exchange_time = self.convert_exchange_timestamp(timestamp_ms);

            // Parse price and quantity
            let price_str_value: std::string::String;
            let price_str = match trade_data["trade_price"].as_str() {
                Some(s) => s,
                None => {
                    price_str_value = trade_data["trade_price"]
                        .as_f64()
                        .map(|p| p.to_string())
                        .unwrap_or_else(|| "0.0".to_string());
                    &price_str_value
                }
            };

            let quantity_str_value: std::string::String;
            let quantity_str = match trade_data["trade_volume"].as_str() {
                Some(s) => s,
                None => {
                    quantity_str_value = trade_data["trade_volume"]
                        .as_f64()
                        .map(|q| q.to_string())
                        .unwrap_or_else(|| "0.0".to_string());
                    &quantity_str_value
                }
            };

            let price =
                rust_decimal::Decimal::from_str(price_str).unwrap_or(rust_decimal::Decimal::ZERO);
            let quantity = rust_decimal::Decimal::from_str(quantity_str)
                .unwrap_or(rust_decimal::Decimal::ZERO);

            // Determine trade direction
            let direction = if trade_data["ask_bid"].as_str().unwrap_or("") == "ASK" {
                OrderSide::Sell
            } else {
                OrderSide::Buy
            };

            // Create Trade object
            let trade = MarketTrade {
                timestamp: self.clock.now(),
                exchange_time_ns: exchange_time,
                price,
                quantity,
                direction,
                instrument_id: instrument_id.clone(),
            };

            trades.push(trade);
        }

        Ok(trades)
    }

    #[inline]
    async fn get_orderbook_snapshot(
        &self,
        symbol: &str,
        depth: Option<u32>,
    ) -> Result<OrderBookSnapshot> {
        // Construct the URL
        let url = format!(
            "{}/orderbook?markets={}",
            self.config.rest_config.base_url, symbol
        );

        // Make the request
        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 orderbooks: Vec<simd_json::OwnedValue> = simd_json::from_slice(&mut bytes_vec)
            .map_err(|e| anyhow!("Failed to parse JSON: {}", e))?;

        // Ensure we have at least one orderbook
        if orderbooks.is_empty() {
            return Err(anyhow!("Empty orderbook response"));
        }

        // Get the first orderbook
        let orderbook = &orderbooks[0];

        // Extract timestamp
        let timestamp_ms = orderbook["timestamp"]
            .as_u64()
            .unwrap_or_else(|| self.clock.raw() / 1_000_000);

        // Convert to nanoseconds using optimized method
        let exchange_time = self.convert_exchange_timestamp(timestamp_ms);

        // Create instrument ID
        let instrument_id =
            rusty_model::instruments::InstrumentId::new(symbol, rusty_model::venues::Venue::Upbit);

        // Create empty OrderBookSnapshot
        let mut order_book_depth =
            OrderBookSnapshot::new_empty(instrument_id, self.clock.raw(), exchange_time);

        // Extract order book units
        let orderbook_units = orderbook["orderbook_units"]
            .as_array()
            .ok_or_else(|| anyhow!("Missing orderbook_units"))?;

        // Limit depth if specified
        let limit = depth.unwrap_or(orderbook_units.len() as u32) as usize;
        let orderbook_units = &orderbook_units[0..std::cmp::min(limit, orderbook_units.len())];

        // Process each level
        for unit in orderbook_units {
            // Extract bid
            if let (Some(bid_price), Some(bid_size)) =
                (unit["bid_price"].as_f64(), unit["bid_size"].as_f64())
            {
                let bid_price_dec = rust_decimal::Decimal::from_f64(bid_price)
                    .unwrap_or(rust_decimal::Decimal::ZERO);
                let bid_size_dec = rust_decimal::Decimal::from_f64(bid_size)
                    .unwrap_or(rust_decimal::Decimal::ZERO);

                order_book_depth.add_bid(bid_price_dec, bid_size_dec);
            }

            // Extract ask
            if let (Some(ask_price), Some(ask_size)) =
                (unit["ask_price"].as_f64(), unit["ask_size"].as_f64())
            {
                let ask_price_dec = rust_decimal::Decimal::from_f64(ask_price)
                    .unwrap_or(rust_decimal::Decimal::ZERO);
                let ask_size_dec = rust_decimal::Decimal::from_f64(ask_size)
                    .unwrap_or(rust_decimal::Decimal::ZERO);

                order_book_depth.add_ask(ask_price_dec, ask_size_dec);
            }
        }

        Ok(order_book_depth)
    }

    #[inline]
    async fn get_realtime_orderbook(&self, symbol: &str) -> Result<SharedSimdOrderBook> {
        // First, get an initial snapshot
        let snapshot = self.get_orderbook_snapshot(symbol, None).await?;

        // Create an OrderBook first, then wrap in SharedOrderBook
        // Convert from book_snapshot::PriceLevel to orderbook::PriceLevel
        let bids: SmallVec<[rusty_model::data::orderbook::PriceLevel; 64]> = snapshot
            .bids
            .iter()
            .map(|level| rusty_model::data::orderbook::PriceLevel::new(level.price, level.quantity))
            .collect();

        let asks: SmallVec<[rusty_model::data::orderbook::PriceLevel; 64]> = snapshot
            .asks
            .iter()
            .map(|level| rusty_model::data::orderbook::PriceLevel::new(level.price, level.quantity))
            .collect();

        let orderbook = OrderBook::new(
            snapshot.instrument_id.symbol.clone(),
            snapshot.timestamp_event,
            snapshot.timestamp_init,
            bids,
            asks,
        );
        let shared_orderbook = SharedSimdOrderBook::from_orderbook(&orderbook);

        // Create a channel for receiving orderbook updates
        let (_, _depth_rx) = mpsc::channel::<OrderbookMessage>(1024);

        // Subscribe to orderbook updates
        let mut symbols = SmallVec::<[String; 8]>::new();
        symbols.push(symbol.into());
        let mut msg_rx = self.subscribe_orderbook(symbols).await?;

        // Spawn a task to update the shared orderbook
        let shared_orderbook_clone = shared_orderbook.clone();
        let clock = self.clock.clone();

        tokio::spawn(async move {
            while let Some(msg) = msg_rx.recv().await {
                // Process the message and update the shared orderbook
                // The OrderbookMessage has the fields directly
                {
                    let data = &msg; // For clarity - now 'data' refers to the message
                    // Extract instrument ID
                    let instrument_id = rusty_model::instruments::InstrumentId::new(
                        data.code.clone(),
                        rusty_model::venues::Venue::Upbit,
                    );

                    // Extract timestamp - timestamp is already u64, not Option<u64>
                    let timestamp_ms = if data.timestamp == 0 {
                        clock.raw() / 1_000_000 // Convert ns to ms as fallback
                    } else {
                        data.timestamp
                    };

                    // Convert to nanoseconds
                    let exchange_time = timestamp::exchange_time_to_nanos(
                        timestamp_ms,
                        TimestampFormat::Milliseconds,
                    );

                    // Update the shared orderbook directly
                    // Upbit sends full depth snapshots, so we can directly update the OrderBook
                    // This is more efficient than creating intermediate OrderBookSnapshot
                    shared_orderbook_clone.write(|ob| {
                        // Update timestamps
                        ob.exchange_timestamp_ns = exchange_time;
                        ob.system_timestamp_ns = clock.raw();

                        // Clear existing levels to apply fresh snapshot
                        ob.bids.clear();
                        ob.asks.clear();

                        // Process orderbook units and directly populate bids/asks
                        // We need to collect and sort the levels since Upbit may not send them sorted
                        let mut bids =
                            SmallVec::<[rusty_model::data::orderbook::PriceLevel; 64]>::new();
                        let mut asks =
                            SmallVec::<[rusty_model::data::orderbook::PriceLevel; 64]>::new();

                        for unit in &data.orderbook_units {
                            // Add bid level (only if size > 0)
                            if unit.bid_size > rust_decimal::Decimal::ZERO {
                                bids.push(rusty_model::data::orderbook::PriceLevel::new(
                                    unit.bid_price,
                                    unit.bid_size,
                                ));
                            }

                            // Add ask level (only if size > 0)
                            if unit.ask_size > rust_decimal::Decimal::ZERO {
                                asks.push(rusty_model::data::orderbook::PriceLevel::new(
                                    unit.ask_price,
                                    unit.ask_size,
                                ));
                            }
                        }

                        // Sort bids in descending order (highest price first)
                        bids.sort_by(|a, b| b.price.cmp(&a.price));

                        // Sort asks in ascending order (lowest price first)
                        asks.sort_by(|a, b| a.price.cmp(&b.price));

                        // Convert and assign the sorted levels to the order book
                        ob.bids = SimdPriceLevels::from_smallvec(&bids);
                        ob.asks = SimdPriceLevels::from_smallvec(&asks);
                    });
                }
            }
        });

        Ok(shared_orderbook)
    }

    #[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!(
                "{}/v1/market/all",
                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> {
        UPBIT_RATE_LIMITS.to_vec()
    }

    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(())
    }
}

/// Simple message handler for Upbit WebSocket messages
#[derive(Debug)]
struct UpbitMessageHandler {
    handler_type: HandlerType,
    clock: Clock,
    stats: Arc<RwLock<ConnectionStats>>,
    trade_tx: Option<mpsc::Sender<TradeMessage>>,
    orderbook_tx: Option<mpsc::Sender<OrderbookMessage>>,
    subscription: Vec<simd_json::OwnedValue>,
}

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

impl UpbitMessageHandler {
    const fn new_trade_handler(
        clock: Clock,
        stats: Arc<RwLock<ConnectionStats>>,
        tx: mpsc::Sender<TradeMessage>,
        subscription: Vec<simd_json::OwnedValue>,
    ) -> Self {
        Self {
            handler_type: HandlerType::Trade,
            clock,
            stats,
            trade_tx: Some(tx),
            orderbook_tx: None,
            subscription,
        }
    }

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

    /// Helper method to parse and send messages with consistent error handling
    ///
    /// This method encapsulates the common pattern of:
    /// 1. Parsing message bytes to a specific type
    /// 2. Checking if sender is available
    /// 3. Attempting to send the message
    /// 4. Logging warnings on send failures
    ///
    /// # Parameters
    /// - `message_bytes`: Mutable reference to message bytes for parsing
    /// - `sender`: Optional reference to the message sender
    /// - `message_context`: Context string for error logging (e.g., "trade", "orderbook")
    async fn parse_and_send_message<T>(
        message_bytes: &mut [u8],
        sender: Option<&mpsc::Sender<T>>,
        message_context: &str,
    ) where
        T: serde::de::DeserializeOwned + Send,
    {
        if let Ok(parsed_message) = simd_json::from_slice::<T>(message_bytes)
            && let Some(tx) = sender
            && let Err(e) = tx.send(parsed_message).await
        {
            log::warn!("Failed to send {message_context} message: {e}");
        }
    }
}

#[async_trait]
impl rusty_common::websocket::MessageHandler for UpbitMessageHandler {
    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() {
            UpbitProvider::update_receive_stats(self.stats.clone(), text.len(), local_time);

            // Parse the message as JSON first
            let mut message_bytes = text.as_bytes().to_vec();
            let json_value: simd_json::OwnedValue = simd_json::from_slice(&mut message_bytes)
                .map_err(|e| {
                    WebSocketError::MessageProcessingError(format!("Failed to parse JSON: {e}"))
                })?;

            // Check if this is a data message
            if let Some(msg_type) = json_value.get("type").and_then(|v| v.as_str()) {
                match self.handler_type {
                    HandlerType::Trade if msg_type == "trade" => {
                        Self::parse_and_send_message::<TradeMessage>(
                            &mut message_bytes,
                            self.trade_tx.as_ref(),
                            "trade",
                        )
                        .await;
                    }
                    HandlerType::Orderbook if msg_type == "orderbook" => {
                        Self::parse_and_send_message::<OrderbookMessage>(
                            &mut message_bytes,
                            self.orderbook_tx.as_ref(),
                            "orderbook",
                        )
                        .await;
                    }
                    _ => {
                        // Handle other message types (errors, heartbeat)
                        log::debug!("Unhandled Upbit message type: {msg_type}");
                    }
                }
            }
        }

        Ok(())
    }

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

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

        // Note: The actual sending would be handled by the WebSocketClient
        // This is just for logging purposes
        Ok(())
    }

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

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