o2-tools 0.1.11

use crate::{
    CallOption,
    call_handler_ext::CallHandlerExt,
    contract_ext::ContractExt,
};
use anyhow::Result;
use fuels::{
    prelude::*,
    tx::StorageSlot,
    types::{
        ContractId,
        Identity,
    },
};

abigen!(
    Contract(
        name = "TradingAccountOracle",
        abi = "artifacts/trade-account-oracle/trade-account-oracle-abi.json"
    ),
    Contract(
        name = "TradingAccount",
        abi = "artifacts/trade-account/trade-account-abi.json"
    ),
    Contract(
        name = "TradingAccountProxy",
        abi = "artifacts/trade-account-proxy/trade-account-proxy-abi.json",
    ),
    Contract(
        name = "TradingAccountProxy0_69_1",
        abi = "artifacts/0.69.1/trade-account-proxy/artifacts/release/trade-account-proxy-abi.json",
    ),
);
pub const TRADE_ACCOUNT_BYTECODE: &[u8] =
    include_bytes!("../artifacts/trade-account/trade-account.bin");
pub const TRADE_ACCOUNT_STORAGE: &[u8] =
    include_bytes!("../artifacts/trade-account/trade-account-storage_slots.json");
pub const TRADE_ACCOUNT_PROXY_BYTECODE: &[u8] =
    include_bytes!("../artifacts/trade-account-proxy/trade-account-proxy.bin");
pub const TRADE_ACCOUNT_PROXY_STORAGE: &[u8] = include_bytes!(
    "../artifacts/trade-account-proxy/trade-account-proxy-storage_slots.json"
);
pub const TRADE_ACCOUNT_ORACLE_BYTECODE: &[u8] =
    include_bytes!("../artifacts/trade-account-oracle/trade-account-oracle.bin");
pub const TRADE_ACCOUNT_ORACLE_STORAGE: &[u8] = include_bytes!(
    "../artifacts/trade-account-oracle/trade-account-oracle-storage_slots.json"
);

/// Configuration for deploying trade account contracts.
/// Contains all bytecode and storage slot information needed for deployment.
#[derive(Clone)]
pub struct TradeAccountDeployConfig {
    /// Bytecode for the trade account implementation contract
    pub trade_account_bytecode: Vec<u8>,
    /// Bytecode for the oracle contract that manages implementation addresses
    pub oracle_bytecode: Vec<u8>,
    /// Bytecode for the proxy contract that delegates to implementations
    pub proxy_bytecode: Vec<u8>,
    /// Storage slots configuration for the trade account contract
    pub trade_account_storage_slots: Vec<StorageSlot>,
    /// Storage slots configuration for the oracle contract
    pub oracle_storage_slots: Vec<StorageSlot>,
    /// Storage slots configuration for the proxy contract
    pub proxy_storage_slots: Vec<StorageSlot>,
    /// Maximum words per blob for large contract deployment
    pub max_words_per_blob: usize,
    /// Salt for contract deployment
    pub salt: Salt,
}

