jup-lend-sdk 0.2.20

use super::errors::ErrorCodes;
use anyhow::Result;

const COEFFICIENT_SIZE_DEBT_FACTOR: u8 = 35;
const EXPONENT_SIZE_DEBT_FACTOR: u8 = 15;
const EXPONENT_MAX_DEBT_FACTOR: u64 = (1 << EXPONENT_SIZE_DEBT_FACTOR) - 1;
const DECIMALS_DEBT_FACTOR: u64 = 16384;
pub const MAX_MASK_DEBT_FACTOR: u64 =
    (1 << (COEFFICIENT_SIZE_DEBT_FACTOR + EXPONENT_SIZE_DEBT_FACTOR)) - 1;

#[allow(dead_code)]
const COEFFICIENT_MAX: u64 = (1 << COEFFICIENT_SIZE_DEBT_FACTOR) - 1;
#[allow(dead_code)]
const COEFFICIENT_MIN: u64 = 1 << (COEFFICIENT_SIZE_DEBT_FACTOR - 1);

pub const PRECISION: u8 = 64;
pub const TWO_POWER_64: u128 = 1 << PRECISION;
const TWO_POWER_69_MINUS_1: u128 = (1 << 69) - 1;

const COEFFICIENT_PLUS_PRECISION: u8 = COEFFICIENT_SIZE_DEBT_FACTOR + PRECISION; // 99
const COEFFICIENT_PLUS_PRECISION_MINUS_1: u8 = COEFFICIENT_PLUS_PRECISION - 1; // 98
const TWO_POWER_COEFFICIENT_PLUS_PRECISION_MINUS_1: u128 =
    (1 << COEFFICIENT_PLUS_PRECISION_MINUS_1) - 1; // (1 << 98) - 1
const TWO_POWER_COEFFICIENT_PLUS_PRECISION_MINUS_1_MINUS_1: u128 =
    (1 << (COEFFICIENT_PLUS_PRECISION_MINUS_1 - 1)) - 1; // (1 << 97) - 1

/// Multiplies a `normal` number with a `big_number1` and then divides by `big_number2`.
///
/// For vault's use case MUST always:
/// - bigNumbers have exponent size 15 bits
/// - bigNumbers have coefficient size 35 bits and have 35th bit always 1
/// - big_number1 (debt factor) always have exponent >= 1 & <= 16384
/// - big_number2 (connection factor) always have exponent >= 1 & <= 32767
/// - big_number2 always >= big_number1
/// - normal is positionRawDebt and is always within 10000 and u64::MAX
///
/// # Returns
/// normal * big_number1 / big_number2
pub fn mul_div_normal(normal: u64, big_number1: u64, big_number2: u64) -> Result<u64> {
    // Handle zero cases early
    if big_number1 == 0 || big_number2 == 0 {
        return Ok(0);
    }

    // Extract exponents from the big numbers
    let exponent1 = big_number1 & EXPONENT_MAX_DEBT_FACTOR;
    let exponent2 = big_number2 & EXPONENT_MAX_DEBT_FACTOR;

    // Calculate net exponent (exponent2 - exponent1)
    if exponent2 < exponent1 {
        return Err(ErrorCodes::BnError.into()); // Should never happen per requirements
    }

    let net_exponent = exponent2 - exponent1;

    if net_exponent < 129 {
        // Extract coefficients
        let coefficient1 = big_number1 >> EXPONENT_SIZE_DEBT_FACTOR;
        let coefficient2 = big_number2 >> EXPONENT_SIZE_DEBT_FACTOR;

        // Calculate (normal * coefficient1) / (coefficient2 << net_exponent)
        // Use u128 for intermediate calculations to prevent overflow
        let numerator: u128 = (normal as u128) * (coefficient1 as u128);
        let denominator: u128 = (coefficient2 as u128) << net_exponent;

        // Check for division by zero
        if denominator == 0 {
            return Err(ErrorCodes::DivisionByZero.into());
        }

        // Calculate result and check for overflow
        let result = numerator / denominator;
        if result > u64::MAX as u128 {
            return Err(ErrorCodes::BnError.into());
        }

        Ok(result as u64)
    } else {
        // If net_exponent >= 129, result will always be 0
        Ok(0)
    }
}

