Contractless/src/verifications/async_funcs/checks/block_balance.rs

use crate::blocks::loans::LoanContractTransaction;
use crate::common::asset_names::loan_collateral_asset_name_or_base;
use crate::common::nft_assets::nft_asset_name;
use crate::common::types::Transaction;
use crate::decode;
use crate::records::balance_sheet::get_wallet_balance::get_balance_with_db;
use crate::records::memory::mempool::{
    get_basecoin_balance, get_coin_balance, signature_exists, BASECOIN,
};
use crate::rpc::commands::transaction_by_txid::request_transaction_by_txid;
use crate::rpc::responses::RpcResponse;
use crate::sled::Db;
use crate::verifications::async_funcs::total_payments::get_total_payments;
use std::collections::{HashMap, HashSet};

#[derive(Clone)]
struct BalanceDebit {
    address: String,
    coin: String,
    amount: u64,
}

struct TransactionBalanceView {
    hash: String,
    signatures: Vec<String>,
    debits: Vec<BalanceDebit>,
    check_saved_txid: bool,
}

pub(crate) struct TransactionBalanceReservation {
    hash: String,
    signatures: Vec<String>,
    debits: HashMap<(String, String), u64>,
    available: HashMap<(String, String), u64>,
    loan_payment: Option<LoanPaymentReservation>,
}

struct LoanPaymentReservation {
    contract_hash: String,
    payment_amount: u64,
    confirmed_paid: u64,
    contract_total: u64,
}

pub struct BlockBalanceTracker {
    transaction_hashes: HashSet<String>,
    transaction_signatures: HashSet<String>,
    reserved: HashMap<(String, String), u64>,
    loan_payments_reserved: HashMap<String, u64>,
}

impl BlockBalanceTracker {
    pub fn new() -> Self {
        Self {
            transaction_hashes: HashSet::new(),
            transaction_signatures: HashSet::new(),
            reserved: HashMap::new(),
            loan_payments_reserved: HashMap::new(),
        }
    }

    pub(crate) fn reserve(
        &mut self,
        reservation: TransactionBalanceReservation,
    ) -> Result<(), String> {
        if !self.transaction_hashes.insert(reservation.hash) {
            return Err("Duplicate transaction hash inside block.".to_string());
        }

        for signature in reservation.signatures {
            if !signature.is_empty() && !self.transaction_signatures.insert(signature) {
                return Err("Duplicate transaction signature inside block.".to_string());
            }
        }

        let loan_payment_update = if let Some(payment) = reservation.loan_payment {
            let already_reserved = self
                .loan_payments_reserved
                .get(&payment.contract_hash)
                .copied()
                .unwrap_or(0);
            let block_total = already_reserved
                .checked_add(payment.payment_amount)
                .ok_or_else(|| "Block loan payment reservation overflowed.".to_string())?;
            let total_after_block = payment
                .confirmed_paid
                .checked_add(block_total)
                .ok_or_else(|| "Block loan repayment total overflowed.".to_string())?;
            if total_after_block > payment.contract_total {
                return Err(format!(
                    "Loan payments inside block exceed remaining contract balance: contract={}",
                    payment.contract_hash
                ));
            }
            Some((payment.contract_hash, block_total))
        } else {
            None
        };

        let mut updates = Vec::new();
        for (key, transaction_debit) in reservation.debits {
            let already_reserved = self.reserved.get(&key).copied().unwrap_or(0);
            let required = already_reserved
                .checked_add(transaction_debit)
                .ok_or_else(|| "Block balance reservation overflowed.".to_string())?;
            let available = reservation.available.get(&key).copied().unwrap_or(0);

            if required > available {
                let (address, coin) = key;
                return Err(format!(
                    "Insufficient block-local funds: address={address} coin={} available={available} required={required}",
                    coin.trim()
                ));
            }

            updates.push((key, required));
        }

        for (key, required) in updates {
            self.reserved.insert(key, required);
        }
        if let Some((contract_hash, block_total)) = loan_payment_update {
            self.loan_payments_reserved
                .insert(contract_hash, block_total);
        }

        Ok(())
    }
}

impl Default for BlockBalanceTracker {
    fn default() -> Self {
        Self::new()
    }
}

