sol-trade-sdk 6.0.2

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?;
    decode_amm_config_account(&account)
}

fn decode_amm_config_account(
    account: &solana_sdk::account::Account,
) -> Result<AmmConfig, anyhow::Error> {
    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"))
    }
}

/// Always include this PDA in collection instructions, even when it is absent on-chain.
pub fn get_creator_fee_share_pda(creator: &Pubkey, amm_config: &Pubkey) -> Pubkey {
    Pubkey::find_program_address(
        &[b"creator_fee_share", creator.as_ref(), amm_config.as_ref()],
        &accounts::RAYDIUM_CPMM,
    )
    .0
}

/// Resolve the current collection-time rate. Only a missing/empty or foreign-owned
/// account falls back to AmmConfig; malformed CPMM accounts are errors.
/// Config and override are fetched together from the same confirmed bank snapshot.
pub async fn fetch_creator_fee_share_rate(
    rpc: &SolanaRpcClient,
    creator: &Pubkey,
    config_address: &Pubkey,
) -> Result<u64, anyhow::Error> {
    let address = get_creator_fee_share_pda(creator, config_address);
    // One bank snapshot avoids mixing config and override values from separate reads.
    let response = rpc
        .get_multiple_accounts_with_commitment(
            &[*config_address, address],
            solana_commitment_config::CommitmentConfig::confirmed(),
        )
        .await?;
    if response.value.len() != 2 {
        return Err(anyhow!("Incomplete Raydium CPMM fee-share snapshot"));
    }
    let config_account = response.value[0]
        .as_ref()
        .ok_or_else(|| anyhow!("Raydium CPMM AmmConfig account is missing"))?;
    let config = decode_amm_config_account(config_account)?;
    resolve_creator_fee_share_rate(&config, creator, config_address, response.value[1].as_ref())
}

/// Resolve a fetched share account with the same missing-account fallback as the program.
/// Zero-lamport observations are closed, even if the source retains old data bytes.
pub fn resolve_creator_fee_share_rate(
    config: &AmmConfig,
    creator: &Pubkey,
    config_address: &Pubkey,
    share_account: Option<&solana_sdk::account::Account>,
) -> Result<u64, anyhow::Error> {
    resolve_creator_fee_share_rate_from_data(
        config,
        creator,
        config_address,
        share_account.map(|account| {
            let data = if account.lamports == 0 { &[][..] } else { account.data.as_slice() };
            (account.owner, data)
        }),
    )
}

/// Borrow subscription bytes without cloning an account on the preparation path.
pub(crate) fn resolve_creator_fee_share_rate_from_data(
    config: &AmmConfig,
    creator: &Pubkey,
    config_address: &Pubkey,
    share_account: Option<(Pubkey, &[u8])>,
) -> Result<u64, anyhow::Error> {
    use crate::instruction::utils::raydium_cpmm_types::creator_fee_share_decode;
    let rate = match share_account {
        Some((owner, data)) if owner == accounts::RAYDIUM_CPMM && !data.is_empty() => {
            let share = creator_fee_share_decode(data)
                .ok_or_else(|| anyhow!("Invalid Raydium CPMM CreatorFeeShare account"))?;
            if share.creator != *creator || share.amm_config != *config_address {
                return Err(anyhow!("CreatorFeeShare creator/config mismatch"));
            }
            share.share_rate
        }
        _ => config.creator_fee_share_rate,
    };
    if rate > 1_000_000 {
        return Err(anyhow!("Creator fee share rate exceeds 1,000,000"));
    }
    Ok(rate)
}

/// Returns (creator payout, protocol share), before any Token-2022 transfer fee.
/// The protocol share rounds down; remaining dust belongs to the creator.
pub fn split_creator_fee(creator_fee: u64, share_rate: u64) -> Result<(u64, u64), anyhow::Error> {
    if share_rate > 1_000_000 {
        return Err(anyhow!("Creator fee share rate exceeds 1,000,000"));
    }
    let protocol = (u128::from(creator_fee) * u128::from(share_rate) / 1_000_000) as u64;
    Ok((creator_fee - protocol, protocol))
}

/// Estimate both creator payouts using a freshly resolved collection-time rate.
/// Pool counters are gross accrued fees; protocol counters also grow on collection.
pub fn estimate_creator_fee_payout(
    pool: &PoolState,
    share_rate: u64,
) -> Result<(u64, u64), anyhow::Error> {
    Ok((
        split_creator_fee(pool.creator_fees_token0, share_rate)?.0,
        split_creator_fee(pool.creator_fees_token1, share_rate)?.0,
    ))
}