o2-tools 0.1.10

use anyhow::Result;
use fuels::{
    prelude::*,
    tx::StorageSlot,
    types::{
        AssetId,
        ContractId,
        Identity,
    },
};

abigen!(
    Contract(
        name = "OrderBook",
        abi = "artifacts/order-book/order-book-abi.json"
    ),
    Contract(
        name = "OrderBookProxy",
        abi = "artifacts/order-book-proxy/order-book-proxy-abi.json"
    ),
    Contract(
        name = "OrderBookWhitelist",
        abi = "artifacts/order-book-whitelist/order-book-whitelist-abi.json",
    ),
    Contract(
        name = "OrderBookBlacklist",
        abi = "artifacts/order-book-blacklist/order-book-blacklist-abi.json",
    ),
);
pub const ORDER_BOOK_BYTECODE: &[u8] =
    include_bytes!("../../../artifacts/order-book/order-book.bin");
pub const ORDER_BOOK_STORAGE: &[u8] =
    include_bytes!("../../../artifacts/order-book/order-book-storage_slots.json");
pub const ORDER_BOOK_PROXY_BYTECODE: &[u8] =
    include_bytes!("../../../artifacts/order-book-proxy/order-book-proxy.bin");
pub const ORDER_BOOK_PROXY_STORAGE: &[u8] = include_bytes!(
    "../../../artifacts/order-book-proxy/order-book-proxy-storage_slots.json"
);
pub const ORDER_BOOK_WHITELIST_BYTECODE: &[u8] =
    include_bytes!("../../../artifacts/order-book-whitelist/order-book-whitelist.bin");
pub const ORDER_BOOK_WHITELIST_STORAGE: &[u8] = include_bytes!(
    "../../../artifacts/order-book-whitelist/order-book-whitelist-storage_slots.json"
);
pub const ORDER_BOOK_BLACKLIST_BYTECODE: &[u8] =
    include_bytes!("../../../artifacts/order-book-blacklist/order-book-blacklist.bin");
pub const ORDER_BOOK_BLACKLIST_STORAGE: &[u8] = include_bytes!(
    "../../../artifacts/order-book-blacklist/order-book-blacklist-storage_slots.json"
);

/// Configuration for deploying OrderBook contracts.
/// Contains bytecode and storage slot information needed for deployment.
#[derive(Clone)]
pub struct OrderBookDeployConfig {
    /// Bytecode for the OrderBook contract
    pub order_book_bytecode: Vec<u8>,
    /// Storage slots configuration for the OrderBook contract
    pub order_book_storage_slots: Vec<StorageSlot>,
    /// Configurables for the OrderBook contract
    pub order_book_configurables: OrderBookConfigurables,
    /// Bytecode for the OrderBookProxy contract
    pub order_book_proxy_bytecode: Vec<u8>,
    /// Storage slots configuration for the OrderBookProxy contract
    pub order_book_proxy_storage_slots: Vec<StorageSlot>,
    /// Configurables for the OrderBookProxy contract
    pub order_book_proxy_configurables: OrderBookProxyConfigurables,
    /// Bytecode for the OrderBookWhitelist contract
    pub order_book_whitelist_bytecode: Vec<u8>,
    /// Storage slots configuration for the OrderBookWhitelist contract
    pub order_book_whitelist_storage_slots: Vec<StorageSlot>,
    /// Bytecode for the OrderBookBlacklist contract
    pub order_book_blacklist_bytecode: Vec<u8>,
    /// Storage slots configuration for the OrderBookBlacklist contract
    pub order_book_blacklist_storage_slots: Vec<StorageSlot>,
    /// Maximum words per blob for deployment
    pub max_words_per_blob: usize,
    /// Owner of the OrderBook proxy
    pub proxy_owner: Option<Identity>,
    /// Owner of the OrderBook
    pub order_book_owner: Option<Identity>,
    /// Salt for contract deployment
    pub salt: Salt,
}

