o2-tools 0.1.10

use fuels::{
    prelude::*,
    types::{
        AssetId,
        Identity,
    },
};

use crate::order_book_deploy::{
    OrderBookDeploy,
    OrderBookDeployConfig,
    OrderBookProxy,
};

// Re-export types from the contract ABI
use crate::{
    order_book_deploy,
    order_book_deploy::{
        OrderArgs,
        OrderBook,
        Side as ContractSide,
        Time,
    },
    trade_account_deploy::CallParams,
};

// Re-export domain types from api-types
pub use o2_api_types::primitives::{
    OrderType,
    Side,
};

pub const DEFAULT_METHOD_GAS: u64 = 1_000_000;

/// Convert a contract OrderType into the shared OrderType.
pub fn order_type_from_contract(order_type: order_book_deploy::OrderType) -> OrderType {
    match order_type {
        order_book_deploy::OrderType::Spot => OrderType::Spot,
        order_book_deploy::OrderType::Limit((price, timestamp)) => {
            OrderType::Limit(price, timestamp.unix as u128)
        }
        order_book_deploy::OrderType::FillOrKill => OrderType::FillOrKill,
        order_book_deploy::OrderType::PostOnly => OrderType::PostOnly,
        order_book_deploy::OrderType::Market => OrderType::Market,
        order_book_deploy::OrderType::BoundedMarket((max_price, min_price)) => {
            OrderType::BoundedMarket {
                max_price,
                min_price,
            }
        }
    }
}

/// Convert a shared OrderType into the contract OrderType.
pub fn order_type_to_contract(order_type: OrderType) -> order_book_deploy::OrderType {
    match order_type {
        OrderType::Spot => order_book_deploy::OrderType::Spot,
        OrderType::Limit(price, timestamp) => order_book_deploy::OrderType::Limit((
            price,
            Time {
                unix: timestamp as u64,
            },
        )),
        OrderType::FillOrKill => order_book_deploy::OrderType::FillOrKill,
        OrderType::PostOnly => order_book_deploy::OrderType::PostOnly,
        OrderType::Market => order_book_deploy::OrderType::Market,
        OrderType::BoundedMarket {
            max_price,
            min_price,
        } => order_book_deploy::OrderType::BoundedMarket((max_price, min_price)),
    }
}

/// Convert a contract Side into the shared Side.
pub fn side_from_contract(order_side: ContractSide) -> Side {
    match order_side {
        ContractSide::Buy => Side::Buy,
        ContractSide::Sell => Side::Sell,
    }
}

/// Convert a shared Side into the contract Side.
pub fn side_to_contract(side: Side) -> ContractSide {
    match side {
        Side::Buy => ContractSide::Buy,
        Side::Sell => ContractSide::Sell,
    }
}

/// Parameters for creating a new order in the order book.
/// Contains all information needed to place a buy or sell order.
#[derive(Debug, Clone)]
pub struct CreateOrderParams {
    /// The price per unit (in quote asset)
    pub price: u64,
    /// The quantity to buy/sell (in base asset)
    pub quantity: u64,
    /// The type of order (Limit or Spot)
    pub order_type: OrderType,
    /// The side of the order (Buy or Sell)
    pub side: Side,
    /// The asset ID of the asset being traded
    pub asset_id: AssetId,
}

impl CreateOrderParams {
    pub fn random(
        side: Side,
        asset_id: AssetId,
        min_price: u64,
        max_price: u64,
        min_quantity: u64,
        max_quantity: u64,
    ) -> Self {
        let price = (rand::random::<u64>() + min_price) % max_price;
        let quantity = (rand::random::<u64>() + min_quantity) % max_quantity;
        let order_type = OrderType::Spot;
        Self {
            price,
            quantity,
            order_type,
            side,
            asset_id,
        }
    }

    /// Converts to the contract's OrderArgs format.
    pub fn to_order_args(&self) -> OrderArgs {
        OrderArgs {
            price: self.price,
            quantity: self.quantity,
            order_type: order_type_to_contract(self.order_type),
        }
    }
}

