//! Interest rate calculations for Drift spot markets.
//!
//! Provides pure-math functions for computing utilization, borrow/deposit rates,
//! and LTV from `SpotMarket` data. No I/O — all inputs are passed by reference.
//!
//! Ports the Drift SDK `calculateInterestRate` piecewise linear model.

use pyra_types::SpotMarket;

use crate::{RateError, RateResult, RATE_PRECISION};

// ─── Constants ───────────────────────────────────────────────────

/// Piecewise segments above optimal utilization.
/// Each entry: (utilization_breakpoint, weight out of 1000).
const BORROW_RATE_SEGMENTS: [(u128, u128); 6] = [
    (850_000, 50),
    (900_000, 100),
    (950_000, 150),
    (990_000, 200),
    (995_000, 250),
    (RATE_PRECISION, 250),
];

const WEIGHTS_DIVISOR: u128 = 1000;

// ─── Token amounts ──────────────────────────────────────────────

/// Compute the precision decrease factor for a market's decimals.
///
/// Returns `10^(19 - decimals)`, which converts scaled balances to token amounts.
fn precision_decrease(decimals: u32) -> RateResult<u128> {
    let exponent = 19u32
        .checked_sub(decimals)
        .ok_or(RateError::InvalidDecimals(decimals))?;
    10u128.checked_pow(exponent).ok_or(RateError::Overflow)
}

/// Calculate total token deposits from the market's deposit_balance and
/// cumulative deposit interest.
pub fn token_deposit_amount(market: &SpotMarket) -> RateResult<u128> {
    let pd = precision_decrease(market.decimals)?;
    market
        .deposit_balance
        .checked_mul(market.cumulative_deposit_interest)
        .ok_or(RateError::Overflow)?
        .checked_div(pd)
        .ok_or(RateError::Overflow)
}

/// Calculate total token borrows from the market's borrow_balance and
/// cumulative borrow interest.
pub fn token_borrow_amount(market: &SpotMarket) -> RateResult<u128> {
    let pd = precision_decrease(market.decimals)?;
    market
        .borrow_balance
        .checked_mul(market.cumulative_borrow_interest)
        .ok_or(RateError::Overflow)?
        .checked_div(pd)
        .ok_or(RateError::Overflow)
}

// ─── Utilization ─────────────────────────────────────────────────

/// Calculate utilization ratio scaled by [`RATE_PRECISION`].
///
/// Returns 0 when both deposits and borrows are zero, and [`RATE_PRECISION`]
/// (100%) when borrows exceed deposits.
pub fn calculate_utilization(market: &SpotMarket) -> RateResult<u128> {
    let deposits = token_deposit_amount(market)?;
    let borrows = token_borrow_amount(market)?;

    if borrows == 0 && deposits == 0 {
        return Ok(0);
    }
    if deposits == 0 {
        return Ok(RATE_PRECISION);
    }

    borrows
        .checked_mul(RATE_PRECISION)
        .ok_or(RateError::Overflow)?
        .checked_div(deposits)
        .ok_or(RateError::Overflow)
}

// ─── Borrow rate ─────────────────────────────────────────────────

