sol-trade-sdk 5.0.5

use crate::{
    common::SolanaRpcClient,
    instruction::utils::raydium_cpmm_types::{
        amm_config_decode, pool_state_decode, AmmConfig, PoolState, AMM_CONFIG_DISCRIMINATOR,
        AMM_CONFIG_SIZE, POOL_STATE_DISCRIMINATOR, POOL_STATE_SIZE,
    },
    trading::core::params::RaydiumCpmmParams,
};
use anyhow::anyhow;
use solana_sdk::pubkey::Pubkey;

/// Constants used as seeds for deriving PDAs (Program Derived Addresses)
pub mod seeds {
    pub const POOL_SEED: &[u8] = b"pool";
    pub const POOL_VAULT_SEED: &[u8] = b"pool_vault";
    pub const OBSERVATION_STATE_SEED: &[u8] = b"observation";
}

/// Constants related to program accounts and authorities
pub mod accounts {
    use solana_sdk::{pubkey, pubkey::Pubkey};
    pub const AUTHORITY: Pubkey = pubkey!("GpMZbSM2GgvTKHJirzeGfMFoaZ8UR2X7F4v8vHTvxFbL");
    pub const RAYDIUM_CPMM: Pubkey = pubkey!("CPMMoo8L3F4NbTegBCKVNunggL7H1ZpdTHKxQB5qKP1C");
    pub const FEE_RATE_DENOMINATOR_VALUE: u128 = 1_000_000;
    pub const TRADE_FEE_RATE: u64 = 2500;
    pub const CREATOR_FEE_RATE: u64 = 0;
    pub const PROTOCOL_FEE_RATE: u64 = 120000;
    pub const FUND_FEE_RATE: u64 = 40000;
    // META
    pub const AUTHORITY_META: solana_sdk::instruction::AccountMeta =
        solana_sdk::instruction::AccountMeta {
            pubkey: AUTHORITY,
            is_signer: false,
            is_writable: false,
        };
}

pub const SWAP_BASE_IN_DISCRIMINATOR: &[u8] = &[143, 190, 90, 218, 196, 30, 51, 222];
pub const SWAP_BASE_OUT_DISCRIMINATOR: &[u8] = &[55, 217, 98, 86, 163, 74, 180, 173];

pub async fn fetch_pool_state(
    rpc: &SolanaRpcClient,
    pool_address: &Pubkey,
) -> Result<PoolState, anyhow::Error> {
    let account = rpc.get_account(pool_address).await?;
    if account.owner != accounts::RAYDIUM_CPMM {
        return Err(anyhow!("Account is not owned by Raydium Cpmm program"));
    }
    let expected_len = 8 + POOL_STATE_SIZE;
    if account.data.len() < expected_len {
        return Err(anyhow!(
            "Raydium CPMM pool account is too short: expected at least {}, got {}",
            expected_len,
            account.data.len()
        ));
    }
    if account.data[..8] != POOL_STATE_DISCRIMINATOR {
        return Err(anyhow!("Account discriminator is not Raydium CPMM PoolState"));
    }
    let pool_state = pool_state_decode(&account.data[8..])
        .ok_or_else(|| anyhow!("Failed to decode pool state"))?;
    Ok(pool_state)
}

pub async fn fetch_amm_config(
    rpc: &SolanaRpcClient,
    config_address: &Pubkey,
) -> Result<AmmConfig, anyhow::Error> {
    let account = rpc.get_account(config_address).await?;
    if account.owner != accounts::RAYDIUM_CPMM {
        return Err(anyhow!("Account is not owned by Raydium CPMM program"));
    }
    let expected_len = 8 + AMM_CONFIG_SIZE;
    if account.data.len() < expected_len {
        return Err(anyhow!(
            "Raydium CPMM config account is too short: expected at least {}, got {}",
            expected_len,
            account.data.len()
        ));
    }
    if account.data[..8] != AMM_CONFIG_DISCRIMINATOR {
        return Err(anyhow!("Account discriminator is not Raydium CPMM AmmConfig"));
    }
    amm_config_decode(&account.data[8..])
        .ok_or_else(|| anyhow!("Failed to decode Raydium CPMM config"))
}

