o2-tools 0.1.11

use crate::{
    CallOption,
    call_handler_ext::CallHandlerExt,
    trade_account_deploy::{
        DeployConfig,
        State as TradingAccountState,
        TRADE_ACCOUNT_PROXY_BYTECODE,
        TradingAccountProxy0_69_1Configurables,
        TradingAccountProxyConfigurables,
    },
};
use anyhow::Result;
use fuels::{
    prelude::*,
    tx::StorageSlot,
    types::{
        Bits256,
        Bytes32,
        ContractId,
        Identity,
    },
};

abigen!(
    Contract(
        name = "TradeAccountRegistry",
        abi = "artifacts/trade-account-registry/trade-account-registry-abi.json"
    ),
    Contract(
        name = "TradeAccountRegistryProxy",
        abi = "artifacts/trade-account-registry-proxy/trade-account-registry-proxy-abi.json"
    )
);

pub const TRADE_ACCOUNT_REGISTER_BYTECODE: &[u8] =
    include_bytes!("../artifacts/trade-account-registry/trade-account-registry.bin");
pub const TRADE_ACCOUNT_REGISTER_STORAGE: &[u8] = include_bytes!(
    "../artifacts/trade-account-registry/trade-account-registry-storage_slots.json"
);
pub const TRADE_ACCOUNT_REGISTER_PROXY_BYTECODE: &[u8] = include_bytes!(
    "../artifacts/trade-account-registry-proxy/trade-account-registry-proxy.bin"
);
pub const TRADE_ACCOUNT_REGISTER_PROXY_STORAGE: &[u8] = include_bytes!(
    "../artifacts/trade-account-registry-proxy/trade-account-registry-proxy-storage_slots.json"
);

/// Configuration for deploying TradeAccountRegistry contracts.
/// Contains bytecode and storage slot information needed for deployment.
#[derive(Clone)]
pub struct TradeAccountRegistryDeployConfig {
    /// Bytecode for the TradeAccountRegistry contract
    pub registry_bytecode: Vec<u8>,
    /// Storage slots configuration for the TradeAccountRegistry contract
    pub registry_storage_slots: Vec<StorageSlot>,
    /// Configurables
    pub registry_config: TradeAccountRegistryConfigurables,
    /// Bytecode for the TradeAccountRegistryProxy contract
    pub registry_proxy_bytecode: Vec<u8>,
    /// Storage slots configuration for the TradeAccountRegistryProxy
    pub registry_proxy_storage_slots: Vec<StorageSlot>,
    /// Register proxy configurations
    pub registry_proxy_config: TradeAccountRegistryProxyConfigurables,
    /// Maximum words per blob for deployment
    pub max_words_per_blob: usize,
    /// Bytecode for the TradeAccountProxy contract
    pub trade_account_proxy_bytecode: Vec<u8>,
    /// Blob id for trading account
    pub trade_account_proxy_root: Bytes32,
    /// The storage slots of the trade account proxy
    pub trade_account_proxy_storage_slots: Vec<StorageSlot>,
    /// Owner for the TradeAccountProxyRegistry contract
    pub proxy_owner: Option<Identity>,
    /// Owner for the TradeAccountRegistry contract
    pub registry_owner: Option<Identity>,
    /// Salt for contract deployment
    pub salt: Salt,
}

pub struct TradeAccountRegistryBlob {
    /// 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 TradeAccountRegistryDeployConfig {
    fn default() -> Self {
        let trade_account_proxy_storage_slots: Vec<StorageSlot> =
            serde_json::from_slice(TRADE_ACCOUNT_REGISTER_PROXY_STORAGE).unwrap();
        let trade_account_proxy_bytecode = Contract::regular(
            TRADE_ACCOUNT_PROXY_BYTECODE.to_vec(),
            Salt::zeroed(),
            trade_account_proxy_storage_slots.clone(),
        );
        let trade_account_proxy_root = trade_account_proxy_bytecode.code_root();

        Self {
            registry_bytecode: TRADE_ACCOUNT_REGISTER_BYTECODE.to_vec(),
            registry_storage_slots: serde_json::from_slice(
                TRADE_ACCOUNT_REGISTER_STORAGE,
            )
            .unwrap(),
            registry_config: TradeAccountRegistryConfigurables::default(),
            registry_proxy_bytecode: TRADE_ACCOUNT_REGISTER_PROXY_BYTECODE.to_vec(),
            registry_proxy_storage_slots: serde_json::from_slice(
                TRADE_ACCOUNT_REGISTER_PROXY_STORAGE,
            )
            .unwrap(),
            registry_proxy_config: TradeAccountRegistryProxyConfigurables::default(),
            max_words_per_blob: 10_000,
            trade_account_proxy_bytecode: TRADE_ACCOUNT_PROXY_BYTECODE.to_vec(),
            trade_account_proxy_root,
            trade_account_proxy_storage_slots,
            proxy_owner: None,
            registry_owner: None,
            salt: Salt::default(),
        }
    }
}

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