/// Calculate annualized borrow rate scaled by [`RATE_PRECISION`].
///
/// Uses a piecewise linear model:
/// - Below optimal utilization: linear ramp from 0 to `optimal_borrow_rate`.
/// - Above optimal: 6 weighted segments distribute `(max_rate - optimal_rate)`.
///
/// The result is floored at `min_borrow_rate` (stored as half-percentage points).
pub fn calculate_borrow_rate(market: &SpotMarket) -> RateResult<u128> {
    let utilization = calculate_utilization(market)?;
    let optimal_util = market.optimal_utilization as u128;
    let optimal_rate = market.optimal_borrow_rate as u128;
    let max_rate = market.max_borrow_rate as u128;

    // min_borrow_rate is stored as a u8 representing half-percentage points.
    // Convert to rate precision: min_borrow_rate * RATE_PRECISION / 200
    // e.g. min_borrow_rate=1 -> 0.5% -> 5000 in rate precision.
    let min_rate = (market.min_borrow_rate as u128)
        .checked_mul(RATE_PRECISION)
        .ok_or(RateError::Overflow)?
        .checked_div(200)
        .ok_or(RateError::Overflow)?;

    let rate = if optimal_util == 0 {
        0u128
    } else if utilization <= optimal_util {
        // Below optimal: linear ramp from 0 to optimal_rate.
        utilization
            .checked_mul(optimal_rate)
            .ok_or(RateError::Overflow)?
            .checked_div(optimal_util)
            .ok_or(RateError::Overflow)?
    } else {
        // Above optimal: piecewise segments distribute (max_rate - optimal_rate)
        let total_extra_rate = max_rate
            .checked_sub(optimal_rate)
            .ok_or(RateError::Overflow)?;

        let mut rate = optimal_rate;
        let mut prev_util = optimal_util;

        for &(breakpoint, weight) in &BORROW_RATE_SEGMENTS {
            let segment_end = if breakpoint > RATE_PRECISION {
                RATE_PRECISION
            } else {
                breakpoint
            };

            // Skip segments that are entirely before optimal_utilization
            if segment_end <= prev_util {
                continue;
            }

            let segment_range = segment_end
                .checked_sub(prev_util)
                .ok_or(RateError::Overflow)?;
            let segment_rate_total = total_extra_rate
                .checked_mul(weight)
                .ok_or(RateError::Overflow)?
                .checked_div(WEIGHTS_DIVISOR)
                .ok_or(RateError::Overflow)?;

            if utilization <= segment_end {
                let partial_util = utilization
                    .checked_sub(prev_util)
                    .ok_or(RateError::Overflow)?;
                let partial_rate = segment_rate_total
                    .checked_mul(partial_util)
                    .ok_or(RateError::Overflow)?
                    .checked_div(segment_range)
                    .ok_or(RateError::Overflow)?;
                rate = rate.checked_add(partial_rate).ok_or(RateError::Overflow)?;
                break;
            } else {
                rate = rate
                    .checked_add(segment_rate_total)
                    .ok_or(RateError::Overflow)?;
                prev_util = segment_end;
            }
        }

        rate
    };

    // Floor at min_borrow_rate
    Ok(if rate > min_rate { rate } else { min_rate })
}

// ─── Deposit rate ────────────────────────────────────────────────

/// Calculate annualized deposit rate scaled by [`RATE_PRECISION`].
///
/// Formula: `borrow_rate * (1 - insurance_fund.total_factor) * utilization`
/// divided by `RATE_PRECISION^2`.
pub fn calculate_deposit_rate(market: &SpotMarket) -> RateResult<u128> {
    let utilization = calculate_utilization(market)?;
    let borrow_rate = calculate_borrow_rate(market)?;
    let insurance_factor = market.insurance_fund.total_factor as u128;

    let factor = RATE_PRECISION
        .checked_sub(insurance_factor)
        .ok_or(RateError::Overflow)?;

    borrow_rate
        .checked_mul(factor)
        .ok_or(RateError::Overflow)?
        .checked_mul(utilization)
        .ok_or(RateError::Overflow)?
        .checked_div(RATE_PRECISION)
        .ok_or(RateError::Overflow)?
        .checked_div(RATE_PRECISION)
        .ok_or(RateError::Overflow)
}

// ─── LTV ─────────────────────────────────────────────────────────

/// Return the loan-to-value ratio from the market's `initial_asset_weight`.
///
/// The raw weight is in Drift weight precision (10,000 = 100%).
/// Returns `f64` for direct compatibility with api-v2 consumers.
pub fn calculate_ltv(market: &SpotMarket) -> f64 {
    market.initial_asset_weight as f64
}

// ─── Convenience helpers ─────────────────────────────────────────