fn add_debit(debits: &mut Vec<BalanceDebit>, address: &str, coin: String, amount: u64) {
    if amount == 0 {
        return;
    }

    debits.push(BalanceDebit {
        address: address.to_string(),
        coin,
        amount,
    });
}

fn add_to_total(
    totals: &mut HashMap<(String, String), u64>,
    debit: &BalanceDebit,
) -> Result<(), String> {
    let key = (debit.address.clone(), debit.coin.clone());
    let current = totals.get(&key).copied().unwrap_or(0);
    let updated = current
        .checked_add(debit.amount)
        .ok_or_else(|| "Block balance reservation overflowed.".to_string())?;
    totals.insert(key, updated);
    Ok(())
}

fn add_amount_to_total(
    totals: &mut HashMap<(String, String), u64>,
    key: (String, String),
    amount: u64,
) -> Result<(), String> {
    let current = totals.get(&key).copied().unwrap_or(0);
    let updated = current
        .checked_add(amount)
        .ok_or_else(|| "Block balance reservation overflowed.".to_string())?;
    totals.insert(key, updated);
    Ok(())
}

async fn transaction_is_pending(view: &TransactionBalanceView) -> Result<bool, String> {
    for signature in &view.signatures {
        if signature.is_empty() {
            continue;
        }

        if signature_exists(signature, &view.hash)
            .await
            .map_err(|err| err.to_string())?
        {
            return Ok(true);
        }
    }

    Ok(false)
}

async fn loan_contract_for_hash(
    db: &Db,
    contract_hash: &str,
) -> Result<LoanContractTransaction, String> {
    let contract_key =
        decode(contract_hash).map_err(|err| format!("Invalid loan contract hash: {err}"))?;
    let RpcResponse::Binary(bytes) = request_transaction_by_txid(db, contract_key).await;

    if bytes.is_empty() || bytes[0] != 7 {
        return Err("Invalid loan contract: referenced transaction is not a loan contract".into());
    }

    LoanContractTransaction::from_bytes(7, &bytes[1..])
        .await
        .map_err(|err| format!("Invalid loan contract: {err}"))
}

fn saved_txid_exists(db: &Db, hash: &str) -> Result<bool, String> {
    let tree = db
        .open_tree("txid")
        .map_err(|err| format!("Failed to open txid tree: {err}"))?;
    let key = decode(hash).map_err(|err| format!("Invalid transaction hash: {err}"))?;

    tree.contains_key(key)
        .map_err(|err| format!("Failed to check txid tree: {err}"))
}

