o2-tools 0.3.21-rc

use crate::{
    CallOption,
    blob_loader,
    contract_ext::ContractExt,
};
use anyhow::Result;
use fuel_core_client::client::FuelClient;
use fuels::{
    core::Configurables,
    prelude::*,
    tx::StorageSlot,
    types::{
        Address,
        AssetId,
        ContractId,
        Identity,
    },
};

abigen!(
    Contract(
        name = "TrialTradingAccountOracle",
        abi = "artifacts/trial-trade-account-oracle/trial-trade-account-oracle-abi.json"
    ),
    Contract(
        name = "TrialTradingAccount",
        abi = "artifacts/trial-trade-account/trial-trade-account-abi.json"
    ),
    Contract(
        name = "TrialTradingAccountProxy",
        abi = "artifacts/trial-trade-account-proxy/trial-trade-account-proxy-abi.json",
    ),
);

pub const TRIAL_TRADE_ACCOUNT_BYTECODE: &[u8] =
    include_bytes!("../artifacts/trial-trade-account/trial-trade-account.bin");
pub const TRIAL_TRADE_ACCOUNT_STORAGE: &[u8] = include_bytes!(
    "../artifacts/trial-trade-account/trial-trade-account-storage_slots.json"
);
pub const TRIAL_TRADE_ACCOUNT_PROXY_BYTECODE: &[u8] = include_bytes!(
    "../artifacts/trial-trade-account-proxy/trial-trade-account-proxy.bin"
);
pub const TRIAL_TRADE_ACCOUNT_PROXY_STORAGE: &[u8] = include_bytes!(
    "../artifacts/trial-trade-account-proxy/trial-trade-account-proxy-storage_slots.json"
);
pub const TRIAL_TRADE_ACCOUNT_ORACLE_BYTECODE: &[u8] = include_bytes!(
    "../artifacts/trial-trade-account-oracle/trial-trade-account-oracle.bin"
);
pub const TRIAL_TRADE_ACCOUNT_ORACLE_STORAGE: &[u8] = include_bytes!(
    "../artifacts/trial-trade-account-oracle/trial-trade-account-oracle-storage_slots.json"
);
pub const DEFAULT_TRIAL_DURATION: u64 = 604800;
pub const DEFAULT_LIQUIDATION_THRESHOLD_AMOUNT: u64 = 0;
pub const DEFAULT_TRIAL_USER_PROFIT_SHARE_BPS: u64 = 0;

#[derive(Clone)]
pub struct TrialTradeAccountDeployConfig {
    pub trial_trade_account_bytecode: Vec<u8>,
    pub oracle_bytecode: Vec<u8>,
    pub oracle_storage_slots: Vec<StorageSlot>,
    pub proxy_bytecode: Vec<u8>,
    pub trial_trade_account_storage_slots: Vec<StorageSlot>,
    pub proxy_storage_slots: Vec<StorageSlot>,
    pub trial_trade_account_config: TrialTradingAccountConfigurables,
    pub proxy_config: TrialTradingAccountProxyConfigurables,
    pub cosigner: Option<Address>,
    pub trial_duration: u64,
    pub liquidation_threshold_amount: u64,
    pub liquidation_threshold_asset_id: AssetId,
    pub trial_allowed_order_book_ids: Vec<ContractId>,
    pub trial_user_profit_share_bps: u64,
    pub trial_platform_payout_identity: Identity,
    pub max_words_per_blob: usize,
    pub salt: Salt,
}

