pyra-margin 0.4.1

use std::collections::HashMap;

use pyra_tokens::{AssetId, Token};
use pyra_types::{KaminoObligation, KaminoReserve};
use solana_pubkey::Pubkey;

use super::balance::{get_kamino_borrow_balance, get_kamino_deposit_balance};
use super::weights::{get_kamino_asset_weight, get_kamino_liability_weight, get_kamino_price};
use crate::common::usdc_base_units_to_cents;
use crate::drift::balance::calculate_value_usdc_base_units;
use crate::error::{MathError, MathResult};

const MARGIN_PRECISION: i128 = 10_000;

/// Position type for Kamino capacity breakdown.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum KaminoPositionType {
    Deposit,
    Borrow,
}

/// Per-position info emitted during capacity calculation.
///
/// Matches Drift's `PositionInfo` structure for downstream consistency.
#[derive(Debug, Clone, PartialEq)]
pub struct KaminoPositionInfo {
    /// Pyra asset ID (from `pyra_tokens`). Use `asset_id.token()` for
    /// reserve address, mint, decimals, etc.
    pub asset_id: AssetId,
    /// Unsigned token balance in base units.
    pub balance: u64,
    pub position_type: KaminoPositionType,
    /// Oracle price in USDC base units (PRICE_PRECISION = 1e6).
    pub price_usdc_base_units: u64,
    /// Asset weight in basis points (10_000 = 100%).
    pub weight_bps: u32,
}

/// Result of Kamino capacity calculation.
///
/// Matches Drift's `CapacityResult` interface.
#[derive(Debug, Clone)]
pub struct KaminoCapacityResult {
    /// Max spendable via liquidating spend (unweighted collateral - slippage - liabilities), in cents.
    pub total_spendable_cents: u64,
    /// Available credit line (LTV-weighted collateral - borrow-factor-adjusted liabilities), in cents.
    pub available_credit_cents: u64,
    /// USDC balance in cents (identified by `usdc_mint`).
    pub usdc_balance_cents: u64,
    /// Total weighted collateral in USDC base units (LTV-weighted).
    pub weighted_collateral_usdc_base_units: u64,
    /// Total weighted liabilities in USDC base units (borrow-factor-adjusted).
    pub weighted_liabilities_usdc_base_units: u64,
    /// Per-position breakdown for downstream use.
    pub position_infos: Vec<KaminoPositionInfo>,
}