impl TradeAccountDeployConfig {
    pub fn old_0_69_1() -> Self {
        pub const TRADE_ACCOUNT_BYTECODE_0_69_1: &[u8] = include_bytes!(
            "../artifacts/0.69.1/trade-account/artifacts/release/trade-account.bin"
        );
        pub const TRADE_ACCOUNT_STORAGE_0_69_1: &[u8] = include_bytes!(
            "../artifacts/0.69.1/trade-account/artifacts/release/trade-account-storage_slots.json"
        );
        pub const TRADE_ACCOUNT_PROXY_BYTECODE_0_69_1: &[u8] = include_bytes!(
            "../artifacts/0.69.1/trade-account-proxy/artifacts/release/trade-account-proxy.bin"
        );
        pub const TRADE_ACCOUNT_PROXY_STORAGE_0_69_1: &[u8] = include_bytes!(
            "../artifacts/0.69.1/trade-account-proxy/artifacts/release/trade-account-proxy-storage_slots.json"
        );
        pub const TRADE_ACCOUNT_ORACLE_BYTECODE_0_69_1: &[u8] = include_bytes!(
            "../artifacts/0.69.1/trade-account-oracle/artifacts/release/trade-account-oracle.bin"
        );
        pub const TRADE_ACCOUNT_ORACLE_STORAGE_0_69_1: &[u8] = include_bytes!(
            "../artifacts/0.69.1/trade-account-oracle/artifacts/release/trade-account-oracle-storage_slots.json"
        );

        Self {
            trade_account_bytecode: TRADE_ACCOUNT_BYTECODE_0_69_1.to_vec(),
            oracle_bytecode: TRADE_ACCOUNT_ORACLE_BYTECODE_0_69_1.to_vec(),
            proxy_bytecode: TRADE_ACCOUNT_PROXY_BYTECODE_0_69_1.to_vec(),
            trade_account_storage_slots: serde_json::from_slice(
                TRADE_ACCOUNT_STORAGE_0_69_1,
            )
            .unwrap(),
            oracle_storage_slots: serde_json::from_slice(
                TRADE_ACCOUNT_ORACLE_STORAGE_0_69_1,
            )
            .unwrap(),
            proxy_storage_slots: serde_json::from_slice(
                TRADE_ACCOUNT_PROXY_STORAGE_0_69_1,
            )
            .unwrap(),
            max_words_per_blob: 10_000,
            salt: Salt::default(),
        }
    }
}

impl Default for TradeAccountDeployConfig {
    fn default() -> Self {
        Self {
            trade_account_bytecode: TRADE_ACCOUNT_BYTECODE.to_vec(),
            oracle_bytecode: TRADE_ACCOUNT_ORACLE_BYTECODE.to_vec(),
            proxy_bytecode: TRADE_ACCOUNT_PROXY_BYTECODE.to_vec(),
            trade_account_storage_slots: serde_json::from_slice(TRADE_ACCOUNT_STORAGE)
                .unwrap(),
            oracle_storage_slots: serde_json::from_slice(TRADE_ACCOUNT_ORACLE_STORAGE)
                .unwrap(),
            proxy_storage_slots: serde_json::from_slice(TRADE_ACCOUNT_PROXY_STORAGE)
                .unwrap(),
            max_words_per_blob: 100_000,
            salt: Salt::default(),
        }
    }
}

/// Result of a complete trade account deployment.
/// Contains all deployed contract instances and their IDs for easy access.
#[derive(Clone)]
pub struct TradeAccountDeploy<W> {
    /// The deployed oracle contract instance
    pub oracle: TradingAccountOracle<W>,
    /// Contract ID of the deployed oracle
    pub oracle_id: ContractId,
    /// Blob ID of the trade account implementation
    pub trade_account_blob_id: BlobId,
    /// The wallet to use for deployment (pays all gas fees)
    pub deployer_wallet: W,
    /// The deployed proxy contract instance
    pub proxy: Option<TradingAccountProxy<W>>,
    /// Contract ID of the deployed proxy
    pub proxy_id: Option<ContractId>,
}