#[derive(Clone, Copy)]
pub struct OrderbookConfig {
    /// Base asset ID for the trading pair
    pub base_asset: AssetId,
    pub base_decimals: u64,
    /// Quote asset ID for the trading pair
    pub quote_asset: AssetId,
    pub quote_decimals: u64,
}

impl OrderbookConfig {
    pub fn get_order_side_asset(&self, order_side: &Side) -> AssetId {
        if order_side == &Side::Buy {
            self.quote_asset
        } else {
            self.base_asset
        }
    }

    pub fn get_order_side_amount(&self, quantity: u64, price: u64, side: &Side) -> u64 {
        if side == &Side::Buy {
            ((quantity as u128 * price as u128) / self.base_decimals as u128)
                .try_into()
                .unwrap()
        } else {
            quantity
        }
    }

    pub fn create_call_params(
        &self,
        params: &CreateOrderParams,
        gas: Option<u64>,
    ) -> CallParams {
        let amount =
            self.get_order_side_amount(params.quantity, params.price, &params.side);
        let asset_id = self.get_order_side_asset(&params.side);
        CallParams::new(amount, asset_id, gas.unwrap_or(u64::MAX))
    }
}

#[derive(Clone)]
pub struct OrderBookManager<W: Account + Clone> {
    /// The OrderBook contract instance
    pub proxy: OrderBookProxy<W>,
    /// The OrderBook contract instance
    pub contract: OrderBook<W>,
    /// The wallet used for creating transactions
    pub gas_payer_wallet: W,
    /// Confguration
    pub config: OrderbookConfig,
}

impl<W: Account + Clone> OrderBookManager<W> {
    /// Creates a new OrderBookManager instance.
    ///
    /// # Arguments
    /// * `contract_id` - The deployed OrderBook contract ID
    /// * `wallet` - The wallet to use for transactions
    /// * `base_asset` - The base asset ID for the trading pair
    /// * `quote_asset` - The quote asset ID for the trading pair
    pub fn new(
        gas_payer_wallet: &W,
        base_decimals: u64,
        quote_decimals: u64,
        order_book_deploy: &OrderBookDeploy<W>,
    ) -> Self {
        let config = OrderbookConfig {
            base_asset: order_book_deploy.base_asset,
            base_decimals,
            quote_asset: order_book_deploy.quote_asset,
            quote_decimals,
        };
        Self {
            proxy: order_book_deploy
                .order_book_proxy
                .clone()
                .with_account(gas_payer_wallet.clone()),
            contract: order_book_deploy
                .order_book
                .clone()
                .with_account(gas_payer_wallet.clone()),
            config,
            gas_payer_wallet: gas_payer_wallet.clone(),
        }
    }

    pub fn create_call_params(
        &self,
        params: &CreateOrderParams,
        gas: Option<u64>,
    ) -> CallParams {
        self.config().create_call_params(params, gas)
    }

    pub fn get_order_side_asset(&self, order_side: &Side) -> AssetId {
        self.config().get_order_side_asset(order_side)
    }

    pub fn get_order_side_amount(&self, quantity: u64, price: u64, side: &Side) -> u64 {
        self.config().get_order_side_amount(quantity, price, side)
    }

    pub async fn balances_of(&self, identity: &Identity) -> anyhow::Result<(u128, u128)> {
        let result = self
            .contract
            .methods()
            .get_settled_balance_of(*identity)
            .simulate(Execution::state_read_only())
            .await?;
        Ok((result.value.0 as u128, result.value.1 as u128))
    }

    pub async fn emit_config(&self) -> anyhow::Result<()> {
        self.contract
            .methods()
            .emit_orderbook_config()
            .call()
            .await?;
        Ok(())
    }

    pub async fn upgrade(
        &self,
        deploy_config: &OrderBookDeployConfig,
    ) -> anyhow::Result<()> {
        let base_asset_id = self.contract.methods().get_base_asset().call().await?.value;
        let quote_asset_id = self
            .contract
            .methods()
            .get_quote_asset()
            .simulate(Execution::state_read_only())
            .await?
            .value;
        let order_book_blob_id = OrderBookDeploy::deploy_order_book_blob(
            &self.gas_payer_wallet,
            base_asset_id,
            quote_asset_id,
            deploy_config,
        )
        .await?;
        self.proxy
            .methods()
            .set_proxy_target(ContractId::new(order_book_blob_id))
            .call()
            .await?;
        Ok(())
    }

