jup-lend-sdk 0.1.2

use std::{rc::Rc, sync::Arc};

use anchor_client::{
    solana_sdk::{commitment_config::CommitmentConfig, signature::Keypair},
    Client, Cluster,
};

use crate::{
    borrow::instructions::VAULTS_PROGRAM_ID,
    programs::vaults::accounts::{Branch, Tick, TickHasDebtArray, VaultState},
};

use super::{
    pda,
    ticks::{get_tick_indices, MIN_TICK, TICK_HAS_DEBT_ARRAY_SIZE},
};

use rayon::iter::{IntoParallelRefIterator, ParallelIterator};

pub fn load_relevant_branches_for_liquidate(
    vault_id: u16,
    vault_state: VaultState,
    cluster: &Cluster,
) -> anyhow::Result<Vec<Branch>> {
    let provider = Client::new_with_options(
        cluster.clone(),
        Rc::new(Keypair::new()),
        CommitmentConfig::confirmed(),
    );
    let program = provider.program(VAULTS_PROGRAM_ID)?;
    let mut branches = vec![];

    let current_branch_id = vault_state.current_branch_id;

    if current_branch_id > 0 {
        if let Ok(current_branch) =
            program.account::<Branch>(pda::get_branch(vault_id, current_branch_id))
        {
            branches.push(current_branch);
        }
    }

    if !branches.is_empty() {
        let mut connected_branch_id = branches[0].connected_branch_id;

        while !branches.iter().any(|b| b.branch_id == connected_branch_id) {
            if let Ok(connected_branch) =
                program.account::<Branch>(pda::get_branch(vault_id, connected_branch_id))
            {
                connected_branch_id = connected_branch.connected_branch_id;
                branches.push(connected_branch);
            } else {
                break;
            }
        }
    }

    if !branches.iter().any(|b| b.branch_id == 0) {
        branches.push(Branch {
            vault_id,
            branch_id: 0,
            ..Default::default()
        });
    }

    Ok(branches)
}

fn find_next_tick_with_debt(
    vault_id: u16,
    start_tick: i32,
    cluster: &Cluster,
) -> anyhow::Result<i32> {
    let provider = Client::new_with_options(
        cluster.clone(),
        Rc::new(Keypair::new()),
        CommitmentConfig::confirmed(),
    );
    let program = provider.program(VAULTS_PROGRAM_ID)?;

    let (array_index, map_index, byte_index, bit_index) = get_tick_indices(start_tick)?;

    let mut current_array_index = array_index;
    let mut current_map_index = map_index;

    let mut tick_has_debt_data = match program
        .account::<TickHasDebtArray>(pda::get_tick_has_debt(vault_id, current_array_index))
    {
        Ok(data) => data,
        Err(_) => return Ok(MIN_TICK),
    };

    tick_has_debt_data.clear_bits(map_index, byte_index, bit_index)?;

    loop {
        let (next_tick, has_next_tick) =
            tick_has_debt_data.fetch_next_top_tick(current_map_index)?;

        if has_next_tick && next_tick != MIN_TICK {
            return Ok(next_tick);
        }

        // No bits found in current array, move to previous array (lower ticks)
        if current_array_index == 0 {
            return Ok(MIN_TICK);
        }

        current_array_index -= 1;
        current_map_index = TICK_HAS_DEBT_ARRAY_SIZE - 1;

        tick_has_debt_data = match program
            .account::<TickHasDebtArray>(pda::get_tick_has_debt(vault_id, current_array_index))
        {
            Ok(data) => data,
            Err(_) => return Ok(MIN_TICK),
        };
    }
}
pub fn load_relevant_ticks_for_liquidate(
    vault_id: u16,
    vault_state: VaultState,
    liquidation_tick: i32,
    cluster: &Cluster,
) -> anyhow::Result<(Vec<Tick>, i32)> {
    let provider = Client::new_with_options(
        cluster.clone(),
        Rc::new(Keypair::new()),
        CommitmentConfig::confirmed(),
    );
    let program = provider.program(VAULTS_PROGRAM_ID)?;

    let mut ticks = vec![];

    let top_tick = vault_state.topmost_tick;

    if top_tick > liquidation_tick {
        if let Ok(current_tick) = program.account::<Tick>(pda::get_tick(vault_id, top_tick)) {
            ticks.push(current_tick);
        }
    }

    let mut next_tick = find_next_tick_with_debt(vault_id, top_tick, &cluster)?;

    if !ticks.is_empty() {
        while next_tick > liquidation_tick && !ticks.iter().any(|t| t.tick == next_tick) {
            if let Ok(next_tick_data) = program.account::<Tick>(pda::get_tick(vault_id, next_tick))
            {
                ticks.push(next_tick_data);
                if let Ok(tick) = find_next_tick_with_debt(vault_id, next_tick, &cluster) {
                    next_tick = tick
                } else {
                    break;
                }
            } else {
                break;
            }
        }
    }

    Ok((ticks, next_tick))
}

pub fn load_relevant_ticks_has_debt_for_liquidate(
    vault_id: u16,
    top_tick: i32,
    next_tick: i32,
    cluster: &Cluster,
) -> anyhow::Result<Vec<TickHasDebtArray>> {
    let provider = Client::new_with_options(
        cluster.clone(),
        Arc::new(Keypair::new()),
        CommitmentConfig::confirmed(),
    );
    let program = provider.program(VAULTS_PROGRAM_ID)?;

    let (top_tick_index, _, _, _) = get_tick_indices(top_tick)?;
    let (next_tick_index, _, _, _) = get_tick_indices(next_tick)?;

    let indices: Vec<u8> = (next_tick_index..=top_tick_index).rev().collect();

    let tick_has_debt_array: Vec<TickHasDebtArray> = indices
        .par_iter()
        .flat_map(|&arr_idx| {
            program.account::<TickHasDebtArray>(pda::get_tick_has_debt(vault_id, arr_idx))
        })
        .collect();

    Ok(tick_has_debt_array)
}