pub struct TradeAccountBlob {
    /// 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> TradeAccountDeploy<W>
where
    W: Account + Clone,
{
    pub fn change_wallet(mut self, wallet: &W) -> Self {
        self.deployer_wallet = wallet.clone();
        self.oracle = self.oracle.with_account(wallet.clone());

        if let Some(proxy) = self.proxy {
            self.proxy = Some(proxy.with_account(wallet.clone()));
        }

        self
    }

    pub async fn from_oracle_id(
        deployer_wallet: &W,
        oracle_id: ContractId,
    ) -> Result<Self>
    where
        W: Account + Clone,
    {
        let oracle: TradingAccountOracle<W> =
            TradingAccountOracle::new(oracle_id, deployer_wallet.clone());
        let trade_account_blob_id = oracle
            .methods()
            .get_trade_account_impl()
            .simulate(Execution::state_read_only())
            .await?
            .value
            .ok_or_else(|| anyhow::anyhow!("Trade account implementation not set"))?
            .into();
        Ok(Self {
            oracle,
            oracle_id,
            trade_account_blob_id,
            deployer_wallet: deployer_wallet.clone(),
            proxy: None,
            proxy_id: None,
        })
    }

    pub fn trade_account_blob_from_config(
        config: &TradeAccountDeployConfig,
    ) -> Result<Blob> {
        let blobs = Contract::regular(
            config.trade_account_bytecode.clone(),
            config.salt,
            config.trade_account_storage_slots.clone(),
        )
        .convert_to_loader(config.max_words_per_blob)?
        .blobs()
        .to_vec();
        let blob = blobs[0].clone();
        Ok(blob)
    }

    pub fn trade_account_proxy_blob_from_config(
        config: &TradeAccountDeployConfig,
    ) -> Result<Blob> {
        let blobs = Contract::regular(
            config.proxy_bytecode.clone(),
            config.salt,
            config.proxy_storage_slots.clone(),
        )
        .convert_to_loader(config.max_words_per_blob)?
        .blobs()
        .to_vec();
        let blob = blobs[0].clone();
        Ok(blob)
    }

    pub async fn trade_account_blob(
        deployer_wallet: &W,
        config: &TradeAccountDeployConfig,
    ) -> Result<TradeAccountBlob> {
        let blob = Self::trade_account_blob_from_config(config)?;
        let blob_id = blob.id();
        let blob_exists = deployer_wallet.try_provider()?.blob_exists(blob_id).await?;

        Ok(TradeAccountBlob {
            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_trade_account_blob(
        deployer_wallet: &W,
        config: &DeployConfig,
    ) -> Result<BlobId> {
        match config {
            DeployConfig::Old0_69_1(config) | DeployConfig::Latest(config) => {
                let trade_account_blob =
                    Self::trade_account_blob(deployer_wallet, config).await?;
                if !trade_account_blob.exists {
                    let mut builder = BlobTransactionBuilder::default()
                        .with_blob(trade_account_blob.blob.clone());

                    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(trade_account_blob.id)
            }
        }
    }

    /// Deploys the oracle contract that manages trade account implementation addresses.
    /// The oracle is initialized with the deployer as owner and the blob ID as the current implementation.
    ///
    /// # Arguments
    /// * `deployer_wallet` - The wallet to use for deployment and set as initial owner
    /// * `trade_account_blob_id` - The blob ID of the trade account implementation
    /// * `config` - Deployment configuration containing oracle bytecode and settings
    ///
    /// # Returns
    /// * `Ok((TradingAccountOracle, ContractId))` - The oracle instance and its contract ID
    /// * `Err(anyhow::Error)` - If deployment or initialization fails
    pub async fn deploy_oracle(
        deployer_wallet: &W,
        trade_account_blob_id: &BlobId,
        config: &DeployConfig,
    ) -> Result<(TradingAccountOracle<W>, ContractId)> {
        match config {
            DeployConfig::Old0_69_1(_) => Err(anyhow::anyhow!(
                "Deployment with old 0.69.1 config is not supported for oracle"
            )),
            DeployConfig::Latest(config) => {
                let contract = Contract::regular(
                    config.oracle_bytecode.clone(),
                    config.salt,
                    config.oracle_storage_slots.clone(),
                )
                .with_configurables(TradingAccountOracleConfigurables::default())
                .with_salt(config.salt);
                let contract_id = contract.contract_id();
                let instance =
                    TradingAccountOracle::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(
                            Identity::Address(deployer_wallet.address()),
                            ContractId::from(*trade_account_blob_id),
                        )
                        .call()
                        .await?;
                }

                Ok((instance, contract_id))
            }
        }
    }

    pub fn trade_account_contract(
        oracle_id: &ContractId,
        owner_identity: &Identity,
        config: &DeployConfig,
    ) -> Result<Contract<fuels::programs::contract::Regular>> {
        match config {
            DeployConfig::Old0_69_1(config) => {
                let configurables = TradingAccountProxy0_69_1Configurables::default()
                    .with_ORACLE_CONTRACT_ID(*oracle_id)?
                    .with_INITIAL_OWNER(State::Initialized(*owner_identity))?;
                let contract = Contract::regular(
                    config.proxy_bytecode.clone(),
                    config.salt,
                    config.proxy_storage_slots.clone(),
                )
                .with_configurables(configurables);
                Ok(contract)
            }
            DeployConfig::Latest(config) => {
                let configurables = TradingAccountProxyConfigurables::default()
                    .with_ORACLE_CONTRACT_ID(*oracle_id)?
                    .with_INITIAL_OWNER(State::Initialized(*owner_identity))?;
                let contract = Contract::regular(
                    config.proxy_bytecode.clone(),
                    config.salt,
                    config.proxy_storage_slots.clone(),
                )
                .with_configurables(configurables);
                Ok(contract)
            }
        }
    }

    /// Deploys the proxy contract that delegates calls to the current implementation.
    /// The proxy is configured with the oracle contract ID and initialized with the specified owner.
    ///
    /// # Arguments
    /// * `deployer_wallet` - The wallet to use for deployment (pays gas fees)
    /// * `owner_wallet` - The wallet to set as the proxy owner
    /// * `oracle_id` - The contract ID of the deployed oracle
    /// * `config` - Deployment configuration containing proxy bytecode and settings
    ///
    /// # Returns
    /// * `Ok((TradingAccountProxy, ContractId))` - The proxy instance and its contract ID
    /// * `Err(anyhow::Error)` - If deployment or initialization fails
    pub async fn deploy_proxy(
        deployer_wallet: &Wallet,
        owner_identity: &Identity,
        oracle_id: ContractId,
        config: &DeployConfig,
        call_option: &CallOption,
    ) -> Result<(TradingAccountProxy<Wallet>, ContractId)>
    where
        W: Account + Clone,
    {
        let contract = Self::trade_account_contract(&oracle_id, owner_identity, config)?;
        let result = deployer_wallet
            .try_provider()?
            .contract_exists(&contract.contract_id())
            .await?;

        let id = if !result {
            match call_option {
                CallOption::AwaitBlock => {
                    contract
                        .deploy(deployer_wallet, TxPolicies::default())
                        .await?
                        .contract_id
                }
                CallOption::AwaitPreconfirmation(ops) => {
                    let result = contract
                        .almost_sync_deploy(
                            deployer_wallet,
                            &ops.data_builder,
                            &ops.utxo_manager,
                            &ops.tx_config,
                        )
                        .await?;

                    // We need to wait for the block to be produced in order to submit next transaction
                    // that will initialize the proxy.
                    if let Some(tx_id) = result.tx_id {
                        deployer_wallet
                            .try_provider()?
                            .client()
                            .await_transaction_commit(&tx_id)
                            .await?;
                    }
                    result.contract_id
                }
            }
        } else {
            contract.contract_id()
        };

        let proxy = TradingAccountProxy::new(id, deployer_wallet.clone());

        let response = proxy
            .methods()
            .proxy_owner()
            .simulate(Execution::state_read_only())
            .await?;

        if response.value == State::Uninitialized {
            let call_handler =
                proxy.methods().initialize().with_contract_ids(&[oracle_id]);

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

        Ok((proxy, id))
    }

    /// Deploys a complete trade account system including implementation blob, oracle, and proxy.
    /// This is the main deployment function that orchestrates the entire process.
    ///
    /// # Arguments
    /// * `deployer_wallet` - The wallet to use for deployment (pays all gas fees)
    /// * `owner_wallet` - The wallet to set as the trade account owner
    /// * `config` - Deployment configuration containing all bytecode and settings
    ///
    /// # Returns
    /// * `Ok(TradeAccountDeploy)` - Complete deployment result with all contract instances
    /// * `Err(anyhow::Error)` - If any part of the deployment fails
    ///
    /// # Process
    /// 1. Deploys the trade account implementation as a blob
    /// 2. Deploys the oracle contract and registers the blob ID
    /// 3. Deploys the proxy contract configured with the oracle
    pub async fn deploy(
        deployer_wallet: &W,
        config: &DeployConfig,
    ) -> Result<TradeAccountDeploy<W>> {
        let trade_account_blob_id =
            Self::deploy_trade_account_blob(deployer_wallet, config).await?;
        let (oracle, oracle_id) =
            Self::deploy_oracle(deployer_wallet, &trade_account_blob_id, config).await?;
        let trade_account_blob_id = oracle
            .methods()
            .get_trade_account_impl()
            .simulate(Execution::state_read_only())
            .await?
            .value
            .unwrap();

        Ok(TradeAccountDeploy {
            oracle,
            oracle_id,
            trade_account_blob_id: trade_account_blob_id.into(),
            deployer_wallet: deployer_wallet.clone(),
            proxy: None,
            proxy_id: None,
        })
    }
}

pub enum DeployConfig {
    // TODO: Remove, when mainnet is upgraded to latest version of the trade account contracts.
    Old0_69_1(TradeAccountDeployConfig),
    Latest(TradeAccountDeployConfig),
}

impl DeployConfig {
    pub fn config(&self) -> &TradeAccountDeployConfig {
        match self {
            DeployConfig::Old0_69_1(config) | DeployConfig::Latest(config) => config,
        }
    }
}

impl TradeAccountDeploy<Wallet> {
    pub async fn deploy_with_account(
        &self,
        owner_identity: &Identity,
        config: &DeployConfig,
        call_option: &CallOption,
    ) -> Result<Self> {
        let (proxy, proxy_id) = Self::deploy_proxy(
            &self.deployer_wallet,
            owner_identity,
            self.oracle_id,
            config,
            call_option,
        )
        .await?;
        Ok(Self {
            proxy: Some(proxy),
            proxy_id: Some(proxy_id),
            oracle: self.oracle.clone(),
            oracle_id: self.oracle_id,
            trade_account_blob_id: self.trade_account_blob_id,
            deployer_wallet: self.deployer_wallet.clone(),
        })
    }
}

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

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

        // Deploy contracts
        let config = DeployConfig::Latest(TradeAccountDeployConfig::default());
        let deployment = TradeAccountDeploy::deploy(&wallet, &config)
            .await
            .unwrap()
            .deploy_with_account(
                &wallet.address().into(),
                &config,
                &CallOption::AwaitBlock,
            )
            .await
            .unwrap();

        // Check if IDs exist by querying the deployed contracts
        let provider = wallet.try_provider().unwrap();

        // Check oracle contract exists
        let oracle_contract_info = provider
            .contract_exists(&deployment.oracle_id)
            .await
            .unwrap();
        assert!(oracle_contract_info, "Oracle contract should exist");

        // Check proxy contract exists
        let proxy_contract_info = provider
            .contract_exists(&deployment.proxy_id.unwrap())
            .await
            .unwrap();
        assert!(proxy_contract_info, "Proxy contract should exist");

        // Verify blob exists by checking the oracle's stored blob ID
        let stored_blob_id = deployment
            .oracle
            .methods()
            .get_trade_account_impl()
            .simulate(Execution::state_read_only())
            .await
            .unwrap()
            .value;

        assert_eq!(
            deployment.trade_account_blob_id,
            BlobId::from(stored_blob_id.unwrap()),
            "Trade account blob ID should match"
        );
    }
}