jup-lend-sdk 0.2.13

use crate::{
    math::{
        bn::{mul_div_big_number, TWO_POWER_64},
        casting::Cast,
        safe_math::SafeMath,
        tick::TickMath,
        u256::safe_multiply_divide,
    },
    programs::vaults::accounts::{Branch, Tick, VaultState},
};
use anchor_lang::zero_copy;
use anyhow::Result;

use super::{
    constants::{
        EXCHANGE_PRICES_PRECISION, FOUR_DECIMALS, INITIAL_BRANCH_DEBT_FACTOR, RATE_OUTPUT_DECIMALS,
        X30,
    },
    errors::ErrorCodes,
};

#[derive(Default)]
#[repr(C, packed)]
#[zero_copy]
pub struct CurrentLiquidity {
    pub debt_remaining: u128, // Debt remaining to liquidate
    pub debt: u128,           // Current liquidatable debt before reaching next check point
    pub col: u128,            // Calculate using debt & ratioCurrent
    pub total_debt_liq: u128, // Total debt liquidated till now
    pub total_col_liq: u128,  // Total collateral liquidated till now
    pub tick: i32,            // Current tick to liquidate
    pub ratio: u128,          // Current ratio to liquidate
    pub tick_status: u8, // if 1 then it's a perfect tick, if 2 that means it's a liquidated tick
    pub ref_tick: i32,   // ref tick to liquidate
    pub ref_ratio: u128, // ratio at ref tick
    pub ref_tick_status: u8, // if 1 then it's a perfect tick, if 2 that means it's a liquidated tick, if 3 that means it's a liquidation threshold
}

impl CurrentLiquidity {
    pub fn check_is_ref_partials_safe_for_tick(
        &self,
        existing_partials: u128,
        partials: u128,
    ) -> Result<()> {
        if existing_partials > 0 && existing_partials >= partials {
            // If refTick is liquidated tick and hence contains partials then checking that
            // current liquidation tick's partial should not be less than last liquidation refTick

            // Not sure if this is even possible to happen but adding checks to avoid it fully
            // If it reverts here then next liquidation on next block should go through fine
            return Err(ErrorCodes::VaultLiquidationReverts.into());
        }

        Ok(())
    }

    pub fn get_final_ratio(&self, col_liquidated: u128, debt_liquidated: u128) -> Result<u128> {
        let final_ratio: u128 = self
            .debt
            .safe_sub(debt_liquidated)?
            .safe_mul(TickMath::ZERO_TICK_SCALED_RATIO)?
            .safe_div(self.col.safe_sub(col_liquidated)?)?;

        Ok(final_ratio)
    }

    pub fn is_perfect_tick(&self) -> bool {
        self.tick_status == 1
    }

    pub fn is_ref_tick_liquidation_threshold(&self) -> bool {
        self.ref_tick_status == 3
    }

    pub fn is_ref_tick_liquidated(&self) -> bool {
        self.ref_tick_status == 2
    }

    pub fn get_debt_from_ratios(&self) -> Result<u128> {
        Ok(self.ref_ratio.safe_mul(self.debt)?.safe_div(self.ratio)?)
    }

    pub fn get_col_from_ratios(&self, col_per_debt: u128) -> Result<u128> {
        // in u128, this would be risky to overflow, use u256 to be safe even though col_per_debt and ref_ratio
        // are inversely proportional to each other and this should not be necessary.
        // col_per_debt can be up to 4.62e25 at oracle precision 1e15
        // ratio can be up to 13002088133096036565414295

        Ok(safe_multiply_divide(
            col_per_debt,
            self.ref_ratio,
            TickMath::ZERO_TICK_SCALED_RATIO,
        )?)
    }
    pub fn get_debt_liquidated(&self, col_per_debt: u128) -> Result<u128> {
        // Calculate the numerator
        let numerator = self
            .debt
            .safe_sub(self.get_debt_from_ratios()?)?
            .safe_mul(10u128.pow(RATE_OUTPUT_DECIMALS))?;

        // Calculate the denominator
        let denominator = 10u128
            .pow(RATE_OUTPUT_DECIMALS)
            .safe_sub(self.get_col_from_ratios(col_per_debt)?)?;

        let debt_liquidated = numerator.safe_div(denominator)?;

        Ok(debt_liquidated)
    }

    fn get_normal_debt_liq(&self, borrow_ex_price: u128) -> Result<u128> {
        Ok(self
            .total_debt_liq
            .safe_mul(borrow_ex_price)?
            .safe_div(EXCHANGE_PRICES_PRECISION)?)
    }

