rusty_ems/exchanges/coinbase/

unified_client.rs

//! Unified Coinbase Trading Client
//!
//! This module provides a unified interface for interacting with Coinbase through multiple protocols:
//! - REST API for order management and account operations
//! - WebSocket for real-time market data and order updates
//! - FIX protocol for institutional-grade trading
//!
//! The unified client coordinates between these protocols to provide optimal performance
//! and reliability for different use cases.

use anyhow::{Result, bail};
use async_trait::async_trait;
use flume::Sender;
use log::{debug, error, info};
use parking_lot::RwLock;
use rusty_common::auth::exchanges::coinbase::CoinbaseAuth;
use rusty_common::utils::id_generation;
use rusty_model::{
    enums::OrderStatus, instruments::InstrumentId, trading_order::Order, venues::Venue,
};
use smallvec::SmallVec;
use smartstring::alias::String as SmartString;
use std::sync::Arc;

use super::websocket_trading::CoinbaseWebsocketTrading as CoinbaseWebSocketTrader;
use super::{
    rest_client::{CoinbaseOrderRequest, CoinbaseRestClient},
    //fix_client::CoinbaseFixClient,
};
use crate::execution_engine::{Exchange as ExchangeTrait, ExecutionReport};
use crate::instrument_registry::{InstrumentRegistry, OrderMetadata};
use rust_decimal::Decimal;

/// Protocol selection for order routing
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum CoinbaseProtocol {
    /// REST API - Reliable, synchronous order placement
    Rest,
    /// WebSocket - Real-time updates, REST for order placement
    WebSocket,
    /// FIX - Low-latency institutional trading
    Fix,
    /// Auto - Automatically select best protocol based on order characteristics
    Auto,
}

/// Unified Coinbase client configuration
#[derive(Debug, Clone)]
pub struct UnifiedClientConfig {
    /// Which protocol to use for order placement
    pub order_protocol: CoinbaseProtocol,
    /// Enable WebSocket for real-time updates
    pub enable_websocket: bool,
    /// Enable FIX protocol (requires institutional access)
    pub enable_fix: bool,
    /// Use sandbox environment
    pub sandbox: bool,
}

impl Default for UnifiedClientConfig {
    fn default() -> Self {
        Self {
            order_protocol: CoinbaseProtocol::Auto,
            enable_websocket: true,
            enable_fix: false,
            sandbox: false,
        }
    }
}

/// Unified Coinbase Trading Client
///
/// Coordinates between REST, WebSocket, and FIX protocols to provide
/// optimal trading experience with automatic failover and protocol selection.
pub struct CoinbaseUnifiedClient {
    /// Shared authentication
    auth: Arc<CoinbaseAuth>,

    /// REST API client (always available)
    rest_client: Arc<CoinbaseRestClient>,

    /// WebSocket client for real-time updates
    ws_client: Option<Arc<CoinbaseWebSocketTrader>>,

    /// FIX client for institutional trading
    //fix_client: Option<Arc<CoinbaseFixClient>>,

    /// Configuration
    config: UnifiedClientConfig,

    /// Connection state
    state: Arc<RwLock<ConnectionState>>,

    /// Instrument registry for order/instrument mapping
    instrument_registry: Arc<dyn InstrumentRegistry>,
}

#[derive(Debug, Default)]
struct ConnectionState {
    rest_connected: bool,
    ws_connected: bool,
    fix_connected: bool,
}

impl CoinbaseUnifiedClient {
    /// Create a new unified client
    pub fn new(
        auth: Arc<CoinbaseAuth>,
        config: UnifiedClientConfig,
        instrument_registry: Arc<dyn InstrumentRegistry>,
    ) -> Result<Self> {
        // Create REST client (always required)
        let rest_client = Arc::new(
            CoinbaseRestClient::new(auth.clone(), config.sandbox)
                .map_err(|e| anyhow::anyhow!("Failed to create REST client: {}", e))?,
        );

        // Create WebSocket client if enabled
        let ws_client = if config.enable_websocket {
            let client = CoinbaseWebSocketTrader::new(
                auth.clone(),
                config.sandbox,
                false, // use_direct
            );
            Some(Arc::new(client))
        } else {
            None
        };

        // Create FIX client if enabled
        // let fix_client = if config.enable_fix {
        //     Some(Arc::new(CoinbaseFixClient::new(auth.clone(), config.sandbox)))
        // } else {
        //     None
        // };

        Ok(Self {
            auth,
            rest_client,
            ws_client,
            //fix_client,
            config,
            state: Arc::new(RwLock::new(ConnectionState::default())),
            instrument_registry,
        })
    }