pub fn get_pool_pda(amm_config: &Pubkey, mint1: &Pubkey, mint2: &Pubkey) -> Option<Pubkey> {
    let seeds: &[&[u8]; 4] =
        &[seeds::POOL_SEED, amm_config.as_ref(), mint1.as_ref(), mint2.as_ref()];
    let program_id: &Pubkey = &accounts::RAYDIUM_CPMM;
    let pda: Option<(Pubkey, u8)> = Pubkey::try_find_program_address(seeds, program_id);
    pda.map(|pubkey| pubkey.0)
}

pub fn get_vault_pda(pool_state: &Pubkey, mint: &Pubkey) -> Option<Pubkey> {
    let seeds: &[&[u8]; 3] = &[seeds::POOL_VAULT_SEED, pool_state.as_ref(), mint.as_ref()];
    let program_id: &Pubkey = &accounts::RAYDIUM_CPMM;
    let pda: Option<(Pubkey, u8)> = Pubkey::try_find_program_address(seeds, program_id);
    pda.map(|pubkey| pubkey.0)
}

pub fn get_observation_state_pda(pool_state: &Pubkey) -> Option<Pubkey> {
    let seeds: &[&[u8]; 2] = &[seeds::OBSERVATION_STATE_SEED, pool_state.as_ref()];
    let program_id: &Pubkey = &accounts::RAYDIUM_CPMM;
    let pda: Option<(Pubkey, u8)> = Pubkey::try_find_program_address(seeds, program_id);
    pda.map(|pubkey| pubkey.0)
}

/// Get the balances of two tokens in the pool
///
/// # Returns
/// Returns token0_balance, token1_balance
pub async fn get_pool_token_balances(
    rpc: &SolanaRpcClient,
    pool_state: &Pubkey,
    token0_mint: &Pubkey,
    token1_mint: &Pubkey,
) -> Result<(u64, u64), anyhow::Error> {
    let token0_vault = get_vault_pda(pool_state, token0_mint)
        .ok_or_else(|| anyhow!("Failed to derive Raydium CPMM token0 vault"))?;
    let token1_vault = get_vault_pda(pool_state, token1_mint)
        .ok_or_else(|| anyhow!("Failed to derive Raydium CPMM token1 vault"))?;
    get_pool_token_balances_from_vaults(rpc, &token0_vault, &token1_vault).await
}

/// Reads pool balances from the vault addresses stored in PoolState.
///
/// Prefer `getMultipleAccounts` over `getTokenAccountBalance` so public RPCs
/// that gate indexed methods (PublicNode/Allnodes) still work.
pub async fn get_pool_token_balances_from_vaults(
    rpc: &SolanaRpcClient,
    token0_vault: &Pubkey,
    token1_vault: &Pubkey,
) -> Result<(u64, u64), anyhow::Error> {
    crate::trading::common::get_multi_token_balances(rpc, token0_vault, token1_vault).await
}

/// Calculate token price (token1/token0)
///
/// # Returns
/// Returns the price of token1 relative to token0
pub async fn calculate_price(
    token0_amount: u64,
    token1_amount: u64,
    mint0_decimals: u8,
    mint1_decimals: u8,
) -> Result<f64, anyhow::Error> {
    if token0_amount == 0 {
        return Err(anyhow!("Token0 balance is zero, cannot calculate price"));
    }
    // Consider decimal precision
    let token0_adjusted = token0_amount as f64 / 10_f64.powi(mint0_decimals as i32);
    let token1_adjusted = token1_amount as f64 / 10_f64.powi(mint1_decimals as i32);
    let price = token1_adjusted / token0_adjusted;
    Ok(price)
}

/// Helper function to get token vault account address
///
/// # Parameters
/// - `pool_state`: Pool state account address
/// - `token_mint`: Token mint address
/// - `protocol_params`: Protocol parameters
///
/// # Returns
/// Returns the corresponding token vault account address
pub fn get_vault_account(
    pool_state: &Pubkey,
    token_mint: &Pubkey,
    protocol_params: &RaydiumCpmmParams,
) -> Result<Pubkey, anyhow::Error> {
    if protocol_params.base_mint == *token_mint && protocol_params.base_vault != Pubkey::default() {
        Ok(protocol_params.base_vault)
    } else if protocol_params.quote_mint == *token_mint
        && protocol_params.quote_vault != Pubkey::default()
    {
        Ok(protocol_params.quote_vault)
    } else {
        get_vault_pda(pool_state, token_mint)
            .ok_or_else(|| anyhow!("Failed to derive Raydium CPMM vault"))
    }
}