    fn get_normal_col_liq(&self, supply_ex_price: u128) -> Result<u128> {
        Ok(self
            .total_col_liq
            .safe_mul(supply_ex_price)?
            .safe_div(EXCHANGE_PRICES_PRECISION)?)
    }

    pub fn get_actual_amounts(
        &self,
        borrow_ex_price: u128,
        supply_ex_price: u128,
        debt_amount: u128,
    ) -> Result<(u128, u128)> {
        let mut actual_debt_amt = self.get_normal_debt_liq(borrow_ex_price)?;
        let mut actual_col_amt = self.get_normal_col_liq(supply_ex_price)?;

        if actual_debt_amt > debt_amount {
            // Calculate new actual_col_amt via ratio
            actual_col_amt = actual_col_amt
                .safe_mul(debt_amount)?
                .safe_div(actual_debt_amt)?;

            actual_debt_amt = debt_amount;
        }

        if actual_debt_amt == 0 {
            return Err(ErrorCodes::VaultInvalidLiquidation.into());
        }

        Ok((actual_debt_amt, actual_col_amt))
    }

    pub fn update_next_iterations_with_ref(&mut self) {
        self.tick = self.ref_tick;
        self.tick_status = self.ref_tick_status;
        self.ratio = self.ref_ratio;
    }

    pub fn get_debt_factor(&self, debt_liquidated: u128) -> Result<u128> {
        // debtFactor = debtFactor * (liquidatableDebt - debtLiquidated) / liquidatableDebt
        // -> debtFactor * leftOverDebt / liquidatableDebt

        let debt_factor = TWO_POWER_64
            .safe_mul(self.debt.safe_sub(debt_liquidated)?.cast()?)?
            .safe_div(self.debt.cast()?)?;

        Ok(debt_factor)
    }

    pub fn reduce_debt_remaining(&mut self, debt_liquidated: u128) -> Result<()> {
        self.debt_remaining = self.debt_remaining.safe_sub(debt_liquidated)?;
        Ok(())
    }

    fn reduce_col(&mut self, col_liquidated: u128) -> Result<()> {
        self.col = self.col.safe_sub(col_liquidated)?;
        Ok(())
    }

    fn increase_total_debt_liq(&mut self, debt_liquidated: u128) -> Result<()> {
        self.total_debt_liq = self.total_debt_liq.safe_add(debt_liquidated)?;
        Ok(())
    }

    fn increase_total_col_liq(&mut self, col_liquidated: u128) -> Result<()> {
        self.total_col_liq = self.total_col_liq.safe_add(col_liquidated)?;
        Ok(())
    }

    fn reduce_debt(&mut self, debt_liquidated: u128) -> Result<()> {
        self.debt = self.debt.safe_sub(debt_liquidated)?;
        Ok(())
    }

    pub fn update_totals(&mut self, debt_liquidated: u128, col_liquidated: u128) -> Result<()> {
        self.reduce_col(col_liquidated)?;
        self.increase_total_col_liq(col_liquidated)?;
        self.increase_total_debt_liq(debt_liquidated)?;
        self.reduce_debt(debt_liquidated)?;

        Ok(())
    }
}

#[derive(Default)]
#[repr(C, packed)]
#[zero_copy]
pub struct TickMemoryVars {
    pub tick: i32,
    pub partials: u128,
}

impl TickMemoryVars {
    pub fn set_partials(&mut self, partials: u128) -> Result<()> {
        if partials == 0 {
            self.partials = 1;
        } else if partials >= X30 {
            self.partials = X30 - 1;
        } else {
            self.partials = partials;
        }

        Ok(())
    }

    pub fn set_tick(&mut self, tick: i32) {
        self.tick = tick;
    }
}

impl Tick {
    pub fn set_fully_liquidated(&mut self) {
        self.set_liquidated(0, 0);
        self.is_fully_liquidated = 1;
    }

    pub fn set_liquidated(&mut self, branch_id: u32, debt_factor: u64) {
        self.raw_debt = 0;
        self.is_liquidated = 1;
        self.debt_factor = debt_factor;
        self.liquidation_branch_id = branch_id;
    }

    pub fn get_raw_debt(&self) -> Result<u128> {
        Ok(self.raw_debt.cast()?)
    }
}

#[derive(Default)]
#[repr(C, packed)]
#[zero_copy]
pub struct BranchState {
    pub minima_tick: i32,          // Minima tick of this branch Y
    pub minima_tick_partials: u32, // Partials of minima tick of branch this is connected to Y
    pub debt_liquidity: u64,       // Debt liquidity at this branch Y

