rustrade-execution 0.10.1

use crate::{
    UnindexedAccountEvent, UnindexedAccountSnapshot,
    balance::AssetBalance,
    client::{ExecutionClient, OrderStatusClient},
    error::{
        ConnectivityError, OrderError, StreamTerminationReason, UnindexedClientError,
        UnindexedOrderError,
    },
    exchange::mock::request::MockExchangeRequest,
    fee::FeeModelConfig,
    fill::SimFillConfig,
    order::{
        Order, OrderKey, UnindexedInactiveOrder, UnindexedOrderKey,
        request::{OrderRequestCancel, OrderRequestOpen, UnindexedOrderResponseCancel},
        state::{Open, OrderState, UnindexedOrderState},
    },
    trade::TradesRead,
};
use chrono::{DateTime, Utc};
use derive_more::Constructor;
use futures::{StreamExt, stream::BoxStream};
use rustrade_instrument::{
    asset::name::AssetNameExchange,
    exchange::ExchangeId,
    instrument::{kind::InstrumentKindDiscriminant, name::InstrumentNameExchange},
};
use serde::{Deserialize, Serialize};
use tokio::sync::{broadcast, mpsc, oneshot};
use tokio_stream::wrappers::BroadcastStream;
use tracing::error;

#[derive(
    Debug, Clone, Eq, PartialEq, Ord, PartialOrd, Hash, Deserialize, Serialize, Constructor,
)]
pub struct MockExecutionConfig {
    pub mocked_exchange: ExchangeId,
    pub initial_state: UnindexedAccountSnapshot,
    pub latency_ms: u64,
    /// Fee model used by the mock exchange to compute trading fees.
    ///
    /// Defaults to [`FeeModelConfig::Zero`]. Use [`FeeModelConfig::Percentage`]
    /// for spot/futures simulation (e.g. 0.1% taker fee).
    #[serde(default)]
    pub fee_model: FeeModelConfig,
    /// Fill model used by the mock exchange to compute execution prices.
    ///
    /// Defaults to [`SimFillConfig::LastPrice`], which fills at the snapshot's `last_price`,
    /// falling back to the side of the book a taker would cross to. Switch to
    /// [`SimFillConfig::BidAsk`] or [`SimFillConfig::Midpoint`] for more realistic spread-cost
    /// simulation when market prices are injected alongside orders.
    #[serde(default)]
    pub fill_model: SimFillConfig,
}

#[derive(Debug, Constructor)]
pub struct MockExecutionClientConfig<FnTime> {
    pub mocked_exchange: ExchangeId,
    pub clock: FnTime,
    pub request_tx: mpsc::UnboundedSender<MockExchangeRequest>,
    pub event_rx: broadcast::Receiver<UnindexedAccountEvent>,
}

impl<FnTime> Clone for MockExecutionClientConfig<FnTime>
where
    FnTime: Clone,
{
    fn clone(&self) -> Self {
        Self {
            mocked_exchange: self.mocked_exchange,
            clock: self.clock.clone(),
            request_tx: self.request_tx.clone(),
            event_rx: self.event_rx.resubscribe(),
        }
    }
}

#[derive(Debug, Constructor)]
pub struct MockExecution<FnTime> {
    pub mocked_exchange: ExchangeId,
    pub clock: FnTime,
    pub request_tx: mpsc::UnboundedSender<MockExchangeRequest>,
    pub event_rx: broadcast::Receiver<UnindexedAccountEvent>,
}

impl<FnTime> Clone for MockExecution<FnTime>
where
    FnTime: Clone,
{
    fn clone(&self) -> Self {
        Self {
            mocked_exchange: self.mocked_exchange,
            clock: self.clock.clone(),
            request_tx: self.request_tx.clone(),
            event_rx: self.event_rx.resubscribe(),
        }
    }
}

impl<FnTime> MockExecution<FnTime>
where
    FnTime: Fn() -> DateTime<Utc>,
{
    pub fn time_request(&self) -> DateTime<Utc> {
        (self.clock)()
    }
}

