pyra_margin/

capacity.rs

use std::cmp;
use std::collections::HashMap;

use pyra_types::{SpotBalanceType, SpotMarket, SpotPosition};

use crate::balance::{calculate_value_usdc_base_units, get_token_balance};
use crate::error::{MathError, MathResult};
use crate::spend_limits::get_remaining_timeframe_limit;
use crate::weights::{calculate_asset_weight, calculate_liability_weight, get_strict_price};

/// USDC has 6 decimals: 1 USDC = 1_000_000 base units = 100 cents.
/// So 1 cent = 10_000 base units.
const USDC_BASE_UNITS_PER_CENT: u64 = 10_000;

const MARGIN_PRECISION: i128 = 10_000;

/// Convert USDC base units to cents.
///
/// 1 USDC = 1_000_000 base units = 100 cents, so divide by 10_000.
pub fn usdc_base_units_to_cents(base_units: u64) -> MathResult<u64> {
    base_units
        .checked_div(USDC_BASE_UNITS_PER_CENT)
        .ok_or(MathError::Overflow)
}

/// Calculate spend limit in cents from vault state.
///
/// Returns `min(transaction_limit, timeframe_remaining, max_cap)` in cents.
///
/// `now_unix` is the current Unix timestamp in seconds (e.g. `Utc::now().timestamp() as u64`).
pub fn calculate_spend_limit_cents(
    vault: &pyra_types::Vault,
    max_transaction_limit_cents: u64,
    now_unix: u64,
) -> MathResult<u64> {
    let timeframe_base_units = get_remaining_timeframe_limit(vault, now_unix);
    let transaction_limit_cents = usdc_base_units_to_cents(vault.spend_limit_per_transaction)?;
    let timeframe_remaining_cents = usdc_base_units_to_cents(timeframe_base_units)?;

    let spend_limit_no_cap = cmp::min(transaction_limit_cents, timeframe_remaining_cents);
    Ok(cmp::min(spend_limit_no_cap, max_transaction_limit_cents))
}

/// Per-position info emitted during capacity calculation.
///
/// Used downstream by liquidating spend jobs to know position sizes and weights.
#[derive(Debug, Clone, PartialEq)]
pub struct PositionInfo {
    pub market_index: u16,
    /// Unsigned token balance in base units.
    pub balance: u64,
    pub position_type: SpotBalanceType,
    pub price_usdc_base_units: u64,
    /// Initial asset weight from the spot market (basis points).
    pub weight_bps: u32,
}

/// Result of spending capacity calculation.
#[derive(Debug, Clone)]
pub struct CapacityResult {
    /// Max spendable via liquidating spend (unweighted collateral - slippage - liabilities), in cents.
    /// Excludes unliquidatable assets from collateral.
    pub total_spendable_cents: u64,
    /// Available credit line (weighted collateral - weighted liabilities), in cents.
    pub available_credit_cents: u64,
    /// USDC balance (market index 0) in cents.
    pub usdc_balance_cents: u64,
    /// Total weighted collateral in USDC base units.
    pub weighted_collateral_usdc_base_units: u64,
    /// Total weighted liabilities in USDC base units.
    pub weighted_liabilities_usdc_base_units: u64,
    /// Per-position breakdown for downstream use.
    pub position_infos: Vec<PositionInfo>,
}