    /// Connect all enabled protocols with report sender
    pub async fn connect_with_report_sender(
        &self,
        report_tx: Sender<ExecutionReport>,
    ) -> Result<()> {
        info!(
            "Connecting Coinbase unified client with protocols: REST={}, WS={}, FIX={}",
            true, self.config.enable_websocket, self.config.enable_fix
        );

        // REST is always connected (stateless)
        self.state.write().rest_connected = true;

        // Connect WebSocket if enabled
        if let Some(ws_client) = &self.ws_client {
            match ws_client.connect(report_tx.clone()).await {
                Ok(()) => {
                    self.state.write().ws_connected = true;
                    info!("WebSocket connected successfully");
                }
                Err(e) => {
                    error!("WebSocket connection failed: {e}");
                    // Continue without WebSocket - REST will handle everything
                }
            }
        }

        // Connect FIX if enabled
        // if let Some(fix_client) = &self.fix_client {
        //     match fix_client.connect(report_tx.clone()).await {
        //         Ok(()) => {
        //             self.state.write().fix_connected = true;
        //             info!("FIX protocol connected successfully");
        //         }
        //         Err(e) => {
        //             error!("FIX connection failed: {e}");
        //             // Continue without FIX
        //         }
        //     }
        // }

        Ok(())
    }

    /// Disconnect all protocols
    pub async fn disconnect(&self) -> Result<()> {
        info!("Disconnecting Coinbase unified client");

        // Disconnect WebSocket
        if let Some(ws_client) = &self.ws_client {
            if let Err(e) = ws_client.disconnect().await {
                error!("WebSocket disconnect error: {e}");
            }
            self.state.write().ws_connected = false;
        }

        // Disconnect FIX
        // if let Some(fix_client) = &self.fix_client {
        //     if let Err(e) = fix_client.disconnect().await {
        //         error!("FIX disconnect error: {e}");
        //     }
        //     self.state.write().fix_connected = false;
        // }

        self.state.write().rest_connected = false;

        Ok(())
    }

    /// Select the best protocol for an order
    fn select_protocol(&self, order: &Order) -> CoinbaseProtocol {
        match self.config.order_protocol {
            CoinbaseProtocol::Auto => {
                let state = self.state.read();

                // Priority: FIX > WebSocket+REST > REST
                // if state.fix_connected && self.fix_client.is_some() {
                //     CoinbaseProtocol::Fix
                // } else
                if state.ws_connected && self.ws_client.is_some() {
                    CoinbaseProtocol::WebSocket
                } else {
                    CoinbaseProtocol::Rest
                }
            }
            protocol => protocol,
        }
    }

    /// Get connection status
    #[must_use]
    pub fn get_connection_status(&self) -> (bool, bool, bool) {
        let state = self.state.read();
        (
            state.rest_connected,
            state.ws_connected,
            state.fix_connected,
        )
    }

    /// Check if any protocol is connected
    #[must_use]
    pub fn is_connected(&self) -> bool {
        let state = self.state.read();
        state.rest_connected || state.ws_connected || state.fix_connected
    }
}