impl<FnTime> ExecutionClient for MockExecution<FnTime>
where
    FnTime: Fn() -> DateTime<Utc> + Clone + Send + Sync,
{
    const EXCHANGE: ExchangeId = ExchangeId::Mock;

    // `MockExchange` models no expiry, funding schedule or contract chain, so `Perpetual`, `Future`
    // and `Option` have no faithful projection onto it. `Spot` and `Cfd` need none: both are
    // positions in an instrument that the mock can open, close and mark without simulating a
    // lifecycle. Kept in step with the projection in `generate_mock_exchange_instruments`.
    const SUPPORTED_KINDS: &'static [InstrumentKindDiscriminant] = &[
        InstrumentKindDiscriminant::Spot,
        InstrumentKindDiscriminant::Cfd,
    ];
    type Config = MockExecutionClientConfig<FnTime>;
    type AccountStream = BoxStream<'static, UnindexedAccountEvent>;

    fn new(config: Self::Config) -> Self {
        Self {
            mocked_exchange: config.mocked_exchange,
            clock: config.clock,
            request_tx: config.request_tx,
            event_rx: config.event_rx,
        }
    }

    async fn account_snapshot(
        &self,
        _: &[AssetNameExchange],
        _: &[InstrumentNameExchange],
    ) -> Result<UnindexedAccountSnapshot, UnindexedClientError> {
        let (response_tx, response_rx) = oneshot::channel();

        self.request_tx
            .send(MockExchangeRequest::fetch_account_snapshot(
                self.time_request(),
                response_tx,
            ))
            .map_err(|_| {
                UnindexedClientError::Connectivity(ConnectivityError::ExchangeOffline(
                    self.mocked_exchange,
                ))
            })?;

        response_rx.await.map_err(|_| {
            UnindexedClientError::Connectivity(ConnectivityError::ExchangeOffline(
                self.mocked_exchange,
            ))
        })
    }

    async fn account_stream(
        &self,
        _: &[AssetNameExchange],
        _: &[InstrumentNameExchange],
    ) -> Result<Self::AccountStream, UnindexedClientError> {
        // `scan` (not `map_while`) so the broadcast-lag terminal is delivered in-band: on lag we
        // yield one StreamTerminated event, set `done`, and the next poll ends the stream — whereas
        // `map_while` would end on lag with a silent EOF and no terminal event.
        let exchange = self.mocked_exchange;
        Ok(BroadcastStream::new(self.event_rx.resubscribe())
            .scan(false, move |done, result| {
                // `futures::StreamExt::scan` ends the stream when the closure resolves to None.
                let item = if *done {
                    // Terminal already emitted on the prior poll — end the stream now.
                    None
                } else {
                    match result {
                        Ok(event) => Some(event),
                        Err(error) => {
                            error!(
                                ?error,
                                "MockExchange Broadcast AccountStream lagged - terminating"
                            );
                            *done = true;
                            Some(UnindexedAccountEvent::stream_terminated(
                                exchange,
                                StreamTerminationReason::Error(
                                    "mock broadcast stream lagged".to_string(),
                                ),
                            ))
                        }
                    }
                };
                futures::future::ready(item)
            })
            .boxed())
    }

    async fn cancel_order(
        &self,
        request: OrderRequestCancel<ExchangeId, &InstrumentNameExchange>,
    ) -> UnindexedOrderResponseCancel {
        let (response_tx, response_rx) = oneshot::channel();

        let key = OrderKey {
            exchange: request.key.exchange,
            instrument: request.key.instrument.clone(),
            strategy: request.key.strategy.clone(),
            cid: request.key.cid.clone(),
        };

        if self
            .request_tx
            .send(MockExchangeRequest::cancel_order(
                self.time_request(),
                response_tx,
                request.into_owned_instrument(),
            ))
            .is_err()
        {
            return UnindexedOrderResponseCancel {
                key,
                state: Err(UnindexedOrderError::Connectivity(
                    ConnectivityError::ExchangeOffline(self.mocked_exchange),
                )),
            };
        }

        match response_rx.await {
            Ok(response) => response,
            Err(_) => UnindexedOrderResponseCancel {
                key,
                state: Err(UnindexedOrderError::Connectivity(
                    ConnectivityError::ExchangeOffline(self.mocked_exchange),
                )),
            },
        }
    }

    async fn open_order(
        &self,
        request: OrderRequestOpen<ExchangeId, &InstrumentNameExchange>,
    ) -> Order<ExchangeId, InstrumentNameExchange, UnindexedOrderState> {
        let (response_tx, response_rx) = oneshot::channel();

        let request = request.into_owned_instrument();

        if self
            .request_tx
            .send(MockExchangeRequest::open_order(
                self.time_request(),
                response_tx,
                // Carries the snapshot the engine sampled when it decided to send this order on
                // `RequestOpen::market`. It is the venue's only price source for a Market order,
                // which carries no limit price of its own. `None` there means nothing sampled
                // one, and the venue rejects rather than guessing.
                request.clone(),
            ))
            .is_err()
        {
            return Order {
                key: request.key,
                side: request.state.side,
                price: request.state.price,
                quantity: request.state.quantity,
                kind: request.state.kind,
                time_in_force: request.state.time_in_force,
                state: OrderState::inactive(OrderError::Connectivity(
                    ConnectivityError::ExchangeOffline(self.mocked_exchange),
                )),
            };
        }

        match response_rx.await {
            Ok(response) => response,
            Err(_) => Order {
                key: request.key,
                side: request.state.side,
                price: request.state.price,
                quantity: request.state.quantity,
                kind: request.state.kind,
                time_in_force: request.state.time_in_force,
                state: OrderState::inactive(OrderError::Connectivity(
                    ConnectivityError::ExchangeOffline(self.mocked_exchange),
                )),
            },
        }
    }

    async fn fetch_balances(
        &self,
        assets: &[AssetNameExchange],
    ) -> Result<Vec<AssetBalance<AssetNameExchange>>, UnindexedClientError> {
        let (response_tx, response_rx) = oneshot::channel();

        self.request_tx
            .send(MockExchangeRequest::fetch_balances(
                self.time_request(),
                assets.to_vec(),
                response_tx,
            ))
            .map_err(|_| {
                UnindexedClientError::Connectivity(ConnectivityError::ExchangeOffline(
                    self.mocked_exchange,
                ))
            })?;

        response_rx.await.map_err(|_| {
            UnindexedClientError::Connectivity(ConnectivityError::ExchangeOffline(
                self.mocked_exchange,
            ))
        })
    }

    async fn fetch_open_orders(
        &self,
        instruments: &[InstrumentNameExchange],
    ) -> Result<Vec<Order<ExchangeId, InstrumentNameExchange, Open>>, UnindexedClientError> {
        let (response_tx, response_rx) = oneshot::channel();

        self.request_tx
            .send(MockExchangeRequest::fetch_orders_open(
                self.time_request(),
                instruments.to_vec(),
                response_tx,
            ))
            .map_err(|_| {
                UnindexedClientError::Connectivity(ConnectivityError::ExchangeOffline(
                    self.mocked_exchange,
                ))
            })?;

        response_rx.await.map_err(|_| {
            UnindexedClientError::Connectivity(ConnectivityError::ExchangeOffline(
                self.mocked_exchange,
            ))
        })
    }

    /// The mock venue holds every trade and answers a span whole, so the read is always
    /// complete.
    async fn fetch_trades(
        &self,
        start: DateTime<Utc>,
        end: DateTime<Utc>,
        instruments: &[InstrumentNameExchange],
    ) -> Result<TradesRead<AssetNameExchange, InstrumentNameExchange>, UnindexedClientError> {
        let (response_tx, response_rx) = oneshot::channel();

        self.request_tx
            .send(MockExchangeRequest::fetch_trades(
                self.time_request(),
                response_tx,
                start,
                end,
            ))
            .map_err(|_| {
                UnindexedClientError::Connectivity(ConnectivityError::ExchangeOffline(
                    self.mocked_exchange,
                ))
            })?;

        let mut trades = response_rx.await.map_err(|_| {
            UnindexedClientError::Connectivity(ConnectivityError::ExchangeOffline(
                self.mocked_exchange,
            ))
        })?;
        if !instruments.is_empty() {
            trades.retain(|trade| instruments.contains(&trade.instrument));
        }
        Ok(TradesRead::complete(trades))
    }
}

