pyra-margin 0.4.4

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

use pyra_tokens::AssetId;
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::drift::balance::calculate_value_usdc_base_units;
use crate::error::{MathError, MathResult};

const MARGIN_PRECISION: i128 = 10_000;
const PRICE_PRECISION: i128 = 1_000_000;

/// Aggregated Kamino margin state across all positions in an obligation.
#[derive(Debug, Clone, Copy)]
pub struct KaminoMarginState {
    /// Total deposit value in USDC base units (unweighted).
    pub total_collateral: i128,
    /// Total LTV-weighted deposit value in USDC base units.
    pub total_weighted_collateral: i128,
    /// Total borrow value in USDC base units (positive, unweighted).
    pub total_liabilities: i128,
    /// Total liquidation-threshold-weighted borrow value in USDC base units.
    pub total_weighted_liabilities: i128,
}

impl KaminoMarginState {
    /// Calculate margin state from a Kamino obligation and reserve data.
    ///
    /// For each deposit: computes market value and applies LTV weight.
    /// For each borrow: computes market value and applies inverse-threshold weight.
    ///
    /// Positions whose reserve is missing from the map are skipped.
    pub fn calculate(
        obligation: &KaminoObligation,
        reserves: &HashMap<Pubkey, KaminoReserve>,
    ) -> MathResult<Self> {
        // 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: i128 = 0;
        let mut total_weighted_collateral: i128 = 0;
        let mut total_liabilities: i128 = 0;
        let mut total_weighted_liabilities: i128 = 0;

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

            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)?;

            total_collateral = total_collateral
                .checked_add(value)
                .ok_or(MathError::Overflow)?;

            let weight = get_kamino_asset_weight(reserve) as i128;
            let weighted = value
                .checked_mul(weight)
                .ok_or(MathError::Overflow)?
                .checked_div(MARGIN_PRECISION)
                .ok_or(MathError::Overflow)?;
            total_weighted_collateral = total_weighted_collateral
                .checked_add(weighted)
                .ok_or(MathError::Overflow)?;
        }

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

            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 liability
            let liability = value.checked_neg().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 factor_adjusted = liability
                .checked_mul(borrow_factor)
                .ok_or(MathError::Overflow)?
                .checked_div(100)
                .ok_or(MathError::Overflow)?;

            total_liabilities = total_liabilities
                .checked_add(factor_adjusted)
                .ok_or(MathError::Overflow)?;

            let weight = get_kamino_liability_weight(reserve)? as i128;
            let weighted = factor_adjusted
                .checked_mul(weight)
                .ok_or(MathError::Overflow)?
                .checked_div(MARGIN_PRECISION)
                .ok_or(MathError::Overflow)?;
            total_weighted_liabilities = total_weighted_liabilities
                .checked_add(weighted)
                .ok_or(MathError::Overflow)?;
        }

        Ok(Self {
            total_collateral,
            total_weighted_collateral,
            total_liabilities,
            total_weighted_liabilities,
        })
    }

    /// Free collateral = weighted collateral - liability, clamped to 0.
    /// Uses `total_liabilities` (not weighted) because Klend only applies
    /// borrow_factor, and using weighted liabilities would double-weight borrows.
    pub fn free_collateral(&self) -> u64 {
        let fc = self
            .total_weighted_collateral
            .saturating_sub(self.total_liabilities)
            .max(0);
        clamp_to_u64(fc)
    }

    /// Credit usage in basis points (0 = no liabilities, 10_000 = 100%).
    /// Capped at 10_000 — an under-collateralized account is at 100% usage.
    /// Returns 0 if collateral is zero or negative.
    pub fn credit_usage_bps(&self) -> MathResult<u64> {
        if self.total_weighted_collateral <= 0 {
            return Ok(0);
        }
        let usage = self
            .total_weighted_liabilities
            .checked_mul(10_000)
            .ok_or(MathError::Overflow)?
            .checked_div(self.total_weighted_collateral)
            .ok_or(MathError::Overflow)?;
        Ok(cmp::min(clamp_to_u64(cmp::max(0, usage)), 10_000))
    }
}

/// Per-reserve withdraw and borrow limits.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct KaminoPositionLimits {
    /// Maximum amount that can be withdrawn (in token base units).
    pub withdraw_limit: u64,
    /// Maximum amount that can be borrowed (in token base units).
    pub borrow_limit: u64,
}

