jup-lend-sdk 0.2.4

use super::{
    errors::ErrorCodes,
    ticks::{
        get_first_tick_for_array_index, get_tick_indices, MIN_TICK, TICK_HAS_DEBT_ARRAY_SIZE,
        TICK_HAS_DEBT_CHILDREN_SIZE, TICK_HAS_DEBT_CHILDREN_SIZE_IN_BITS, TOTAL_INDICES_NEEDED,
    },
    VaultProgram,
};
use crate::{
    math::{casting::Cast, safe_math::SafeMath, tick::TickMath},
    programs::vaults::accounts::{Tick, TickHasDebtArray},
};
use anchor_lang::prelude::AccountMeta;
use anyhow::Result;

impl TickHasDebtArray {
    pub fn has_bits(&self, map_id: usize) -> bool {
        self.tick_has_debt[map_id]
            .children_bits
            .iter()
            .any(|&x| x != 0)
    }

    pub fn clear_bits_for_tick(&mut self, tick: i32) -> Result<()> {
        let (_, map_index, byte_index, bit_index) = get_tick_indices(tick)?;

        self.clear_bits(map_index, byte_index, bit_index)?;
        Ok(())
    }

    // Create a mask for the current byte that will:
    // - Keep all bits lower than bit_index (these are smaller ticks)
    // - Clear the current bit and all higher bits (these are the current and larger ticks)
    pub fn clear_bits(
        &mut self,
        map_index: usize,
        byte_index: usize,
        bit_index: usize,
    ) -> Result<()> {
        // Create a working copy of the current map's bitmap
        let mut bitmap: [u8; TICK_HAS_DEBT_CHILDREN_SIZE] =
            self.tick_has_debt[map_index].children_bits;

        // Clear the current tick's bit and all higher bits in the current byte
        if bit_index > 0 {
            // Create mask to keep only bits lower than current bit
            let mask = (1 << bit_index) - 1;
            bitmap[byte_index] &= mask;
        } else {
            // If bit_index is 0, clear the entire byte
            bitmap[byte_index] = 0;
        }

        // Clear all bytes with higher indices (representing higher ticks)
        for i in (byte_index + 1)..TICK_HAS_DEBT_CHILDREN_SIZE {
            bitmap[i] = 0;
        }

        self.tick_has_debt[map_index].children_bits = bitmap;

        Ok(())
    }

    pub fn fetch_next_top_tick(&mut self, mut map_index: usize) -> Result<(i32, bool)> {
        // Search for the next tick with debt
        loop {
            if self.has_bits(map_index.cast()?) {
                let (next_tick, has_next_tick) = self.get_next_tick(map_index.cast()?)?;

                if has_next_tick {
                    return Ok((next_tick, true));
                }
            }

            // No bits found in current map, move to the previous map (lower ticks)
            if map_index == 0 {
                if self.index == 0 {
                    // No more maps to check in this array
                    return Ok((MIN_TICK, true));
                } else {
                    return Ok((MIN_TICK, false));
                }
            }

            map_index -= 1;
        }
    }

    fn get_most_significant_bit(&self, map_id: usize, byte_idx: usize) -> u32 {
        let bits: u8 = self.tick_has_debt[map_id].children_bits[byte_idx];

        // Find most significant bit in the byte
        bits.leading_zeros()
    }

    fn get_next_tick(&self, map_index: usize) -> Result<(i32, bool)> {
        for byte_idx in (0..TICK_HAS_DEBT_CHILDREN_SIZE).rev() {
            if self.tick_has_debt[map_index].children_bits[byte_idx] != 0 {
                // Find the highest set bit in this byte
                let leading_zeros = self.get_most_significant_bit(map_index, byte_idx);
                // 7 - leading_zeros ensures we count from the left of the byte
                let bit_pos = 7 - leading_zeros as usize;

                // Calculate the tick within the map (0-255)
                let tick_within_map = byte_idx * 8 + bit_pos;

                // Calculate the actual tick value
                let map_first_tick = self.get_first_tick_for_map_index(map_index)?;
                return Ok((map_first_tick + tick_within_map as i32, true));
            }
        }

        Ok((MIN_TICK, false))
    }

    fn get_first_tick_for_map_index(&self, map_index: usize) -> Result<i32> {
        get_first_tick_for_map_in_array(self.index, map_index.cast()?)
    }
}

/// Get the first tick for a given map_index within a specific TickHasDebtArray
pub fn get_first_tick_for_map_in_array(array_index: u8, map_index: u8) -> Result<i32> {
    if array_index >= TOTAL_INDICES_NEEDED as u8 || map_index >= TICK_HAS_DEBT_ARRAY_SIZE as u8 {
        return Err(ErrorCodes::VaultTickHasDebtOutOfRange.into());
    }

    // Each array covers 2048 ticks, each map within array covers 256 ticks
    let array_first_tick = get_first_tick_for_array_index(array_index)?;

    let map_first_tick = array_first_tick
        .safe_add((map_index as i32).safe_mul(TICK_HAS_DEBT_CHILDREN_SIZE_IN_BITS as i32)?)?;

    Ok(map_first_tick)
}