/// Calculate Kamino spending capacity from an obligation and reserve data.
///
/// Matches Drift's `calculate_capacity` interface:
/// - **total_spendable_cents**: max amount for liquidating spends (collateral - slippage - liabilities)
/// - **available_credit_cents**: credit line (LTV-weighted collateral - borrow-factor-adjusted liabilities)
/// - **usdc_balance_cents**: direct USDC holdings
///
/// Reserves are looked up by pubkey in the `reserves` map. Positions whose
/// reserve is missing from the map are skipped. Positions whose reserve has
/// no matching `pyra_tokens` entry are also skipped.
pub fn calculate_kamino_capacity(
    obligation: &KaminoObligation,
    reserves: &HashMap<Pubkey, KaminoReserve>,
    max_slippage_bps: u64,
) -> MathResult<KaminoCapacityResult> {
    // Elevation group support not implemented. When elevation_group > 0,
    // LTV/liquidation_threshold come from the elevation group config (not
    // reserve.config) and borrow_factor becomes 1.0. Fail loudly rather
    // than silently computing wrong values.
    if obligation.elevation_group != 0 {
        return Err(MathError::Overflow);
    }

    let mut total_collateral_usdc_base_units: u64 = 0;
    let mut total_liabilities_usdc_base_units: u64 = 0;
    let mut total_weighted_collateral: u64 = 0;
    let mut total_weighted_liabilities: u64 = 0;
    let mut usdc_balance_base_units: u64 = 0;
    let mut position_infos: Vec<KaminoPositionInfo> = Vec::new();

    for deposit in &obligation.deposits {
        if deposit.deposited_amount == 0 {
            continue;
        }
        let Some(reserve) = reserves.get(&deposit.deposit_reserve) else {
            continue;
        };
        let Some(token) = Token::find_by_kamino_reserve(&deposit.deposit_reserve) else {
            continue;
        };
        let asset_id = token.asset_id;

        let balance = get_kamino_deposit_balance(deposit, reserve)?;
        let price = get_kamino_price(reserve)?;
        let price_u64 = u64::try_from(price).map_err(|_| MathError::Overflow)?;
        let decimals = u32::try_from(reserve.liquidity.mint_decimals).map_err(|_| MathError::Overflow)?;
        let value = calculate_value_usdc_base_units(balance, price_u64, decimals)?;
        let value_u64 = u64::try_from(value).map_err(|_| MathError::Overflow)?;

        // Accumulate unweighted collateral
        total_collateral_usdc_base_units = total_collateral_usdc_base_units
            .checked_add(value_u64)
            .ok_or(MathError::Overflow)?;

        // LTV-weighted value
        let asset_weight = get_kamino_asset_weight(reserve) as i128;
        let weighted = value
            .checked_mul(asset_weight)
            .ok_or(MathError::Overflow)?
            .checked_div(MARGIN_PRECISION)
            .ok_or(MathError::Overflow)?;
        let weighted_u64 = u64::try_from(weighted).map_err(|_| MathError::Overflow)?;
        total_weighted_collateral = total_weighted_collateral
            .checked_add(weighted_u64)
            .ok_or(MathError::Overflow)?;

        // Track USDC balance
        if token.mint == pyra_tokens::USDC.mint && usdc_balance_base_units == 0 {
            let balance_u64 = u64::try_from(balance).map_err(|_| MathError::Overflow)?;
            usdc_balance_base_units = balance_u64;
        }

        // Store position info
        let balance_u64 = u64::try_from(balance.unsigned_abs()).map_err(|_| MathError::Overflow)?;
        let weight_bps = u32::try_from(get_kamino_asset_weight(reserve)).map_err(|_| MathError::Overflow)?;
        position_infos.push(KaminoPositionInfo {
            asset_id,
            balance: balance_u64,
            position_type: KaminoPositionType::Deposit,
            price_usdc_base_units: price_u64,
            weight_bps,
        });
    }

    for borrow in &obligation.borrows {
        if borrow.borrowed_amount_sf == 0 {
            continue;
        }
        let Some(reserve) = reserves.get(&borrow.borrow_reserve) else {
            continue;
        };
        let Some(token) = Token::find_by_kamino_reserve(&borrow.borrow_reserve) else {
            continue;
        };
        let asset_id = token.asset_id;

        let balance = get_kamino_borrow_balance(borrow, reserve)?;
        let price = get_kamino_price(reserve)?;
        let price_u64 = u64::try_from(price).map_err(|_| MathError::Overflow)?;
        let decimals = u32::try_from(reserve.liquidity.mint_decimals).map_err(|_| MathError::Overflow)?;
        let value = calculate_value_usdc_base_units(balance, price_u64, decimals)?;

        // value is negative (borrow convention), negate to get positive
        let raw_borrow = value.checked_neg().ok_or(MathError::Overflow)?;
        let raw_borrow_u64 = u64::try_from(raw_borrow).map_err(|_| MathError::Overflow)?;
        total_liabilities_usdc_base_units = total_liabilities_usdc_base_units
            .checked_add(raw_borrow_u64)
            .ok_or(MathError::Overflow)?;

        // Apply borrow factor: max(1.0, borrow_factor_pct/100)
        let borrow_factor = i128::from(std::cmp::max(100u64, reserve.config.borrow_factor_pct));
        let adjusted_borrow = raw_borrow
            .checked_mul(borrow_factor)
            .ok_or(MathError::Overflow)?
            .checked_div(100)
            .ok_or(MathError::Overflow)?;
        let adjusted_u64 = u64::try_from(adjusted_borrow).map_err(|_| MathError::Overflow)?;
        total_weighted_liabilities = total_weighted_liabilities
            .checked_add(adjusted_u64)
            .ok_or(MathError::Overflow)?;

        // Store position info
        let balance_u64 = u64::try_from(balance.unsigned_abs()).map_err(|_| MathError::Overflow)?;
        let weight_bps = u32::try_from(get_kamino_liability_weight(reserve)?).map_err(|_| MathError::Overflow)?;
        position_infos.push(KaminoPositionInfo {
            asset_id,
            balance: balance_u64,
            position_type: KaminoPositionType::Borrow,
            price_usdc_base_units: price_u64,
            weight_bps,
        });
    }

    // Available credit = weighted collateral - weighted liabilities
    let available_credit_base_units = total_weighted_collateral
        .saturating_sub(total_weighted_liabilities);
    let available_credit_cents = usdc_base_units_to_cents(available_credit_base_units)?;

    // 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(KaminoCapacityResult {
        total_spendable_cents,
        available_credit_cents,
        usdc_balance_cents,
        weighted_collateral_usdc_base_units: total_weighted_collateral,
        weighted_liabilities_usdc_base_units: total_weighted_liabilities,
        position_infos,
    })
}