    pub async fn accumulated_fees(&self) -> anyhow::Result<(u64, u64)> {
        let fees = self
            .contract
            .methods()
            .current_fees()
            .simulate(Execution::state_read_only())
            .await?
            .value;
        Ok(fees)
    }

    pub async fn get_whitelist_id(&self) -> anyhow::Result<Option<ContractId>> {
        let result = self
            .contract
            .methods()
            .get_whitelist_id()
            .simulate(Execution::state_read_only())
            .await?;
        Ok(result.value)
    }

    pub async fn get_blacklist_id(&self) -> anyhow::Result<Option<ContractId>> {
        let result = self
            .contract
            .methods()
            .get_blacklist_id()
            .simulate(Execution::state_read_only())
            .await?;
        Ok(result.value)
    }

    pub fn base_asset(&self) -> AssetId {
        self.config().base_asset
    }

    pub fn base_decimals(&self) -> u64 {
        self.config().base_decimals
    }

    pub fn quote_asset(&self) -> AssetId {
        self.config().quote_asset
    }

    pub fn quote_decimals(&self) -> u64 {
        self.config().quote_decimals
    }

    pub fn config(&self) -> OrderbookConfig {
        self.config
    }

    pub async fn is_paused(&self) -> anyhow::Result<bool> {
        let result = self
            .contract
            .methods()
            .is_paused()
            .simulate(Execution::state_read_only())
            .await?;
        Ok(result.value)
    }
}

#[cfg(test)]
mod tests {
    use crate::{
        helpers::get_asset_balance,
        order_book_deploy::{
            OrderBookDeployConfig,
            OrderCreatedEvent,
            OrderMatchedEvent,
        },
    };

    use super::*;
    use fuels::test_helpers::{
        WalletsConfig,
        launch_custom_provider_and_get_wallets,
    };