impl Default for TrialTradeAccountDeployConfig {
    fn default() -> Self {
        Self {
            trial_trade_account_bytecode: TRIAL_TRADE_ACCOUNT_BYTECODE.to_vec(),
            oracle_bytecode: TRIAL_TRADE_ACCOUNT_ORACLE_BYTECODE.to_vec(),
            oracle_storage_slots: serde_json::from_slice(
                TRIAL_TRADE_ACCOUNT_ORACLE_STORAGE,
            )
            .unwrap(),
            proxy_bytecode: TRIAL_TRADE_ACCOUNT_PROXY_BYTECODE.to_vec(),
            trial_trade_account_storage_slots: serde_json::from_slice(
                TRIAL_TRADE_ACCOUNT_STORAGE,
            )
            .unwrap(),
            proxy_storage_slots: serde_json::from_slice(
                TRIAL_TRADE_ACCOUNT_PROXY_STORAGE,
            )
            .unwrap(),
            trial_trade_account_config: TrialTradingAccountConfigurables::default(),
            proxy_config: TrialTradingAccountProxyConfigurables::default(),
            cosigner: None,
            trial_duration: DEFAULT_TRIAL_DURATION,
            liquidation_threshold_amount: DEFAULT_LIQUIDATION_THRESHOLD_AMOUNT,
            liquidation_threshold_asset_id: AssetId::zeroed(),
            trial_allowed_order_book_ids: Vec::new(),
            trial_user_profit_share_bps: DEFAULT_TRIAL_USER_PROFIT_SHARE_BPS,
            trial_platform_payout_identity: Identity::Address(Address::zeroed()),
            max_words_per_blob: 10_000,
            salt: Salt::default(),
        }
    }
}

impl TrialTradeAccountDeployConfig {
    pub fn with_oracle_id(mut self, oracle_id: ContractId) -> Result<Self> {
        self.trial_trade_account_config = self
            .trial_trade_account_config
            .with_ORACLE_CONTRACT_ID(oracle_id)?;
        self.proxy_config = self.proxy_config.with_ORACLE_CONTRACT_ID(oracle_id)?;
        Ok(self)
    }

    pub fn with_cosigner(mut self, cosigner: Address) -> Self {
        self.cosigner = Some(cosigner);
        self
    }

    pub fn with_trial_duration(mut self, trial_duration: u64) -> Self {
        self.trial_duration = trial_duration;
        self
    }

    pub fn with_liquidation_threshold_amount(
        mut self,
        liquidation_threshold_amount: u64,
    ) -> Self {
        self.liquidation_threshold_amount = liquidation_threshold_amount;
        self
    }

    pub fn with_liquidation_threshold_asset_id(
        mut self,
        liquidation_threshold_asset_id: AssetId,
    ) -> Self {
        self.liquidation_threshold_asset_id = liquidation_threshold_asset_id;
        self
    }

    pub fn with_trial_allowed_order_book_ids(
        mut self,
        trial_allowed_order_book_ids: Vec<ContractId>,
    ) -> Self {
        self.trial_allowed_order_book_ids = trial_allowed_order_book_ids;
        self
    }

    pub fn with_trial_user_profit_share_bps(
        mut self,
        trial_user_profit_share_bps: u64,
    ) -> Self {
        self.trial_user_profit_share_bps = trial_user_profit_share_bps;
        self
    }

    pub fn with_trial_platform_payout_identity(
        mut self,
        trial_platform_payout_identity: Identity,
    ) -> Self {
        self.trial_platform_payout_identity = trial_platform_payout_identity;
        self
    }

    pub fn with_initial_session_id(mut self, session_id: Identity) -> Result<Self> {
        self.proxy_config = self.proxy_config.with_INITIAL_SESSION_ID(session_id)?;
        Ok(self)
    }

    pub fn with_trial_trade_account_registry(
        mut self,
        registry_id: ContractId,
    ) -> Result<Self> {
        self.proxy_config = self
            .proxy_config
            .with_TRIAL_TRADE_ACCOUNT_REGISTRY(registry_id)?;
        Ok(self)
    }
}

#[derive(Clone)]
pub struct TrialTradeAccountDeploy<W> {
    pub oracle: TrialTradingAccountOracle<W>,
    pub oracle_id: ContractId,
    pub trial_trade_account_blob_id: BlobId,
    pub deployer_wallet: W,
    pub proxy: Option<TrialTradingAccountProxy<W>>,
    pub proxy_id: Option<ContractId>,
}

pub struct TrialTradeAccountBlob {
    pub id: BlobId,
    pub exists: bool,
    pub blob: Blob,
    pub data_blobs: Vec<Blob>,
}

impl<W> TrialTradeAccountDeploy<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> {
        let oracle = TrialTradingAccountOracle::new(oracle_id, deployer_wallet.clone());
        let trial_trade_account_blob_id = oracle
            .methods()
            .get_trial_account_impl()
            .simulate(Execution::state_read_only())
            .await?
            .value
            .ok_or_else(|| anyhow::anyhow!("Trial trade account implementation not set"))?
            .into();