/// Calculate spending capacity from Drift spot positions.
///
/// This is the core calculation used for card transaction authorization.
/// It computes:
/// - **total_spendable**: max amount for liquidating spends (collateral minus slippage minus liabilities,
///   excluding unliquidatable assets)
/// - **available_credit**: credit line from weighted margin (weighted collateral minus weighted liabilities)
/// - **usdc_balance**: direct USDC holdings
///
/// Positions in `unliquidatable_market_indices` are excluded from unweighted collateral
/// (affecting `total_spendable`) but still included in weighted calculations
/// (affecting `available_credit`).
///
/// Positions whose market index is missing from `spot_market_map` or `price_map` are skipped.
pub fn calculate_capacity(
    spot_positions: &[SpotPosition],
    spot_market_map: &HashMap<u16, SpotMarket>,
    price_map: &HashMap<u16, u64>,
    unliquidatable_market_indices: &[u16],
    max_slippage_bps: u64,
) -> MathResult<CapacityResult> {
    let mut total_collateral_usdc_base_units: u64 = 0;
    let mut total_liabilities_usdc_base_units: u64 = 0;

    let mut total_weighted_collateral_usdc_base_units: u64 = 0;
    let mut total_weighted_liabilities_usdc_base_units: u64 = 0;

    let mut usdc_balance_base_units: u64 = 0;

    let mut position_infos: Vec<PositionInfo> = Vec::new();

    for position in spot_positions {
        let market_index = position.market_index;

        let Some(spot_market) = spot_market_map.get(&market_index) else {
            continue;
        };
        let Some(price_usdc_base_units) = price_map.get(&market_index).copied() else {
            continue;
        };

        // Step 1: Calculate token balance and USDC value
        let token_balance_base_units = get_token_balance(position, spot_market)?;

        let is_asset = token_balance_base_units >= 0;
        let twap5min = spot_market
            .historical_oracle_data
            .last_oracle_price_twap5min;
        let strict_price = get_strict_price(price_usdc_base_units, twap5min, is_asset);

        let value_usdc_base_units = calculate_value_usdc_base_units(
            token_balance_base_units,
            strict_price,
            spot_market.decimals,
        )?;

        // Accumulate unweighted totals (excluding unliquidatable collateral)
        let is_unliquidatable_collateral =
            unliquidatable_market_indices.contains(&market_index) && value_usdc_base_units > 0;
        if !is_unliquidatable_collateral {
            update_running_totals(
                &mut total_collateral_usdc_base_units,
                &mut total_liabilities_usdc_base_units,
                value_usdc_base_units,
            )?;
        }

        // Step 2: Apply IMF-adjusted weights
        let token_amount_unsigned = token_balance_base_units.unsigned_abs();
        let weight_bps = if is_asset {
            calculate_asset_weight(token_amount_unsigned, price_usdc_base_units, spot_market)?
                as i128
        } else {
            calculate_liability_weight(token_amount_unsigned, spot_market)? as i128
        };
        let weighted_value_usdc_base_units = value_usdc_base_units
            .checked_mul(weight_bps)
            .ok_or(MathError::Overflow)?
            .checked_div(MARGIN_PRECISION)
            .ok_or(MathError::Overflow)?;

        update_running_totals(
            &mut total_weighted_collateral_usdc_base_units,
            &mut total_weighted_liabilities_usdc_base_units,
            weighted_value_usdc_base_units,
        )?;

        // Step 3: Track USDC balance (market index 0)
        if market_index == 0 && usdc_balance_base_units == 0 && token_balance_base_units > 0 {
            usdc_balance_base_units =
                u64::try_from(token_balance_base_units).map_err(|_| MathError::Overflow)?;
        }

        // Step 4: Store position info
        let token_balance_unsigned =
            u64::try_from(token_balance_base_units.unsigned_abs()).map_err(|_| MathError::Overflow)?;
        position_infos.push(PositionInfo {
            market_index,
            balance: token_balance_unsigned,
            position_type: position.balance_type.clone(),
            price_usdc_base_units,
            weight_bps: spot_market.initial_asset_weight,
        });
    }

    // Step 5: Available credit = weighted collateral - weighted liabilities
    let available_credit_base_units = total_weighted_collateral_usdc_base_units
        .saturating_sub(total_weighted_liabilities_usdc_base_units);
    let available_credit_cents = usdc_base_units_to_cents(available_credit_base_units)?;

    // Step 6: Total spendable = collateral - slippage - liabilities (for liquidating spends)
    let max_slippage_usdc_base_units = total_collateral_usdc_base_units
        .checked_mul(max_slippage_bps)
        .ok_or(MathError::Overflow)?
        .checked_div(10_000)
        .ok_or(MathError::Overflow)?;
    let total_spendable_base_units = total_collateral_usdc_base_units
        .saturating_sub(max_slippage_usdc_base_units)
        .saturating_sub(total_liabilities_usdc_base_units);
    let total_spendable_cents = usdc_base_units_to_cents(total_spendable_base_units)?;

    let usdc_balance_cents = usdc_base_units_to_cents(usdc_balance_base_units)?;

    Ok(CapacityResult {
        total_spendable_cents,
        available_credit_cents,
        usdc_balance_cents,
        weighted_collateral_usdc_base_units: total_weighted_collateral_usdc_base_units,
        weighted_liabilities_usdc_base_units: total_weighted_liabilities_usdc_base_units,
        position_infos,
    })
}

