wp-solana-test-utils 0.1.1

//! Raydium CLMM online fixture setup (requires RPC)

use std::str::FromStr;

use anyhow::Result;
use solana_client::nonblocking::rpc_client::RpcClient;
use solana_sdk::{account::Account, pubkey::Pubkey};
use wp_solana_raydium_clmm_client::{
    generated::accounts::PoolState,
    pda::{get_tick_array_address, get_tick_array_start_tick_index},
};
use wp_solana_test_core::{
    new_test_context, set_account, svm_builders::build_spl_token_account, TestContext,
};

use super::{
    offline::RaydiumClmmAccounts,
    types::{metadata_program_id, raydium_clmm_program_id},
};
use crate::internal::{fetch_and_add_program, setup_mint_and_fund};

/// Tick array size in Raydium CLMM (number of ticks per array)
const TICK_ARRAY_SIZE: i32 = 60;

/// Setup Raydium CLMM test fixture with LiteSVM, pool, USDC mint, and WSOL
/// mint.
///
/// This function sets up a test environment by fetching live state from RPC,
/// including:
/// - [`TestContext`] with LiteSVM + mock RPC + funded payer
/// - Raydium CLMM program
/// - Pool state account
/// - Token mints (USDC, WSOL)
/// - Token accounts for the payer
/// - Tick arrays around the current tick
pub async fn setup_raydium_clmm_fixture_online(
    pool_address: Option<Pubkey>,
) -> Result<(TestContext, RaydiumClmmAccounts)> {
    let ctx = new_test_context()?;

    // Get Raydium CLMM program ID
    let program_id = raydium_clmm_program_id();
    // Get Metaplex Metadata program ID
    let meta_id = metadata_program_id();

    // Create RPC client (use environment variable or default to mainnet)
    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);

    // Fetch and add both programs in parallel
    tracing::info!("Fetching Raydium CLMM and Metaplex Metadata programs from RPC in parallel");
    let (clmm_result, meta_result) = tokio::join!(
        fetch_and_add_program(&ctx, &client, program_id, "target/deploy/raydium-clmm.so",),
        fetch_and_add_program(&ctx, &client, meta_id, "target/deploy/metaplex-metadata.so",),
    );
    clmm_result?;
    meta_result?;

    // Use provided pool address or default
    let pool_address = pool_address.unwrap_or_else(|| {
        Pubkey::from_str("3ucNos4NbumPLZNWztqGHNFFgkHeRMBQAVemeeomsUxv")
            .unwrap_or_else(|_| Pubkey::new_unique())
    });

    // USDC mint address on Solana mainnet
    const USDC_MINT: &str = "EPjFWdd5AufqSSqeM2qN1xzybapC8G4wEGGkZwyTDt1v";
    let usdc_mint = Pubkey::from_str(USDC_MINT)
        .map_err(|e| anyhow::anyhow!("Invalid USDC mint address: {}", e))?;

    // WSOL mint address on Solana mainnet
    const WSOL_MINT: &str = "So11111111111111111111111111111111111111112";
    let wsol_mint = Pubkey::from_str(WSOL_MINT)
        .map_err(|e| anyhow::anyhow!("Invalid WSOL mint address: {}", e))?;

    // Fetch pool account, USDC mint, and WSOL mint in parallel
    tracing::info!(
        %pool_address,
        "Fetching pool, USDC mint, and WSOL mint accounts in parallel"
    );
    let (pool_account_result, usdc_mint_account_result, wsol_mint_account_result) = tokio::join!(
        client.get_account(&pool_address),
        client.get_account(&usdc_mint),
        client.get_account(&wsol_mint),
    );

    let pool_account =
        pool_account_result.map_err(|e| anyhow::anyhow!("Failed to fetch pool account: {}", e))?;
    let usdc_mint_account = usdc_mint_account_result
        .map_err(|e| anyhow::anyhow!("Failed to fetch USDC mint account: {}", e))?;
    let wsol_mint_account = wsol_mint_account_result
        .map_err(|e| anyhow::anyhow!("Failed to fetch WSOL mint account: {}", e))?;

    let pool_account_data = pool_account.data.clone();

    // Parse pool state to get current_tick, tick_spacing, and token vaults
    let pool_state = PoolState::from_bytes(pool_account_data.as_slice())
        .map_err(|e| anyhow::anyhow!("Failed to parse Raydium CLMM pool account: {:?}", e))?;

    // Add pool account to LiteSVM
    set_account(&ctx, pool_address, pool_account)?;
    tracing::info!("Pool account added successfully");

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

    // Setup USDC mint, override authority, and fund payer
    tracing::info!("Setting up USDC mint and token account");
    let usdc_amount: u64 = 1_000_000_000_000_000;
    let payer_usdc_ata =
        setup_mint_and_fund(&ctx, usdc_mint, usdc_mint_account, false, usdc_amount)?;

    // Setup WSOL mint and fund payer
    tracing::info!("Setting up WSOL mint and token account");
    let wsol_amount: u64 = 10_000_000_000; // 10 SOL
    let payer_wsol_ata =
        setup_mint_and_fund(&ctx, wsol_mint, wsol_mint_account, true, wsol_amount)?;

    tracing::info!(
        usdc_mint = %usdc_mint,
        payer_usdc_ata = %payer_usdc_ata,
        wsol_mint = %wsol_mint,
        payer_wsol_ata = %payer_wsol_ata,
        "Setup complete"
    );

    // Extract pool state information
    let current_tick = pool_state.tick_current;
    let tick_spacing = pool_state.tick_spacing;
    let token_vault0 = pool_state.token_vault0;
    let token_vault1 = pool_state.token_vault1;
    let rewards = pool_state.reward_infos;

    // Fetch token vault accounts in parallel
    tracing::info!("Fetching pool token vault accounts in parallel");
    let (vault0_result, vault1_result) =
        tokio::join!(client.get_account(&token_vault0), client.get_account(&token_vault1),);

    // Setup token vault 0
    match vault0_result {
        Ok(vault_account) => {
            set_account(&ctx, token_vault0, vault_account)?;
            tracing::info!(vault = %token_vault0, "Token vault 0 added");
        }
        Err(e) => {
            tracing::warn!(
                error = %e,
                vault = %token_vault0,
                "Failed to fetch token vault 0; creating fallback"
            );
            let fallback = build_spl_token_account(pool_state.token_mint0, pool_address, 0);
            set_account(&ctx, token_vault0, fallback)?;
            tracing::info!(
                vault = %token_vault0,
                "Fallback token vault 0 created"
            );
        }
    }

    // Setup token vault 1
    match vault1_result {
        Ok(vault_account) => {
            set_account(&ctx, token_vault1, vault_account)?;
            tracing::info!(vault = %token_vault1, "Token vault 1 added");
        }
        Err(e) => {
            tracing::warn!(
                error = %e,
                vault = %token_vault1,
                "Failed to fetch token vault 1; creating fallback"
            );
            let fallback = build_spl_token_account(pool_state.token_mint1, pool_address, 0);
            set_account(&ctx, token_vault1, fallback)?;
            tracing::info!(
                vault = %token_vault1,
                "Fallback token vault 1 created"
            );
        }
    }

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

    for (index, reward_info) in rewards.iter().enumerate() {
        if reward_info.reward_state == 0 {
            continue;
        }

        let reward_mint = reward_info.token_mint;
        let reward_vault = reward_info.token_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,
                    vault = %reward_vault,
                    "Failed to fetch reward vault; creating fallback"
                );
                let fallback = build_spl_token_account(reward_mint, pool_address, 0);
                set_account(&ctx, reward_vault, fallback)?;
                tracing::info!(
                    index,
                    vault = %reward_vault,
                    "Fallback reward vault created"
                );
            }
        }

        if reward_mint != pool_state.token_mint0 && reward_mint != pool_state.token_mint1 {
            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"
                    );
                }
            }
        }
    }

    // Get tick array data for the pool
    tracing::info!("Fetching tick array data");

    let tick_array_step = TICK_ARRAY_SIZE * tick_spacing as i32;
    let tick_array_start_index = get_tick_array_start_tick_index(current_tick, tick_spacing);

    let tick_array_indices = [
        tick_array_start_index - tick_array_step * 2,
        tick_array_start_index - tick_array_step,
        tick_array_start_index,
        tick_array_start_index + tick_array_step,
        tick_array_start_index + tick_array_step * 2,
    ];

    let tick_array_addresses: Vec<(i32, Pubkey)> = tick_array_indices
        .iter()
        .map(|&idx| (idx, get_tick_array_address(&pool_address, idx)))
        .collect();

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

    tracing::info!("Fetching all tick arrays in parallel");
    let tick_array_results = futures::future::join_all(tick_array_futures).await;

    for ((tick_array_index, tick_array_address), result) in
        tick_array_addresses.iter().zip(tick_array_results)
    {
        tracing::debug!(
            address = %tick_array_address,
            "Tick array address"
        );

        match result {
            Ok(tick_array_account) => {
                tracing::debug!(tick_array_index, "Fetched tick array from mainnet");

                set_account(&ctx, *tick_array_address, tick_array_account)?;
                tracing::info!(
                    address = %tick_array_address,
                    tick_array_index,
                    "Tick array added"
                );
            }
            Err(e) => {
                const TICK_ARRAY_ACCOUNT_SIZE: usize = 10240;

                tracing::warn!(
                    tick_array_index,
                    error = %e,
                    "Tick array not found on mainnet"
                );
                tracing::debug!(tick_array_index, "Creating empty tick array account");

                let min_rent = {
                    let svm = ctx.lock_svm();
                    svm.minimum_balance_for_rent_exemption(TICK_ARRAY_ACCOUNT_SIZE)
                };
                let empty_data = vec![0u8; TICK_ARRAY_ACCOUNT_SIZE];

                set_account(
                    &ctx,
                    *tick_array_address,
                    Account {
                        lamports: min_rent,
                        data: empty_data,
                        owner: program_id,
                        executable: false,
                        rent_epoch: 0,
                    },
                )?;

                tracing::info!(
                    address = %tick_array_address,
                    tick_array_index,
                    size = TICK_ARRAY_ACCOUNT_SIZE,
                    rent = min_rent,
                    "Empty tick array account created"
                );
            }
        }
    }

    tracing::info!("Tick array fetching complete");

    Ok((
        ctx,
        RaydiumClmmAccounts {
            program_id,
            metadata_program_id: meta_id,
            pool_address,
            usdc_mint,
            wsol_mint,
            payer_usdc_ata,
            payer_wsol_ata,
            tick_spacing,
        },
    ))
}