    #[tokio::test]
    async fn test_order_book_manager() {
        let base_asset = AssetId::from([1; 32]);
        let quote_asset = AssetId::from([2; 32]);
        let initial_balance = 1_000_000_000_000_000u64;

        // Start fuel-core
        let mut wallets = launch_custom_provider_and_get_wallets(
            WalletsConfig::new_multiple_assets(
                3,
                vec![
                    AssetConfig {
                        id: AssetId::default(),
                        num_coins: 1,
                        coin_amount: initial_balance,
                    },
                    AssetConfig {
                        id: quote_asset,
                        num_coins: 1,
                        coin_amount: initial_balance,
                    },
                    AssetConfig {
                        id: base_asset,
                        num_coins: 1,
                        coin_amount: initial_balance,
                    },
                ],
            ),
            None,
            None,
        )
        .await
        .unwrap();
        let deployer_wallet = wallets.pop().unwrap();
        let maker_wallet = wallets.pop().unwrap();
        let taker_wallet = wallets.pop().unwrap();

        // Deploy OrderBook
        let mut config = OrderBookDeployConfig::default();

        config.order_book_configurables = config
            .order_book_configurables
            .with_MAKER_FEE(0.into())
            .unwrap()
            .with_TAKER_FEE(0.into())
            .unwrap();

        let deployment =
            OrderBookDeploy::deploy(&deployer_wallet, base_asset, quote_asset, &config)
                .await
                .unwrap();

        // Check if contract exists
        let provider = deployer_wallet.try_provider().unwrap();
        let contract_exists = provider
            .contract_exists(&deployment.contract_id)
            .await
            .unwrap();
        assert!(contract_exists, "OrderBook contract should exist");

        // Verify configuration
        assert_eq!(deployment.base_asset, base_asset);
        assert_eq!(deployment.quote_asset, quote_asset);

        let order_book_manager = OrderBookManager::new(
            &deployer_wallet,
            10u64.pow(9),
            10u64.pow(9),
            &deployment,
        );
        let maker_order_params = CreateOrderParams {
            price: 1_000_000_000,
            quantity: 2_000_000_000,
            order_type: OrderType::Spot,
            side: Side::Buy,
            asset_id: quote_asset,
        };
        let maker_call_params =
            order_book_manager.create_call_params(&maker_order_params, None);
        let maker_order_book_instance = order_book_manager
            .contract
            .clone()
            .with_account(maker_wallet.clone());
        let result = maker_order_book_instance
            .methods()
            .create_order(maker_order_params.to_order_args())
            .with_tx_policies(TxPolicies::default())
            .call_params(CallParameters::new(
                maker_call_params.coins,
                maker_call_params.asset_id,
                maker_call_params.gas,
            ))
            .unwrap()
            .call()
            .await
            .unwrap();
        let maker_order_created_events =
            result.decode_logs_with_type::<OrderCreatedEvent>().unwrap();
        let maker_order_created_event = maker_order_created_events.first().unwrap();
        let taker_order_params = CreateOrderParams {
            price: 1_000_000_000,
            quantity: 1_000_000_000,
            order_type: OrderType::Spot,
            side: Side::Sell,
            asset_id: base_asset,
        };
        let taker_call_params =
            order_book_manager.create_call_params(&taker_order_params, None);
        let result = order_book_manager
            .contract
            .clone()
            .with_account(taker_wallet.clone())
            .methods()
            .create_order(taker_order_params.to_order_args())
            .with_tx_policies(TxPolicies::default())
            .call_params(CallParameters::new(
                taker_call_params.coins,
                taker_call_params.asset_id,
                taker_call_params.gas,
            ))
            .unwrap()
            .with_variable_output_policy(VariableOutputPolicy::Exactly(10))
            .call()
            .await
            .unwrap();

        let maker_balances = maker_wallet.get_balances().await.unwrap();
        let taker_balances = taker_wallet.get_balances().await.unwrap();
        let (maker_balance_base, maker_balance_quote) = order_book_manager
            .balances_of(&Identity::Address(maker_wallet.address()))
            .await
            .unwrap();
        let (taker_balance_base, taker_balance_quote) = order_book_manager
            .balances_of(&Identity::Address(taker_wallet.address()))
            .await
            .unwrap();

        assert_eq!(
            get_asset_balance(&maker_balances, &base_asset) + maker_balance_base,
            initial_balance as u128 + taker_order_params.quantity as u128
        );
        assert_eq!(
            get_asset_balance(&maker_balances, &quote_asset) + maker_balance_quote,
            initial_balance as u128 - maker_call_params.coins as u128
        );
        assert_eq!(
            get_asset_balance(&taker_balances, &base_asset) + taker_balance_base,
            initial_balance as u128 - taker_order_params.quantity as u128
        );
        assert_eq!(
            get_asset_balance(&taker_balances, &quote_asset) + taker_balance_quote,
            initial_balance as u128 + 1_000_000_000u128
        );

        // Log receipts
        let matches = result.decode_logs_with_type::<OrderMatchedEvent>().unwrap();
        // Only one match
        assert_eq!(matches.len(), 1);
        let match_event = matches.first().unwrap();
        assert_eq!(match_event.price, maker_order_params.price);
        assert_eq!(match_event.quantity, taker_order_params.quantity);

        let _ = maker_order_book_instance
            .methods()
            .settle_balances(vec![
                Identity::Address(maker_wallet.address()),
                Identity::Address(taker_wallet.address()),
            ])
            .with_variable_output_policy(VariableOutputPolicy::Exactly(5))
            .call()
            .await
            .unwrap();
        let maker_quote_balance_before_cancel =
            get_asset_balance(&maker_wallet.get_balances().await.unwrap(), &quote_asset);
        let cancel_result = maker_order_book_instance
            .methods()
            .cancel_order(maker_order_created_event.order_id)
            .with_tx_policies(TxPolicies::default())
            .with_variable_output_policy(VariableOutputPolicy::Exactly(10))
            .call()
            .await
            .unwrap();
        let maker_quote_balance_after =
            get_asset_balance(&maker_wallet.get_balances().await.unwrap(), &quote_asset);

        assert!(cancel_result.value);
        assert_eq!(
            maker_quote_balance_after,
            maker_quote_balance_before_cancel + 1_000_000_000u128
        );
    }
}