#[cfg(test)]
#[allow(
    clippy::unwrap_used,
    clippy::expect_used,
    clippy::panic,
    clippy::arithmetic_side_effects
)]
mod tests {
    use super::*;
    use pyra_types::{
        KaminoBigFractionBytes, KaminoLastUpdate, KaminoObligationCollateral,
        KaminoObligationLiquidity, KaminoReserveCollateral, KaminoReserveConfig,
        KaminoReserveLiquidity,
    };

    const FRACTION_ONE: u128 = 1 << 60;

    /// Real USDC kamino reserve address from pyra_tokens.
    fn usdc_reserve_key() -> Pubkey {
        pyra_tokens::USDC.kamino_reserve.unwrap()
    }

    fn rate_to_bsf(rate: u128) -> KaminoBigFractionBytes {
        KaminoBigFractionBytes {
            value: [rate as u64, (rate >> 64) as u64, 0, 0],
        }
    }

    fn make_usdc_reserve(
        ltv: u8,
        liq_threshold: u8,
        total_available: u64,
        collateral_supply: u64,
    ) -> KaminoReserve {
        let price_sf = 1u128 << 60; // $1
        KaminoReserve {
            lending_market: Pubkey::default(),
            liquidity: KaminoReserveLiquidity {
                mint_pubkey: pyra_tokens::USDC.mint,
                supply_vault: Pubkey::default(),
                fee_vault: Pubkey::default(),
                total_available_amount: total_available,
                borrowed_amount_sf: 0,
                cumulative_borrow_rate_bsf: rate_to_bsf(FRACTION_ONE),
                mint_decimals: 6,
                market_price_sf: price_sf,
                accumulated_protocol_fees_sf: 0,
                accumulated_referrer_fees_sf: 0,
                pending_referrer_fees_sf: 0,
                token_program: Pubkey::default(),
            },
            collateral: KaminoReserveCollateral {
                mint_pubkey: Pubkey::default(),
                supply_vault: Pubkey::default(),
                mint_total_supply: collateral_supply,
            },
            config: KaminoReserveConfig {
                loan_to_value_pct: ltv,
                liquidation_threshold_pct: liq_threshold,
                protocol_take_rate_pct: 0,
                protocol_liquidation_fee_pct: 0,
                borrow_factor_pct: 100,
                deposit_limit: 0,
                borrow_limit: 0,
                fees: Default::default(),
                borrow_rate_curve: Default::default(),
                deposit_withdrawal_cap: Default::default(),
                debt_withdrawal_cap: Default::default(),
                elevation_groups: [0; 20],
            },
            last_update: KaminoLastUpdate::default(),
        }
    }

    fn empty_obligation() -> KaminoObligation {
        KaminoObligation {
            lending_market: Pubkey::default(),
            owner: Pubkey::default(),
            deposits: vec![],
            borrows: vec![],
            deposited_value_sf: 0,
            borrowed_assets_market_value_sf: 0,
            allowed_borrow_value_sf: 0,
            unhealthy_borrow_value_sf: 0,
            has_debt: 0,
            elevation_group: 0,
            last_update: KaminoLastUpdate::default(),
        }
    }

    #[test]
    fn empty_obligation_zero_capacity() {
        let result = calculate_kamino_capacity(&empty_obligation(), &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 key = usdc_reserve_key();
        let reserve = make_usdc_reserve(100, 100, 10_000_000_000, 10_000_000_000);
        let reserves: HashMap<Pubkey, KaminoReserve> = [(key, reserve)].into();

        let mut obligation = empty_obligation();
        obligation.deposits.push(KaminoObligationCollateral {
            deposit_reserve: key,
            deposited_amount: 100_000_000, // 100 USDC
            market_value_sf: 0,
        });

        let result = calculate_kamino_capacity(&obligation, &reserves, 0).unwrap();
        assert_eq!(result.total_spendable_cents, 10_000);
        assert_eq!(result.available_credit_cents, 10_000);
        assert_eq!(result.usdc_balance_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].asset_id, pyra_tokens::USDC.asset_id);
        assert_eq!(result.position_infos[0].position_type, KaminoPositionType::Deposit);
    }

