scuriolus 0.1.0

mod data;

use super::clock::Clock;
use super::market::{mexc_base::Kline, Balance, KlinesParams, Market};
use super::order::{Crypto, Order, OrderDraft, OrderStatus};
use anyhow::Error;
use data::Data;

const DB_ADDRESS: &str = "db/main/";

/// The [`Core`]  is responsible for interfacing with a given [`Market`]'s API and tracking all orders placed by the strategies.
#[derive(Debug)]
pub struct Core<M: Market> {
    name: String,
    market: M,
    data: Data,
    clock: Box<dyn Clock>,
}

impl<M: Market> Core<M> {
    pub async fn new(name: String, market: M) -> Result<Self, Error> {
        Self::new_with_db_address(name, market, DB_ADDRESS.to_string()).await
    }

    async fn new_with_db_address(
        name: String,
        market: M,
        db_address: String,
    ) -> Result<Self, Error> {
        let data = Data::load_data(name.clone(), db_address).await?;
        let clock = market.clock();

        Ok(Core {
            name,
            market,
            data,
            clock,
        })
    }

    /// Send an order to the market and save it to the database.
    ///
    /// First, it creates an `Order` from the given `OrderDraft` and sends it to the market.
    /// If the sending is successful, it sets the market's ID of the order and saves the order
    /// to the database. If the sending fails, it sets the status of the order to `Failed` and logs
    /// the error.
    ///
    /// # Arguments
    ///
    /// * `draft` - The draft of the order to be sent.
    ///
    /// # Returns
    ///
    /// * `Result<String, Error>` - On success, returns the ID of the order as a string.
    pub async fn send_order(&self, draft: OrderDraft) -> Result<String, Error> {
        let mut order = Order::new(
            draft.asset,
            draft.currency,
            draft.side,
            draft.order_type,
            draft.amount,
            draft.price,
        )?;

        match self.market.send_order(&order).await {
            Ok(id) => {
                order.set_market_id(id);
            }
            Err(e) => {
                order.set_status(OrderStatus::Failed);
                tracing::error!(
                    "Core {}: Failed to send order {}: {:#?}",
                    self.name,
                    order.id(),
                    e
                );
            }
        }

        let id = order.id().clone();

        self.data.add_order(order).await.unwrap();

        //TODO if order at market price : update soon

        Ok(id)
    }

    pub fn clock(&self) -> Box<dyn Clock> {
        self.clock.clone_box()
    }

    /// Get the list of opened orders on the market and update their status with the information from the market.
    ///
    /// # Returns
    ///
    /// Returns the list of updated orders.
    pub async fn update_opened_orders(&self) -> Result<Vec<Order>, Error> {
        let mut orders: Vec<Order> = self.data.opened_orders().await;
        for order in &mut orders {
            self.update_order(order).await?;
        }
        Ok(orders)
    }

    /// Cancel all opened orders on the market and update their status accordingly.
    ///
    /// # Returns
    ///
    /// Returns the list of orders with their updated statuses after being canceled.
    pub async fn cancel_opened_orders(&self) -> Result<Vec<Order>, Error> {
        let mut orders: Vec<Order> = self.data.opened_orders().await;
        for order in &mut orders {
            self.cancel_order(order).await?;
        }
        Ok(orders)
    }

    /// Get an order by its id from the database.
    ///
    /// # Arguments
    ///
    /// * `id` - The id of the order to retrieve.
    ///
    /// # Returns
    ///
    /// Returns the order with the given id if it exists
    pub async fn get_order(&self, id: &str) -> Option<Order> {
        self.data.get_order(id).await
    }

    /// Returns the list of opened orders from the database.
    ///
    /// # Returns
    ///
    /// Returns a vector of opened orders.
    pub async fn get_opened_orders(&self) -> Vec<Order> {
        self.data.opened_orders().await
    }

    /// Updates an order's status and details based on the market's response.
    ///
    /// This function sends an update request for the specified order to the market. If the order's
    /// status or details have changed after the update, it updates the local order accordingly and saves the changes in the
    /// database.
    ///
    /// # Arguments
    ///
    /// * `order` - A mutable reference to the `Order` object to be updated.
    ///
    /// # Returns
    ///
    /// Returns `Ok(true)` if the order's status or details changed
    pub async fn update_order(&self, order: &mut Order) -> Result<bool, Error> {
        tracing::debug!("Updating order {}", order.id());
        let query_output = self.market.update_order(order).await?;
        if *order.status() != query_output.status.into()
            || *order.price() != Some(query_output.price)
            || *order.executed_qty() != query_output.executed_quantity
            || *order.cummulative_quote_qty() != query_output.cummulative_quote_quantity
        {
            order.set_status(query_output.status.into());
            if *order.price() != Some(query_output.price) {
                order.set_price(query_output.price)?;
            }
            order.set_executed(
                query_output.executed_quantity,
                query_output.cummulative_quote_quantity,
            );

            self.data.update_order(order.clone()).await?;
            Ok(true)
        } else {
            Ok(false)
        }
    }