/// Compute all Drift rates for a market, returned as `(deposit_rate, borrow_rate, ltv)` decimals.
///
/// Deposit and borrow rates are fractions (e.g. `0.08` = 8%).
/// LTV is returned as the raw `initial_asset_weight` cast to `f64`.
pub fn compute_drift_rates(market: &SpotMarket) -> RateResult<(f64, f64, f64)> {
    let deposit_rate = crate::rate_to_decimal(calculate_deposit_rate(market)?);
    let borrow_rate = crate::rate_to_decimal(calculate_borrow_rate(market)?);
    let ltv = calculate_ltv(market);
    Ok((deposit_rate, borrow_rate, ltv))
}

// ─── Tests ───────────────────────────────────────────────────────

#[cfg(test)]
mod tests {
    use super::*;
    use pyra_types::{HistoricalOracleData, InsuranceFund};

    fn make_market(
        deposit_balance: u128,
        borrow_balance: u128,
        optimal_utilization: u32,
        optimal_borrow_rate: u32,
        max_borrow_rate: u32,
    ) -> SpotMarket {
        SpotMarket {
            pubkey: vec![],
            market_index: 0,
            initial_asset_weight: 8000,
            initial_liability_weight: 12_000,
            decimals: 6,
            cumulative_deposit_interest: 10_000_000_000, // 10^10
            cumulative_borrow_interest: 10_000_000_000,  // 10^10
            deposit_balance,
            borrow_balance,
            optimal_utilization,
            optimal_borrow_rate,
            max_borrow_rate,
            min_borrow_rate: 0,
            imf_factor: 0,
            scale_initial_asset_weight_start: 0,
            insurance_fund: InsuranceFund { total_factor: 0 },
            historical_oracle_data: HistoricalOracleData::default(),
            oracle: None,
        }
    }

    // ─── Token amounts ───────────────────────────────────────────

    #[test]
    fn token_deposit_amount_basic() {
        let market = make_market(2_000_000, 0, 700_000, 80_000, 1_000_000);
        // deposit_balance * cumulative_interest / precision_decrease
        // = 2_000_000 * 10^10 / 10^(19-6) = 2_000_000 * 10^10 / 10^13 = 2000
        let deposits = token_deposit_amount(&market);
        assert!(deposits.is_ok());
        assert_eq!(deposits.unwrap_or(0), 2000);
    }

    #[test]
    fn token_borrow_amount_basic() {
        let market = make_market(0, 1_000_000, 700_000, 80_000, 1_000_000);
        // 1_000_000 * 10^10 / 10^13 = 1000
        let borrows = token_borrow_amount(&market);
        assert!(borrows.is_ok());
        assert_eq!(borrows.unwrap_or(0), 1000);
    }

    #[test]
    fn decimals_exceeding_19_returns_error() {
        let mut market = make_market(1_000_000, 500_000, 700_000, 80_000, 1_000_000);
        market.decimals = 20;
        assert!(calculate_utilization(&market).is_err());
    }

    // ─── Utilization ─────────────────────────────────────────────

    #[test]
    fn utilization_zero_both() {
        let market = make_market(0, 0, 700_000, 80_000, 1_000_000);
        assert_eq!(calculate_utilization(&market).unwrap_or(1), 0);
    }

    #[test]
    fn utilization_zero_deposits() {
        // Use a large enough borrow_balance so that integer division yields > 0.
        // With decimals=6, precision_decrease=10^13, cumulative=10^10:
        // token_borrow = 10_000 * 10^10 / 10^13 = 10. deposits=0 -> RATE_PRECISION.
        let market = make_market(0, 10_000, 700_000, 80_000, 1_000_000);
        assert_eq!(calculate_utilization(&market).unwrap_or(0), RATE_PRECISION);
    }

    #[test]
    fn utilization_fifty_percent() {
        let market = make_market(2_000_000, 1_000_000, 700_000, 80_000, 1_000_000);
        assert_eq!(calculate_utilization(&market).unwrap_or(0), 500_000);
    }

    #[test]
    fn utilization_full() {
        let market = make_market(1_000_000, 1_000_000, 700_000, 80_000, 1_000_000);
        assert_eq!(calculate_utilization(&market).unwrap_or(0), RATE_PRECISION);
    }

