frag.rs

mod idl;
use anchor_lang::{AnchorDeserialize, prelude::AccountInfo};
use anyhow::{anyhow, Result};
use anyhow_ext::Context;
use idl::*;
use rust_decimal::Decimal;
use sanctum_lst_list::{
    PoolInfo::{SPool, SanctumSpl, SanctumSplMulti, Spl},
    SanctumLst, SanctumLstList, SplPoolAccounts,
};
use solana_client::nonblocking::rpc_client::RpcClient;
use solana_sdk::{
    account::Account,
    native_token::LAMPORTS_PER_SOL,
    program_pack::Pack,
    pubkey,
    pubkey::Pubkey
};
use spl_token::state::Mint;
use s_sol_val_calc_prog_aggregate::{
    KnownLstSolValCalc, LidoLstSolValCalc, LstSolValCalc, MarinadeLstSolValCalc,
    MutableLstSolValCalc, SanctumSplLstSolValCalc, SanctumSplMultiLstSolValCalc,
    SplLstSolValCalc, SplLstSolValCalcInitKeys, WsolLstSolValCalc
};
use std::{
    collections::HashMap,
    sync::{Arc, atomic::AtomicU64}
};
use crate::{get_epoch, inf_quote, to_lst, utils::account_utils};

const NSOL_POOL: Pubkey = pubkey!("AkbZvKxUAxMKz92FF7g5k2YLCJftg8SnYEPWdmZTt3mp");
const FRAGSOL_FUND: Pubkey = pubkey!("3TK9fNePM4qdKC4dwvDe8Bamv14prDqdVfuANxPeiryb");
const INF_MINT: Pubkey = pubkey!("5oVNBeEEQvYi1cX3ir8Dx5n1P7pdxydbGF2X4TxVusJm");

pub async fn simulate_nsol_live_price(client: &RpcClient) -> Result<Decimal> {
    let s = fetch_nsol_pool_account(client).await?;
    let start_epoch = client.get_epoch_info().await?.epoch;
    let SanctumLstList {
        sanctum_lst_list: lsts,
    } = SanctumLstList::load();
    let mut total_weight = Decimal::ZERO;
    let mint_account = client.get_account(&s.normalized_token_mint);
    let mut futs = Vec::new();
    for i in 0..s.supported_tokens.len() {
        let lst_mint = &s.supported_tokens[i].mint;
        futs.push(calc_from_lst(client, to_lst(&lsts, lst_mint)));
    }
    let prices = futures::future::try_join_all(futs).await?;
    let mint_account = mint_account.await?;
    for (i, lst) in s.supported_tokens.iter().enumerate() {
        let price = prices[i];
        total_weight += Decimal::from(lst.locked_amount) * price;
    }

    let nsol_mint = Mint::unpack_from_slice(&mint_account.data)?;
    let end_epoch = client.get_epoch_info().await?.epoch;
    if start_epoch != end_epoch {
        return Err(anyhow::anyhow!("Epoch changed during fragsol calculation"));
    }
    Ok(total_weight / Decimal::from(nsol_mint.supply))
}

pub async fn get_latest_fragsol_price(rpc: &RpcClient) -> Result<Decimal> {
    let fund = fetch_fragsol_fund_account(rpc).await?;
    // let epoch_schedule = rpc
        // .get_epoch_schedule()
        // .await
        // .context("failed to fetch epoch schedule")?;
    // let current_epoch = rpc
        // .get_epoch_info()
        // .await
        // .context("failed to fetch epoch info")?
        // .epoch;
    // let updated_epoch = epoch_schedule.get_epoch(fund.receipt_token_value_updated_slot);
    // if current_epoch != updated_epoch {
        // return Err(anyhow::anyhow!("fragsol token account needs epoch refresh"));
    // }
    Ok(Decimal::from_i128_with_scale(
        fund.one_receipt_token_as_sol.into(),
        9,
    ))
}

async fn fetch_nsol_pool_account(rpc: &RpcClient) -> Result<NormalizedTokenPoolAccount> {
    let account_data = rpc.get_account_data(&NSOL_POOL).await?;
    let mut account_data_slice = &account_data[8..];
    NormalizedTokenPoolAccount::deserialize(&mut account_data_slice)
        .context("failed to deserialize nsol pool account")
}

async fn fetch_fragsol_fund_account(rpc: &RpcClient) -> Result<FundAccount> {
    let account_data = rpc
        .get_account_data(&FRAGSOL_FUND)
        .await
        .context("failed to fetch fragsol fund account")?;
    let mut account_data_slice_without_discriminator = &account_data[8..];
    FundAccount::deserialize(&mut account_data_slice_without_discriminator)
        .context("failed to deserialize fragsol fund account")
}

fn to_account_info<'a>(key: &'a Pubkey, account: &'a Account) -> AccountInfo<'a> {
    account_utils::to_account_info_from_owned(key, account.clone())
}