    // For non-merged branches
    pub debt_factor: u64, // Debt factor coefficient, 35 coefficient | 15 exponent, Y

    // For all branches
    pub connected_branch_id: u32, // Branch's ID with which this branch is connected Y
    pub connected_minima_tick: i32, // Minima tick of branch this is connected to
}

#[derive(Default)]
#[repr(C, packed)]
#[zero_copy]
pub struct BranchMemoryVars {
    pub id: u32,
    pub data: BranchState,
    pub debt_factor: u64,
    pub minima_tick: i32,
    pub base_branch_data: BranchState,
}

impl BranchMemoryVars {
    pub fn reset_branch_data(&mut self) {
        self.data = BranchState::default();
        self.data.minima_tick = TickMath::COLD_TICK;
        self.data.connected_minima_tick = TickMath::COLD_TICK;
    }

    pub fn update_branch_to_base_branch(&mut self) {
        self.id = self.data.connected_branch_id;
        self.data = self.base_branch_data;
        self.minima_tick = self.base_branch_data.connected_minima_tick;
    }

    pub fn update_branch_debt_factor(&mut self, debt_factor: u128) -> Result<()> {
        self.debt_factor = mul_div_big_number(self.debt_factor, debt_factor.cast()?)?;
        Ok(())
    }

    pub fn set_branch_data(&mut self, branch: &Branch) -> Result<()> {
        self.data = BranchState {
            minima_tick: branch.minima_tick,
            minima_tick_partials: branch.minima_tick_partials,
            debt_liquidity: branch.debt_liquidity,
            debt_factor: branch.debt_factor,
            connected_branch_id: branch.connected_branch_id,
            connected_minima_tick: branch.connected_minima_tick,
        };

        Ok(())
    }

    pub fn set_branch_data_in_memory(&mut self, branch: &Branch) -> Result<()> {
        self.id = branch.branch_id;

        self.set_branch_data(branch)?;

        self.debt_factor = self.data.debt_factor;

        if self.debt_factor == 0 {
            self.debt_factor = INITIAL_BRANCH_DEBT_FACTOR.cast()?;
        }

        self.minima_tick = self.data.connected_minima_tick;

        Ok(())
    }

    pub fn set_base_branch_data(&mut self, branch: &Branch) -> Result<()> {
        self.base_branch_data = BranchState {
            minima_tick: branch.minima_tick,
            minima_tick_partials: branch.minima_tick_partials,
            debt_liquidity: branch.debt_liquidity,
            debt_factor: branch.debt_factor,
            connected_branch_id: branch.connected_branch_id,
            connected_minima_tick: branch.connected_minima_tick,
        };

        Ok(())
    }

    pub fn get_current_ratio_from_minima_tick(&self) -> Result<u128> {
        // Calculate ratios
        let ratio = TickMath::get_ratio_at_tick(self.minima_tick)?;
        let ratio_one_less = ratio
            .safe_mul(FOUR_DECIMALS)?
            .safe_div(TickMath::TICK_SPACING)?;

        let length = ratio.safe_sub(ratio_one_less)?;

        // Calculate current ratio with partials
        let current_ratio = ratio_one_less.safe_add(
            length
                .safe_mul(self.data.minima_tick_partials.cast()?)?
                .safe_div(X30)?,
        )?;

        Ok(current_ratio)
    }
}

impl VaultState {
    fn bump_total_branch_id(&mut self) {
        self.total_branch_id += 1;
    }

    pub fn reset_top_tick(&mut self) {
        self.topmost_tick = TickMath::COLD_TICK;
    }

    pub fn update_branch_info_by_one(&mut self) {
        // increment total branches by 1
        self.bump_total_branch_id();

        // reset branch liquidated as new branch is initialized
        self.reset_branch_liquidated();

        // set current branch id to total branches
        self.current_branch_id = self.total_branch_id;
    }

    pub fn reset_branch_liquidated(&mut self) {
        self.branch_liquidated = 0;
    }

    pub fn get_tick_status(&self) -> u8 {
        if self.is_branch_liquidated() {
            2
        } else {
            1
        }
    }

    pub fn is_branch_liquidated(&self) -> bool {
        self.branch_liquidated == 1
    }

    pub fn get_top_tick(&self) -> i32 {
        self.topmost_tick
    }

    pub fn get_total_borrow(&self) -> Result<u128> {
        Ok(self.total_borrow.cast()?)
    }

    pub fn reduce_total_borrow(&mut self, amount: u128) -> Result<()> {
        self.total_borrow = self.total_borrow.safe_sub(amount.cast()?)?;
        Ok(())
    }