    #[test]
    fn deposit_and_borrow() {
        let key = usdc_reserve_key();
        let reserve = make_usdc_reserve(100, 100, 10_000_000_000, 10_000_000_000);
        let reserves: HashMap<Pubkey, KaminoReserve> = [(key, reserve)].into();

        let mut obligation = empty_obligation();
        obligation.deposits.push(KaminoObligationCollateral {
            deposit_reserve: key,
            deposited_amount: 100_000_000,
            market_value_sf: 0,
        });
        obligation.borrows.push(KaminoObligationLiquidity {
            borrow_reserve: key,
            cumulative_borrow_rate_bsf: rate_to_bsf(FRACTION_ONE),
            borrowed_amount_sf: 50_000_000u128 * FRACTION_ONE,
            market_value_sf: 0,
            borrow_factor_adjusted_market_value_sf: 0,
        });

        let result = calculate_kamino_capacity(&obligation, &reserves, 0).unwrap();
        assert_eq!(result.total_spendable_cents, 5_000);
        assert_eq!(result.available_credit_cents, 5_000);
        assert_eq!(result.usdc_balance_cents, 10_000);
        assert_eq!(result.position_infos.len(), 2);
    }

    #[test]
    fn ltv_weight_applied() {
        let key = usdc_reserve_key();
        // 80% LTV — weighted collateral should be 80% of unweighted
        let reserve = make_usdc_reserve(80, 85, 10_000_000_000, 10_000_000_000);
        let reserves: HashMap<Pubkey, KaminoReserve> = [(key, reserve)].into();

        let mut obligation = empty_obligation();
        obligation.deposits.push(KaminoObligationCollateral {
            deposit_reserve: key,
            deposited_amount: 100_000_000, // 100 USDC
            market_value_sf: 0,
        });

        let result = calculate_kamino_capacity(&obligation, &reserves, 0).unwrap();
        // Unweighted = $100 = 10_000 cents
        assert_eq!(result.total_spendable_cents, 10_000);
        // LTV-weighted = $80 = 8_000 cents
        assert_eq!(result.available_credit_cents, 8_000);
        assert_eq!(result.weighted_collateral_usdc_base_units, 80_000_000);
        assert_eq!(result.position_infos[0].weight_bps, 8_000);
    }

    #[test]
    fn at_limit_spendable_zero() {
        let key = usdc_reserve_key();
        let reserve = make_usdc_reserve(100, 100, 10_000_000_000, 10_000_000_000);
        let reserves: HashMap<Pubkey, KaminoReserve> = [(key, reserve)].into();

        let mut obligation = empty_obligation();
        obligation.deposits.push(KaminoObligationCollateral {
            deposit_reserve: key,
            deposited_amount: 100_000_000,
            market_value_sf: 0,
        });
        obligation.borrows.push(KaminoObligationLiquidity {
            borrow_reserve: key,
            cumulative_borrow_rate_bsf: rate_to_bsf(FRACTION_ONE),
            borrowed_amount_sf: 100_000_000u128 * FRACTION_ONE,
            market_value_sf: 0,
            borrow_factor_adjusted_market_value_sf: 0,
        });

        let result = calculate_kamino_capacity(&obligation, &reserves, 0).unwrap();
        assert_eq!(result.total_spendable_cents, 0);
        assert_eq!(result.available_credit_cents, 0);
    }

    #[test]
    fn borrow_factor_increases_effective_borrow() {
        let key = usdc_reserve_key();
        let mut reserve = make_usdc_reserve(100, 100, 10_000_000_000, 10_000_000_000);
        reserve.config.borrow_factor_pct = 150;
        let reserves: HashMap<Pubkey, KaminoReserve> = [(key, reserve)].into();

        let mut obligation = empty_obligation();
        obligation.deposits.push(KaminoObligationCollateral {
            deposit_reserve: key,
            deposited_amount: 100_000_000,
            market_value_sf: 0,
        });
        obligation.borrows.push(KaminoObligationLiquidity {
            borrow_reserve: key,
            cumulative_borrow_rate_bsf: rate_to_bsf(FRACTION_ONE),
            borrowed_amount_sf: 50_000_000u128 * FRACTION_ONE,
            market_value_sf: 0,
            borrow_factor_adjusted_market_value_sf: 0,
        });

        let result = calculate_kamino_capacity(&obligation, &reserves, 0).unwrap();
        assert_eq!(result.total_spendable_cents, 5_000);
        assert_eq!(result.available_credit_cents, 2_500);
        assert_eq!(result.weighted_liabilities_usdc_base_units, 75_000_000);
    }