/// Multiplies a `big_number` with normal `number1` and then divides by `TWO_POWER_64`.
///
/// For vault's use case (calculating new branch debt factor after liquidation):
/// - number1 is debtFactor, initialized as TWO_POWER_64 and reduced from there
/// - big_number is branch debt factor, which starts with specific values and reduces
/// - big_number must have exponent size 15 bits and be >= 1 & <= 16384
/// - big_number must have coefficient size 35 bits and have 35th bit always 1
///
/// # Returns
/// big_number * number1 / TWO_POWER_64
pub fn mul_div_big_number(big_number: u64, number1: u128) -> Result<u64> {
    // Handle zero case early
    if big_number == 0 {
        return Ok(0);
    }

    // Extract coefficient from big_number
    let coefficient = big_number >> EXPONENT_SIZE_DEBT_FACTOR;
    let exponent = big_number & EXPONENT_MAX_DEBT_FACTOR;

    // Calculate result numerator: big_number coefficient * normal number
    let result_numerator: u128 = (coefficient as u128) * number1;

    // Find the most significant bit position
    let mut diff: u8;
    if result_numerator > TWO_POWER_COEFFICIENT_PLUS_PRECISION_MINUS_1 {
        diff = COEFFICIENT_PLUS_PRECISION;
    } else if result_numerator > TWO_POWER_COEFFICIENT_PLUS_PRECISION_MINUS_1_MINUS_1 {
        diff = COEFFICIENT_PLUS_PRECISION_MINUS_1;
    } else {
        diff = most_significant_bit(result_numerator);
    }

    // Calculate difference in bits to make the result_numerator 35 bits again
    diff = diff.saturating_sub(COEFFICIENT_SIZE_DEBT_FACTOR);

    // Shift result_numerator by the difference
    let adjusted_coefficient = (result_numerator >> diff) as u64;

    // Calculate new exponent
    let result_exponent = exponent.saturating_add(diff as u64);

    // Divide by TWO_POWER_64 by reducing exponent by 64
    if result_exponent > PRECISION as u64 {
        let final_exponent = result_exponent - PRECISION as u64;

        // Check that we don't exceed the exponent max
        if final_exponent > EXPONENT_MAX_DEBT_FACTOR {
            return Err(ErrorCodes::BnError.into());
        }

        // Combine coefficient and exponent
        Ok((adjusted_coefficient << EXPONENT_SIZE_DEBT_FACTOR) | final_exponent)
    } else {
        // If we would underflow the exponent, this is an error case
        // Debt factor should never become a BigNumber with exponent <= 0
        Err(ErrorCodes::BnError.into())
    }
}

/// Multiplies a `big_number1` with another `big_number2`.
///
/// For vault's use case (calculating connection factor of merged branches):
/// - bigNumbers must have exponent size 15 bits and be >= 1 & <= 32767
/// - bigNumber must have coefficient size 35 bits and have 35th bit always 1
/// - Sum of exponents should be > 16384
///
/// # Returns
/// BigNumber format with coefficient and exponent
pub fn mul_big_number(big_number1: u64, big_number2: u64) -> Result<u64> {
    // Extract coefficients and exponents
    let coefficient1: u64 = big_number1 >> EXPONENT_SIZE_DEBT_FACTOR;
    let coefficient2: u64 = big_number2 >> EXPONENT_SIZE_DEBT_FACTOR;
    let exponent1: u64 = big_number1 & EXPONENT_MAX_DEBT_FACTOR;
    let exponent2: u64 = big_number2 & EXPONENT_MAX_DEBT_FACTOR;

    // Calculate result coefficient
    // res coefficient at max can be 34359738367 * 34359738367 = 1180591620648691826689 (X35 * X35 fits in 70 bits)
    let res_coefficient: u128 = (coefficient1 as u128) * (coefficient2 as u128);

    // Determine overflow length based on result size
    let overflow_len = if res_coefficient > TWO_POWER_69_MINUS_1 {
        COEFFICIENT_SIZE_DEBT_FACTOR as u64
    } else {
        (COEFFICIENT_SIZE_DEBT_FACTOR - 1) as u64
    };

    // Adjust coefficient to fit in 35 bits
    let adjusted_coefficient = (res_coefficient >> overflow_len) as u64;

    // Calculate result exponent
    let res_exponent = exponent1 + exponent2 + overflow_len;

    // Check for exponent underflow
    if res_exponent < DECIMALS_DEBT_FACTOR {
        return Err(ErrorCodes::BnError.into());
    }

    // Adjust exponent
    let final_exponent = res_exponent - DECIMALS_DEBT_FACTOR;

    // Check for exponent overflow
    if final_exponent > EXPONENT_MAX_DEBT_FACTOR {
        // If exponent exceeds max, user is ~100% liquidated
        return Ok(MAX_MASK_DEBT_FACTOR);
    }

    // Combine coefficient and exponent
    Ok((adjusted_coefficient << EXPONENT_SIZE_DEBT_FACTOR) | final_exponent)
}