async fn calc_from_calculator(
    client: &RpcClient,
    mut calc: KnownLstSolValCalc,
) -> Result<Decimal> {
    // Get account keys needed for calculation
    let accounts = calc.get_accounts_to_update();

    // Fetch all accounts in a single RPC call
    let account_datas: Vec<Account> = client
        .get_multiple_accounts(&accounts)
        .await
        .map_err(|_| anyhow!("Failed to get accounts"))?
        .into_iter()
        .flatten()
        .collect();

    if accounts.len() != account_datas.len() {
        return Err(anyhow!("Account data mismatch, failed to fetch all needed calc accounts"));
    }

    // Create account info map
    let mut data_map = HashMap::new();
    for i in 0..accounts.len() {
        let key = &accounts[i];
        let account_info = to_account_info(key, &account_datas[i]);
        data_map.insert(*key, account_info);
    }

    // Update calculator with account data
    calc.update(&data_map)
        .map_err(|_| anyhow!("Failed to update calc"))?;

    // Calculate LST to SOL conversion
    let range = calc
        .lst_to_sol(LAMPORTS_PER_SOL)
        .map_err(|_| anyhow!("Failed to calculate lst_to_sol"))?;

    // Verify range consistency (allow up to 1 lamport for rounding)
    if range.get_max() - range.get_min() > 1 {
        return Err(anyhow!("SanctumCalcError: Unexpected range"));
    }

    // Return as decimal
    Ok(Decimal::from(range.get_max()) / Decimal::from(LAMPORTS_PER_SOL))
}

async fn calc_from_lst(client: &RpcClient, lst: Option<&SanctumLst>) -> Result<Decimal> {
    let lst = lst.ok_or(anyhow!("SanctumLstNotFound"))?;
    let epoch = get_epoch(client).await?;

    // Create the appropriate calculator based on pool type
    let calc = match &lst.pool {
        // SPL stake pool with multiple validators
        SanctumSplMulti(SplPoolAccounts { pool, .. }) => {
            // Create calculator with initialized values
            let mut stake_pool_calc = SanctumSplMultiLstSolValCalc::from_keys(
                SplLstSolValCalcInitKeys {
                    lst_mint: lst.mint,
                    stake_pool_addr: *pool,
                },
                epoch,
            );

            // Since we just created it, we know its structure, so directly set epoch
            if let Some(c) = &stake_pool_calc.0.calc {
                stake_pool_calc.0.shared_current_epoch = Arc::new(AtomicU64::new(c.last_update_epoch));
            }

            KnownLstSolValCalc::SanctumSplMulti(stake_pool_calc)
        }

        // Sanctum SPL stake pool
        SanctumSpl(SplPoolAccounts { pool, .. }) => {
            // Create calculator with initialized values
            let mut stake_pool_calc = SanctumSplLstSolValCalc::from_keys(
                SplLstSolValCalcInitKeys {
                    lst_mint: lst.mint,
                    stake_pool_addr: *pool,
                },
                epoch,
            );

            // Since we just created it, we know its structure, so directly set epoch
            if let Some(c) = &stake_pool_calc.0.calc {
                stake_pool_calc.0.shared_current_epoch = Arc::new(AtomicU64::new(c.last_update_epoch));
            }

            KnownLstSolValCalc::SanctumSpl(stake_pool_calc)
        }

        // Standard SPL stake pool
        Spl(SplPoolAccounts { pool, .. }) => {
            // Create calculator with initialized values
            let mut stake_pool_calc = SplLstSolValCalc::from_keys(
                SplLstSolValCalcInitKeys {
                    lst_mint: lst.mint,
                    stake_pool_addr: *pool,
                },
                epoch,
            );

            // Since we just created it, we know its structure, so directly set epoch
            if let Some(c) = &stake_pool_calc.calc {
                stake_pool_calc.shared_current_epoch = Arc::new(AtomicU64::new(c.last_update_epoch));
            }

            KnownLstSolValCalc::Spl(stake_pool_calc)
        }

        // Marinade LST
        sanctum_lst_list::PoolInfo::Marinade => {
            MarinadeLstSolValCalc::default().into()
        }

        // Lido LST
        sanctum_lst_list::PoolInfo::Lido => {
            // For Lido we won't have computed_in_epoch until after update,
            // so we'll return a standard KnownLstSolValCalc with the epoch set
            LidoLstSolValCalc {
                shared_current_epoch: epoch,
                calc: None,
            }.into()
        }

        // Reserve pool (WSOL)
        sanctum_lst_list::PoolInfo::ReservePool => {
            WsolLstSolValCalc {}.into()
        }

        // Sanctum pool
        SPool(_) => {
            if lst.mint != INF_MINT {
                return Err(anyhow!("InvalidSpool"));
            }
            return inf_quote(client).await;
        }
    };

    // Calculate the LST price using the configured calculator
    calc_from_calculator(client, calc).await
}

#[cfg(test)]
mod tests {
    use super::*;

    #[tokio::test]
    async fn test_simulate_fragsol_redemption() {
        let client = RpcClient::new("https://api.mainnet-beta.solana.com".to_string());
        println!("{:#?}", get_latest_fragsol_price(&client).await.unwrap());
    }

    #[tokio::test]
    async fn test_simulate_nsol_price() {
        let client = RpcClient::new("https://api.mainnet-beta.solana.com".to_string());
        println!("{:#?}", simulate_nsol_live_price(&client).await.unwrap());
    }
}