    #[test]
    fn slippage_reduces_spendable() {
        let key = usdc_reserve_key();
        let reserve = make_usdc_reserve(100, 100, 10_000_000_000, 10_000_000_000);
        let reserves: HashMap<Pubkey, KaminoReserve> = [(key, reserve)].into();

        let mut obligation = empty_obligation();
        obligation.deposits.push(KaminoObligationCollateral {
            deposit_reserve: key,
            deposited_amount: 100_000_000,
            market_value_sf: 0,
        });

        let result = calculate_kamino_capacity(&obligation, &reserves, 1_000).unwrap();
        assert_eq!(result.total_spendable_cents, 9_000);
        assert_eq!(result.available_credit_cents, 10_000);
    }

    #[test]
    fn elevation_group_nonzero_errors() {
        let mut obligation = empty_obligation();
        obligation.elevation_group = 1;
        assert!(calculate_kamino_capacity(&obligation, &HashMap::new(), 0).is_err());
    }

    #[test]
    fn missing_reserve_skipped() {
        let unknown_key = Pubkey::new_unique();
        let mut obligation = empty_obligation();
        obligation.deposits.push(KaminoObligationCollateral {
            deposit_reserve: unknown_key,
            deposited_amount: 100_000_000,
            market_value_sf: 0,
        });

        let result = calculate_kamino_capacity(&obligation, &HashMap::new(), 0).unwrap();
        assert_eq!(result.total_spendable_cents, 0);
        assert!(result.position_infos.is_empty());
    }

    #[test]
    fn unknown_token_reserve_skipped() {
        // Reserve exists in HashMap but has no matching pyra_tokens entry
        let unknown_key = Pubkey::new_unique();
        let reserve = make_usdc_reserve(100, 100, 10_000_000_000, 10_000_000_000);
        let reserves: HashMap<Pubkey, KaminoReserve> = [(unknown_key, reserve)].into();

        let mut obligation = empty_obligation();
        obligation.deposits.push(KaminoObligationCollateral {
            deposit_reserve: unknown_key,
            deposited_amount: 100_000_000,
            market_value_sf: 0,
        });

        let result = calculate_kamino_capacity(&obligation, &reserves, 0).unwrap();
        assert_eq!(result.total_spendable_cents, 0);
        assert!(result.position_infos.is_empty());
    }

    #[test]
    fn position_infos_correct_types() {
        let key = usdc_reserve_key();
        let reserve = make_usdc_reserve(100, 100, 10_000_000_000, 10_000_000_000);
        let reserves: HashMap<Pubkey, KaminoReserve> = [(key, reserve)].into();

        let mut obligation = empty_obligation();
        obligation.deposits.push(KaminoObligationCollateral {
            deposit_reserve: key,
            deposited_amount: 100_000_000,
            market_value_sf: 0,
        });
        obligation.borrows.push(KaminoObligationLiquidity {
            borrow_reserve: key,
            cumulative_borrow_rate_bsf: rate_to_bsf(FRACTION_ONE),
            borrowed_amount_sf: 30_000_000u128 * FRACTION_ONE,
            market_value_sf: 0,
            borrow_factor_adjusted_market_value_sf: 0,
        });

        let result = calculate_kamino_capacity(&obligation, &reserves, 0).unwrap();
        assert_eq!(result.position_infos.len(), 2);
        assert_eq!(result.position_infos[0].asset_id, pyra_tokens::USDC.asset_id);
        assert_eq!(result.position_infos[0].position_type, KaminoPositionType::Deposit);
        assert_eq!(result.position_infos[0].balance, 100_000_000);
        assert_eq!(result.position_infos[0].weight_bps, 10_000);
        assert_eq!(result.position_infos[1].asset_id, pyra_tokens::USDC.asset_id);
        assert_eq!(result.position_infos[1].position_type, KaminoPositionType::Borrow);
        assert_eq!(result.position_infos[1].balance, 30_000_000);
        assert_eq!(result.position_infos[1].weight_bps, 10_000);
    }
}