/// Calculate per-reserve withdraw and borrow limits.
///
/// - `withdraw_limit`: how much collateral can be removed while staying above LTV.
/// - `borrow_limit`: total max outflow for this reserve. Equals `withdraw_limit` (what can
///   be removed from the deposit) plus `max_additional` (new debt capacity from remaining
///   free collateral after removing this position). This gives Pyra the total token amount
///   a user could extract, whether by withdrawing collateral or borrowing against other
///   positions.
///
/// `asset_id` identifies the token (from `pyra_tokens`). The reserve address is
/// looked up via `asset_id.token()`.
pub fn calculate_kamino_position_limits(
    margin_state: &KaminoMarginState,
    reserve: &KaminoReserve,
    asset_id: AssetId,
    obligation: &KaminoObligation,
) -> MathResult<KaminoPositionLimits> {
    let reserve_address = asset_id
        .token()
        .kamino_reserve
        .ok_or(MathError::Overflow)?;
    let price = get_kamino_price(reserve)?;
    if price == 0 {
        return Ok(KaminoPositionLimits {
            withdraw_limit: 0,
            borrow_limit: 0,
        });
    }

    let free_collateral = i128::from(margin_state.free_collateral());
    let asset_weight = get_kamino_asset_weight(reserve) as i128;
    let decimals = u32::try_from(reserve.liquidity.mint_decimals).map_err(|_| MathError::Overflow)?;
    let (numerator_scale, denominator_scale) = decimal_scale(decimals)?;

    // Find deposit balance for this reserve
    let deposit_balance = obligation
        .deposits
        .iter()
        .find(|d| d.deposit_reserve == reserve_address)
        .map(|d| get_kamino_deposit_balance(d, reserve))
        .transpose()?
        .unwrap_or(0);
    let deposit_balance_u64 = clamp_to_u64(cmp::max(0, deposit_balance));

    // Max withdraw: how much can be withdrawn while staying above LTV.
    // No liabilities or zero asset weight -> can withdraw full deposit.
    let withdraw_limit = if asset_weight == 0 || margin_state.total_weighted_liabilities == 0 {
        deposit_balance_u64
    } else {
        // Floor division for safety: underestimates max withdraw (conservative)
        let limit = free_collateral
            .checked_mul(MARGIN_PRECISION)
            .and_then(|v| v.checked_div(asset_weight))
            .and_then(|v| v.checked_mul(PRICE_PRECISION))
            .and_then(|v| v.checked_div(price))
            .and_then(|v| v.checked_mul(numerator_scale as i128))
            .and_then(|v| v.checked_div(denominator_scale as i128))
            .ok_or(MathError::Overflow)?;

        cmp::min(deposit_balance_u64, clamp_to_u64(cmp::max(0, limit)))
    };

    // Max borrow: withdraw_limit + max additional liability from remaining free collateral.
    // Subtract this position's weighted deposit value from free collateral first.
    let free_collateral_after = if deposit_balance > 0 {
        let price_u64 = u64::try_from(price).map_err(|_| MathError::Overflow)?;
        let position_value =
            calculate_value_usdc_base_units(deposit_balance, price_u64, decimals)?;
        let weighted = position_value
            .checked_mul(asset_weight)
            .and_then(|v| v.checked_div(MARGIN_PRECISION))
            .ok_or(MathError::Overflow)?;
        cmp::max(0, free_collateral.saturating_sub(weighted))
    } else {
        free_collateral
    };

    let liability_weight = get_kamino_liability_weight(reserve)? as i128;
    let max_additional = free_collateral_after
        .checked_mul(MARGIN_PRECISION)
        .and_then(|v| v.checked_div(liability_weight))
        .and_then(|v| v.checked_mul(PRICE_PRECISION))
        .and_then(|v| v.checked_div(price))
        .and_then(|v| v.checked_mul(numerator_scale as i128))
        .and_then(|v| v.checked_div(denominator_scale as i128))
        .ok_or(MathError::Overflow)?;

    let borrow_limit = (withdraw_limit as i128)
        .checked_add(cmp::max(0, max_additional))
        .ok_or(MathError::Overflow)?;

    Ok(KaminoPositionLimits {
        withdraw_limit,
        borrow_limit: clamp_to_u64(cmp::max(0, borrow_limit)),
    })
}

/// Compute decimal scaling factors for converting between token precision and USDC (6 decimals).
fn decimal_scale(token_decimals: u32) -> MathResult<(u32, u32)> {
    if token_decimals > 6 {
        let numerator = 10u32
            .checked_pow(
                token_decimals
                    .checked_sub(6)
                    .ok_or(MathError::Overflow)?,
            )
            .ok_or(MathError::Overflow)?;
        Ok((numerator, 1))
    } else {
        let denominator = 10u32
            .checked_pow(6u32.checked_sub(token_decimals).ok_or(MathError::Overflow)?)
            .ok_or(MathError::Overflow)?;
        Ok((1, denominator))
    }
}