    // ─── Borrow rate ─────────────────────────────────────────────

    #[test]
    fn borrow_rate_below_optimal() {
        let market = make_market(2_000_000, 1_000_000, 700_000, 80_000, 1_000_000);
        let rate = calculate_borrow_rate(&market).unwrap_or(0) as i128;
        // 50% utilization, optimal at 70%. Rate = 500000 * 80000 / 700000 = ~57142
        assert!(rate == 57142 || rate == 57143);
    }

    #[test]
    fn borrow_rate_at_optimal() {
        let market = make_market(10_000_000, 7_000_000, 700_000, 80_000, 1_000_000);
        assert_eq!(calculate_borrow_rate(&market).unwrap_or(0), 80_000);
    }

    #[test]
    fn borrow_rate_above_optimal() {
        let market = make_market(10_000_000, 9_000_000, 700_000, 80_000, 1_000_000);
        let rate = calculate_borrow_rate(&market).unwrap_or(0);
        // 90% utilization > 70% optimal, so rate should be > optimal_borrow_rate
        assert!(rate > 80_000);
    }

    #[test]
    fn borrow_rate_zero_optimal_util() {
        let market = make_market(1_000_000, 500_000, 0, 80_000, 1_000_000);
        // optimal_util=0 -> rate=0, but min_rate=0 too, so 0
        assert_eq!(calculate_borrow_rate(&market).unwrap_or(1), 0);
    }

    #[test]
    fn borrow_rate_min_rate_floor() {
        let mut market = make_market(1_000_000, 1, 700_000, 80_000, 1_000_000);
        market.min_borrow_rate = 2; // 1% = 10_000 in rate precision
        let rate = calculate_borrow_rate(&market).unwrap_or(0);
        assert!(rate >= 10_000);
    }

    #[test]
    fn borrow_rate_high_optimal_utilization() {
        // optimal_utilization = 900_000 (90%), above first segment breakpoint (850_000).
        let market = make_market(10_000_000, 9_500_000, 900_000, 80_000, 2_000_000);
        let rate = calculate_borrow_rate(&market);
        assert!(
            rate.is_ok(),
            "Should not underflow with high optimal_utilization"
        );
        let r = rate.unwrap_or(0);
        assert!(r > 80_000);
    }

    // ─── Deposit rate ────────────────────────────────────────────

    #[test]
    fn deposit_rate_basic() {
        let market = make_market(2_000_000, 1_000_000, 700_000, 80_000, 1_000_000);
        let deposit_rate = calculate_deposit_rate(&market).unwrap_or(0);
        // borrow_rate ~= 57142, no insurance, 50% util
        // deposit_rate = 57142 * 1000000 * 500000 / (1e6 * 1e6) = ~28571
        assert!(deposit_rate > 27000 && deposit_rate < 30000);
    }

    #[test]
    fn deposit_rate_with_insurance() {
        let mut market = make_market(2_000_000, 1_000_000, 700_000, 80_000, 1_000_000);
        market.insurance_fund.total_factor = 50_000; // 5%
        let deposit_rate = calculate_deposit_rate(&market).unwrap_or(0) as i128;
        // deposit_rate = borrow_rate * 0.95 * 0.5 = ~27142
        assert!(deposit_rate > 27000 && deposit_rate < 28000);
    }

    #[test]
    fn deposit_rate_zero_utilization() {
        let market = make_market(0, 0, 700_000, 80_000, 1_000_000);
        assert_eq!(calculate_deposit_rate(&market).unwrap_or(1), 0);
    }

    // ─── LTV ─────────────────────────────────────────────────────

    #[test]
    fn ltv_from_asset_weight() {
        let market = make_market(1, 0, 700_000, 80_000, 1_000_000);
        assert!((calculate_ltv(&market) - 8000.0).abs() < f64::EPSILON);
    }

    // ─── Helpers ─────────────────────────────────────────────────

