raydium-launchlab-sdk 0.1.6

use anyhow::Result;
use borsh::BorshSerialize;
use raydium_launchlab::accounts::GlobalConfig;
use raydium_launchlab::{
    accounts::PoolState, instructions, math, pda, PROGRAM_ID, SYSTEM_PROGRAM_ID, TOKEN_PROGRAM_ID,
};
use solana_client::rpc_client::RpcClient;
use solana_sdk::{
    commitment_config::CommitmentConfig,
    instruction::{AccountMeta, Instruction},
    pubkey::Pubkey,
    signature::{Keypair, Signer},
    system_instruction, system_program,
    transaction::Transaction,
};
use spl_associated_token_account::{
    get_associated_token_address, instruction::create_associated_token_account_idempotent,
};
use spl_token::instruction as token_instruction;
use std::str::FromStr;

const WSOL_MINT: &str = "So11111111111111111111111111111111111111112";
const USD1_MINT: &str = "USD1ttGY1N17NEEHLmELoaybftRBUSErhqYiQzvEmuB";
const USDC_MINT: &str = "EPjFWdd5AufqSSqeM2qN1xzybapC8G4wEGGkZwyTDt1v";
const USDC: &str = "EPjFWdd5AufqSSqeM2qN1xzybapC8G4wEGGkZwyTDt1v";

pub fn get_creator_fee_vault_pda(creator: &Pubkey, quote_mint: &Pubkey) -> (Pubkey, u8) {
    Pubkey::find_program_address(&[creator.as_ref(), quote_mint.as_ref()], &PROGRAM_ID)
}

pub fn get_platform_fee_vault_pda(platform_config: &Pubkey, quote_mint: &Pubkey) -> (Pubkey, u8) {
    Pubkey::find_program_address(
        &[platform_config.as_ref(), quote_mint.as_ref()],
        &PROGRAM_ID,
    )
}

fn main() -> Result<()> {
    let rpc_url = "https://api.mainnet-beta.solana.com";
    let client = RpcClient::new_with_commitment(rpc_url, CommitmentConfig::confirmed());

    let private_key_base58 = "";
    let payer = if private_key_base58.is_empty() {
        Keypair::new()
    } else {
        Keypair::from_base58_string(private_key_base58)
    };
    println!("Payer pubkey: {}", payer.pubkey());

    let base_mint = Pubkey::from_str("AmTben3tShkszjRPtiWhSPpafHJHPCKygxrwRkuZN92i")?;

    let quote_amount = 1_000_000_0; // if wsol then 9 decimals if usd1 or usdc then 6 decimals

    let (instructions, pool_state) =
        build_buy_instruction(&client, &payer, &base_mint, quote_amount, 0)?;
    println!("instructions: {:?}", instructions);
    println!("pool_state: {:?}", pool_state);
    let recent_blockhash = client.get_latest_blockhash()?;
    let tx = Transaction::new_signed_with_payer(
        &instructions,
        Some(&payer.pubkey()),
        &[&payer],
        recent_blockhash,
    );

    println!(
        "\nSending Buy transaction with {} instructions...",
        instructions.len()
    );
    // COMMENTING OUT SEND for safety in example, or keep it inside match if user wants to run

    match client.send_and_confirm_transaction(&tx) {
        Ok(signature) => {
            println!("\n✅ Buy Transaction confirmed!");
            println!("Signature: {}", signature);
            println!("Explorer: https://solscan.io/tx/{}", signature);

            // NOW SEND ATOMIC BUY AND SELL
            println!("\n--- Sending Atomic Buy + Sell Transaction ---");

            let estimated_tokens_received = math::calculate_buy_amount(
                quote_amount,
                pool_state.virtual_base as u128,
                pool_state.virtual_quote as u128,
                pool_state.real_base as u128,
                pool_state.real_quote as u128,
                100_000_000,
                0,
                100, // 1% slippage
            );

            let estimated_sol_received = math::calculate_sell_amount(
                estimated_tokens_received,
                pool_state.virtual_base as u128,
                pool_state.virtual_quote as u128,
                pool_state.real_base as u128,
                pool_state.real_quote as u128,
                100_000_000,
                0,
                100, // 1% slippage
            );
            println!(
                "Estimated SOL received from Sell: {} ({:.2} SOL)",
                estimated_sol_received,
                estimated_sol_received as f64 / 1_000_000_000.0
            );

            // Generate Sell Instructions
            let sell_instructions = build_sell_instructions(
                &client,
                &payer,
                &base_mint,
                estimated_tokens_received,
                estimated_sol_received,
            )?;

            // Combine Instructions: Buy + Sell
            let mut combined_instructions = sell_instructions;

            // Get fresh blockhash for combined tx
            let fresh_blockhash = client.get_latest_blockhash()?;
            let tx_combined = Transaction::new_signed_with_payer(
                &combined_instructions,
                Some(&payer.pubkey()),
                &[&payer],
                fresh_blockhash,
            );

            println!(
                "\nSending Combined Transaction ({} instructions)...",
                combined_instructions.len()
            );
            match client.send_and_confirm_transaction(&tx_combined) {
                Ok(sig) => {
                    println!("✅ Combined Buy+Sell Transaction confirmed!");
                    println!("Signature: {}", sig);
                }
                Err(e) => {
                    println!("❌ Combined Transaction failed: {:?}", e);
                }
            }
        }
        Err(e) => {
            println!("❌ Buy Transaction failed: {:?}", e);
        }
    }

    // For verifying compilation, I will just call the math functions to ensure they link
    let estimated_tokens_received = math::calculate_buy_amount(
        quote_amount,
        pool_state.virtual_base as u128,
        pool_state.virtual_quote as u128,
        pool_state.real_base as u128,
        pool_state.real_quote as u128,
        100_000_000,
        0,
        100, // 1% slippage
    );
    println!("Estimated tokens: {}", estimated_tokens_received);

    let estimated_sol_received = math::calculate_sell_amount(
        estimated_tokens_received,
        pool_state.virtual_base as u128,
        pool_state.virtual_quote as u128,
        pool_state.real_base as u128,
        pool_state.real_quote as u128,
        100_000_000,
        0,
        100, // 1% slippage
    );
    println!("Estimated SOL back: {}", estimated_sol_received);

    Ok(())
}