/// Clamp an i128 value to u64 range.
fn clamp_to_u64(value: i128) -> u64 {
    if value < 0 {
        0
    } else if value > u64::MAX as i128 {
        u64::MAX
    } else {
        value as u64
    }
}

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

    const FRACTION_ONE: u128 = 1 << 60;

    fn usdc_asset_id() -> AssetId {
        pyra_tokens::USDC.asset_id
    }

    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_reserve(
        ltv: u8,
        liq_threshold: u8,
        price_usd: u128,
        decimals: u64,
        total_available: u64,
        collateral_supply: u64,
    ) -> KaminoReserve {
        let price_sf = price_usd << 60;
        KaminoReserve {
            lending_market: Pubkey::default(),
            liquidity: KaminoReserveLiquidity {
                mint_pubkey: Pubkey::default(),
                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: decimals,
                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(),
        }
    }

    // --- KaminoMarginState tests ---

    #[test]
    fn empty_margin_state() {
        let state =
            KaminoMarginState::calculate(&empty_obligation(), &HashMap::new()).unwrap();
        assert_eq!(state.total_collateral, 0);
        assert_eq!(state.total_weighted_collateral, 0);
        assert_eq!(state.total_liabilities, 0);
        assert_eq!(state.total_weighted_liabilities, 0);
        assert_eq!(state.free_collateral(), 0);
        assert_eq!(state.credit_usage_bps().unwrap(), 0);
    }

    #[test]
    fn margin_state_single_deposit() {
        let reserve_key = Pubkey::new_unique();
        let reserve = make_reserve(100, 100, 1, 6, 10_000_000_000, 10_000_000_000);
        let reserves: HashMap<Pubkey, KaminoReserve> = [(reserve_key, reserve)].into();

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

        let state = KaminoMarginState::calculate(&obligation, &reserves).unwrap();
        assert_eq!(state.total_collateral, 10_000_000);
        assert_eq!(state.total_weighted_collateral, 10_000_000);
        assert_eq!(state.total_liabilities, 0);
        assert_eq!(state.free_collateral(), 10_000_000);
        assert_eq!(state.credit_usage_bps().unwrap(), 0);
    }

    #[test]
    fn margin_state_deposit_and_borrow() {
        let reserve_key = Pubkey::new_unique();
        let reserve = make_reserve(100, 100, 1, 6, 10_000_000_000, 10_000_000_000);
        let reserves: HashMap<Pubkey, KaminoReserve> = [(reserve_key, reserve)].into();

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

        let state = KaminoMarginState::calculate(&obligation, &reserves).unwrap();
        assert_eq!(state.total_collateral, 10_000_000);
        assert_eq!(state.total_liabilities, 5_000_000);
        assert_eq!(state.free_collateral(), 5_000_000);
        assert_eq!(state.credit_usage_bps().unwrap(), 5_000); // 50%
    }

    #[test]
    fn free_collateral_clamped_to_zero() {
        let state = KaminoMarginState {
            total_collateral: 5_000_000,
            total_weighted_collateral: 5_000_000,
            total_liabilities: 10_000_000,
            total_weighted_liabilities: 10_000_000,
        };
        assert_eq!(state.free_collateral(), 0);
        assert_eq!(state.credit_usage_bps().unwrap(), 10_000); // 100% capped
    }

    #[test]
    fn credit_usage_over_100_capped() {
        let state = KaminoMarginState {
            total_collateral: 5_000_000,
            total_weighted_collateral: 5_000_000,
            total_liabilities: 20_000_000,
            total_weighted_liabilities: 20_000_000,
        };
        assert_eq!(state.credit_usage_bps().unwrap(), 10_000);
    }

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

    // --- Position limits tests ---

    #[test]
    fn withdraw_limit_no_borrows() {
        let key = usdc_reserve_key();
        let reserve = make_reserve(100, 100, 1, 6, 10_000_000_000, 10_000_000_000);

        let state = KaminoMarginState {
            total_collateral: 10_000_000,
            total_weighted_collateral: 10_000_000,
            total_liabilities: 0,
            total_weighted_liabilities: 0,
        };

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

        let limits = calculate_kamino_position_limits(
            &state,
            &reserve,
            usdc_asset_id(),
            &obligation,
        )
        .unwrap();
        assert_eq!(limits.withdraw_limit, 10_000_000);
    }

    #[test]
    fn withdraw_limit_with_borrows() {
        let key = usdc_reserve_key();
        let reserve = make_reserve(100, 100, 1, 6, 10_000_000_000, 10_000_000_000);

        let state = KaminoMarginState {
            total_collateral: 10_000_000,
            total_weighted_collateral: 10_000_000,
            total_liabilities: 5_000_000,
            total_weighted_liabilities: 5_000_000,
        };

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

        let limits = calculate_kamino_position_limits(
            &state,
            &reserve,
            usdc_asset_id(),
            &obligation,
        )
        .unwrap();
        assert_eq!(limits.withdraw_limit, 5_000_000);
    }

    #[test]
    fn borrow_limit_basic() {
        let key = usdc_reserve_key();
        let reserve = make_reserve(100, 100, 1, 6, 10_000_000_000, 10_000_000_000);

        let state = KaminoMarginState {
            total_collateral: 10_000_000,
            total_weighted_collateral: 10_000_000,
            total_liabilities: 0,
            total_weighted_liabilities: 0,
        };

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

        let limits = calculate_kamino_position_limits(
            &state,
            &reserve,
            usdc_asset_id(),
            &obligation,
        )
        .unwrap();
        assert_eq!(limits.withdraw_limit, 10_000_000);
        assert_eq!(limits.borrow_limit, 10_000_000);
    }

    #[test]
    fn borrow_limit_with_other_collateral() {
        let usdc_reserve = make_reserve(100, 100, 1, 6, 10_000_000_000, 10_000_000_000);

        // Other collateral provides $80 weighted collateral
        let state = KaminoMarginState {
            total_collateral: 100_000_000,
            total_weighted_collateral: 80_000_000,
            total_liabilities: 0,
            total_weighted_liabilities: 0,
        };

        // No USDC deposit — user wants to borrow USDC against other collateral
        let obligation = empty_obligation();

        let limits = calculate_kamino_position_limits(
            &state,
            &usdc_reserve,
            usdc_asset_id(),
            &obligation,
        )
        .unwrap();
        assert_eq!(limits.withdraw_limit, 0);
        assert_eq!(limits.borrow_limit, 80_000_000);
    }

    #[test]
    fn borrow_factor_increases_weighted_liability() {
        let reserve_key = Pubkey::new_unique();
        let mut reserve = make_reserve(100, 100, 1, 6, 10_000_000_000, 10_000_000_000);
        reserve.config.borrow_factor_pct = 150;

        let mut obligation = empty_obligation();
        obligation.deposits.push(KaminoObligationCollateral {
            deposit_reserve: reserve_key,
            deposited_amount: 100_000_000,
            market_value_sf: 0,
        });
        obligation.borrows.push(KaminoObligationLiquidity {
            borrow_reserve: 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 reserves: HashMap<Pubkey, KaminoReserve> = [(reserve_key, reserve)].into();
        let state = KaminoMarginState::calculate(&obligation, &reserves).unwrap();

        assert_eq!(state.total_liabilities, 75_000_000);
    }

    #[test]
    fn zero_price_returns_zero_limits() {
        let mut reserve = make_reserve(80, 85, 0, 6, 10_000_000_000, 10_000_000_000);
        reserve.liquidity.market_price_sf = 0;

        let state = KaminoMarginState {
            total_collateral: 10_000_000,
            total_weighted_collateral: 8_000_000,
            total_liabilities: 0,
            total_weighted_liabilities: 0,
        };

        let obligation = empty_obligation();
        let limits = calculate_kamino_position_limits(
            &state,
            &reserve,
            usdc_asset_id(),
            &obligation,
        )
        .unwrap();
        assert_eq!(limits.withdraw_limit, 0);
        assert_eq!(limits.borrow_limit, 0);
    }

    #[test]
    fn invalid_asset_id_errors() {
        // Asset ID 9999 doesn't exist — AssetId::new rejects it
        assert!(AssetId::new(9999).is_err());
    }

    // --- decimal_scale tests ---

    #[test]
    fn decimal_scale_usdc() {
        assert_eq!(decimal_scale(6).unwrap(), (1, 1));
    }

    #[test]
    fn decimal_scale_sol() {
        assert_eq!(decimal_scale(9).unwrap(), (1_000, 1));
    }

    #[test]
    fn decimal_scale_small() {
        assert_eq!(decimal_scale(4).unwrap(), (1, 100));
    }
}