    #[test]
    fn compute_drift_rates_basic() {
        let market = make_market(2_000_000, 1_000_000, 700_000, 80_000, 1_000_000);
        let result = compute_drift_rates(&market);
        assert!(result.is_ok());
        let (deposit, borrow, ltv) = result.unwrap_or((0.0, 0.0, 0.0));
        assert!(deposit > 0.0);
        assert!(borrow > 0.0);
        assert!((ltv - 8000.0).abs() < f64::EPSILON);
    }

    // ─── Proptest ────────────────────────────────────────────────

    mod proptests {
        use super::*;
        use proptest::prelude::*;

        prop_compose! {
            fn arb_market()(
                deposit_balance in 0u128..=1_000_000_000_000,
                borrow_balance in 0u128..=1_000_000_000_000,
                optimal_utilization in 0u32..=1_000_000,
                optimal_borrow_rate in 0u32..=10_000_000,
                max_borrow_rate in 0u32..=100_000_000,
                min_borrow_rate in 0u8..=200,
                insurance_factor in 0u32..=999_999,
                decimals in 0u32..=19,
            ) -> SpotMarket {
                SpotMarket {
                    pubkey: vec![],
                    market_index: 0,
                    initial_asset_weight: 8000,
                    initial_liability_weight: 12_000,
                    decimals,
                    cumulative_deposit_interest: 10_000_000_000,
                    cumulative_borrow_interest: 10_000_000_000,
                    deposit_balance,
                    borrow_balance,
                    optimal_utilization,
                    optimal_borrow_rate,
                    max_borrow_rate,
                    min_borrow_rate,
                    imf_factor: 0,
                    scale_initial_asset_weight_start: 0,
                    insurance_fund: InsuranceFund { total_factor: insurance_factor },
                    historical_oracle_data: HistoricalOracleData::default(),
                    oracle: None,
                }
            }
        }

        proptest! {
            #[test]
            fn utilization_never_panics(market in arb_market()) {
                let _ = calculate_utilization(&market);
            }

            #[test]
            fn utilization_zero_when_no_activity(market in arb_market()) {
                // When both deposit and borrow token amounts are zero, utilization = 0
                if let (Ok(d), Ok(b)) = (token_deposit_amount(&market), token_borrow_amount(&market)) {
                    if d == 0 && b == 0 {
                        prop_assert_eq!(calculate_utilization(&market).unwrap_or(1), 0);
                    }
                }
            }

            #[test]
            fn borrow_rate_never_panics(market in arb_market()) {
                let _ = calculate_borrow_rate(&market);
            }

            #[test]
            fn borrow_rate_geq_min_rate(market in arb_market()) {
                if let Ok(rate) = calculate_borrow_rate(&market) {
                    let min_rate = (market.min_borrow_rate as u128)
                        .checked_mul(RATE_PRECISION)
                        .and_then(|v| v.checked_div(200))
                        .unwrap_or(0);
                    prop_assert!(
                        rate >= min_rate,
                        "rate {} < min_rate {} (min_borrow_rate={})",
                        rate, min_rate, market.min_borrow_rate
                    );
                }
            }

            #[test]
            fn deposit_rate_never_panics(market in arb_market()) {
                let _ = calculate_deposit_rate(&market);
            }

            #[test]
            fn deposit_rate_leq_borrow_rate_when_util_bounded(market in arb_market()) {
                // deposit_rate <= borrow_rate only holds when utilization <= 100%.
                // When borrows > deposits, utilization exceeds 100% and the
                // formula can produce deposit_rate > borrow_rate.
                if let (Ok(util), Ok(deposit), Ok(borrow)) = (
                    calculate_utilization(&market),
                    calculate_deposit_rate(&market),
                    calculate_borrow_rate(&market),
                ) {
                    if util <= RATE_PRECISION {
                        prop_assert!(
                            deposit <= borrow,
                            "deposit {} > borrow {} (util {})",
                            deposit, borrow, util
                        );
                    }
                }
            }

            #[test]
            fn rate_to_decimal_non_negative(rate in 0u128..=100_000_000u128) {
                let decimal = crate::rate_to_decimal(rate);
                prop_assert!(decimal >= 0.0);
            }
        }
    }
}