/// Accumulate a signed value into positive/negative running totals.
fn update_running_totals(
    total_positive: &mut u64,
    total_negative: &mut u64,
    value: i128,
) -> MathResult<()> {
    let value_unsigned = u64::try_from(value.unsigned_abs()).map_err(|_| MathError::Overflow)?;

    if value >= 0 {
        *total_positive = total_positive
            .checked_add(value_unsigned)
            .ok_or(MathError::Overflow)?;
    } else {
        *total_negative = total_negative
            .checked_add(value_unsigned)
            .ok_or(MathError::Overflow)?;
    }

    Ok(())
}

#[cfg(test)]
#[allow(
    clippy::unwrap_used,
    clippy::expect_used,
    clippy::panic,
    clippy::arithmetic_side_effects
)]
mod tests {
    use super::*;
    use pyra_types::{HistoricalOracleData, InsuranceFund, Vault};

    fn make_spot_market_with_twap(
        market_index: u16,
        decimals: u32,
        initial_asset_weight: u32,
        initial_liability_weight: u32,
        twap5min: i64,
    ) -> SpotMarket {
        let precision_decrease = 10u128.pow(19u32.saturating_sub(decimals));
        SpotMarket {
            pubkey: vec![],
            market_index,
            initial_asset_weight,
            initial_liability_weight,
            imf_factor: 0,
            scale_initial_asset_weight_start: 0,
            decimals,
            cumulative_deposit_interest: precision_decrease,
            cumulative_borrow_interest: precision_decrease,
            deposit_balance: 0,
            borrow_balance: 0,
            optimal_utilization: 0,
            optimal_borrow_rate: 0,
            max_borrow_rate: 0,
            min_borrow_rate: 0,
            insurance_fund: InsuranceFund::default(),
            historical_oracle_data: HistoricalOracleData {
                last_oracle_price_twap5min: twap5min,
            },
            oracle: None,
        }
    }

    /// Convenience: creates a spot market where TWAP matches the oracle price.
    fn make_spot_market(
        market_index: u16,
        decimals: u32,
        initial_asset_weight: u32,
        initial_liability_weight: u32,
        oracle_price: u64,
    ) -> SpotMarket {
        make_spot_market_with_twap(
            market_index,
            decimals,
            initial_asset_weight,
            initial_liability_weight,
            oracle_price as i64,
        )
    }

    fn make_position(market_index: u16, scaled_balance: u64, is_deposit: bool) -> SpotPosition {
        SpotPosition {
            market_index,
            scaled_balance,
            balance_type: if is_deposit {
                SpotBalanceType::Deposit
            } else {
                SpotBalanceType::Borrow
            },
            ..Default::default()
        }
    }

    fn make_vault(
        spend_limit_per_transaction: u64,
        spend_limit_per_timeframe: u64,
        remaining_spend_limit_per_timeframe: u64,
        next_timeframe_reset_timestamp: u64,
        timeframe_in_seconds: u64,
    ) -> Vault {
        Vault {
            owner: vec![0; 32],
            bump: 0,
            spend_limit_per_transaction,
            spend_limit_per_timeframe,
            remaining_spend_limit_per_timeframe,
            next_timeframe_reset_timestamp,
            timeframe_in_seconds,
        }
    }

    // --- usdc_base_units_to_cents ---

    #[test]
    fn cents_basic() {
        assert_eq!(usdc_base_units_to_cents(1_000_000).unwrap(), 100); // 1 USDC = 100 cents
    }