    /// Cancels an order and updates its status and details based on the market's response.
    ///
    /// This function sends a cancellation request for the specified order to the market. If the order's
    /// status or details have changed after the cancellation, it updates the local order accordingly and saves the changes in the
    /// database.
    ///
    /// # Arguments
    ///
    /// * `order` - A mutable reference to the `Order` object to be canceled.
    ///
    /// # Returns
    ///
    /// Returns `Ok(true)` if the order's status or details changed (by update meanwhile or cancellation)
    pub async fn cancel_order(&self, order: &mut Order) -> Result<bool, Error> {
        tracing::debug!("Canceling order {}", order.id());
        let cancel_output = self.market.cancel_order(order).await?;
        if *order.status() != cancel_output.status.into() {
            order.set_status(cancel_output.status.into());
            if *order.price() != Some(cancel_output.price) {
                order.set_price(cancel_output.price)?;
            }
            if *order.executed_qty() != cancel_output.executed_quantity
                || *order.cummulative_quote_qty() != cancel_output.cummulative_quote_quantity
            {
                order.set_executed(
                    cancel_output.executed_quantity,
                    cancel_output.cummulative_quote_quantity,
                );
            }
            self.data.update_order(order.clone()).await?;
            Ok(true)
        } else {
            Ok(false)
        }
    }

    /// Get klines from the market.
    ///
    /// # Arguments
    ///
    /// * `params`: The parameters to query the klines with.
    ///
    /// # Returns
    ///
    /// Returns the list of klines that were requested.
    pub async fn get_klines(&self, params: KlinesParams) -> Result<Vec<Kline>, Error> {
        Ok(self.market.klines(params).await?)
    }

    /// Get the current balance of a crypto on the market.
    ///
    /// # Arguments
    ///
    /// * `crypto`: The crypto to get the balance of.
    ///
    /// # Returns
    ///
    /// Returns the balance of the given crypto.
    pub async fn get_balance(&self, crypto: Crypto) -> Result<Balance, Error> {
        Ok(self.market.get_balance(crypto).await?)
    }
}

#[cfg(test)]
mod tests {
    use std::future::ready;

    use chrono::Utc;
    use mexc_rs::spot::v3::cancel_order::CancelOrderOutput;
    use mexc_rs::spot::v3::klines::Kline;
    use mexc_rs::spot::v3::{enums::OrderStatus, query_order::QueryOrderOutput};
    use rust_decimal::Decimal;
    use serial_test::serial;

    use super::super::{
        clock::{CheatClockFactory, RunningCheatClockFactory},
        core::Core,
        market::{mexc_base::mexc_enums, KlinesParams, MockMarket},
        order::{Amount, Crypto, OrderDraft, Quantity},
    };

    const DB_ADDRESS: &str = "db/test/";

    #[tokio::test]
    #[serial]
    async fn reload_and_update_orders() {
        let mut market = MockMarket::new();

        market
            .expect_send_order()
            .once()
            .returning(|_| Box::pin(ready(Ok("123456".to_string()))));
        market.expect_clock().once().returning(|| {
            let (remote, clock) = RunningCheatClockFactory {}.get_clock(Utc::now()).unwrap();
            remote.not_blocking();
            clock
        });

        let id = {
            let core =
                Core::new_with_db_address("test".to_string(), market, DB_ADDRESS.to_string())
                    .await
                    .unwrap();
            core.data.clear().await;

            let id = core
                .send_order(OrderDraft {
                    asset: Crypto::BTC,
                    currency: Crypto::USDT,
                    side: mexc_enums::OrderSide::Buy,
                    order_type: mexc_enums::OrderType::Limit,
                    amount: Quantity::Asset(Amount::ONE),
                    price: Some(Amount::from(10000)),
                })
                .await
                .unwrap();

            drop(core);

            id
        };

        // wait for db to fully close
        tokio::time::sleep(std::time::Duration::from_millis(100)).await;

        let mut market = MockMarket::new();

        market.expect_update_order().once().returning(|order| {
            Box::pin(ready(Ok(QueryOrderOutput {
                symbol: "BTCUSDT".to_string(),
                original_client_order_id: None,
                order_id: order.id().to_string(),
                client_order_id: None,
                price: order.price().unwrap(),
                original_quantity: order.quantity().get_amount(),
                executed_quantity: *order.executed_qty(),
                cummulative_quote_quantity: *order.cummulative_quote_qty(),
                status: OrderStatus::New,
                time_in_force: None,
                order_type: *order.order_type(),
                side: *order.side(),
                stop_price: order.price().unwrap(),
                time: Utc::now(),
                update_time: Utc::now(),
                is_working: true,
            })))
        });

        market
            .expect_clock()
            .once()
            .returning(|| RunningCheatClockFactory {}.get_clock(Utc::now()).unwrap().1);

        let core = Core::new_with_db_address("test".to_string(), market, DB_ADDRESS.to_string())
            .await
            .unwrap();

        let updated_orders = core.update_opened_orders().await.unwrap();

        assert_eq!(updated_orders.len(), 1);
        assert_eq!(core.get_order(&id).await.unwrap(), updated_orders[0]);
    }

