o2-tools 0.1.11

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

abigen!(
    Contract(
        name = "OrderBookRegistry",
        abi = "artifacts/order-book-registry/order-book-registry-abi.json"
    ),
    Contract(
        name = "OrderBookRegistryProxy",
        abi = "artifacts/order-book-registry-proxy/order-book-registry-proxy-abi.json"
    )
);

pub const ORDER_BOOK_REGISTER_BYTECODE: &[u8] =
    include_bytes!("../artifacts/order-book-registry/order-book-registry.bin");
pub const ORDER_BOOK_REGISTER_STORAGE: &[u8] = include_bytes!(
    "../artifacts/order-book-registry/order-book-registry-storage_slots.json"
);
pub const ORDER_BOOK_REGISTER_PROXY_BYTECODE: &[u8] = include_bytes!(
    "../artifacts/order-book-registry-proxy/order-book-registry-proxy.bin"
);
pub const ORDER_BOOK_REGISTER_PROXY_STORAGE: &[u8] = include_bytes!(
    "../artifacts/order-book-registry-proxy/order-book-registry-proxy-storage_slots.json"
);

/// Configuration for deploying OrderBookRegistry contracts.
/// Contains bytecode and storage slot information needed for deployment.
#[derive(Clone)]
pub struct OrderBookRegistryDeployConfig {
    /// Bytecode for the OrderBookRegistry contract
    pub registry_bytecode: Vec<u8>,
    /// Storage slots configuration for the OrderBookRegistry contract
    pub registry_storage_slots: Vec<StorageSlot>,
    /// Configurables for the OrderBookRegistry contract
    pub registry_configurables: OrderBookRegistryConfigurables,
    /// Bytecode for the OrderBookRegistryProxy contract
    pub registry_proxy_bytecode: Vec<u8>,
    /// Storage slots configuration for the OrderBookRegistryProxy
    pub registry_proxy_storage_slots: Vec<StorageSlot>,
    /// Register proxy configurations
    pub registry_proxy_config: OrderBookRegistryProxyConfigurables,
    /// Maximum words per blob for deployment
    pub max_words_per_blob: usize,
    /// Owner for the OrderBookProxyRegistry contract
    pub proxy_owner: Option<Identity>,
    /// Owner for the OrderBookRegistry contract
    pub registry_owner: Option<Identity>,
    /// Salt for contract deployment
    pub salt: Salt,
}

impl Default for OrderBookRegistryDeployConfig {
    fn default() -> Self {
        Self {
            registry_bytecode: ORDER_BOOK_REGISTER_BYTECODE.to_vec(),
            registry_storage_slots: serde_json::from_slice(ORDER_BOOK_REGISTER_STORAGE)
                .unwrap(),
            registry_configurables: OrderBookRegistryConfigurables::default(),
            registry_proxy_bytecode: ORDER_BOOK_REGISTER_PROXY_BYTECODE.to_vec(),
            registry_proxy_storage_slots: serde_json::from_slice(
                ORDER_BOOK_REGISTER_PROXY_STORAGE,
            )
            .unwrap(),
            registry_proxy_config: OrderBookRegistryProxyConfigurables::default(),
            max_words_per_blob: 10_000,
            proxy_owner: None,
            registry_owner: None,
            salt: Salt::default(),
        }
    }
}