    #[test]
    fn cents_zero() {
        assert_eq!(usdc_base_units_to_cents(0).unwrap(), 0);
    }

    #[test]
    fn cents_sub_cent_truncates() {
        assert_eq!(usdc_base_units_to_cents(9_999).unwrap(), 0); // < 1 cent
        assert_eq!(usdc_base_units_to_cents(10_000).unwrap(), 1); // exactly 1 cent
    }

    // --- calculate_spend_limit_cents ---

    const NOW: u64 = 1_700_000_000;

    #[test]
    fn spend_limit_basic() {
        // 10 USDC per tx, 50 USDC per timeframe, 30 USDC remaining, max cap 1000 cents
        let vault = make_vault(
            10_000_000,  // 10 USDC per tx
            50_000_000,  // 50 USDC per timeframe
            30_000_000,  // 30 USDC remaining
            NOW + 3600,  // active timeframe
            86_400,
        );
        let limit = calculate_spend_limit_cents(&vault, 10_000, NOW).unwrap();
        // min(1000 cents tx, 3000 cents remaining, 10000 cents cap) = 1000
        assert_eq!(limit, 1000);
    }

    #[test]
    fn spend_limit_expired_timeframe_uses_full() {
        let vault = make_vault(
            100_000_000, // 100 USDC per tx
            50_000_000,  // 50 USDC per timeframe
            10_000_000,  // 10 USDC remaining (ignored because expired)
            NOW - 100,   // expired
            86_400,
        );
        let limit = calculate_spend_limit_cents(&vault, 100_000, NOW).unwrap();
        // min(10000 cents tx, 5000 cents full timeframe, 100000 cap) = 5000
        assert_eq!(limit, 5000);
    }

    #[test]
    fn spend_limit_capped_by_max() {
        let vault = make_vault(
            1_000_000_000, // 1000 USDC per tx
            1_000_000_000, // 1000 USDC per timeframe
            1_000_000_000, // 1000 USDC remaining
            NOW + 3600,
            86_400,
        );
        let limit = calculate_spend_limit_cents(&vault, 500, NOW).unwrap();
        assert_eq!(limit, 500);
    }

    #[test]
    fn spend_limit_zero_timeframe() {
        let vault = make_vault(10_000_000, 50_000_000, 30_000_000, NOW + 3600, 0);
        let limit = calculate_spend_limit_cents(&vault, 10_000, NOW).unwrap();
        // timeframe_in_seconds == 0 -> remaining = 0 -> limit = 0
        assert_eq!(limit, 0);
    }

    // --- calculate_capacity ---

    #[test]
    fn empty_positions() {
        let result = calculate_capacity(&[], &HashMap::new(), &HashMap::new(), &[], 0).unwrap();
        assert_eq!(result.total_spendable_cents, 0);
        assert_eq!(result.available_credit_cents, 0);
        assert_eq!(result.usdc_balance_cents, 0);
        assert_eq!(result.weighted_collateral_usdc_base_units, 0);
        assert_eq!(result.weighted_liabilities_usdc_base_units, 0);
        assert!(result.position_infos.is_empty());
    }

    #[test]
    fn single_usdc_deposit() {
        let usdc = make_spot_market(0, 6, 10_000, 10_000, 1_000_000);
        let positions = vec![make_position(0, 100_000_000, true)]; // 100 USDC

        let mut markets = HashMap::new();
        markets.insert(0, usdc);
        let mut prices = HashMap::new();
        prices.insert(0, 1_000_000u64);

        let result = calculate_capacity(&positions, &markets, &prices, &[], 0).unwrap();

        assert_eq!(result.usdc_balance_cents, 10_000); // $100
        assert_eq!(result.total_spendable_cents, 10_000);
        assert_eq!(result.available_credit_cents, 10_000);
        assert_eq!(result.weighted_collateral_usdc_base_units, 100_000_000);
        assert_eq!(result.weighted_liabilities_usdc_base_units, 0);
        assert_eq!(result.position_infos.len(), 1);
        assert_eq!(result.position_infos[0].market_index, 0);
    }