/// Divides a `big_number1` by `big_number2`.
///
/// For vault's use case (calculating connectionFactor):
/// - Numbers must have exponent size 15 bits and be >= 1 & <= 16384
/// - Numbers must have coefficient size 35 bits and have 35th bit always 1
/// - Numbers must never be 0
///
/// # Returns
/// BigNumber format with coefficient and exponent
pub fn div_big_number(big_number1: u64, big_number2: u64) -> Result<u64> {
    // Handle zero cases early
    if big_number1 == 0 {
        return Ok(0);
    }
    if big_number2 == 0 {
        return Err(ErrorCodes::DivisionByZero.into());
    }

    // Extract coefficients and exponents
    let coefficient1 = big_number1 >> EXPONENT_SIZE_DEBT_FACTOR;
    let coefficient2 = big_number2 >> EXPONENT_SIZE_DEBT_FACTOR;
    let exponent1 = big_number1 & EXPONENT_MAX_DEBT_FACTOR;
    let exponent2 = big_number2 & EXPONENT_MAX_DEBT_FACTOR;

    // Check for division by zero coefficient
    if coefficient2 == 0 {
        return Err(ErrorCodes::DivisionByZero.into());
    }

    // Calculate result coefficient: (coefficient1 << PRECISION) / coefficient2
    let res_coefficient: u128 =
        ((coefficient1 as u128) << PRECISION as u128) / (coefficient2 as u128);

    // Determine overflow length
    let overflow_len = if (res_coefficient >> PRECISION as u128) == 1 {
        (PRECISION + 1) as u64
    } else {
        PRECISION as u64
    };

    // Adjust overflow length
    let adjusted_overflow_len = overflow_len - COEFFICIENT_SIZE_DEBT_FACTOR as u64;

    // Adjust coefficient to fit in 35 bits
    let adjusted_coefficient = (res_coefficient >> adjusted_overflow_len) as u64;

    // Calculate result exponent components
    let addition_part = exponent1 + DECIMALS_DEBT_FACTOR + adjusted_overflow_len;
    let subtraction_part = exponent2 + PRECISION as u64;

    // Check if addition part is greater than subtraction part
    if addition_part > subtraction_part {
        let final_exponent = addition_part - subtraction_part;

        // Check that we don't exceed the exponent max
        if final_exponent > EXPONENT_MAX_DEBT_FACTOR {
            return Err(ErrorCodes::BnError.into());
        }

        // Combine coefficient and exponent
        Ok((adjusted_coefficient << EXPONENT_SIZE_DEBT_FACTOR) | final_exponent)
    } else {
        // If we would underflow the exponent, this is an error case
        // Connection factor should never become a BigNumber with exponent <= 0
        Err(ErrorCodes::BnError.into())
    }
}

/// Gets the most significant bit position of a number (1-indexed)
/// Returns 0 for input 0, otherwise returns the position of the highest set bit
pub fn most_significant_bit(normal: u128) -> u8 {
    if normal == 0 {
        return 0;
    }

    // Use built-in leading_zeros for accuracy and performance
    // leading_zeros counts from the left, so MSB position is 128 - leading_zeros
    // But we want 1-indexed position, so it's 128 - leading_zeros
    128 - (normal.leading_zeros() as u8)
}

#[allow(dead_code)]
/// Helper function to create a big number from coefficient and exponent
fn create_big_number(coefficient: u64, exponent: u64) -> u64 {
    (coefficient << EXPONENT_SIZE_DEBT_FACTOR) | exponent
}

#[allow(dead_code)]
/// Helper function to extract coefficient from big number
fn extract_coefficient(big_number: u64) -> u64 {
    big_number >> EXPONENT_SIZE_DEBT_FACTOR
}

#[allow(dead_code)]
/// Helper function to extract exponent from big number
fn extract_exponent(big_number: u64) -> u64 {
    big_number & EXPONENT_MAX_DEBT_FACTOR
}