pyra-margin 0.4.2

use pyra_types::{
    KaminoBigFractionBytes, KaminoObligationCollateral, KaminoObligationLiquidity, KaminoReserve,
    KAMINO_FRACTION_SCALE,
};

use crate::error::{MathError, MathResult};

/// Extract the lower 128 bits of a BigFractionBytes as u128.
///
/// BigFractionBytes stores a 256-bit value as `[u64; 4]` in little-endian order.
/// For cumulative borrow rates (starting at ~2^60 and growing slowly), the
/// lower 128 bits contain all significant data.
fn bigfraction_lower_u128(bsf: &KaminoBigFractionBytes) -> u128 {
    // Safe: cumulative borrow rate starts at ~2^60 and grows slowly.
    // Would need ~2^68x compounding to overflow into upper 128 bits.
    debug_assert!(
        bsf.value[2] == 0 && bsf.value[3] == 0,
        "cumulative borrow rate overflow: upper 128 bits are non-zero"
    );
    let lo = bsf.value[0] as u128;
    let hi = bsf.value[1] as u128;
    hi.checked_shl(64).unwrap_or(0) | lo
}

/// Compute `(a * b) / c` without u128 overflow, using split multiplication.
///
/// Returns `Err(Overflow)` if the final result exceeds u128 or `c` is zero.
fn checked_mul_div(a: u128, b: u128, c: u128) -> MathResult<u128> {
    if c == 0 {
        return Err(MathError::Overflow);
    }
    // Fast path: direct multiplication fits
    if let Some(product) = a.checked_mul(b) {
        return product.checked_div(c).ok_or(MathError::Overflow);
    }
    // Slow path: split a = (a/c)*c + (a%c), so a*b/c = (a/c)*b + (a%c)*b/c
    let q = a.checked_div(c).ok_or(MathError::Overflow)?;
    let r = a.checked_rem(c).ok_or(MathError::Overflow)?;
    let term1 = q.checked_mul(b).ok_or(MathError::Overflow)?;
    let term2 = if let Some(rb) = r.checked_mul(b) {
        rb.checked_div(c).ok_or(MathError::Overflow)?
    } else {
        // r*b overflows u128. Split b: r*(b/c) + r*(b%c)/c
        let bq = b.checked_div(c).ok_or(MathError::Overflow)?;
        let br = b.checked_rem(c).ok_or(MathError::Overflow)?;
        let t1 = r.checked_mul(bq).ok_or(MathError::Overflow)?;
        // r < c and br < c, so r*br < c^2 which fits u128 when c < 2^64
        let t2 = r
            .checked_mul(br)
            .ok_or(MathError::Overflow)?
            .checked_div(c)
            .ok_or(MathError::Overflow)?;
        t1.checked_add(t2).ok_or(MathError::Overflow)?
    };
    term1.checked_add(term2).ok_or(MathError::Overflow)
}

/// Compute total liquidity supply in a Kamino reserve (base units).
///
/// Matches Klend's `ReserveLiquidity::total_supply()`:
/// `total_supply = available + floor(borrowed_sf / 2^60)
///                 - floor(protocol_fees_sf / 2^60)
///                 - floor(referrer_fees_sf / 2^60)
///                 - floor(pending_referrer_fees_sf / 2^60)`
///
/// NOTE: Each sf field is independently truncated via `/ 2^60` before summing.
/// Klend keeps everything in Fraction precision and truncates once at the end.
/// This can introduce up to 4 lamports of error, which is acceptable for
/// off-chain margin calculations.
fn total_liquidity_base_units(reserve: &KaminoReserve) -> MathResult<u128> {
    let liq = &reserve.liquidity;
    let borrowed_base = liq
        .borrowed_amount_sf
        .checked_div(KAMINO_FRACTION_SCALE)
        .ok_or(MathError::Overflow)?;
    let protocol_fees = liq
        .accumulated_protocol_fees_sf
        .checked_div(KAMINO_FRACTION_SCALE)
        .ok_or(MathError::Overflow)?;
    let referrer_fees = liq
        .accumulated_referrer_fees_sf
        .checked_div(KAMINO_FRACTION_SCALE)
        .ok_or(MathError::Overflow)?;
    let pending_fees = liq
        .pending_referrer_fees_sf
        .checked_div(KAMINO_FRACTION_SCALE)
        .ok_or(MathError::Overflow)?;

    (liq.total_available_amount as u128)
        .checked_add(borrowed_base)
        .ok_or(MathError::Overflow)?
        .checked_sub(protocol_fees)
        .ok_or(MathError::Overflow)?
        .checked_sub(referrer_fees)
        .ok_or(MathError::Overflow)?
        .checked_sub(pending_fees)
        .ok_or(MathError::Overflow)
}

