Contractless/src/bin/lookup_remote_balance.rs

use contractless::common::cli_prompts::{prompt_hidden_nonempty, prompt_wallet_path};
use contractless::common::network_startup::get_connections;
use contractless::env;
use contractless::from_str;
use contractless::json;
use contractless::read_to_string;
use contractless::records::memory::response_channels::generate_uid;
use contractless::standalone_tools::connections::handshake;
use contractless::standalone_tools::vanity_resolver::resolve_wallet_address_input;
use contractless::tilde;
use contractless::wallets::structures::Wallet;
use contractless::Value;
use rustyline::completion::FilenameCompleter;
use rustyline::error::ReadlineError;
use rustyline::{history::DefaultHistory, CompletionType, Config, Editor};
use rustyline_derive::Completer;
use rustyline_derive::Helper as RustyHelper;
use rustyline_derive::Highlighter as RustyHighlighter;
use rustyline_derive::Hinter as RustyHinter;
use rustyline_derive::Validator as RustyValidator;

fn format_balance(balance: u64) -> String {
    // Balance values are atomic units and display with 8 decimal places.
    let whole = balance / 100_000_000;
    let fractional = balance % 100_000_000;
    format!("{whole}.{fractional:08}")
}

fn extract_address(contents: &str) -> Result<String, String> {
    // Address files may be plain text or saved wallet JSON.
    let trimmed = contents.trim();
    if trimmed.is_empty() {
        return Err("Address file is empty".to_string());
    }

    if trimmed.starts_with('{') {
        let value: Value =
            from_str(trimmed).map_err(|e| format!("Failed to parse wallet JSON: {e}"))?;
        let address = value
            .get("short_address")
            .and_then(|v| v.as_str())
            .or_else(|| value.get("vanity_address").and_then(|v| v.as_str()))
            .ok_or_else(|| "Wallet JSON does not contain a usable address field".to_string())?;
        return Ok(address.trim().to_string());
    }

    Ok(trimmed.to_string())
}

#[derive(RustyHelper, Completer, RustyHinter, RustyHighlighter, RustyValidator)]
struct PathHelper {
    #[rustyline(Completer)]
    completer: FilenameCompleter,
}

fn prompt_for_path(prompt: &str) -> Result<String, String> {
    // Rustyline gives the interactive prompt filesystem completion.
    let config = Config::builder()
        .completion_type(CompletionType::List)
        .build();
    let mut editor =
        Editor::<PathHelper, DefaultHistory>::with_config(config).map_err(|e| e.to_string())?;
    editor.set_helper(Some(PathHelper {
        completer: FilenameCompleter::new(),
    }));

    match editor.readline(prompt) {
        Ok(line) => Ok(line.trim().to_string()),
        Err(ReadlineError::Interrupted) | Err(ReadlineError::Eof) => {
            Err("Input cancelled".to_string())
        }
        Err(err) => Err(format!("Failed to read address file path: {err}")),
    }
}

#[tokio::main]
async fn main() {
    // Command 23 asks a peer for all asset balances for one wallet address.
    let hashmap_key = generate_uid();

    let rpc_command = 23;

    // Accept an address file path as an arg or prompt interactively.
    let args: Vec<String> = env::args().collect();

    if args.len() > 1 && args.len() != 2 {
        println!("Usage: ./address_balance_lookup <address_file>");
        return;
    }

    let address_file_path = if args.len() == 2 {
        args[1].clone()
    } else {
        match prompt_for_path("Please enter the path to the file containing the wallet address: ") {
            Ok(path) => path,
            Err(err) => {
                eprintln!("{err}");
                return;
            }
        }
    };

    // Extract the lookup address from the selected file before opening a peer connection.
    let expanded_path = tilde(&address_file_path).to_string();
    let address_contents = read_to_string(&expanded_path)
        .await
        .expect("Failed to read address file");
    let wallet_address = match extract_address(&address_contents) {
        Ok(address) => address,
        Err(err) => {
            eprintln!("{err}");
            return;
        }
    };
    let wallet_path = prompt_wallet_path().await;
    let encryption_key = prompt_hidden_nonempty(
        "What is your wallet decryption key? ",
        "Wallet key cannot be empty. Please try again.",
    )
    .await;
    let wallet = match Wallet::try_obtain_wallet(encryption_key, Some(&wallet_path)).await {
        Ok(wallet) => wallet,
        Err(err) => {
            eprintln!("Wallet decryption failed: {err}");
            return;
        }
    };
    let wallet_address = match resolve_wallet_address_input(&wallet_address, &wallet).await {
        Ok(address) => address,
        Err(err) => {
            eprintln!("wallet address is not valid: {err}");
            return;
        }
    };

    let json = wallet_address;
    // Try each configured peer until one returns a parsable balance response or text error.
    let connections = get_connections().await;

    let mut connected: bool = false;
    for conn in connections {
        if connected {
            break;
        }
        let socket_address = conn.parse().expect("Failed to parse the socket address");
        let result = handshake::connect_and_handshake(
            socket_address,
            json.clone(),
            rpc_command,
            handshake::HandshakeWallet::WalletParts {
                public_key: wallet.saved.public_key.clone(),
                private_key: wallet.saved.private_key.clone(),
            },
            hashmap_key,
        )
        .await;
        match result {
            Ok(response) => {
                if response.is_empty() {
                    connected = true;
                    println!(
                        "{}",
                        serde_json::to_string_pretty(&json!({ "balances": [] })).unwrap()
                    );
                    break;
                }
                // Balance rows are 15 bytes asset name, 4 bytes NFT series, and 8 bytes balance.
                if response.len() % 27 == 0 {
                    let mut balances = Vec::new();
                    for entry in response.chunks_exact(27) {
                        let name_bytes = &entry[..15];
                        let nft_series = u32::from_le_bytes(entry[15..19].try_into().unwrap());
                        let balance = u64::from_le_bytes(entry[19..27].try_into().unwrap());
                        let asset_name = String::from_utf8_lossy(name_bytes).trim_end().to_string();
                        let formatted_balance = format_balance(balance);
                        balances.push(json!({
                            "asset": asset_name,
                            "nft_series": nft_series,
                            "balance": formatted_balance
                        }));
                    }
                    println!(
                        "{}",
                        serde_json::to_string_pretty(&json!({ "balances": balances })).unwrap()
                    );
                    connected = true;
                } else {
                    let response_text = String::from_utf8_lossy(&response);
                    let trimmed = response_text.trim();
                    if !trimmed.is_empty() {
                        println!("{trimmed}");
                        connected = true;
                    }
                }
            }
            Err(_) => connected = false,
        }
    }
    if !connected {
        eprintln!("failed to connect");
    }
}