/// Answered from the venue's own ledger: an order it filled, cancelled or retired at its deadline
/// is reported, and one still open or never held is omitted. The venue rejects an order only on
/// arrival, in the response to its open request, so it never reports `OpenFailed`.
impl<FnTime> OrderStatusClient for MockExecution<FnTime>
where
    FnTime: Fn() -> DateTime<Utc> + Clone + Send + Sync,
{
    async fn fetch_ended_orders(
        &self,
        orders: &[UnindexedOrderKey],
    ) -> Result<Vec<UnindexedInactiveOrder>, UnindexedClientError> {
        let (response_tx, response_rx) = oneshot::channel();

        self.request_tx
            .send(MockExchangeRequest::fetch_orders_ended(
                self.time_request(),
                orders.to_vec(),
                response_tx,
            ))
            .map_err(|_| {
                UnindexedClientError::Connectivity(ConnectivityError::ExchangeOffline(
                    self.mocked_exchange,
                ))
            })?;

        response_rx.await.map_err(|_| {
            UnindexedClientError::Connectivity(ConnectivityError::ExchangeOffline(
                self.mocked_exchange,
            ))
        })
    }
}

#[cfg(test)]
#[allow(clippy::unwrap_used, clippy::expect_used)] // Test code: panics on bad input are acceptable
mod tests {
    use super::*;
    use crate::{AccountEventKind, balance::Balance};
    use rust_decimal::Decimal;
    use rustrade_integration::collection::snapshot::Snapshot;