/// Convert cToken deposit amount to underlying liquidity base units.
///
/// Uses the collateral exchange rate from the reserve:
/// `liquidity = deposited_amount * total_liquidity / collateral_supply`
///
/// Reference: `CollateralExchangeRate::collateral_to_liquidity()` in Klend
/// `state/reserve.rs`.
///
/// Returns a positive i128 (deposits are always non-negative).
pub fn get_kamino_deposit_balance(
    collateral: &KaminoObligationCollateral,
    reserve: &KaminoReserve,
) -> MathResult<i128> {
    let deposited = collateral.deposited_amount as u128;
    if deposited == 0 {
        return Ok(0);
    }

    let collateral_supply = reserve.collateral.mint_total_supply as u128;
    if collateral_supply == 0 {
        // Initial 1:1 exchange rate when no collateral has been minted
        return i128::try_from(deposited).map_err(|_| MathError::Overflow);
    }

    let total_liquidity = total_liquidity_base_units(reserve)?;
    if total_liquidity == 0 {
        return Ok(0);
    }

    let result = checked_mul_div(deposited, total_liquidity, collateral_supply)?;
    i128::try_from(result).map_err(|_| MathError::Overflow)
}

/// Convert `borrowed_amount_sf` to base units with interest accrual.
///
/// 1. Converts from Fraction scale: `base = borrowed_amount_sf >> 60`
/// 2. Applies interest accrual if reserve rate > obligation rate:
///    `accrued = base * reserve_rate / obligation_rate`
///
/// Reference: `ObligationLiquidity::accrue_interest()` and
/// `calculate_amount_with_accrued_interest()` in Klend `state/obligation.rs`.
///
/// Returns a **negative** i128 (matching Drift convention where borrows are negative).
pub fn get_kamino_borrow_balance(
    liquidity: &KaminoObligationLiquidity,
    reserve: &KaminoReserve,
) -> MathResult<i128> {
    let borrowed_sf = liquidity.borrowed_amount_sf;
    if borrowed_sf == 0 {
        return Ok(0);
    }

    // Convert from Fraction scale to base units.
    // NOTE: Klend does the multiplication in full sf precision before truncating
    // (multiply sf values first for 256-bit precision, then floor). We truncate
    // first then multiply, which can lose up to 1 lamport. Acceptable for
    // off-chain margin/display use.
    let base_amount = borrowed_sf
        .checked_div(KAMINO_FRACTION_SCALE)
        .ok_or(MathError::Overflow)?;

    // Apply interest accrual: amount * new_rate / old_rate
    let obligation_rate = bigfraction_lower_u128(&liquidity.cumulative_borrow_rate_bsf);
    let reserve_rate = bigfraction_lower_u128(&reserve.liquidity.cumulative_borrow_rate_bsf);

    let accrued = if obligation_rate == 0 || reserve_rate <= obligation_rate {
        base_amount
    } else {
        checked_mul_div(base_amount, reserve_rate, obligation_rate)?
    };

    // Negate: borrows are negative by convention
    let signed = i128::try_from(accrued).map_err(|_| MathError::Overflow)?;
    Ok(signed.saturating_neg())
}