impl<W> TradeAccountRegistryManager<W>
where
    W: Account + Clone,
{
    pub fn change_wallet(mut self, wallet: &W) -> Self {
        self.wallet = wallet.clone();
        self.registry = self.registry.with_account(wallet.clone());

        self
    }

    pub fn new(deployer_wallet: W, contract_id: ContractId) -> Self {
        let proxy = TradeAccountRegistryProxy::new(contract_id, deployer_wallet.clone());
        let registry = TradeAccountRegistry::new(contract_id, deployer_wallet.clone());

        Self {
            registry_proxy: proxy,
            registry,
            contract_id,
            wallet: deployer_wallet,
        }
    }

    pub async fn register_blob(
        deployer_wallet: &W,
        trade_account_oracle_id: ContractId,
        trade_account_proxy_blob_id: BlobId,
        config: &TradeAccountRegistryDeployConfig,
    ) -> Result<TradeAccountRegistryBlob> {
        let registry_owner = config
            .registry_owner
            .unwrap_or(Identity::Address(deployer_wallet.address()));
        // Configure the contract with the trading pair assets
        let configurables = config
            .registry_config
            .clone()
            .with_ORACLE_CONTRACT_ID(trade_account_oracle_id)?
            .with_TRADE_ACCOUNT_PROXY_ROOT(Bits256(*config.trade_account_proxy_root))?
            .with_DEFAULT_TRADE_ACCOUNT_PROXY(ContractId::from(
                trade_account_proxy_blob_id,
            ))?
            .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(TradeAccountRegistryBlob {
            id: blob_id,
            exists: blob_exists,
            blob: blob.clone(),
        })
    }

    pub async fn register_proxy_blob(
        deployer_wallet: &W,
        config: &TradeAccountRegistryDeployConfig,
    ) -> Result<TradeAccountRegistryBlob> {
        let blobs = Contract::regular(
            config.trade_account_proxy_bytecode.clone(),
            Salt::zeroed(),
            config.trade_account_proxy_storage_slots.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(TradeAccountRegistryBlob {
            id: blob_id,
            exists: blob_exists,
            blob: blob.clone(),
        })
    }

    pub async fn deploy_register_blob(
        deployer_wallet: &W,
        trade_account_oracle_id: ContractId,
        config: &TradeAccountRegistryDeployConfig,
    ) -> Result<BlobId> {
        let proxy_blob = Self::register_proxy_blob(deployer_wallet, config).await?;
        if !proxy_blob.exists {
            let mut builder =
                BlobTransactionBuilder::default().with_blob(proxy_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)?;
        }

        let register_blob = Self::register_blob(
            deployer_wallet,
            trade_account_oracle_id,
            proxy_blob.id,
            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: &TradeAccountRegistryDeployConfig,
    ) -> Result<(TradeAccountRegistryProxy<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 = TradeAccountRegistryProxy::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 a TradeAccountRegistry contract.
    ///
    /// # Arguments
    /// * `deployer_wallet` - The wallet to use for deployment (pays gas fees)
    /// * `config` - Deployment configuration containing bytecode and settings
    /// * `oracle_id` - The oracle contract ID to configure
    ///
    /// # Returns
    /// * `Ok(TradeAccountRegistryDeploy)` - Complete deployment result with contract instance
    /// * `Err(anyhow::Error)` - If deployment or initialization fails
    ///
    /// # Process
    /// 1. Deploys the TradeAccountRegistry contract
    /// 2. Initializes the contract with the specified owner
    pub async fn deploy(
        deployer_wallet: &W,
        trade_account_oracle_id: ContractId,
        config: &TradeAccountRegistryDeployConfig,
    ) -> Result<TradeAccountRegistryManager<W>> {
        let register_blob_id =
            Self::deploy_register_blob(deployer_wallet, trade_account_oracle_id, config)
                .await?;
        let (proxy, requires_initialization) =
            TradeAccountRegistryManager::deploy_register_proxy(
                deployer_wallet,
                &register_blob_id,
                config,
            )
            .await?;
        let register_deploy = TradeAccountRegistryManager::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)
    }

    /// Upgrades the TradeAccountRegistry contract with new bytecode.
    ///
    /// # Arguments
    /// * `config` - New deployment configuration containing updated bytecode and settings
    ///
    /// # Returns
    /// * `Ok(BlobId)` - The ID of the new blob containing the upgraded contract
    /// * `Err(anyhow::Error)` - If the upgrade fails
    ///
    /// # Process
    /// 1. Deploys new blob with updated bytecode
    /// 2. Updates the proxy to point to the new blob
    pub async fn upgrade(
        &self,
        trade_account_oracle_id: ContractId,
        config: &TradeAccountRegistryDeployConfig,
    ) -> Result<BlobId> {
        // Deploy new blob with updated bytecode
        let new_blob_id =
            Self::deploy_register_blob(&self.wallet, trade_account_oracle_id, 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_trade_account(&self, owner: Identity) -> Result<Option<ContractId>> {
        Ok(self
            .registry
            .methods()
            .get_trade_account(owner)
            .simulate(Execution::state_read_only())
            .await?
            .value)
    }

    pub async fn oracle_id(&self) -> Result<ContractId> {
        Ok(self
            .registry
            .methods()
            .get_oracle_id()
            .simulate(Execution::state_read_only())
            .await?
            .value)
    }
}

impl TradeAccountRegistryManager<Wallet> {
    pub async fn register_trade_account(
        &self,
        owner_identity: Identity,
        trade_account_id: ContractId,
        trade_account_oracle_id: ContractId,
        call_option: &CallOption,
        deploy_config: &DeployConfig,
    ) -> Result<(ContractId, u64)> {
        let configurables: Vec<fuels::core::Configurable> = match deploy_config {
            DeployConfig::Old0_69_1(_) => {
                TradingAccountProxy0_69_1Configurables::default()
                    .with_ORACLE_CONTRACT_ID(trade_account_oracle_id)?
                    .with_INITIAL_OWNER(TradingAccountState::Initialized(owner_identity))?
                    .into()
            }
            DeployConfig::Latest(_) => TradingAccountProxyConfigurables::default()
                .with_ORACLE_CONTRACT_ID(trade_account_oracle_id)?
                .with_INITIAL_OWNER(TradingAccountState::Initialized(owner_identity))?
                .into(),
        };
        let call_handler = self
            .registry
            .methods()
            .register_contract(
                trade_account_id,
                Some(
                    configurables
                        .into_iter()
                        .map(|c| (c.offset, c.data))
                        .collect(),
                ),
            )
            .with_contract_ids(&[trade_account_id, trade_account_oracle_id]);

        let result = match call_option {
            CallOption::AwaitBlock => call_handler.call().await?,
            CallOption::AwaitPreconfirmation(ops) => {
                call_handler
                    .almost_sync_call(
                        &ops.data_builder,
                        &ops.utxo_manager,
                        &ops.tx_config,
                    )
                    .await?
                    .tx_status?
            }
        };

        match result.value {
            Ok(_) => Ok((trade_account_id, result.tx_status.total_gas)),
            Err(e) => Err(anyhow::anyhow!("Failed to register trade account: {:?}", e)),
        }
    }
}

#[cfg(test)]
mod tests_trade_account_registry {
    use super::*;
    use crate::{
        trade_account::TradeAccountManager,
        trade_account_deploy::{
            DeployConfig,
            TradeAccountDeploy,
            TradeAccountDeployConfig,
        },
    };
    use fuels::test_helpers::{
        WalletsConfig,
        launch_custom_provider_and_get_wallets,
    };

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

        let deploy_config = DeployConfig::Latest(TradeAccountDeployConfig::default());
        let trade_account_deploy =
            TradeAccountDeploy::deploy(&deployer_wallet, &deploy_config)
                .await
                .unwrap();
        let user1_identity = Identity::Address(user1_wallet.address());

        // Test register of trade account for user 1
        let user1_trade_account_deployed = trade_account_deploy
            .deploy_with_account(&user1_identity, &deploy_config, &CallOption::AwaitBlock)
            .await
            .unwrap();
        let user1_trade_account = TradeAccountManager::new(
            &user1_wallet,
            user1_trade_account_deployed.proxy_id.unwrap(),
        )
        .await
        .unwrap();

        let trade_account_registry = TradeAccountRegistryManager::deploy(
            &deployer_wallet,
            trade_account_deploy.oracle_id,
            &Default::default(),
        )
        .await
        .unwrap();

        trade_account_registry
            .register_trade_account(
                user1_identity,
                user1_trade_account.contract_id(),
                trade_account_deploy.oracle_id,
                &CallOption::AwaitBlock,
                &deploy_config,
            )
            .await
            .unwrap();

        let result = trade_account_registry
            .get_trade_account(user1_identity)
            .await
            .unwrap();
        assert_eq!(
            result,
            Some(user1_trade_account.contract_id()),
            "Trade account should register successfully"
        );

        // Upgrate trade account regsitry
        let current_blob_id = trade_account_registry
            .registry_proxy
            .methods()
            .proxy_target()
            .simulate(Execution::state_read_only())
            .await
            .unwrap()
            .value;
        trade_account_registry
            .upgrade(
                trade_account_deploy.oracle_id,
                &TradeAccountRegistryDeployConfig {
                    trade_account_proxy_root: Bytes32::new([1; 32]),
                    ..Default::default()
                },
            )
            .await
            .unwrap();
        let upgraded_blob_id = trade_account_registry
            .registry_proxy
            .methods()
            .proxy_target()
            .simulate(Execution::state_read_only())
            .await
            .unwrap()
            .value;
        assert_ne!(
            current_blob_id, upgraded_blob_id,
            "Blob id should be updated after upgrade"
        );

        // Test register of trade account for user 2 after upgrade
        let result = trade_account_registry
            .get_trade_account(user1_identity)
            .await
            .unwrap();
        assert_eq!(
            result,
            Some(user1_trade_account.contract_id()),
            "Trade account registered should keep working after upgrade"
        );
    }

    #[tokio::test]
    async fn test_trade_account_registration_on_registry() {
        // 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 user_wallet = wallets.pop().unwrap();

        let deploy_config = DeployConfig::Latest(TradeAccountDeployConfig::default());
        let trade_account_deploy =
            TradeAccountDeploy::deploy(&deployer_wallet, &deploy_config)
                .await
                .unwrap();
        let owner_identity = Identity::Address(user_wallet.address());
        let trade_account_deployed = trade_account_deploy
            .deploy_with_account(&owner_identity, &deploy_config, &CallOption::AwaitBlock)
            .await
            .unwrap();
        let trade_account = TradeAccountManager::new(
            &user_wallet,
            trade_account_deployed.proxy_id.unwrap(),
        )
        .await
        .unwrap();

        let trade_account_registry = TradeAccountRegistryManager::deploy(
            &deployer_wallet,
            trade_account_deploy.oracle_id,
            &Default::default(),
        )
        .await
        .unwrap();

        trade_account_registry
            .register_trade_account(
                owner_identity,
                trade_account.contract_id(),
                trade_account_deploy.oracle_id,
                &CallOption::AwaitBlock,
                &deploy_config,
            )
            .await
            .unwrap();

        let result = trade_account_registry
            .get_trade_account(owner_identity)
            .await
            .unwrap();
        assert_eq!(result, Some(trade_account.contract_id()));

        trade_account_registry
            .register_trade_account(
                owner_identity,
                trade_account.contract_id(),
                trade_account_deploy.oracle_id,
                &CallOption::AwaitBlock,
                &deploy_config,
            )
            .await
            .expect_err("Should fail if already register");
    }
}