async fn transaction_balance_view(
    db: &Db,
    transaction: &Transaction,
) -> Result<TransactionBalanceView, String> {
    match transaction {
        Transaction::Genesis(tx) => Ok(TransactionBalanceView {
            hash: tx.unsigned.hash().await,
            signatures: Vec::new(),
            debits: Vec::new(),
            check_saved_txid: false,
        }),
        Transaction::Rewards(tx) => Ok(TransactionBalanceView {
            hash: tx.unsigned.hash().await,
            signatures: Vec::new(),
            debits: Vec::new(),
            check_saved_txid: true,
        }),
        Transaction::Transfer(tx) => {
            let transfer = &tx.unsigned_transfer;
            let mut debits = Vec::new();
            let asset = nft_asset_name(&transfer.coin, transfer.nft_series);

            // The sender pays the transferred asset plus the base-coin fee.
            add_debit(&mut debits, &transfer.sender, asset, transfer.value);
            add_debit(
                &mut debits,
                &transfer.sender,
                BASECOIN.clone(),
                transfer.txfee,
            );

            Ok(TransactionBalanceView {
                hash: transfer.hash().await,
                signatures: vec![tx.signature.clone()],
                debits,
                check_saved_txid: true,
            })
        }
        Transaction::Burn(tx) => {
            let burn = &tx.unsigned_burn;
            let mut debits = Vec::new();
            let asset = nft_asset_name(&burn.coin, burn.nft_series);

            add_debit(&mut debits, &burn.address, asset, burn.value);
            add_debit(&mut debits, &burn.address, BASECOIN.clone(), burn.txfee);

            Ok(TransactionBalanceView {
                hash: burn.hash().await,
                signatures: vec![tx.signature.clone()],
                debits,
                check_saved_txid: true,
            })
        }
        Transaction::Token(tx) => {
            let token = &tx.unsigned_create_token;
            let mut debits = Vec::new();

            add_debit(&mut debits, &token.creator, BASECOIN.clone(), token.txfee);

            Ok(TransactionBalanceView {
                hash: token.hash().await,
                signatures: vec![tx.signature.clone()],
                debits,
                check_saved_txid: true,
            })
        }
        Transaction::IssueToken(tx) => {
            let token = &tx.unsigned_issue_token;
            let mut debits = Vec::new();

            add_debit(&mut debits, &token.creator, BASECOIN.clone(), token.txfee);

            Ok(TransactionBalanceView {
                hash: token.hash().await,
                signatures: vec![tx.signature.clone()],
                debits,
                check_saved_txid: true,
            })
        }
        Transaction::Nft(tx) => {
            let nft = &tx.unsigned_create_nft;
            let mut debits = Vec::new();

            add_debit(&mut debits, &nft.creator, BASECOIN.clone(), nft.txfee);

            Ok(TransactionBalanceView {
                hash: nft.hash().await,
                signatures: vec![tx.signature.clone()],
                debits,
                check_saved_txid: true,
            })
        }
        Transaction::Marketing(tx) => {
            let marketing = &tx.unsigned_marketing;
            let mut debits = Vec::new();

            add_debit(
                &mut debits,
                &marketing.advertiser,
                BASECOIN.clone(),
                marketing.txfee,
            );

            Ok(TransactionBalanceView {
                hash: marketing.hash().await,
                signatures: vec![tx.signature.clone()],
                debits,
                check_saved_txid: true,
            })
        }
        Transaction::Swap(tx) => {
            let swap = &tx.unsigned_swap;
            let mut debits = Vec::new();
            let asset1 = nft_asset_name(&swap.ticker1, swap.nft_series1);
            let asset2 = nft_asset_name(&swap.ticker2, swap.nft_series2);
            let sender1_asset_total = swap
                .value1
                .checked_add(swap.tip1)
                .ok_or_else(|| "Swap sender1 debit overflowed.".to_string())?;
            let sender2_asset_total = swap
                .value2
                .checked_add(swap.tip2)
                .ok_or_else(|| "Swap sender2 debit overflowed.".to_string())?;

            add_debit(&mut debits, &swap.sender1, asset1, sender1_asset_total);
            add_debit(&mut debits, &swap.sender1, BASECOIN.clone(), swap.txfee1);
            add_debit(&mut debits, &swap.sender2, asset2, sender2_asset_total);
            add_debit(&mut debits, &swap.sender2, BASECOIN.clone(), swap.txfee2);

            Ok(TransactionBalanceView {
                hash: swap.hash().await,
                signatures: vec![tx.signature1.clone(), tx.signature2.clone()],
                debits,
                check_saved_txid: true,
            })
        }
        Transaction::Lender(tx) => {
            let loan = &tx.unsigned_loan_contract;
            let mut debits = Vec::new();
            let collateral = loan_collateral_asset_name_or_base(&loan.collateral, &BASECOIN)
                .ok_or_else(|| "Invalid loan collateral asset.".to_string())?;

            add_debit(
                &mut debits,
                &loan.lender,
                loan.loan_coin.clone(),
                loan.loan_amount,
            );
            add_debit(&mut debits, &loan.lender, BASECOIN.clone(), loan.txfee);
            add_debit(
                &mut debits,
                &loan.borrower,
                collateral,
                loan.collateral_amount,
            );

            Ok(TransactionBalanceView {
                hash: loan.hash().await,
                signatures: vec![tx.signature1.clone(), tx.signature2.clone()],
                debits,
                check_saved_txid: true,
            })
        }
        Transaction::Borrower(tx) => {
            let payment = &tx.unsigned_contract_payment;
            let contract = loan_contract_for_hash(db, &payment.contract_hash).await?;
            let mut debits = Vec::new();
            let payment_total = payment
                .payback_amount
                .checked_add(payment.tip)
                .ok_or_else(|| "Loan payment debit overflowed.".to_string())?;

            add_debit(
                &mut debits,
                &payment.address,
                contract.unsigned_loan_contract.loan_coin,
                payment_total,
            );
            add_debit(
                &mut debits,
                &payment.address,
                BASECOIN.clone(),
                payment.txfee,
            );

            Ok(TransactionBalanceView {
                hash: payment.hash().await,
                signatures: vec![tx.signature.clone()],
                debits,
                check_saved_txid: true,
            })
        }
        Transaction::Collateral(tx) => {
            let collateral = &tx.unsigned_collateral_claim;
            let mut debits = Vec::new();

            add_debit(
                &mut debits,
                &collateral.address,
                BASECOIN.clone(),
                collateral.txfee,
            );

            Ok(TransactionBalanceView {
                hash: collateral.hash().await,
                signatures: vec![tx.signature.clone()],
                debits,
                check_saved_txid: true,
            })
        }
        Transaction::Vanity(tx) => {
            let vanity = &tx.unsigned_vanity_address;
            let mut debits = Vec::new();

            add_debit(&mut debits, &vanity.address, BASECOIN.clone(), vanity.txfee);

            Ok(TransactionBalanceView {
                hash: vanity.hash().await,
                signatures: vec![tx.signature.clone()],
                debits,
                check_saved_txid: true,
            })
        }
        Transaction::StorageKey(tx) => {
            let storage_key = &tx.unsigned_storagekey;
            let mut debits = Vec::new();

            add_debit(
                &mut debits,
                &storage_key.address,
                BASECOIN.clone(),
                storage_key.txfee,
            );

            Ok(TransactionBalanceView {
                hash: storage_key.hash().await,
                signatures: vec![tx.signature.clone()],
                debits,
                check_saved_txid: true,
            })
        }
        Transaction::StorageBool(tx) => {
            let storage = &tx.unsigned_storagebool;
            let mut debits = Vec::new();
            add_debit(
                &mut debits,
                &storage.address,
                BASECOIN.clone(),
                storage.txfee,
            );
            Ok(TransactionBalanceView {
                hash: storage.hash().await,
                signatures: vec![tx.signature.clone()],
                debits,
                check_saved_txid: true,
            })
        }
        Transaction::StorageU8(tx) => {
            let storage = &tx.unsigned_storageu8;
            let mut debits = Vec::new();
            add_debit(
                &mut debits,
                &storage.address,
                BASECOIN.clone(),
                storage.txfee,
            );
            Ok(TransactionBalanceView {
                hash: storage.hash().await,
                signatures: vec![tx.signature.clone()],
                debits,
                check_saved_txid: true,
            })
        }
        Transaction::StorageU16(tx) => {
            let storage = &tx.unsigned_storageu16;
            let mut debits = Vec::new();
            add_debit(
                &mut debits,
                &storage.address,
                BASECOIN.clone(),
                storage.txfee,
            );
            Ok(TransactionBalanceView {
                hash: storage.hash().await,
                signatures: vec![tx.signature.clone()],
                debits,
                check_saved_txid: true,
            })
        }
        Transaction::StorageU32(tx) => {
            let storage = &tx.unsigned_storageu32;
            let mut debits = Vec::new();
            add_debit(
                &mut debits,
                &storage.address,
                BASECOIN.clone(),
                storage.txfee,
            );
            Ok(TransactionBalanceView {
                hash: storage.hash().await,
                signatures: vec![tx.signature.clone()],
                debits,
                check_saved_txid: true,
            })
        }
        Transaction::StorageU64(tx) => {
            let storage = &tx.unsigned_storageu64;
            let mut debits = Vec::new();
            add_debit(
                &mut debits,
                &storage.address,
                BASECOIN.clone(),
                storage.txfee,
            );
            Ok(TransactionBalanceView {
                hash: storage.hash().await,
                signatures: vec![tx.signature.clone()],
                debits,
                check_saved_txid: true,
            })
        }
        Transaction::StorageU128(tx) => {
            let storage = &tx.unsigned_storageu128;
            let mut debits = Vec::new();
            add_debit(
                &mut debits,
                &storage.address,
                BASECOIN.clone(),
                storage.txfee,
            );
            Ok(TransactionBalanceView {
                hash: storage.hash().await,
                signatures: vec![tx.signature.clone()],
                debits,
                check_saved_txid: true,
            })
        }
        Transaction::StorageString(tx) => {
            let storage = &tx.unsigned_storage;
            let mut debits = Vec::new();
            add_debit(
                &mut debits,
                &storage.address,
                BASECOIN.clone(),
                storage.txfee,
            );
            Ok(TransactionBalanceView {
                hash: storage.hash().await,
                signatures: vec![tx.signature.clone()],
                debits,
                check_saved_txid: true,
            })
        }
    }
}