#[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::{
        KaminoLastUpdate, KaminoReserveCollateral, KaminoReserveConfig, KaminoReserveLiquidity,
    };
    use solana_pubkey::Pubkey;

    const FRACTION_ONE: u128 = 1 << 60;

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

    fn make_reserve(
        total_available: u64,
        borrowed_sf: u128,
        collateral_supply: u64,
        mint_decimals: u64,
    ) -> KaminoReserve {
        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: borrowed_sf,
                cumulative_borrow_rate_bsf: rate_to_bsf(FRACTION_ONE),
                mint_decimals,
                market_price_sf: 0,
                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: 80,
                liquidation_threshold_pct: 85,
                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 make_collateral(deposited_amount: u64) -> KaminoObligationCollateral {
        KaminoObligationCollateral {
            deposit_reserve: Pubkey::default(),
            deposited_amount,
            market_value_sf: 0,
        }
    }

    fn make_liquidity(borrowed_sf: u128, obligation_rate: u128) -> KaminoObligationLiquidity {
        KaminoObligationLiquidity {
            borrow_reserve: Pubkey::default(),
            cumulative_borrow_rate_bsf: rate_to_bsf(obligation_rate),
            borrowed_amount_sf: borrowed_sf,
            market_value_sf: 0,
            borrow_factor_adjusted_market_value_sf: 0,
        }
    }

    // --- get_kamino_deposit_balance tests ---

    #[test]
    fn deposit_zero_returns_zero() {
        let reserve = make_reserve(1_000_000, 0, 1_000_000, 6);
        let collateral = make_collateral(0);
        assert_eq!(get_kamino_deposit_balance(&collateral, &reserve).unwrap(), 0);
    }

    #[test]
    fn deposit_one_to_one_exchange_rate() {
        // total_liquidity = 1_000_000, collateral_supply = 1_000_000 -> 1:1 rate
        let reserve = make_reserve(1_000_000, 0, 1_000_000, 6);
        let collateral = make_collateral(500_000);
        assert_eq!(
            get_kamino_deposit_balance(&collateral, &reserve).unwrap(),
            500_000
        );
    }

    #[test]
    fn deposit_exchange_rate_with_interest() {
        // total_liquidity = 1_100_000 (1M available + 100K borrowed)
        // collateral_supply = 1_000_000 -> rate = 1.1
        let borrowed_sf = 100_000u128 * FRACTION_ONE;
        let reserve = make_reserve(1_000_000, borrowed_sf, 1_000_000, 6);
        let collateral = make_collateral(1_000_000);
        assert_eq!(
            get_kamino_deposit_balance(&collateral, &reserve).unwrap(),
            1_100_000 // 1M cTokens * 1.1 = 1.1M liquidity
        );
    }

    #[test]
    fn deposit_zero_collateral_supply_uses_one_to_one() {
        let reserve = make_reserve(1_000_000, 0, 0, 6);
        let collateral = make_collateral(500_000);
        assert_eq!(
            get_kamino_deposit_balance(&collateral, &reserve).unwrap(),
            500_000
        );
    }

    #[test]
    fn deposit_exchange_rate_with_fees_subtracted() {
        // 1M available + 200K borrowed = 1.2M gross
        // Fees: 50K protocol + 30K referrer + 20K pending = 100K total
        // Net total_liquidity = 1.2M - 100K = 1.1M
        // collateral_supply = 1M -> rate = 1.1
        let borrowed_sf = 200_000u128 * FRACTION_ONE;
        let mut reserve = make_reserve(1_000_000, borrowed_sf, 1_000_000, 6);
        reserve.liquidity.accumulated_protocol_fees_sf = 50_000u128 * FRACTION_ONE;
        reserve.liquidity.accumulated_referrer_fees_sf = 30_000u128 * FRACTION_ONE;
        reserve.liquidity.pending_referrer_fees_sf = 20_000u128 * FRACTION_ONE;

        let collateral = make_collateral(1_000_000);
        assert_eq!(
            get_kamino_deposit_balance(&collateral, &reserve).unwrap(),
            1_100_000
        );
    }

    #[test]
    fn deposit_zero_total_liquidity_returns_zero() {
        let reserve = make_reserve(0, 0, 1_000_000, 6);
        let collateral = make_collateral(500_000);
        assert_eq!(get_kamino_deposit_balance(&collateral, &reserve).unwrap(), 0);
    }

    // --- get_kamino_borrow_balance tests ---

    #[test]
    fn borrow_zero_returns_zero() {
        let reserve = make_reserve(0, 0, 0, 6);
        let liquidity = make_liquidity(0, FRACTION_ONE);
        assert_eq!(get_kamino_borrow_balance(&liquidity, &reserve).unwrap(), 0);
    }

    #[test]
    fn borrow_returns_negative() {
        let borrowed_sf = 1_000_000u128 * FRACTION_ONE;
        let reserve = make_reserve(0, 0, 0, 6);
        let liquidity = make_liquidity(borrowed_sf, FRACTION_ONE);
        assert_eq!(
            get_kamino_borrow_balance(&liquidity, &reserve).unwrap(),
            -1_000_000
        );
    }

    #[test]
    fn borrow_with_interest_accrual() {
        // Borrowed 1M at rate 1.0, reserve now at 1.1 -> accrued ~ 1.1M
        // Use exact 2x rate to avoid fixed-point rounding
        let borrowed_sf = 1_000_000u128 * FRACTION_ONE;
        let obligation_rate = FRACTION_ONE;
        let reserve_rate = FRACTION_ONE * 2; // 2.0x (exact in fixed-point)

        let mut reserve = make_reserve(0, 0, 0, 6);
        reserve.liquidity.cumulative_borrow_rate_bsf = rate_to_bsf(reserve_rate);

        let liquidity = make_liquidity(borrowed_sf, obligation_rate);
        assert_eq!(
            get_kamino_borrow_balance(&liquidity, &reserve).unwrap(),
            -2_000_000
        );

        // Also verify 1.1x with rounding tolerance
        let rate_1_1 = FRACTION_ONE + FRACTION_ONE / 10;
        let mut reserve2 = make_reserve(0, 0, 0, 6);
        reserve2.liquidity.cumulative_borrow_rate_bsf = rate_to_bsf(rate_1_1);

        let liquidity2 = make_liquidity(borrowed_sf, obligation_rate);
        let balance = get_kamino_borrow_balance(&liquidity2, &reserve2).unwrap();
        // Floor division may lose 1 unit
        assert!(balance == -1_100_000 || balance == -1_099_999);
    }

    #[test]
    fn borrow_no_accrual_when_rates_equal() {
        let borrowed_sf = 500_000u128 * FRACTION_ONE;
        let rate = FRACTION_ONE;

        let mut reserve = make_reserve(0, 0, 0, 6);
        reserve.liquidity.cumulative_borrow_rate_bsf = rate_to_bsf(rate);

        let liquidity = make_liquidity(borrowed_sf, rate);
        assert_eq!(
            get_kamino_borrow_balance(&liquidity, &reserve).unwrap(),
            -500_000
        );
    }

    #[test]
    fn borrow_no_accrual_when_obligation_rate_zero() {
        let borrowed_sf = 500_000u128 * FRACTION_ONE;
        let reserve = make_reserve(0, 0, 0, 6);
        let liquidity = make_liquidity(borrowed_sf, 0);
        assert_eq!(
            get_kamino_borrow_balance(&liquidity, &reserve).unwrap(),
            -500_000
        );
    }

    // --- checked_mul_div tests ---

    #[test]
    fn mul_div_basic() {
        assert_eq!(checked_mul_div(100, 200, 50).unwrap(), 400);
    }

    #[test]
    fn mul_div_zero_divisor() {
        assert!(checked_mul_div(100, 200, 0).is_err());
    }

    #[test]
    fn mul_div_large_values() {
        // Test the overflow split path
        let a = u128::MAX / 2;
        let b = 3u128;
        let c = 2u128;
        let result = checked_mul_div(a, b, c).unwrap();
        let expected = (a / c) * b + (a % c) * b / c;
        assert_eq!(result, expected);
    }
}