pub struct OrderBookBlob {
    /// The ID of the deployed blob
    pub id: BlobId,
    /// Whether the blob already exists
    pub exists: bool,
    /// The blob data containing the contract bytecode
    pub blob: Blob,
}

impl Default for OrderBookDeployConfig {
    fn default() -> Self {
        Self {
            order_book_bytecode: ORDER_BOOK_BYTECODE.to_vec(),
            order_book_storage_slots: serde_json::from_slice(ORDER_BOOK_STORAGE).unwrap(),
            order_book_whitelist_bytecode: ORDER_BOOK_WHITELIST_BYTECODE.to_vec(),
            order_book_configurables: OrderBookConfigurables::default(),
            order_book_proxy_bytecode: ORDER_BOOK_PROXY_BYTECODE.to_vec(),
            order_book_proxy_storage_slots: serde_json::from_slice(
                ORDER_BOOK_PROXY_STORAGE,
            )
            .unwrap(),
            order_book_proxy_configurables: OrderBookProxyConfigurables::default(),
            order_book_whitelist_storage_slots: serde_json::from_slice(
                ORDER_BOOK_WHITELIST_STORAGE,
            )
            .unwrap(),
            order_book_blacklist_bytecode: ORDER_BOOK_BLACKLIST_BYTECODE.to_vec(),
            order_book_blacklist_storage_slots: serde_json::from_slice(
                ORDER_BOOK_BLACKLIST_STORAGE,
            )
            .unwrap(),
            max_words_per_blob: 100_000,
            proxy_owner: None,
            order_book_owner: None,
            salt: Salt::default(),
        }
    }
}

impl OrderBookDeployConfig {
    pub fn with_configurables(order_book_configurables: OrderBookConfigurables) -> Self {
        Self {
            order_book_configurables,
            ..OrderBookDeployConfig::default()
        }
    }
}

/// Result of an OrderBook deployment.
/// Contains the deployed contract instance and configuration.
#[derive(Clone)]
pub struct OrderBookDeploy<W: Account + Clone> {
    /// The deployed OrderBook contract instance
    pub order_book_proxy: OrderBookProxy<W>,
    /// The deployed OrderBook contract instance
    pub order_book: OrderBook<W>,
    /// Contract ID of the deployed OrderBook
    pub contract_id: ContractId,
    /// Base asset ID for the trading pair
    pub base_asset: AssetId,
    /// Quote asset ID for the trading pair
    pub quote_asset: AssetId,
    /// The deployed trade account whitelist contract instance
    pub whitelist: Option<OrderBookWhitelist<W>>,
    /// Contract ID of the deployed trade account whitelist
    pub whitelist_id: Option<ContractId>,
}

impl<W: Account + Clone> OrderBookDeploy<W> {
    pub fn new(
        w: W,
        contract_id: ContractId,
        base_asset: AssetId,
        quote_asset: AssetId,
    ) -> Self {
        let order_book = OrderBook::new(contract_id, w.clone());
        let order_book_proxy = OrderBookProxy::new(contract_id, w.clone());
        Self {
            order_book,
            order_book_proxy,
            contract_id,
            base_asset,
            quote_asset,
            whitelist: None,
            whitelist_id: None,
        }
    }

    pub async fn initialize(&self) -> Result<()> {
        // Initialize the OrderBookProxy with the configured owner via configurables
        self.order_book_proxy
            .methods()
            .initialize_proxy()
            .call()
            .await?;
        // Initialize the OrderBook contract with the specified owner
        self.order_book.methods().initialize().call().await?;
        Ok(())
    }