    #[test]
    fn deposit_and_borrow() {
        let usdc = make_spot_market(0, 6, 10_000, 10_000, 1_000_000);
        let positions = vec![
            make_position(0, 100_000_000, true), // 100 USDC deposit
            make_position(0, 50_000_000, false),  // 50 USDC borrow
        ];

        let mut markets = HashMap::new();
        markets.insert(0, usdc);
        let mut prices = HashMap::new();
        prices.insert(0, 1_000_000u64);

        let result = calculate_capacity(&positions, &markets, &prices, &[], 0).unwrap();

        assert_eq!(result.usdc_balance_cents, 10_000); // 100 USDC deposit
        assert_eq!(result.total_spendable_cents, 5_000); // 100 - 50 = 50 USDC
        assert_eq!(result.available_credit_cents, 5_000);
    }

    #[test]
    fn unliquidatable_excluded_from_spendable() {
        // Market 4 is unliquidatable
        let usdc = make_spot_market(0, 6, 10_000, 10_000, 1_000_000);
        let weth = make_spot_market(4, 9, 8_000, 12_000, 100_000_000);

        let positions = vec![
            make_position(0, 10_000_000, true),    // 10 USDC
            make_position(4, 1_000_000_000, true),  // 1 wETH at $100
        ];

        let mut markets = HashMap::new();
        markets.insert(0, usdc);
        markets.insert(4, weth);
        let mut prices = HashMap::new();
        prices.insert(0, 1_000_000u64);
        prices.insert(4, 100_000_000u64); // $100

        let unliquidatable = vec![4u16];
        let result =
            calculate_capacity(&positions, &markets, &prices, &unliquidatable, 0).unwrap();

        // total_spendable: only USDC (10M base = 1000 cents), wETH excluded
        assert_eq!(result.total_spendable_cents, 1_000);
        // available_credit: includes wETH weighted (100M * 80% = 80M) + USDC (10M * 100% = 10M) = 90M = 9000 cents
        assert_eq!(result.available_credit_cents, 9_000);
        assert_eq!(result.position_infos.len(), 2);
    }

    #[test]
    fn slippage_reduces_spendable() {
        let usdc = make_spot_market(0, 6, 10_000, 10_000, 1_000_000);
        let positions = vec![make_position(0, 100_000_000, true)]; // 100 USDC

        let mut markets = HashMap::new();
        markets.insert(0, usdc);
        let mut prices = HashMap::new();
        prices.insert(0, 1_000_000u64);

        // 10% slippage = 1000 bps
        let result = calculate_capacity(&positions, &markets, &prices, &[], 1_000).unwrap();

        // 100 USDC collateral, 10% slippage = 10 USDC, spendable = 90 USDC = 9000 cents
        assert_eq!(result.total_spendable_cents, 9_000);
        // available_credit not affected by slippage
        assert_eq!(result.available_credit_cents, 10_000);
    }

    #[test]
    fn missing_market_skipped() {
        let positions = vec![make_position(5, 1_000_000, true)];

        let result =
            calculate_capacity(&positions, &HashMap::new(), &HashMap::new(), &[], 0).unwrap();

        assert_eq!(result.total_spendable_cents, 0);
        assert!(result.position_infos.is_empty());
    }

    #[test]
    fn missing_price_skipped() {
        let usdc = make_spot_market(0, 6, 10_000, 10_000, 1_000_000);
        let positions = vec![make_position(0, 1_000_000, true)];

        let mut markets = HashMap::new();
        markets.insert(0, usdc);

        let result =
            calculate_capacity(&positions, &markets, &HashMap::new(), &[], 0).unwrap();

        assert_eq!(result.total_spendable_cents, 0);
        assert!(result.position_infos.is_empty());
    }