        Ok(Self {
            oracle,
            oracle_id,
            trial_trade_account_blob_id,
            deployer_wallet: deployer_wallet.clone(),
            proxy: None,
            proxy_id: None,
        })
    }

    pub fn trial_trade_account_blob_from_config(
        config: &TrialTradeAccountDeployConfig,
    ) -> Result<(Vec<Blob>, Blob)> {
        blob_loader::build_loader_blobs(
            config.trial_trade_account_bytecode.clone(),
            config.salt,
            config.trial_trade_account_storage_slots.clone(),
            config.trial_trade_account_config.clone(),
            config.max_words_per_blob,
        )
    }

    pub fn trial_trade_account_proxy_blob_from_config(
        config: &TrialTradeAccountDeployConfig,
    ) -> Result<(Vec<Blob>, Blob)> {
        blob_loader::build_loader_blobs(
            config.proxy_bytecode.clone(),
            config.salt,
            config.proxy_storage_slots.clone(),
            Configurables::from(config.proxy_config.clone()),
            config.max_words_per_blob,
        )
    }

    pub async fn trial_trade_account_blob(
        deployer_wallet: &W,
        config: &TrialTradeAccountDeployConfig,
    ) -> Result<TrialTradeAccountBlob> {
        let (data_blobs, loader_blob) =
            Self::trial_trade_account_blob_from_config(config)?;
        let loader_blob_id = loader_blob.id();
        let loader_blob_exists = deployer_wallet
            .try_provider()?
            .blob_exists(loader_blob_id)
            .await?;

        Ok(TrialTradeAccountBlob {
            id: loader_blob_id,
            exists: loader_blob_exists,
            blob: loader_blob,
            data_blobs,
        })
    }

    pub async fn deploy_trial_trade_account_blob(
        deployer_wallet: &W,
        config: &DeployConfig,
    ) -> Result<BlobId> {
        match config {
            DeployConfig::Latest(config) => {
                let trial_trade_account_blob =
                    Self::trial_trade_account_blob(deployer_wallet, config).await?;
                blob_loader::upload_loader_blobs(
                    deployer_wallet,
                    trial_trade_account_blob.data_blobs,
                    trial_trade_account_blob.blob,
                )
                .await
            }
        }
    }

    /// Deploy the dedicated trial oracle (salted, so the id is deterministic
    /// per config) and initialize it with the deployer as owner. Idempotent:
    /// an already-deployed oracle is loaded instead.
    pub async fn deploy_oracle(
        deployer_wallet: &W,
        config: &DeployConfig,
    ) -> Result<(TrialTradingAccountOracle<W>, ContractId)> {
        match config {
            DeployConfig::Latest(config) => {
                let contract = Contract::regular(
                    config.oracle_bytecode.clone(),
                    config.salt,
                    config.oracle_storage_slots.clone(),
                )
                .with_configurables(TrialTradingAccountOracleConfigurables::default())
                .with_salt(config.salt);
                let contract_id = contract.contract_id();
                let instance =
                    TrialTradingAccountOracle::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()))
                        .call()
                        .await?;
                }

                Ok((instance, contract_id))
            }
        }
    }

    /// Deploy the full trial oracle stack: the trial oracle itself, the trial
    /// implementation blob, and the oracle state (implementation id and,
    /// when configured, the cosigner).
    pub async fn deploy(
        deployer_wallet: &W,
        config: &DeployConfig,
    ) -> Result<TrialTradeAccountDeploy<W>> {
        let (_, oracle_id) = Self::deploy_oracle(deployer_wallet, config).await?;
        Self::deploy_to_oracle(deployer_wallet, oracle_id, config).await
    }

    pub async fn deploy_to_oracle(
        deployer_wallet: &W,
        oracle_id: ContractId,
        config: &DeployConfig,
    ) -> Result<TrialTradeAccountDeploy<W>> {
        let config = match config {
            DeployConfig::Latest(config) => {
                DeployConfig::Latest(config.clone().with_oracle_id(oracle_id)?)
            }
        };
        let trial_trade_account_blob_id =
            Self::deploy_trial_trade_account_blob(deployer_wallet, &config).await?;
        let oracle = TrialTradingAccountOracle::new(oracle_id, deployer_wallet.clone());
        let implementation = ContractId::from(trial_trade_account_blob_id);
        let current_impl = oracle
            .methods()
            .get_trial_account_impl()
            .simulate(Execution::state_read_only())
            .await?
            .value;

        if current_impl != Some(implementation) {
            oracle
                .methods()
                .set_trial_account_impl(implementation)
                .call()
                .await?;
        }

        if let Some(cosigner) = config.config().cosigner {
            let current_cosigner = oracle
                .methods()
                .get_cosigner()
                .simulate(Execution::state_read_only())
                .await?
                .value;

            if current_cosigner != Some(cosigner) {
                oracle.methods().set_cosigner(cosigner).call().await?;
            }
        }

        Ok(TrialTradeAccountDeploy {
            oracle,
            oracle_id,
            trial_trade_account_blob_id,
            deployer_wallet: deployer_wallet.clone(),
            proxy: None,
            proxy_id: None,
        })
    }
}

