wp-solana-test-utils 0.1.1

//! Meteora DLMM online fixture setup (requires RPC)

use std::str::FromStr;

use anyhow::Result;
use solana_client::nonblocking::rpc_client::RpcClient;
use solana_sdk::pubkey::Pubkey;
use wp_solana_meteora_dlmm_client::{
    generated::accounts::LbPair,
    pda::{
        bin_id_to_bin_array_index, derive_bin_array_address, derive_bin_array_bitmap_extension,
        get_event_authority,
    },
};
use wp_solana_test_core::{new_test_context, set_account, TestContext};

use super::{offline::MeteoraDlmmAccounts, types::meteora_dlmm_program_id};
use crate::internal::{fetch_and_add_program, setup_mint_and_fund};

/// Setup Meteora DLMM test fixture with LiteSVM, LB pair, and token mints.
///
/// This function sets up a test environment by fetching live state from RPC.
pub async fn setup_meteora_dlmm_fixture_online(
    lb_pair_address: Option<Pubkey>,
) -> Result<(TestContext, MeteoraDlmmAccounts)> {
    let ctx = new_test_context()?;

    let program_id = meteora_dlmm_program_id();

    let rpc_url = std::env::var("SOLANA_RPC_URL")
        .unwrap_or_else(|_| "https://api.mainnet-beta.solana.com".to_string());
    let client = RpcClient::new(rpc_url);

    tracing::info!("Fetching Meteora DLMM program from RPC");
    fetch_and_add_program(&ctx, &client, program_id, "target/deploy/meteora-dlmm.so").await?;

    let lb_pair_address = lb_pair_address.unwrap_or_else(|| {
        Pubkey::from_str("5rCf1DM8LjKTw4YqhnoLcngyZYeNnQqztScTogYHAS6")
            .unwrap_or_else(|_| Pubkey::new_unique())
    });

    tracing::info!(
        %lb_pair_address,
        "Fetching LB pair account from RPC"
    );
    let lb_pair_account = client
        .get_account(&lb_pair_address)
        .await
        .map_err(|e| anyhow::anyhow!("Failed to fetch LB pair account: {}", e))?;

    let lb_pair_account_data = lb_pair_account.data.clone();

    let lb_pair = LbPair::from_bytes(lb_pair_account_data.as_slice())
        .map_err(|e| anyhow::anyhow!("Failed to parse Meteora DLMM LB pair account: {:?}", e))?;

    let token_x_mint = lb_pair.token_x_mint;
    let token_y_mint = lb_pair.token_y_mint;
    let bin_step = lb_pair.bin_step;
    let reserve_x = lb_pair.reserve_x;
    let reserve_y = lb_pair.reserve_y;
    let reward_infos = lb_pair.reward_infos;
    let active_id = lb_pair.active_id;

    tracing::info!("Fetching token X and token Y mint accounts in parallel");
    let (token_x_mint_account_result, token_y_mint_account_result) =
        tokio::join!(client.get_account(&token_x_mint), client.get_account(&token_y_mint),);

    let token_x_mint_account = token_x_mint_account_result
        .map_err(|e| anyhow::anyhow!("Failed to fetch token X mint account: {}", e))?;
    let token_y_mint_account = token_y_mint_account_result
        .map_err(|e| anyhow::anyhow!("Failed to fetch token Y mint account: {}", e))?;

    set_account(&ctx, lb_pair_address, lb_pair_account)?;
    tracing::info!("LB pair account added successfully");

    tracing::info!(
        program_id = %program_id,
        %lb_pair_address,
        token_x_mint = %token_x_mint,
        token_y_mint = %token_y_mint,
        "Program is ready for execution"
    );

    // Setup token mints, override authorities, and fund payer
    let wsol_mint = spl_token::native_mint::id();
    let spl_token_amount: u64 = 1_000_000_000_000_000;
    let wsol_amount: u64 = 10_000_000_000; // 10 SOL

    tracing::info!("Setting up token X mint and token account");
    let is_token_x_native = token_x_mint == wsol_mint;
    let token_x_amount = if is_token_x_native { wsol_amount } else { spl_token_amount };
    let payer_token_x_ata = setup_mint_and_fund(
        &ctx,
        token_x_mint,
        token_x_mint_account,
        is_token_x_native,
        token_x_amount,
    )?;

    tracing::info!("Setting up token Y mint and token account");
    let is_token_y_native = token_y_mint == wsol_mint;
    let token_y_amount = if is_token_y_native { wsol_amount } else { spl_token_amount };
    let payer_token_y_ata = setup_mint_and_fund(
        &ctx,
        token_y_mint,
        token_y_mint_account,
        is_token_y_native,
        token_y_amount,
    )?;

    // Fetch reserve accounts in parallel
    tracing::info!("Fetching reserve accounts in parallel");
    let (reserve_x_result, reserve_y_result) =
        tokio::join!(client.get_account(&reserve_x), client.get_account(&reserve_y),);

    match reserve_x_result {
        Ok(reserve_account) => {
            set_account(&ctx, reserve_x, reserve_account)?;
            tracing::info!(reserve = %reserve_x, "Reserve X added");
        }
        Err(e) => {
            tracing::warn!(error = %e, "Failed to fetch reserve X");
        }
    }

    match reserve_y_result {
        Ok(reserve_account) => {
            set_account(&ctx, reserve_y, reserve_account)?;
            tracing::info!(reserve = %reserve_y, "Reserve Y added");
        }
        Err(e) => {
            tracing::warn!(error = %e, "Failed to fetch reserve Y");
        }
    }

    // Setup reward vault accounts
    tracing::info!("Setting up reward vault accounts");

    for (index, reward_info) in reward_infos.iter().enumerate() {
        if reward_info.mint == Pubkey::default() {
            continue;
        }

        let reward_mint = reward_info.mint;
        let reward_vault = reward_info.vault;

        tracing::info!(
            index,
            mint = %reward_mint,
            vault = %reward_vault,
            "Setting up reward"
        );

        match client.get_account(&reward_vault).await {
            Ok(vault_account) => {
                set_account(&ctx, reward_vault, vault_account)?;
                tracing::info!(
                    index,
                    vault = %reward_vault,
                    "Reward vault added"
                );
            }
            Err(e) => {
                tracing::warn!(
                    index,
                    error = %e,
                    "Failed to fetch reward vault"
                );
            }
        }

        if reward_mint != token_x_mint && reward_mint != token_y_mint {
            match client.get_account(&reward_mint).await {
                Ok(mint_account) => {
                    set_account(&ctx, reward_mint, mint_account)?;
                    tracing::info!(
                        index,
                        mint = %reward_mint,
                        "Reward mint added"
                    );
                }
                Err(e) => {
                    tracing::warn!(
                        index,
                        error = %e,
                        "Failed to fetch reward mint"
                    );
                }
            }
        }
    }

    // Setup bin arrays around the active bin
    tracing::info!("Setting up bin arrays");

    let active_bin_array_index = bin_id_to_bin_array_index(active_id)
        .map_err(|e| anyhow::anyhow!("Failed to calculate bin array index: {:?}", e))?;

    let bin_array_range = 3i64;
    let min_bin_array_index = active_bin_array_index - bin_array_range;
    let max_bin_array_index = active_bin_array_index + bin_array_range;

    tracing::info!(min_bin_array_index, max_bin_array_index, "Fetching bin arrays");

    let bin_array_addresses: Vec<(i64, Pubkey)> = (min_bin_array_index..=max_bin_array_index)
        .map(|idx| {
            let (addr, _) = derive_bin_array_address(&lb_pair_address, idx);
            (idx, addr)
        })
        .collect();

    let (bin_array_bitmap_extension_address, _) =
        derive_bin_array_bitmap_extension(&lb_pair_address);
    let event_authority_address = get_event_authority();

    let bin_array_futures: Vec<_> =
        bin_array_addresses.iter().map(|(_, addr)| client.get_account(addr)).collect();

    tracing::info!("Fetching bin arrays, bitmap extension, and event authority");
    let (bin_array_results, bitmap_result, event_authority_result) = tokio::join!(
        futures::future::join_all(bin_array_futures),
        client.get_account(&bin_array_bitmap_extension_address),
        client.get_account(&event_authority_address),
    );

    for ((idx, bin_array_address), result) in bin_array_addresses.iter().zip(bin_array_results) {
        match result {
            Ok(bin_array_account) => {
                set_account(&ctx, *bin_array_address, bin_array_account)?;
                tracing::info!(
                    index = idx,
                    address = %bin_array_address,
                    "Bin array added"
                );
            }
            Err(e) => {
                tracing::warn!(
                    index = idx,
                    error = %e,
                    "Bin array not found (may need initialization)"
                );
            }
        }
    }

    // Setup bin array bitmap extension
    tracing::info!("Setting up bin array bitmap extension");
    match bitmap_result {
        Ok(bitmap_account) => {
            set_account(&ctx, bin_array_bitmap_extension_address, bitmap_account)?;
            tracing::info!(
                address = %bin_array_bitmap_extension_address,
                "Bin array bitmap extension added"
            );
        }
        Err(e) => {
            tracing::warn!(
                error = %e,
                "Bin array bitmap extension not found"
            );
        }
    }

    // Setup event authority
    tracing::info!("Setting up event authority");
    match event_authority_result {
        Ok(event_authority_account) => {
            set_account(&ctx, event_authority_address, event_authority_account)?;
            tracing::info!(
                address = %event_authority_address,
                "Event authority added"
            );
        }
        Err(e) => {
            tracing::warn!(
                error = %e,
                "Event authority not found"
            );
        }
    }

    tracing::info!(
        token_x_mint = %token_x_mint,
        payer_token_x_ata = %payer_token_x_ata,
        token_y_mint = %token_y_mint,
        payer_token_y_ata = %payer_token_y_ata,
        reserve_x = %reserve_x,
        reserve_y = %reserve_y,
        "Setup complete"
    );

    Ok((
        ctx,
        MeteoraDlmmAccounts {
            program_id,
            lb_pair_address,
            token_x_mint,
            token_y_mint,
            payer_token_x_ata,
            payer_token_y_ata,
            bin_step,
        },
    ))
}