/// Result of an OrderBookRegistry deployment.
/// Contains the deployed contract instance and configuration.
#[derive(Clone)]
pub struct OrderBookRegistryManager<W> {
    /// The OrderBookRegistryProxy contract instance
    pub registry_proxy: OrderBookRegistryProxy<W>,
    /// The deployed OrderBookRegistry contract instance
    pub registry: OrderBookRegistry<W>,
    /// Contract ID of the deployed OrderBookRegistry
    pub contract_id: ContractId,
    /// The wallet used for deployment
    pub deployer_wallet: W,
}

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<W> OrderBookRegistryManager<W>
where
    W: Account + Clone,
{
    pub fn new(deployer_wallet: W, contract_id: ContractId) -> Self {
        let proxy = OrderBookRegistryProxy::new(contract_id, deployer_wallet.clone());
        let registry = OrderBookRegistry::new(contract_id, deployer_wallet.clone());
        Self {
            registry_proxy: proxy,
            registry,
            contract_id,
            deployer_wallet,
        }
    }

    pub async fn register_blob(
        deployer_wallet: &W,
        config: &OrderBookRegistryDeployConfig,
    ) -> Result<OrderBookBlob> {
        let registry_owner = config
            .registry_owner
            .unwrap_or(Identity::Address(deployer_wallet.address()));
        let configurables = config
            .registry_configurables
            .clone()
            .with_INITIAL_OWNER(State::Initialized(registry_owner))?;
        let blobs = Contract::regular(
            config.registry_bytecode.clone(),
            config.salt,
            config.registry_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(),
        })
    }

    pub async fn deploy_register_blob(
        deployer_wallet: &W,
        config: &OrderBookRegistryDeployConfig,
    ) -> Result<BlobId> {
        let register_blob = Self::register_blob(deployer_wallet, config).await?;

        if !register_blob.exists {
            let mut builder =
                BlobTransactionBuilder::default().with_blob(register_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(register_blob.id)
    }

    pub async fn deploy_register_proxy(
        deployer_wallet: &W,
        registry_blob_id: &BlobId,
        config: &OrderBookRegistryDeployConfig,
    ) -> Result<(OrderBookRegistryProxy<W>, bool)> {
        let proxy_owner = config
            .proxy_owner
            .unwrap_or(Identity::Address(deployer_wallet.address()));
        let configurables = config
            .registry_proxy_config
            .clone()
            .with_INITIAL_OWNER(State::Initialized(proxy_owner))?
            .with_INITIAL_TARGET(ContractId::new(*registry_blob_id))?;
        let contract = Contract::regular(
            config.registry_proxy_bytecode.clone(),
            config.salt,
            config.registry_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 proxy = OrderBookRegistryProxy::new(contract_id, deployer_wallet.clone());

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

        Ok((proxy, requires_initialization))
    }

    /// Deploys an OrderBookRegistry contract.
    ///
    /// # Arguments
    /// * `deployer_wallet` - The wallet to use for deployment (pays gas fees)
    /// * `config` - Deployment configuration containing bytecode and settings
    ///
    /// # Returns
    /// * `Ok(OrderBookRegistryDeploy)` - Complete deployment result with contract instance
    /// * `Err(anyhow::Error)` - If deployment or initialization fails
    ///
    /// # Process
    /// 1. Deploys the OrderBookRegistry contract
    /// 2. Initializes the contract with the specified owner
    pub async fn deploy(
        deployer_wallet: &W,
        config: &OrderBookRegistryDeployConfig,
    ) -> Result<OrderBookRegistryManager<W>> {
        let register_blob_id =
            OrderBookRegistryManager::deploy_register_blob(deployer_wallet, config)
                .await?;
        let (proxy, requires_initialization) =
            OrderBookRegistryManager::deploy_register_proxy(
                deployer_wallet,
                &register_blob_id,
                config,
            )
            .await?;
        let register_deploy =
            OrderBookRegistryManager::new(deployer_wallet.clone(), proxy.contract_id());

        // If the contract is already deployed, we skip deployment
        if requires_initialization {
            // Initialize the proxy contract
            register_deploy
                .registry_proxy
                .methods()
                .initialize_proxy()
                .call()
                .await?;
            // Initialize the register with the specified owner
            register_deploy
                .registry
                .methods()
                .initialize()
                .call()
                .await?;
        }

        Ok(register_deploy)
    }

    pub async fn upgrade(
        &self,
        config: &OrderBookRegistryDeployConfig,
    ) -> Result<BlobId> {
        // Deploy new blob with updated bytecode
        let new_blob_id =
            Self::deploy_register_blob(&self.deployer_wallet, config).await?;

        // Update the proxy to point to the new blob
        self.registry_proxy
            .methods()
            .set_proxy_target(ContractId::new(new_blob_id))
            .call()
            .await?;

        Ok(new_blob_id)
    }

    pub async fn get_order_book(
        &self,
        market_id: MarketId,
    ) -> Result<Option<ContractId>> {
        Ok(self
            .registry
            .methods()
            .get_order_book(market_id)
            .simulate(Execution::state_read_only())
            .await?
            .value)
    }

    pub async fn register_order_book(
        &self,
        market_id: MarketId,
        order_book_id: ContractId,
    ) -> Result<ContractId> {
        let _ = self
            .registry
            .methods()
            .register_order_book(order_book_id, market_id)
            .call()
            .await?;
        Ok(order_book_id)
    }
}

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

    #[tokio::test]
    async fn test_order_book_registry_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();
        let owner_wallet = wallets.pop().unwrap();

        // Deploy Register with separate owner
        let config = OrderBookRegistryDeployConfig {
            proxy_owner: Some(Identity::Address(owner_wallet.address())),
            registry_owner: Some(Identity::Address(owner_wallet.address())),
            ..Default::default()
        };
        let deployment = OrderBookRegistryManager::deploy(&deployer_wallet, &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, "Register contract should exist");

        // Verify owner is set correctly
        let contract_owner = deployment
            .registry
            .methods()
            .owner()
            .simulate(Execution::state_read_only())
            .await
            .unwrap()
            .value;

        match contract_owner {
            State::Initialized(identity) => match identity {
                Identity::Address(address) => {
                    assert_eq!(address, owner_wallet.address(), "Owner should match");
                }
                _ => panic!("Owner should be an address"),
            },
            _ => panic!("Owner should be initialized"),
        }
    }

    #[tokio::test]
    async fn test_order_book_registry_register_order_book() {
        // 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 Register with separate owner
        let config = OrderBookRegistryDeployConfig::default();
        let order_book_registry_deploy =
            OrderBookRegistryManager::deploy(&deployer_wallet, &config)
                .await
                .unwrap();
        let order_book_registry_manager = OrderBookRegistryManager::new(
            deployer_wallet.clone(),
            order_book_registry_deploy.contract_id,
        );

        let market_id = MarketId {
            base_asset: AssetId::new([1; 32]),
            quote_asset: AssetId::new([2; 32]),
        };
        let order_book_id = ContractId::new([3; 32]);

        order_book_registry_manager
            .register_order_book(market_id.clone(), order_book_id)
            .await
            .unwrap();
        let order_book_id_result = order_book_registry_manager
            .get_order_book(market_id.clone())
            .await
            .unwrap();
        assert_eq!(
            order_book_id_result,
            Some(order_book_id),
            "Should return the correct order book id"
        );
    }
}