async fn pending_reserved_amount(db: &Db, address: &str, coin: &str) -> Result<u64, String> {
    if coin.trim().eq_ignore_ascii_case(BASECOIN.trim()) {
        get_basecoin_balance(db, address)
            .await
            .map_err(|err| err.to_string())
    } else {
        get_coin_balance(db, address, coin)
            .await
            .map_err(|err| err.to_string())
    }
}

pub(crate) async fn prepare_transaction_balance_reservation(
    db: &Db,
    transaction: &Transaction,
) -> Result<TransactionBalanceReservation, String> {
    let view = transaction_balance_view(db, transaction).await?;
    if view.check_saved_txid && saved_txid_exists(db, &view.hash)? {
        return Err("This transaction already exists.".to_string());
    }

    let already_reserved_in_mempool = transaction_is_pending(&view).await?;
    let loan_payment = if let Transaction::Borrower(tx) = transaction {
        let payment = &tx.unsigned_contract_payment;
        let contract = loan_contract_for_hash(db, &payment.contract_hash).await?;
        let contract_total = contract
            .unsigned_loan_contract
            .payment_amount
            .checked_mul(contract.unsigned_loan_contract.payment_number as u64)
            .ok_or_else(|| "Loan repayment total overflowed.".to_string())?;
        Some(LoanPaymentReservation {
            contract_hash: payment.contract_hash.clone(),
            payment_amount: payment.payback_amount,
            confirmed_paid: get_total_payments(db, &payment.contract_hash).await,
            contract_total,
        })
    } else {
        None
    };
    let mut debits = HashMap::new();
    let mut available = HashMap::new();

    // Aggregate each transaction by address/asset first. Base-coin transfers,
    // for example, can debit the same balance for both value and fee.
    for debit in &view.debits {
        add_to_total(&mut debits, debit)?;
    }

    for ((address, coin), transaction_debit) in &debits {
        let confirmed_balance = get_balance_with_db(db, address, coin)
            .await
            .map_err(|err| err.to_string())?;
        let mempool_reserved = pending_reserved_amount(db, address, coin).await?;
        let included_reserved = if already_reserved_in_mempool {
            *transaction_debit
        } else {
            0
        };
        let external_mempool_reserved = mempool_reserved.saturating_sub(included_reserved);
        let spendable = confirmed_balance.saturating_sub(external_mempool_reserved);

        add_amount_to_total(&mut available, (address.clone(), coin.clone()), spendable)?;
    }

    Ok(TransactionBalanceReservation {
        hash: view.hash,
        signatures: view.signatures,
        debits,
        available,
        loan_payment,
    })
}