impl CoinbaseUnifiedClient {
    /// Convert Order to `CoinbaseOrderRequest`
    fn order_to_coinbase_request(order: &Order) -> Result<CoinbaseOrderRequest> {
        use super::rest_client::{
            CoinbaseLimitOrder, CoinbaseMarketOrder, CoinbaseOrderConfiguration,
            CoinbaseOrderRequest, CoinbaseRestClient, CoinbaseStopLimitOrder,
        };
        use rusty_model::enums::OrderType;

        let client_order_id = order.id.to_string().into();
        let product_id = order.symbol.clone();
        let side = CoinbaseRestClient::order_side_to_string(order.side);

        let order_configuration = match order.order_type {
            OrderType::Market => CoinbaseOrderConfiguration::MarketIoc {
                market_market_ioc: CoinbaseMarketOrder {
                    base_size: Some(order.quantity.to_string().into()),
                    quote_size: None,
                },
            },
            OrderType::Limit => {
                let limit_price = order
                    .price
                    .ok_or_else(|| anyhow::anyhow!("Limit order requires price"))?
                    .to_string()
                    .into();

                CoinbaseOrderConfiguration::LimitGtc {
                    limit_limit_gtc: CoinbaseLimitOrder {
                        base_size: order.quantity.to_string().into(),
                        limit_price,
                        post_only: None,
                    },
                }
            }
            OrderType::StopLimit => {
                let limit_price = order
                    .price
                    .ok_or_else(|| anyhow::anyhow!("Stop limit order requires price"))?
                    .to_string()
                    .into();
                let stop_price = order
                    .stop_price
                    .ok_or_else(|| anyhow::anyhow!("Stop limit order requires stop price"))?
                    .to_string()
                    .into();

                // Determine stop direction based on side and current price
                let stop_direction = if order.side == rusty_model::enums::OrderSide::Buy {
                    "STOP_DIRECTION_STOP_UP".into()
                } else {
                    "STOP_DIRECTION_STOP_DOWN".into()
                };

                CoinbaseOrderConfiguration::StopLimitGtc {
                    stop_limit_stop_limit_gtc: CoinbaseStopLimitOrder {
                        base_size: order.quantity.to_string().into(),
                        limit_price,
                        stop_price,
                        stop_direction,
                    },
                }
            }
            _ => bail!("Unsupported order type: {:?}", order.order_type),
        };

        Ok(CoinbaseOrderRequest {
            client_order_id,
            product_id,
            side,
            order_configuration,
        })
    }

    /// Place order with proper error handling and report generation
    async fn place_order_internal(
        &self,
        order: Order,
        report_tx: Sender<ExecutionReport>,
    ) -> Result<()> {
        // Convert Order to CoinbaseOrderRequest
        let coinbase_request = Self::order_to_coinbase_request(&order)?;

        // Store order details before moving
        let order_id = order.id;
        let instrument_id = InstrumentId::new(order.symbol.clone(), order.venue);
        let quantity = order.quantity;

        // Place order via REST API
        match self.rest_client.place_order(&coinbase_request).await {
            Ok(response) => {
                if response.success {
                    // Cache order mapping for future lookups
                    let metadata = OrderMetadata {
                        size: Some(quantity),
                        order_type: None, // Will be determined from order_configuration
                        side: Some(coinbase_request.side.clone()),
                    };
                    self.instrument_registry.cache_order_mapping(
                        &response.order_id,
                        instrument_id.clone(),
                        Some(metadata),
                    );

                    // Send acceptance report
                    let report = ExecutionReport {
                        id: id_generation::generate_ack_id(&order_id.to_string()),
                        order_id: response.order_id.clone(),
                        exchange_timestamp: rusty_common::time::get_epoch_timestamp_ns(),
                        system_timestamp: rusty_common::time::get_epoch_timestamp_ns(),
                        instrument_id,
                        status: OrderStatus::New,
                        filled_quantity: Decimal::ZERO,
                        remaining_quantity: quantity,
                        execution_price: None,
                        reject_reason: None,
                        exchange_execution_id: Some(response.order_id),
                        is_final: false,
                    };

                    if let Err(e) = report_tx.send_async(report).await {
                        error!("Failed to send acceptance report: {e}");
                    }
                } else {
                    // Send rejection report
                    let report = ExecutionReport {
                        id: id_generation::generate_rejection_id(&order_id.to_string()),
                        order_id: coinbase_request.client_order_id.clone(),
                        exchange_timestamp: rusty_common::time::get_epoch_timestamp_ns(),
                        system_timestamp: rusty_common::time::get_epoch_timestamp_ns(),
                        instrument_id,
                        status: OrderStatus::Rejected,
                        filled_quantity: Decimal::ZERO,
                        remaining_quantity: quantity,
                        execution_price: None,
                        reject_reason: Some(response.failure_reason),
                        exchange_execution_id: None,
                        is_final: true,
                    };

                    if let Err(e) = report_tx.send_async(report).await {
                        error!("Failed to send rejection report: {e}");
                    }
                }
                Ok(())
            }
            Err(e) => {
                // Send error report
                let report = ExecutionReport {
                    id: id_generation::generate_report_id("error", &order_id.to_string()),
                    order_id: coinbase_request.client_order_id,
                    exchange_timestamp: 0,
                    system_timestamp: rusty_common::time::get_epoch_timestamp_ns(),
                    instrument_id,
                    status: OrderStatus::Rejected,
                    filled_quantity: Decimal::ZERO,
                    remaining_quantity: quantity,
                    execution_price: None,
                    reject_reason: Some(e.to_string().into()),
                    exchange_execution_id: None,
                    is_final: true,
                };

                if let Err(send_err) = report_tx.send_async(report).await {
                    error!("Failed to send error report: {send_err}");
                }

                Err(e)
            }
        }
    }

