use anyhow::{anyhow, Result};
use clap::Args;
use fuels::{
accounts::{wallet::Wallet, ViewOnlyAccount},
prelude::*,
};
use std::{
cmp::max,
collections::{BTreeMap, HashMap},
path::Path,
};
use url::Url;
use crate::{
account::{
derive_account, derive_and_cache_addresses, print_balance, print_balance_empty,
read_cached_addresses, verify_address_and_update_cache,
},
format::List,
utils::load_wallet,
DEFAULT_CACHE_ACCOUNTS,
};
#[derive(Debug, Args)]
#[group(skip)]
pub struct Balance {
#[clap(flatten)]
pub(crate) account: crate::account::Balance,
#[clap(long)]
pub(crate) accounts: bool,
}
pub enum AccountVerification {
No,
Yes(String),
}
pub type AccountBalances = Vec<HashMap<String, u64>>;
pub type AccountsMap = BTreeMap<usize, Bech32Address>;
pub fn collect_accounts_with_verification(
wallet_path: &Path,
verification: AccountVerification,
) -> Result<AccountsMap> {
let wallet = load_wallet(wallet_path)?;
let mut addresses = read_cached_addresses(&wallet.crypto.ciphertext)?;
if let AccountVerification::Yes(password) = verification {
for (&ix, addr) in addresses.iter_mut() {
let account = derive_account(wallet_path, ix, &password)?;
if verify_address_and_update_cache(ix, &account, addr, &wallet.crypto.ciphertext)? {
*addr = account.address().clone();
}
}
}
Ok(addresses)
}
pub fn get_derived_accounts(
wallet_path: &Path,
unverified: bool,
target_accounts: Option<usize>,
) -> Result<AccountsMap> {
let wallet = load_wallet(wallet_path)?;
let addresses = if unverified {
read_cached_addresses(&wallet.crypto.ciphertext)?
} else {
BTreeMap::new()
};
let target_accounts = target_accounts.unwrap_or(1);
if !unverified || addresses.len() < target_accounts {
let prompt = "Please enter your wallet password to verify accounts: ";
let password = rpassword::prompt_password(prompt)?;
let phrase_recovered = eth_keystore::decrypt_key(wallet_path, password)?;
let phrase = String::from_utf8(phrase_recovered)?;
let range = 0..max(target_accounts, DEFAULT_CACHE_ACCOUNTS);
derive_and_cache_addresses(&wallet, &phrase, range)
} else {
Ok(addresses)
}
}
pub fn print_account_balances(accounts_map: &AccountsMap, account_balances: &AccountBalances) {
let mut list = List::default();
list.add_newline();
for (ix, balance) in accounts_map.keys().zip(account_balances) {
let balance: BTreeMap<_, _> = balance
.iter()
.map(|(id, &val)| (id.clone(), u128::from(val)))
.collect();
if balance.is_empty() {
continue;
}
list.add_seperator();
list.add(format!("Account {ix}"), accounts_map[ix].to_string());
list.add_newline();
for (asset_id, amount) in balance {
list.add("Asset ID", asset_id);
list.add("Amount", amount.to_string());
}
list.add_seperator();
}
println!("{}", list.to_string());
}
pub(crate) async fn list_account_balances(
node_url: &Url,
addresses: &BTreeMap<usize, Bech32Address>,
) -> Result<(Vec<HashMap<String, u64>>, BTreeMap<String, u128>)> {
println!("Connecting to {node_url}");
let provider = Provider::connect(&node_url).await?;
println!("Fetching and summing balances of the following accounts:");
for (ix, addr) in addresses {
println!(" {ix:>3}: {addr}");
}
let accounts: Vec<_> = addresses
.values()
.map(|addr| Wallet::from_address(addr.clone(), Some(provider.clone())))
.collect();
let account_balances =
futures::future::try_join_all(accounts.iter().map(|acc| acc.get_balances())).await?;
let mut total_balance = BTreeMap::default();
for acc_bal in &account_balances {
for (asset_id, amt) in acc_bal {
let entry = total_balance.entry(asset_id.clone()).or_insert(0u128);
*entry = entry.checked_add(u128::from(*amt)).ok_or_else(|| {
anyhow!("Failed to display balance for asset {asset_id}: Value out of range.")
})?;
}
}
Ok((account_balances, total_balance))
}
pub async fn cli(wallet_path: &Path, balance: &Balance) -> Result<()> {
let verification = if !balance.account.unverified.unverified {
let prompt = "Please enter your wallet password to verify accounts: ";
let password = rpassword::prompt_password(prompt)?;
AccountVerification::Yes(password)
} else {
AccountVerification::No
};
let addresses = collect_accounts_with_verification(wallet_path, verification)?;
let node_url = &balance.account.node_url;
let (account_balances, total_balance) = list_account_balances(node_url, &addresses).await?;
if balance.accounts {
print_account_balances(&addresses, &account_balances);
}
println!("\nTotal:");
if total_balance.is_empty() {
print_balance_empty(&balance.account.node_url);
} else {
print_balance(&total_balance);
}
Ok(())
}