Contractless/src/bin/lookup_remote_balance.rs

use blockchain::common::cli_prompts::prompt_hidden_nonempty;
use blockchain::common::network_startup::get_connections;
use blockchain::env;
use blockchain::from_str;
use blockchain::read_to_string;
use blockchain::records::memory::response_channels::generate_uid;
use blockchain::standalone_tools::connections::handshake;
use blockchain::tilde;
use blockchain::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('{') {
        // Prefer short_address for balance lookups, but accept long_address too.
        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("long_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 encryption_key = prompt_hidden_nonempty(
        "What is your wallet decryption key? ",
        "Wallet key cannot be empty. Please try again.",
    )
    .await;

    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::WalletKey(encryption_key.clone()),
            hashmap_key,
        )
        .await;
        match result {
            Ok(response) => {
                if response.is_empty() {
                    connected = true;
                    println!("[]");
                    break;
                }
                // Balance rows are 15 bytes asset name, 4 bytes NFT series, and 8 bytes balance.
                if response.len() % 27 == 0 {
                    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);
                        if nft_series > 0 {
                            println!("{asset_name}:{nft_series} = {formatted_balance}");
                        } else {
                            println!("{asset_name} = {formatted_balance}");
                        }
                    }
                    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");
    }
}