fn build_buy_instruction(
    client: &RpcClient,
    payer: &Keypair,
    base_mint: &Pubkey,
    quote_amount: u64,
    minimum_amount_out: u64,
) -> Result<(Vec<Instruction>, PoolState)> {
    let quote_mint_usd1 = Pubkey::from_str(USD1_MINT)?;
    let quote_mint_wsol = Pubkey::from_str(WSOL_MINT)?;
    let quote_mint_usdc = Pubkey::from_str(USDC_MINT)?;

    let (pool_state_usd1, _) = pda::get_pool_state_pda(base_mint, &quote_mint_usd1);
    let (pool_state_wsol, _) = pda::get_pool_state_pda(base_mint, &quote_mint_wsol);
    let (pool_state_usdc, _) = pda::get_pool_state_pda(base_mint, &quote_mint_usdc);

    let (pool_state_address, quote_mint, is_wsol_pool) =
        if client.get_account_data(&pool_state_usd1).is_ok() {
            println!("Found USD1 pool: {}", pool_state_usd1);
            (pool_state_usd1, quote_mint_usd1, false)
        } else if client.get_account_data(&pool_state_wsol).is_ok() {
            println!("Found WSOL pool: {}", pool_state_wsol);
            (pool_state_wsol, quote_mint_wsol, true)
        } else if client.get_account_data(&pool_state_usdc).is_ok() {
            println!("Found USDC pool: {}", pool_state_usdc);
            (pool_state_usdc, quote_mint_usdc, false)
        } else {
            anyhow::bail!("No pool found for this token");
        };

    println!("is_wsol_pool = {}", is_wsol_pool);

    let mut instructions = Vec::new();

    let pool_state_data = client.get_account_data(&pool_state_address)?;
    let pool_state = PoolState::try_from_bytes(&pool_state_data)?;

    let actual_minimum = if minimum_amount_out == 0 {
        math::calculate_buy_amount(
            quote_amount as u64,
            pool_state.virtual_base as u128,
            pool_state.virtual_quote as u128,
            pool_state.real_base as u128,
            pool_state.real_quote as u128,
            100_000_000, // 1%
            0,
            100, // 1%
        )
    } else {
        minimum_amount_out
    };

    println!("Actual minimum amount out: {}", actual_minimum / 1000000);

    let base_uses_token_2022 = (pool_state.token_program_flag & 0x01) != 0;
    let quote_uses_token_2022 = (pool_state.token_program_flag & 0x02) != 0;

    let base_token_program = if base_uses_token_2022 {
        Pubkey::from_str("TokenzQdBNbLqP5VEhdkAS6EPFLC1PHnBqCXEpPxuEb")?
    } else {
        TOKEN_PROGRAM_ID
    };
    let quote_token_program = if quote_uses_token_2022 {
        Pubkey::from_str("TokenzQdBNbLqP5VEhdkAS6EPFLC1PHnBqCXEpPxuEb")?
    } else {
        TOKEN_PROGRAM_ID
    };

    let user_base_ata = if base_uses_token_2022 {
        spl_associated_token_account::get_associated_token_address_with_program_id(
            &payer.pubkey(),
            base_mint,
            &base_token_program,
        )
    } else {
        get_associated_token_address(&payer.pubkey(), base_mint)
    };
    let user_quote_ata = get_associated_token_address(&payer.pubkey(), &quote_mint);

    instructions.push(create_associated_token_account_idempotent(
        &payer.pubkey(),
        &payer.pubkey(),
        base_mint,
        &base_token_program,
    ));

    instructions.push(create_associated_token_account_idempotent(
        &payer.pubkey(),
        &payer.pubkey(),
        &quote_mint,
        &quote_token_program,
    ));

    if is_wsol_pool {
        instructions.push(system_instruction::transfer(
            &payer.pubkey(),
            &user_quote_ata,
            quote_amount,
        ));
        instructions.push(token_instruction::sync_native(
            &TOKEN_PROGRAM_ID,
            &user_quote_ata,
        )?);
    }

    let (base_vault, _) = pda::get_pool_vault_pda(&pool_state_address, base_mint);
    let (quote_vault, _) = pda::get_pool_vault_pda(&pool_state_address, &quote_mint);
    let (creator_fee_vault, _) = pda::get_creator_fee_vault_pda(&pool_state.creator, &quote_mint);
    let (platform_fee_vault, _) =
        pda::get_platform_fee_vault_pda(&pool_state.platform_config, &quote_mint);

    let mut buy_ix = instructions::buy_exact_in(
        &payer.pubkey(),
        &pool_state.global_config,
        &pool_state.platform_config,
        &pool_state_address,
        &user_base_ata,
        &user_quote_ata,
        &base_vault,
        &quote_vault,
        base_mint,
        &quote_mint,
        &base_token_program,
        quote_amount,
        actual_minimum,
        0,
    );

    buy_ix.accounts[12] = AccountMeta::new_readonly(quote_token_program, false);

    buy_ix
        .accounts
        .push(AccountMeta::new_readonly(SYSTEM_PROGRAM_ID, false));
    buy_ix
        .accounts
        .push(AccountMeta::new(platform_fee_vault, false));
    buy_ix
        .accounts
        .push(AccountMeta::new(creator_fee_vault, false));

    instructions.push(buy_ix);

    Ok((instructions, pool_state))
}