    /// Cancel order with proper report generation
    async fn cancel_order_internal(
        &self,
        order_id: SmartString,
        report_tx: Sender<ExecutionReport>,
    ) -> Result<()> {
        match self.rest_client.cancel_order(&order_id).await {
            Ok(response) => {
                if response.success {
                    // Lookup instrument from cached mapping
                    let instrument_id = self
                        .instrument_registry
                        .lookup_by_order_id(&order_id)
                        .unwrap_or_else(|| {
                            // Fallback: create instrument with normalized symbol
                            let normalized_symbol = self
                                .instrument_registry
                                .normalize_symbol(&order_id, Venue::Coinbase);
                            InstrumentId::new(normalized_symbol, Venue::Coinbase)
                        });

                    // Remove mapping since order is cancelled
                    self.instrument_registry.remove_mapping(&order_id);

                    // Send cancellation report
                    let report = ExecutionReport {
                        id: id_generation::generate_cancel_id(&order_id),
                        order_id: order_id.clone(),
                        exchange_timestamp: rusty_common::time::get_epoch_timestamp_ns(),
                        system_timestamp: rusty_common::time::get_epoch_timestamp_ns(),
                        instrument_id,
                        status: OrderStatus::Cancelled,
                        filled_quantity: Decimal::ZERO,
                        remaining_quantity: Decimal::ZERO,
                        execution_price: None,
                        reject_reason: None,
                        exchange_execution_id: None,
                        is_final: true,
                    };

                    if let Err(e) = report_tx.send_async(report).await {
                        error!("Failed to send cancellation report: {e}");
                    }
                } else {
                    // Lookup instrument from cached mapping
                    let instrument_id = self
                        .instrument_registry
                        .lookup_by_order_id(&order_id)
                        .unwrap_or_else(|| {
                            // Fallback: create instrument with normalized symbol
                            let normalized_symbol = self
                                .instrument_registry
                                .normalize_symbol(&order_id, Venue::Coinbase);
                            InstrumentId::new(normalized_symbol, Venue::Coinbase)
                        });

                    // Cancellation failed
                    let report = ExecutionReport {
                        id: id_generation::generate_report_id("cancel_fail", order_id.as_ref()),
                        order_id: order_id.clone(),
                        exchange_timestamp: rusty_common::time::get_epoch_timestamp_ns(),
                        system_timestamp: rusty_common::time::get_epoch_timestamp_ns(),
                        instrument_id,
                        status: OrderStatus::Rejected,
                        filled_quantity: Decimal::ZERO,
                        remaining_quantity: Decimal::ZERO,
                        execution_price: None,
                        reject_reason: Some(response.failure_reason),
                        exchange_execution_id: None,
                        is_final: false,
                    };

                    if let Err(e) = report_tx.send_async(report).await {
                        error!("Failed to send cancel failure report: {e}");
                    }
                }
                Ok(())
            }
            Err(e) => Err(e),
        }
    }
}