    pub async fn deploy_order_book_blacklist(
        deployer_wallet: &W,
        owner_identity: &Identity,
        config: &OrderBookDeployConfig,
    ) -> Result<OrderBookBlacklist<W>> {
        let contract = Contract::regular(
            config.order_book_blacklist_bytecode.clone(),
            config.salt,
            config.order_book_blacklist_storage_slots.clone(),
        )
        .with_configurables(
            OrderBookBlacklistConfigurables::default()
                .with_INITIAL_OWNER(State::Initialized(*owner_identity))?,
        )
        .with_salt(config.salt);
        let contract_id = contract.contract_id();
        let instance = OrderBookBlacklist::new(contract_id, deployer_wallet.clone());
        let contract_exists = deployer_wallet
            .try_provider()?
            .contract_exists(&contract_id)
            .await?;
        if !contract_exists {
            contract
                .deploy(deployer_wallet, TxPolicies::default())
                .await?;
            instance.methods().initialize().call().await?;
        }
        Ok(instance)
    }

    pub async fn deploy_order_book_whitelist(
        deployer_wallet: &W,
        owner_identity: &Identity,
        config: &OrderBookDeployConfig,
    ) -> Result<OrderBookWhitelist<W>> {
        let contract = Contract::regular(
            config.order_book_whitelist_bytecode.clone(),
            config.salt,
            config.order_book_whitelist_storage_slots.clone(),
        )
        .with_configurables(
            OrderBookWhitelistConfigurables::default()
                .with_INITIAL_OWNER(State::Initialized(*owner_identity))?,
        )
        .with_salt(config.salt);
        let contract_id = contract.contract_id();
        let instance = OrderBookWhitelist::new(contract_id, deployer_wallet.clone());
        let contract_exists = deployer_wallet
            .try_provider()?
            .contract_exists(&contract_id)
            .await?;

        if !contract_exists {
            contract
                .deploy(deployer_wallet, TxPolicies::default())
                .await?;
            instance.methods().initialize().call().await?;
        }

        Ok(instance)
    }

    pub async fn order_book_blob(
        deployer_wallet: &W,
        base_asset: AssetId,
        quote_asset: AssetId,
        config: &OrderBookDeployConfig,
    ) -> Result<OrderBookBlob> {
        let order_book_owner = config
            .order_book_owner
            .unwrap_or(Identity::Address(deployer_wallet.address()));
        // Configure the contract with the trading pair assets
        let configurables = config
            .order_book_configurables
            .clone()
            .with_BASE_ASSET(base_asset)?
            .with_QUOTE_ASSET(quote_asset)?
            .with_INITIAL_OWNER(State::Initialized(order_book_owner))?;
        let blobs = Contract::regular(
            config.order_book_bytecode.clone(),
            config.salt,
            config.order_book_storage_slots.clone(),
        )
        .with_configurables(configurables.clone())
        .convert_to_loader(config.max_words_per_blob)?
        .blobs()
        .to_vec();
        let blob = blobs[0].clone();
        let blob_id = blob.id();
        let blob_exists = deployer_wallet.try_provider()?.blob_exists(blob_id).await?;

        Ok(OrderBookBlob {
            id: blob_id,
            exists: blob_exists,
            blob: blob.clone(),
        })
    }

    /// Deploys the trade account implementation as a blob.
    /// Large contracts are deployed as blobs to handle size limitations.
    ///
    /// # Arguments
    /// * `deployer_wallet` - The wallet to use for deployment (pays gas fees)
    /// * `config` - Deployment configuration containing bytecode and settings
    ///
    /// # Returns
    /// * `Ok(BlobId)` - The ID of the deployed blob
    /// * `Err(anyhow::Error)` - If deployment fails
    pub async fn deploy_order_book_blob(
        deployer_wallet: &W,
        base_asset: AssetId,
        quote_asset: AssetId,
        config: &OrderBookDeployConfig,
    ) -> Result<BlobId> {
        let order_book_blob =
            Self::order_book_blob(deployer_wallet, base_asset, quote_asset, config)
                .await?;
        if !order_book_blob.exists {
            let mut builder =
                BlobTransactionBuilder::default().with_blob(order_book_blob.blob);
            deployer_wallet.adjust_for_fee(&mut builder, 0).await?;
            deployer_wallet.add_witnesses(&mut builder)?;

            let tx = builder.build(&deployer_wallet.try_provider()?).await?;

            deployer_wallet
                .try_provider()?
                .send_transaction_and_await_commit(tx)
                .await?
                .check(None)?;
        }
        Ok(order_book_blob.id)
    }