#[cfg(test)]
mod tests {
    use super::{BlockBalanceTracker, LoanPaymentReservation, TransactionBalanceReservation};
    use std::collections::HashMap;

    fn payment_reservation(
        hash: &str,
        contract_hash: &str,
        payment_amount: u64,
        confirmed_paid: u64,
        contract_total: u64,
    ) -> TransactionBalanceReservation {
        TransactionBalanceReservation {
            hash: hash.to_string(),
            signatures: vec![format!("signature-{hash}")],
            debits: HashMap::new(),
            available: HashMap::new(),
            loan_payment: Some(LoanPaymentReservation {
                contract_hash: contract_hash.to_string(),
                payment_amount,
                confirmed_paid,
                contract_total,
            }),
        }
    }

    #[test]
    fn cumulative_block_payments_cannot_exceed_contract_total() {
        let mut tracker = BlockBalanceTracker::new();

        tracker
            .reserve(payment_reservation("one", "contract", 50, 20, 100))
            .unwrap();
        tracker
            .reserve(payment_reservation("two", "contract", 30, 20, 100))
            .unwrap();

        let error = tracker
            .reserve(payment_reservation("three", "contract", 1, 20, 100))
            .unwrap_err();
        assert!(error.contains("exceed remaining contract balance"));
    }
}