    fn filler_event() -> UnindexedAccountEvent {
        // Content is irrelevant — these events are overwritten by the lag before they can be
        // delivered; they exist only to overflow the broadcast buffer.
        UnindexedAccountEvent::new(
            ExchangeId::Mock,
            AccountEventKind::BalanceSnapshot(Snapshot::new(AssetBalance {
                asset: AssetNameExchange::new("btc"),
                balance: Balance::new(Decimal::ONE, Decimal::ONE),
                time_exchange: Utc::now(),
            })),
        )
    }

    /// `fetch_trades` asks the mock exchange for its span and keeps the requested instruments'
    /// trades, reading the span whole.
    #[tokio::test]
    async fn fetch_trades_reads_the_span_for_the_requested_instruments() {
        use crate::{
            order::id::{OrderId, StrategyId},
            trade::{AssetFees, Trade, TradeId},
        };
        use rustrade_instrument::Side;

        let (_event_tx, event_rx) = broadcast::channel::<UnindexedAccountEvent>(1);
        let (request_tx, mut request_rx) = mpsc::unbounded_channel();
        let client = MockExecution::new(ExchangeId::Mock, Utc::now, request_tx, event_rx);
        let start = Utc::now() - chrono::TimeDelta::hours(1);
        let end = Utc::now();

        let trade = move |instrument: &str| {
            Trade::new(
                TradeId::new(instrument),
                OrderId::new("order"),
                InstrumentNameExchange::new(instrument),
                StrategyId::unknown(),
                start,
                Side::Buy,
                Decimal::ONE,
                Decimal::ONE,
                None,
                AssetFees::new(AssetNameExchange::new("usdt"), Decimal::ZERO, None),
            )
        };
        let exchange = tokio::spawn(async move {
            let request = request_rx.recv().await.expect("a request");
            let crate::exchange::mock::request::MockExchangeRequestKind::FetchTrades {
                response_tx,
                start: asked_start,
                end: asked_end,
            } = request.kind
            else {
                panic!("expected FetchTrades, got {:?}", request.kind);
            };
            assert_eq!((asked_start, asked_end), (start, end));
            response_tx
                .send(vec![trade("btc_usdt"), trade("eth_usdt")])
                .expect("the client awaits the answer");
        });

        let read = client
            .fetch_trades(start, end, &[InstrumentNameExchange::new("eth_usdt")])
            .await
            .unwrap();
        exchange.await.unwrap();
        assert_eq!(read, TradesRead::complete(vec![trade("eth_usdt")]));
    }

    /// Broadcast lag is a terminal stream death: the account stream must deliver an in-band
    /// `StreamTerminated(Error)` (not a silent EOF) and then end.
    #[tokio::test]
    async fn account_stream_emits_stream_terminated_on_broadcast_lag() {
        // Capacity 1 so a burst of sends with no reader immediately overflows → lag.
        let (event_tx, event_rx) = broadcast::channel::<UnindexedAccountEvent>(1);
        let (request_tx, _request_rx) = mpsc::unbounded_channel();
        // `MockExecution` derives `Constructor`, so its inherent `new` takes the four fields
        // directly (the `ExecutionClient::new(config)` trait method is shadowed here).
        let client = MockExecution::new(ExchangeId::Mock, Utc::now, request_tx, event_rx);

        let mut stream = client.account_stream(&[], &[]).await.unwrap();

        // Overflow the buffer before the stream's receiver reads anything → force a lag.
        for _ in 0..3 {
            event_tx.send(filler_event()).unwrap();
        }

        // The lag is surfaced in-band as the terminal event...
        let event = stream.next().await.expect("expected a terminal event");
        assert!(
            matches!(
                event.kind,
                AccountEventKind::StreamTerminated(StreamTerminationReason::Error(_))
            ),
            "expected StreamTerminated(Error) on broadcast lag, got {:?}",
            event.kind
        );
        assert_eq!(event.exchange, ExchangeId::Mock);

        // ...after which the stream ends.
        assert!(
            stream.next().await.is_none(),
            "stream must end after the terminal event"
        );
    }
}