    #[test]
    fn multi_position_with_unliquidatable_and_slippage() {
        let usdc = make_spot_market(0, 6, 10_000, 10_000, 1_000_000);
        let m4 = make_spot_market(4, 9, 8_000, 12_000, 200_000_000); // unliquidatable
        let m5 = make_spot_market(5, 9, 8_000, 12_000, 100_000_000);
        let m6 = make_spot_market(6, 6, 10_000, 10_000, 1_000_000);

        let positions = vec![
            make_position(0, 50_000_000, true),     // 50 USDC
            make_position(4, 1_000_000_000, true),   // 1 token @ $200 (unliquidatable)
            make_position(5, 500_000_000, true),     // 0.5 token @ $100
            make_position(6, 20_000_000, false),     // 20 USDC-like borrow
        ];

        let mut markets = HashMap::new();
        markets.insert(0, usdc);
        markets.insert(4, m4);
        markets.insert(5, m5);
        markets.insert(6, m6);
        let mut prices = HashMap::new();
        prices.insert(0, 1_000_000u64);
        prices.insert(4, 200_000_000u64);
        prices.insert(5, 100_000_000u64);
        prices.insert(6, 1_000_000u64);

        let unliquidatable = vec![4u16, 32u16];
        let result =
            calculate_capacity(&positions, &markets, &prices, &unliquidatable, 500).unwrap();

        // Unweighted collateral (excluding market 4): 50M (USDC) + 50M (m5: 0.5 * $100) = 100M
        // Unweighted liabilities: 20M (m6 borrow)
        // Slippage: 100M * 500/10000 = 5M
        // total_spendable = 100M - 5M - 20M = 75M base = 7500 cents
        assert_eq!(result.total_spendable_cents, 7_500);

        // Weighted collateral: 50M*100% + 200M*80% + 50M*80% = 50M + 160M + 40M = 250M
        // Weighted liabilities: 20M*100% = 20M
        // available_credit = 250M - 20M = 230M base = 23000 cents
        assert_eq!(result.available_credit_cents, 23_000);
        assert_eq!(result.usdc_balance_cents, 5_000);
        assert_eq!(result.position_infos.len(), 4);
    }

    // --- update_running_totals ---

    #[test]
    fn running_totals_positive() {
        let mut pos = 0u64;
        let mut neg = 0u64;
        update_running_totals(&mut pos, &mut neg, 100).unwrap();
        assert_eq!(pos, 100);
        assert_eq!(neg, 0);
    }

    #[test]
    fn running_totals_negative() {
        let mut pos = 0u64;
        let mut neg = 0u64;
        update_running_totals(&mut pos, &mut neg, -50).unwrap();
        assert_eq!(pos, 0);
        assert_eq!(neg, 50);
    }

    #[test]
    fn running_totals_accumulate() {
        let mut pos = 10u64;
        let mut neg = 5u64;
        update_running_totals(&mut pos, &mut neg, 20).unwrap();
        update_running_totals(&mut pos, &mut neg, -15).unwrap();
        assert_eq!(pos, 30);
        assert_eq!(neg, 20);
    }
}

#[cfg(test)]
#[allow(
    clippy::unwrap_used,
    clippy::expect_used,
    clippy::panic,
    clippy::arithmetic_side_effects
)]
mod proptests {
    use super::*;
    use proptest::prelude::*;

    proptest! {
        #[test]
        fn usdc_cents_never_exceeds_base_units(base_units in 0u64..=u64::MAX) {
            let cents = usdc_base_units_to_cents(base_units).unwrap();
            prop_assert!(cents <= base_units, "cents {} > base_units {}", cents, base_units);
        }

        #[test]
        fn spendable_le_collateral_minus_liabilities(
            collateral_base in 0u64..=1_000_000_000_000u64,
            liabilities_base in 0u64..=500_000_000_000u64,
        ) {
            // Spendable should never exceed collateral - liabilities (without slippage)
            let collateral_cents = usdc_base_units_to_cents(collateral_base).unwrap();
            let liabilities_cents = usdc_base_units_to_cents(liabilities_base).unwrap();
            let max_possible = collateral_cents.saturating_sub(liabilities_cents);

            // Since we're testing the formula directly, we can verify the invariant
            let spendable_base = collateral_base.saturating_sub(liabilities_base);
            let spendable_cents = usdc_base_units_to_cents(spendable_base).unwrap();
            prop_assert!(spendable_cents <= max_possible + 1, "rounding violation");
        }
    }
}