    pub fn reduce_total_supply(&mut self, amount: u128) -> Result<()> {
        self.total_supply = self.total_supply.safe_sub(amount.cast()?)?;
        Ok(())
    }

    pub fn reset_absorbed_amounts(&mut self) {
        self.absorbed_debt_amount = 0;
        self.absorbed_col_amount = 0;
    }

    pub fn absorb_dust_amount_for_liquidate(
        &mut self,
        current_data: &mut CurrentLiquidity,
    ) -> Result<()> {
        let absorbed_debt: u128 = self.absorbed_debt_amount;
        let absorbed_col: u128 = self.absorbed_col_amount;

        if absorbed_debt > current_data.debt_remaining {
            // Removing collateral in equal proportion as debt
            current_data.total_col_liq = absorbed_col
                .safe_mul(current_data.debt_remaining)?
                .safe_div(absorbed_debt)?;

            // Update absorbed amounts
            self.absorbed_col_amount = absorbed_col.safe_sub(current_data.total_col_liq)?;

            // Update debt
            current_data.total_debt_liq = current_data.debt_remaining;
            self.absorbed_debt_amount = absorbed_debt.safe_sub(current_data.debt_remaining)?;

            current_data.debt_remaining = 0;
        } else {
            // Clean out all absorbed debt and collateral
            self.reset_absorbed_amounts();

            current_data.debt_remaining = current_data.debt_remaining.safe_sub(absorbed_debt)?;
            current_data.total_debt_liq = absorbed_debt;
            current_data.total_col_liq = absorbed_col;
        }

        Ok(())
    }

    pub fn set_branch_liquidated(&mut self) {
        self.branch_liquidated = 1;
    }
    pub fn update_state_at_liq_end(&mut self, tick: i32, branch_id: u32) -> Result<()> {
        self.set_branch_liquidated();
        self.topmost_tick = tick;
        self.current_branch_id = branch_id;

        Ok(())
    }
}

pub enum BranchStatus {
    NotLiquidated,
    Liquidated,
    Merged,
    Closed,
}

impl Branch {
    pub fn set_state_after_absorb(&mut self, branch_data: &BranchState) {
        self.set_status_as_closed();

        self.minima_tick = branch_data.minima_tick;
        self.minima_tick_partials = branch_data.minima_tick_partials;
        self.debt_liquidity = branch_data.debt_liquidity;
        self.debt_factor = branch_data.debt_factor;
        self.connected_branch_id = branch_data.connected_branch_id;
        self.connected_minima_tick = branch_data.connected_minima_tick;
    }

    pub fn set_status_as_closed(&mut self) {
        self.status = BranchStatus::Closed as u8;
    }

    pub fn set_status_as_merged(&mut self) {
        self.status = BranchStatus::Merged as u8;
    }

    pub fn merge_with_base_branch(&mut self, connection_factor: u64) -> Result<()> {
        self.set_status_as_merged();

        // set new connectionFactor or debtFactor
        self.debt_factor = connection_factor;

        // deleting debt / partials / minima tick of current branch
        self.minima_tick = TickMath::COLD_TICK;
        self.minima_tick_partials = 0;
        self.debt_liquidity = 0;

        Ok(())
    }

    pub fn set_connections(&mut self, branch_id: u32, minima_tick: i32) -> Result<()> {
        self.connected_branch_id = branch_id;
        self.connected_minima_tick = minima_tick;

        Ok(())
    }

    pub fn reset_branch_data(&mut self) {
        self.status = 0;
        self.minima_tick = TickMath::COLD_TICK;
        self.minima_tick_partials = 0;
        self.debt_liquidity = 0;
        self.debt_factor = 0;
        self.connected_branch_id = 0;
        self.connected_minima_tick = TickMath::COLD_TICK;
    }

    pub fn set_status_as_liquidated(&mut self) {
        self.status = BranchStatus::Liquidated as u8;
    }

    pub fn set_branch_debt(&mut self, new_debt_liquidity_raw: u128) -> Result<()> {
        self.debt_liquidity = new_debt_liquidity_raw.cast()?;
        Ok(())
    }

    pub fn update_state_at_liq_end(
        &mut self,
        tick: i32,
        partials: u128,
        debt: u128,
        debt_factor: u128,
    ) -> Result<()> {
        self.set_status_as_liquidated();
        self.minima_tick = tick;
        self.minima_tick_partials = partials.cast()?;
        self.set_branch_debt(debt)?;
        self.debt_factor = debt_factor.cast()?;

        Ok(())
    }
}