    pub async fn deploy_order_book_proxy(
        deployer_wallet: &W,
        order_book_blob_id: &BlobId,
        config: &OrderBookDeployConfig,
    ) -> Result<(OrderBookProxy<W>, bool)> {
        let proxy_owner = config
            .proxy_owner
            .unwrap_or(Identity::Address(deployer_wallet.address()));
        let blob_id = ContractId::new(*order_book_blob_id);
        let configurables = config
            .order_book_proxy_configurables
            .clone()
            .with_INITIAL_TARGET(blob_id)?
            .with_INITIAL_OWNER(State::Initialized(proxy_owner))?;
        let contract = Contract::regular(
            config.order_book_proxy_bytecode.clone(),
            config.salt,
            config.order_book_proxy_storage_slots.clone(),
        )
        .with_configurables(configurables);

        let contract_id = contract.contract_id();
        let already_deployed = deployer_wallet
            .try_provider()?
            .contract_exists(&contract_id)
            .await?;
        let order_book_proxy = OrderBookProxy::new(contract_id, deployer_wallet.clone());

        if !already_deployed {
            contract
                .deploy(deployer_wallet, TxPolicies::default())
                .await?;
        }
        let requires_initialization = !already_deployed;

        Ok((order_book_proxy, requires_initialization))
    }

    /// Deploys an OrderBook contract with the specified trading pair.
    ///
    /// # Arguments
    /// * `deployer_wallet` - The wallet to use for deployment (pays gas fees)
    /// * `owner` - The identity to set as the OrderBook owner
    /// * `base_asset` - The base asset ID for the trading pair
    /// * `quote_asset` - The quote asset ID for the trading pair
    /// * `config` - Deployment configuration containing bytecode and settings
    ///
    /// # Returns
    /// * `Ok(OrderBookDeploy)` - Complete deployment result with contract instance
    /// * `Err(anyhow::Error)` - If deployment or initialization fails
    ///
    /// # Process
    /// 1. Deploys the OrderBook contract with asset configurables
    /// 2. Initializes the contract with the specified owner
    pub async fn deploy(
        deployer_wallet: &W,
        base_asset: AssetId,
        quote_asset: AssetId,
        config: &OrderBookDeployConfig,
    ) -> Result<OrderBookDeploy<W>> {
        // Configure the contract with the trading pair assets
        let (deploy, requires_initialization) = Self::deploy_without_initialization(
            deployer_wallet,
            base_asset,
            quote_asset,
            config,
        )
        .await?;

        if requires_initialization {
            // Initialize the OrderBookProxy with the configured owner via configurables
            deploy.initialize().await?;
        }

        Ok(deploy)
    }