    #[tokio::test]
    #[serial]
    async fn cancel_orders() {
        let mut market = MockMarket::new();

        market
            .expect_send_order()
            .once()
            .returning(|_| Box::pin(ready(Ok("123456".to_string()))));

        market.expect_cancel_order().once().returning(|order| {
            Box::pin(ready(Ok(CancelOrderOutput {
                symbol: "BTCUSDT".to_string(),
                original_client_order_id: None,
                order_id: order.id().to_string(),
                client_order_id: None,
                price: order.price().unwrap(),
                original_quantity: order.quantity().get_amount(),
                executed_quantity: *order.executed_qty(),
                cummulative_quote_quantity: *order.cummulative_quote_qty(),
                status: OrderStatus::Canceled,
                time_in_force: None,
                order_type: *order.order_type(),
                side: *order.side(),
            })))
        });
        market.expect_clock().once().returning(|| {
            let (remote, clock) = RunningCheatClockFactory {}.get_clock(Utc::now()).unwrap();
            remote.not_blocking();
            clock
        });

        let core = Core::new_with_db_address("test".to_string(), market, DB_ADDRESS.to_string())
            .await
            .unwrap();
        core.data.clear().await;

        let id = core
            .send_order(OrderDraft {
                asset: Crypto::BTC,
                currency: Crypto::USDT,
                side: mexc_enums::OrderSide::Buy,
                order_type: mexc_enums::OrderType::Limit,
                amount: Quantity::Asset(Amount::ONE),
                price: Some(Amount::from(10000)),
            })
            .await
            .unwrap();

        let canceled_orders = core.cancel_opened_orders().await.unwrap();

        assert_eq!(canceled_orders.len(), 1);
        assert_eq!(core.get_order(&id).await.unwrap(), canceled_orders[0]);
    }

    #[tokio::test]
    #[serial]
    async fn get_klines() {
        let mut market = MockMarket::new();

        market.expect_klines().once().returning(|_| {
            Box::pin(ready(Ok(vec![Kline {
                open_time: Utc::now(),
                open: Decimal::ZERO,
                high: Decimal::from(42),
                low: Decimal::ZERO,
                close: Decimal::ONE,
                volume: Decimal::ONE,
                close_time: Utc::now(),
                quote_asset_volume: Decimal::ONE,
            }])))
        });
        market.expect_clock().once().returning(|| {
            let (remote, clock) = RunningCheatClockFactory {}.get_clock(Utc::now()).unwrap();
            remote.not_blocking();
            clock
        });

        let core = Core::new_with_db_address("test".to_string(), market, DB_ADDRESS.to_string())
            .await
            .unwrap();

        let klines = core
            .get_klines(KlinesParams {
                symbol: "BTCUSDT".to_string(),
                interval: mexc_enums::KlineInterval::OneMinute,
                start_time: None,
                end_time: None,
                limit: None,
            })
            .await
            .unwrap();

        let kline = // no clone
            Kline {
                open_time: Utc::now(),
                open: Decimal::ZERO,
                high: Decimal::from(42),
                low: Decimal::ZERO,
                close: Decimal::ONE,
                volume: Decimal::ONE,
                close_time: Utc::now(),
                quote_asset_volume: Decimal::ONE,
            };

        assert_eq!(klines.len(), 1);
        assert_eq!(klines[0].high, kline.high);
    }
}