Contractless/src/records/balance_sheet/operations.rs

use crate::fs;
use crate::io::{self, stdout, Read, Seek, Write};
use crate::records::balance_sheet::pathing::{address_root_path, balance_file_path};
use crate::records::wallet_registry::resolve_canonical_registered_short_address;
use crate::sled::Db;
use crate::OpenOptions;
use crate::Path;

fn prune_empty_balance_directories(address: &str, file_path: &Path) {
    // Remove empty nested balance directories after the last asset file for an
    // address has been deleted.
    let Some(address_root) = address_root_path(address) else {
        return;
    };
    let mut current = file_path.parent();

    while let Some(dir) = current {
        if dir == address_root.as_path() || dir.starts_with(&address_root) {
            if fs::remove_dir(dir).is_ok() {
                current = dir.parent();
            } else {
                break;
            }
        } else {
            break;
        }
    }
}

pub fn balance_sheet_operation(
    address: &str,
    balance: u64,
    coin_type: &str,
    operand: &str,
) -> Result<(), io::Error> {
    // Reject invalid wallet keys before touching the balance-sheet path.
    if address_root_path(address).is_none() {
        return Err(io::Error::new(
            io::ErrorKind::InvalidInput,
            "Invalid wallet address for balance-sheet operation",
        ));
    }
    let file_path = balance_file_path(address, coin_type);
    let file_exists = Path::new(&file_path).exists();

    if let Some(parent_dir) = file_path.parent() {
        // Balance files are nested by address and asset, so parent folders may
        // need to be created before the actual wallet.bal file can be opened.
        fs::create_dir_all(parent_dir).map_err(|e| {
            eprintln!("Error creating directory: {e}");
            e
        })?;
    }

    // Open the balance file in place so the existing 8-byte value can be
    // replaced without changing the path.
    let mut file = OpenOptions::new()
        .read(true)
        .write(true)
        .create(true)
        .truncate(false)
        .open(&file_path)
        .map_err(|e| {
            eprintln!("Error opening or creating file: {e}");
            e
        })?;

    let mut buffer = [0; 8];
    let mut file_balance = if file_exists {
        // Existing balances are stored as a single little-endian u64.
        file.read_exact(&mut buffer).map_err(|e| {
            eprintln!("Error reading file balance_sheet address {address}: {e}");
            e
        })?;
        u64::from_le_bytes(buffer)
    } else {
        0
    };

    // Apply the requested delta while preventing an unsigned underflow.
    match operand {
        "addition" => file_balance += balance,
        "subtraction" => {
            if balance > file_balance {
                return Err(io::Error::new(
                    io::ErrorKind::InvalidInput,
                    "Invalid operation: balance too large",
                ));
            }
            file_balance -= balance;
        }
        _ => {
            return Err(io::Error::new(
                io::ErrorKind::InvalidInput,
                "Invalid operand",
            ))
        }
    }

    // Replace the stored balance from byte zero.
    file.seek(std::io::SeekFrom::Start(0)).map_err(|e| {
        eprintln!("Error seeking file: {e}");
        e
    })?;
    file.write_all(&file_balance.to_le_bytes()).map_err(|e| {
        eprintln!("Error writing file: {e}");
        e
    })?;

    stdout().flush().map_err(|e| {
        eprintln!("Error flushing stdout: {e}");
        e
    })?;

    // Keep nonzero balances on disk; zero balances are removed so empty asset
    // folders do not accumulate forever.
    if file_balance > 0 {
        file.set_len(8).map_err(|e| {
            eprintln!("Error truncating file: {e}");
            e
        })?;
    } else {
        fs::remove_file(&file_path).map_err(|e| {
            eprintln!("Error removing file: {e}");
            e
        })?;
        prune_empty_balance_directories(address, &file_path);
    }

    Ok(())
}

pub fn balance_sheet_operation_with_db(
    db: &Db,
    address: &str,
    balance: u64,
    coin_type: &str,
    operand: &str,
) -> Result<(), io::Error> {
    // Vanity or alternate registered addresses resolve to the canonical short
    // address before the filesystem balance is updated.
    let canonical_address = resolve_canonical_registered_short_address(db, address)
        .map_err(|err| io::Error::other(format!("Wallet registry lookup failed: {err}")))?
        .unwrap_or_else(|| address.to_string());

    balance_sheet_operation(&canonical_address, balance, coin_type, operand)
}