impl TrialTradeAccountDeploy<Wallet> {
    pub fn trial_trade_account_contract(
        oracle_id: &ContractId,
        config: &DeployConfig,
    ) -> Result<Contract<fuels::programs::contract::Regular>> {
        match config {
            DeployConfig::Latest(config) => {
                let configurables = config
                    .proxy_config
                    .clone()
                    .with_ORACLE_CONTRACT_ID(*oracle_id)?;
                let contract = Contract::regular(
                    config.proxy_bytecode.clone(),
                    config.salt,
                    config.proxy_storage_slots.clone(),
                )
                .with_configurables(configurables);
                Ok(contract)
            }
        }
    }

    pub async fn deploy_proxy(
        deployer_wallet: &Wallet,
        oracle_id: ContractId,
        config: &DeployConfig,
        call_option: &CallOption,
        dry_run_client: Option<&FuelClient>,
        submit_clients: &[FuelClient],
    ) -> Result<(TrialTradingAccountProxy<Wallet>, ContractId, Option<u64>)> {
        let contract = Self::trial_trade_account_contract(&oracle_id, config)?;
        let already_deployed = deployer_wallet
            .try_provider()?
            .contract_exists(&contract.contract_id())
            .await?;

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

                    let mut total_gas = None;
                    if let Some(tx_id) = result.tx_id {
                        let status = deployer_wallet
                            .try_provider()?
                            .client()
                            .await_transaction_commit(&tx_id)
                            .await?;
                        if let fuel_core_client::client::types::TransactionStatus::Success {
                            total_gas: gas,
                            ..
                        } = status
                        {
                            total_gas = Some(gas);
                        }
                    }
                    (result.contract_id, total_gas)
                }
            }
        } else {
            (contract.contract_id(), None)
        };

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

        let _ = dry_run_client;

        Ok((proxy, id, total_gas))
    }

    /// Deploys the trial account proxy. Returns the deploy handle and, when a
    /// fresh preconfirmation deploy happened, the deploy transaction's total
    /// gas (None when the contract already existed or gas is unavailable).
    pub async fn deploy_with_account(
        &self,
        config: &DeployConfig,
        call_option: &CallOption,
        dry_run_client: Option<&FuelClient>,
        submit_clients: &[FuelClient],
    ) -> Result<(Self, Option<u64>)> {
        let (proxy, proxy_id, total_gas) = Self::deploy_proxy(
            &self.deployer_wallet,
            self.oracle_id,
            config,
            call_option,
            dry_run_client,
            submit_clients,
        )
        .await?;
        Ok((
            Self {
                proxy: Some(proxy),
                proxy_id: Some(proxy_id),
                oracle: self.oracle.clone(),
                oracle_id: self.oracle_id,
                trial_trade_account_blob_id: self.trial_trade_account_blob_id,
                deployer_wallet: self.deployer_wallet.clone(),
            },
            total_gas,
        ))
    }
}

#[derive(Clone)]
pub enum DeployConfig {
    Latest(TrialTradeAccountDeployConfig),
}

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