#[async_trait]
impl ExchangeTrait for CoinbaseUnifiedClient {
    async fn place_order(&self, order: Order, report_tx: Sender<ExecutionReport>) -> Result<()> {
        let protocol = self.select_protocol(&order);
        debug!("Placing order {} via {:?} protocol", order.id, protocol);

        match protocol {
            CoinbaseProtocol::Rest => self.place_order_internal(order, report_tx).await,
            CoinbaseProtocol::WebSocket => {
                // WebSocket uses REST for order placement
                if let Some(ws_client) = &self.ws_client {
                    ws_client.place_order(order, report_tx).await
                } else {
                    // Fallback to REST
                    self.place_order_internal(order, report_tx).await
                }
            }
            CoinbaseProtocol::Fix => {
                // if let Some(fix_client) = &self.fix_client {
                //     fix_client.place_order(order, report_tx).await
                // } else {
                // Fallback to REST
                self.place_order_internal(order, report_tx).await
                // }
            }
            CoinbaseProtocol::Auto => {
                // This should never happen as select_protocol handles Auto
                self.place_order_internal(order, report_tx).await
            }
        }
    }

    async fn cancel_order(
        &self,
        order_id: SmartString,
        report_tx: Sender<ExecutionReport>,
    ) -> Result<()> {
        // Use the same protocol selection logic
        let protocol = self.config.order_protocol;
        debug!("Cancelling order {} via {:?} protocol", order_id, protocol);

        match protocol {
            CoinbaseProtocol::Rest | CoinbaseProtocol::Auto => {
                self.cancel_order_internal(order_id, report_tx).await
            }
            CoinbaseProtocol::WebSocket => {
                if let Some(ws_client) = &self.ws_client {
                    ws_client.cancel_order(order_id, report_tx).await
                } else {
                    self.cancel_order_internal(order_id, report_tx).await
                }
            }
            CoinbaseProtocol::Fix => {
                // if let Some(fix_client) = &self.fix_client {
                //     fix_client.cancel_order(order_id, report_tx).await
                // } else {
                self.cancel_order_internal(order_id, report_tx).await
                // }
            }
        }
    }

    async fn modify_order(
        &self,
        order_id: SmartString,
        new_quantity: Option<rust_decimal::Decimal>,
        new_price: Option<rust_decimal::Decimal>,
        report_tx: Sender<ExecutionReport>,
    ) -> Result<()> {
        // Coinbase Advanced Trade doesn't support order modification
        // Must cancel and replace
        bail!("Order modification not supported by Coinbase. Please cancel and place a new order.")
    }

    async fn get_order_status(&self, order_id: &str) -> Result<OrderStatus> {
        // Always use REST for order status queries
        let order = self.rest_client.get_order(order_id).await?;
        Ok(CoinbaseRestClient::string_to_order_status(&order.status))
    }

    /// Get the exchange venue
    fn venue(&self) -> Venue {
        Venue::Coinbase
    }

    /// Cancel all orders
    async fn cancel_all_orders(
        &self,
        instrument_id: Option<InstrumentId>,
        report_tx: Sender<ExecutionReport>,
    ) -> Result<()> {
        // Use REST API to cancel all orders
        // TODO: Implement filtering by instrument_id
        bail!("Cancel all orders not yet implemented for Coinbase unified client")
    }

    /// Connect to the exchange
    async fn connect(&self, report_sender: Sender<ExecutionReport>) -> Result<()> {
        // Use the proper connect_with_report_sender method
        self.connect_with_report_sender(report_sender).await
    }