fn build_sell_instructions(
    client: &RpcClient,
    payer: &Keypair,
    base_mint: &Pubkey,
    total_tokens_to_sell: u64,
    minimum_amount_out: u64,
) -> anyhow::Result<Vec<Instruction>> {
    let quote_mint_usd1 = Pubkey::from_str(USD1_MINT)?;
    let quote_mint_wsol = Pubkey::from_str(WSOL_MINT)?;
    let quote_mint_usdc = Pubkey::from_str(USDC)?;

    let (pool_state_usd1, _) = pda::get_pool_state_pda(base_mint, &quote_mint_usd1);
    let (pool_state_wsol, _) = pda::get_pool_state_pda(base_mint, &quote_mint_wsol);
    let (pool_state_usdc, _) = pda::get_pool_state_pda(base_mint, &quote_mint_usdc);

    let (pool_state_address, quote_mint, is_wsol_pool) =
        if client.get_account_data(&pool_state_usd1).is_ok() {
            println!("Found USD1 pool for selling: {}", pool_state_usd1);
            (pool_state_usd1, quote_mint_usd1, false)
        } else if client.get_account_data(&pool_state_wsol).is_ok() {
            println!("Found WSOL pool for selling: {}", pool_state_wsol);
            (pool_state_wsol, quote_mint_wsol, true)
        } else if client.get_account_data(&pool_state_usdc).is_ok() {
            println!("Found USDC pool for selling: {}", pool_state_usdc);
            (pool_state_usdc, quote_mint_usdc, false)
        } else {
            anyhow::bail!("No pool found for this token");
        };

    let pool_state_data = client.get_account_data(&pool_state_address)?;
    let pool_state = PoolState::try_from_bytes(&pool_state_data)?;

    let base_uses_token_2022 = (pool_state.token_program_flag & 0x01) != 0;
    let quote_uses_token_2022 = (pool_state.token_program_flag & 0x02) != 0;

    let base_token_program = if base_uses_token_2022 {
        Pubkey::from_str("TokenzQdBNbLqP5VEhdkAS6EPFLC1PHnBqCXEpPxuEb")?
    } else {
        TOKEN_PROGRAM_ID
    };
    let quote_token_program = if quote_uses_token_2022 {
        Pubkey::from_str("TokenzQdBNbLqP5VEhdkAS6EPFLC1PHnBqCXEpPxuEb")?
    } else {
        TOKEN_PROGRAM_ID
    };

    let estimated_out = math::calculate_sell_amount(
        total_tokens_to_sell,
        pool_state.virtual_base as u128,
        pool_state.virtual_quote as u128,
        pool_state.real_base as u128,
        pool_state.real_quote as u128,
        100_000_000,
        0,
        100,
    );

    println!(
        "Estimated output: {} (decimals={}, min_out={})",
        estimated_out, pool_state.quote_decimals, minimum_amount_out
    );

    let user_base_token = if base_uses_token_2022 {
        spl_associated_token_account::get_associated_token_address_with_program_id(
            &payer.pubkey(),
            base_mint,
            &base_token_program,
        )
    } else {
        get_associated_token_address(&payer.pubkey(), base_mint)
    };

    let user_quote_token = if quote_uses_token_2022 {
        spl_associated_token_account::get_associated_token_address_with_program_id(
            &payer.pubkey(),
            &quote_mint,
            &quote_token_program,
        )
    } else {
        get_associated_token_address(&payer.pubkey(), &quote_mint)
    };

    let (base_vault, _) = pda::get_pool_vault_pda(&pool_state_address, base_mint);
    let (quote_vault, _) = pda::get_pool_vault_pda(&pool_state_address, &quote_mint);
    let (creator_fee_vault, _) = pda::get_creator_fee_vault_pda(&pool_state.creator, &quote_mint);
    let (platform_fee_vault, _) =
        pda::get_platform_fee_vault_pda(&pool_state.platform_config, &quote_mint);

    let mut sell_ix = instructions::sell_exact_in(
        &payer.pubkey(),
        &pool_state.global_config,
        &pool_state.platform_config,
        &pool_state_address,
        &user_base_token,
        &user_quote_token,
        &base_vault,
        &quote_vault,
        base_mint,
        &quote_mint,
        &base_token_program,
        total_tokens_to_sell,
        minimum_amount_out,
        0,
    );

    sell_ix.accounts[12] = AccountMeta::new_readonly(quote_token_program, false);

    sell_ix
        .accounts
        .push(AccountMeta::new_readonly(system_program::id(), false));
    sell_ix
        .accounts
        .push(AccountMeta::new(platform_fee_vault, false));
    sell_ix
        .accounts
        .push(AccountMeta::new(creator_fee_vault, false));

    let mut instructions = vec![sell_ix];
    if is_wsol_pool {
        instructions.push(token_instruction::close_account(
            &TOKEN_PROGRAM_ID,
            &user_quote_token,
            &payer.pubkey(),
            &payer.pubkey(),
            &[],
        )?);
    }
    Ok(instructions)
}