    /// Deploys an OrderBook contract like [`Self::deploy`], but without initializing it.
    pub async fn deploy_without_initialization(
        deployer_wallet: &W,
        base_asset: AssetId,
        quote_asset: AssetId,
        config: &OrderBookDeployConfig,
    ) -> Result<(OrderBookDeploy<W>, bool)> {
        let order_book_blob_id = Self::deploy_order_book_blob(
            deployer_wallet,
            base_asset,
            quote_asset,
            config,
        )
        .await?;
        let (deploy_order_book_proxy, requires_initialization) =
            Self::deploy_order_book_proxy(deployer_wallet, &order_book_blob_id, config)
                .await?;
        let contract_id = deploy_order_book_proxy.contract_id();
        let deploy = OrderBookDeploy::new(
            deployer_wallet.clone(),
            contract_id,
            base_asset,
            quote_asset,
        );

        Ok((deploy, requires_initialization))
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use fuels::test_helpers::{
        WalletsConfig,
        launch_custom_provider_and_get_wallets,
    };

    #[tokio::test]
    async fn test_deploy_order_book_blob() {
        // Start fuel-core
        let mut wallets = launch_custom_provider_and_get_wallets(
            WalletsConfig::new(Some(2), Some(1), Some(1_000_000_000)),
            None,
            None,
        )
        .await
        .unwrap();
        let deployer_wallet = wallets.pop().unwrap();
        let contract_owner = wallets.pop().unwrap();

        // Deploy OrderBook
        let base_asset = AssetId::from([1u8; 32]);
        let quote_asset = AssetId::from([2u8; 32]);
        let config = OrderBookDeployConfig {
            proxy_owner: Some(Identity::Address(contract_owner.address())),
            order_book_owner: Some(Identity::Address(contract_owner.address())),
            ..Default::default()
        };

        let order_book_blob_id = OrderBookDeploy::deploy_order_book_blob(
            &deployer_wallet,
            base_asset,
            quote_asset,
            &config,
        )
        .await
        .unwrap();
        let (order_book_deploy, requires_initialization) =
            OrderBookDeploy::deploy_without_initialization(
                &deployer_wallet,
                base_asset,
                quote_asset,
                &config,
            )
            .await
            .unwrap();
        let order_book = order_book_deploy.order_book;
        let order_book_proxy = order_book_deploy.order_book_proxy;

        if requires_initialization {
            order_book_proxy
                .methods()
                .initialize_proxy()
                .call()
                .await
                .unwrap();
            order_book.methods().initialize().call().await.unwrap();
        }

        assert_eq!(
            order_book_proxy
                .methods()
                .proxy_owner()
                .simulate(Execution::state_read_only())
                .await
                .unwrap()
                .value,
            State::Initialized(Identity::Address(contract_owner.address()))
        );
        assert_eq!(
            order_book_proxy
                .methods()
                .proxy_target()
                .simulate(Execution::state_read_only())
                .await
                .unwrap()
                .value,
            Some(ContractId::new(order_book_blob_id))
        );
        assert_eq!(
            order_book
                .methods()
                .owner()
                .simulate(Execution::state_read_only())
                .await
                .unwrap()
                .value,
            State::Initialized(Identity::Address(contract_owner.address()))
        );
        assert_eq!(
            order_book
                .methods()
                .get_base_asset()
                .simulate(Execution::state_read_only())
                .await
                .unwrap()
                .value,
            base_asset,
        );
        assert_eq!(
            order_book
                .methods()
                .get_quote_asset()
                .simulate(Execution::state_read_only())
                .await
                .unwrap()
                .value,
            quote_asset,
        );

        // Test upgrading contract configuration
        let new_base_asset = AssetId::from([3u8; 32]);
        let order_book_blob_id = OrderBookDeploy::deploy_order_book_blob(
            &deployer_wallet,
            new_base_asset,
            quote_asset,
            &config,
        )
        .await
        .unwrap();

        order_book_proxy
            .with_account(contract_owner)
            .methods()
            .set_proxy_target(ContractId::new(order_book_blob_id))
            .call()
            .await
            .unwrap();
        assert_eq!(
            order_book
                .methods()
                .get_base_asset()
                .simulate(Execution::state_read_only())
                .await
                .unwrap()
                .value,
            new_base_asset
        );
    }

    #[tokio::test]
    async fn test_order_book_deployment() {
        // Start fuel-core
        let mut wallets = launch_custom_provider_and_get_wallets(
            WalletsConfig::new(Some(2), Some(1), Some(1_000_000_000)),
            None,
            None,
        )
        .await
        .unwrap();
        let deployer_wallet = wallets.pop().unwrap();

        // Deploy OrderBook
        let base_asset = AssetId::new([1u8; 32]);
        let quote_asset = AssetId::new([2u8; 32]);
        let config = OrderBookDeployConfig::default();

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