    /// Disconnect from the exchange
    async fn disconnect(&self) -> Result<()> {
        // Call the internal disconnect method
        info!("Disconnecting Coinbase unified client");

        // Disconnect WebSocket
        if let Some(ws_client) = &self.ws_client {
            if let Err(e) = ws_client.disconnect().await {
                error!("WebSocket disconnect error: {e}");
            }
            self.state.write().ws_connected = false;
        }

        // Disconnect FIX
        // if let Some(fix_client) = &self.fix_client {
        //     if let Err(e) = fix_client.disconnect().await {
        //         error!("FIX disconnect error: {e}");
        //     }
        //     self.state.write().fix_connected = false;
        // }

        self.state.write().rest_connected = false;

        Ok(())
    }

    /// Check if connected
    async fn is_connected(&self) -> bool {
        let state = self.state.read();
        state.rest_connected || state.ws_connected || state.fix_connected
    }

    /// Get available instruments
    async fn get_instruments(&self) -> Result<SmallVec<[InstrumentId; 32]>> {
        // Use REST API to get products/instruments
        let products = self.rest_client.get_products().await?;
        let instruments: SmallVec<[InstrumentId; 32]> = products
            .into_iter()
            .filter(|p| p.status == "online" && !p.trading_disabled)
            .map(|p| InstrumentId::new(p.product_id, Venue::Coinbase))
            .collect();
        Ok(instruments)
    }

    /// Send FIX message (not supported for unified client)
    async fn send_fix_message(&self, _message: Vec<u8>) -> Result<()> {
        bail!("FIX message sending not supported through unified client")
    }

    /// Receive FIX message (not supported for unified client)
    async fn receive_fix_message(&self) -> Result<Vec<u8>> {
        bail!("FIX message receiving not supported through unified client")
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use rusty_common::auth::exchanges::coinbase::CoinbaseAuth;
    use std::str::FromStr;

    #[test]
    fn test_unified_client_creation() {
        let auth = Arc::new(CoinbaseAuth::new_hmac(
            "test_key".into(),
            "test_secret".into(),
        ));

        let config = UnifiedClientConfig {
            order_protocol: CoinbaseProtocol::Auto,
            enable_websocket: true,
            enable_fix: false,
            sandbox: true,
        };

        let registry = crate::instrument_registry::create_shared_registry();
        let client = CoinbaseUnifiedClient::new(auth, config, registry);
        assert!(client.is_ok());

        let client = client.unwrap();
        assert!(!client.is_connected());
    }

    #[test]
    fn test_protocol_selection() {
        let auth = Arc::new(CoinbaseAuth::new_hmac(
            "test_key".into(),
            "test_secret".into(),
        ));

        let config = UnifiedClientConfig {
            order_protocol: CoinbaseProtocol::Rest,
            enable_websocket: true,
            enable_fix: false,
            sandbox: true,
        };

        let registry = crate::instrument_registry::create_shared_registry();
        let client = CoinbaseUnifiedClient::new(auth, config, registry).unwrap();

        let order = Order {
            id: rusty_model::types::OrderId::new(),
            client_id: rusty_model::types::ClientId::new("test"),
            symbol: "BTC-USD".into(),
            side: rusty_model::enums::OrderSide::Buy,
            order_type: rusty_model::enums::OrderType::Limit,
            quantity: rust_decimal::Decimal::from_str("0.01").unwrap(),
            price: Some(rust_decimal::Decimal::from_str("50000").unwrap()),
            stop_price: None,
            exchange_order_id: None,
            venue: rusty_model::venues::Venue::Coinbase,
            filled_quantity: rust_decimal::Decimal::ZERO,
            average_fill_price: None,
            status: OrderStatus::New,
            creation_time_ns: rusty_common::time::get_epoch_timestamp_ns(),
            update_time_ns: rusty_common::time::get_epoch_timestamp_ns(),
            time_in_force: rusty_model::enums::TimeInForce::GTC,
            metadata: simd_json::json!(null),
        };

        let protocol = client.select_protocol(&order);
        assert_eq!(protocol, CoinbaseProtocol::Rest);
    }
}