pub async fn get_tick_has_debt_from_remaining_accounts_liquidate(
    remaining_accounts: &Vec<AccountMeta>,
    remaining_accounts_indices: &Vec<u8>,
    program: &VaultProgram,
) -> Result<Vec<TickHasDebtArray>> {
    // remaining_accounts_indices[0] is oracle sources length
    // remaining_accounts_indices[1] is branches length
    // remaining_accounts_indices[2] is ticks length
    // remaining_accounts_indices[3] is tick has debt length

    let tick_has_debt_length: usize = remaining_accounts_indices[3].cast::<usize>()?;

    let start_index: usize = remaining_accounts_indices[0].cast::<usize>()?
        + remaining_accounts_indices[1].cast::<usize>()?
        + remaining_accounts_indices[2].cast::<usize>()?;

    let end_index: usize = start_index + tick_has_debt_length;

    Ok(futures::future::try_join_all(
        remaining_accounts
            .iter()
            .take(end_index)
            .skip(start_index)
            .map(|account| program.account::<TickHasDebtArray>(account.pubkey)),
    )
    .await?)
}

pub fn fetch_next_tick_absorb(
    tick_has_debts: &mut Vec<TickHasDebtArray>,
    ticks: &mut Vec<Tick>,
    current_tick: i32,
    max_tick: i32,
) -> Result<(i32, u128, u128)> {
    let (mut array_index, mut map_index, _, _) = get_tick_indices(current_tick)?;

    let mut current_tick_has_debt = tick_has_debts
        .get_mut(array_index as usize)
        .ok_or(ErrorCodes::VaultTickHasDebtNotFound)?;

    if current_tick_has_debt.index != array_index {
        return Err(ErrorCodes::VaultTickHasDebtIndexMismatch.into());
    }

    let mut col_absorbed: u128 = 0;
    let mut debt_absorbed: u128 = 0;

    // For last user remaining in vault there could be a lot of loop iterations
    // Chances of this to happen is extremely low (like ~0%)
    loop {
        let (next_tick, has_next_tick) = current_tick_has_debt.fetch_next_top_tick(map_index)?;

        if has_next_tick {
            if next_tick > max_tick {
                let tick_data = ticks
                    .iter_mut()
                    .find(|t| t.tick == next_tick)
                    .ok_or(ErrorCodes::VaultTickNotFound)?;

                let tick_debt = tick_data.get_raw_debt()?;
                let ratio = TickMath::get_ratio_at_tick(next_tick)?;

                debt_absorbed = debt_absorbed.safe_add(tick_debt)?;
                col_absorbed = col_absorbed.safe_add(
                    tick_debt
                        .safe_mul(TickMath::ZERO_TICK_SCALED_RATIO)?
                        .safe_div(ratio)?,
                )?;

                tick_data.set_fully_liquidated();
                current_tick_has_debt.clear_bits_for_tick(next_tick)?;
            } else {
                return Ok((next_tick, col_absorbed, debt_absorbed));
            }
        } else {
            array_index -= 1;
            map_index = TICK_HAS_DEBT_ARRAY_SIZE - 1;
            current_tick_has_debt = tick_has_debts
                .get_mut(array_index as usize)
                .ok_or(ErrorCodes::VaultTickHasDebtNotFound)?;
        }
    }
}

pub fn fetch_next_tick_liquidate(
    tick_has_debts: &mut Vec<TickHasDebtArray>,
    current_tick: i32,
    liquidation_tick: i32,
    clear_bits: bool,
) -> Result<i32> {
    let (mut array_index, mut map_index, byte_index, bit_index) = get_tick_indices(current_tick)?;

    let mut current_tick_has_debt = tick_has_debts
        .get_mut(array_index as usize)
        .ok_or(ErrorCodes::VaultTickHasDebtNotFound)?;

    if clear_bits {
        current_tick_has_debt.clear_bits(map_index, byte_index, bit_index)?;
    }

    // For last user remaining in vault there could be a lot of loop iterations
    // Chances of this to happen is extremely low (like ~0%)
    loop {
        let (next_tick, has_next_tick) = current_tick_has_debt.fetch_next_top_tick(map_index)?;

        if has_next_tick {
            return Ok(next_tick);
        } else {
            if current_tick_has_debt.get_first_tick_for_map_index(map_index)? < liquidation_tick {
                return Ok(TickMath::COLD_TICK);
            }

            array_index -= 1;
            map_index = TICK_HAS_DEBT_ARRAY_SIZE - 1;
            current_tick_has_debt = tick_has_debts
                .get_mut(array_index as usize)
                .ok_or(ErrorCodes::VaultTickHasDebtNotFound)?;
        }
    }
}