simulator-client 0.7.2

use std::{collections::BTreeMap, time::Duration};

use simulator_api::{AccountData, BinaryEncoding, EncodedBinary};
use solana_address::Address;
use solana_client::nonblocking::rpc_client::RpcClient;
use thiserror::Error;
use tokio_retry::{
    Retry,
    strategy::{ExponentialBackoff, jitter},
};
use tracing::warn;

use crate::error::err_chain;

const MODIFY_RPC_RETRIES: usize = 10;
// `ExponentialBackoff::from_millis(base)` multiplies `current` by `base` each
// step, so naive `from_millis(500)` blows up to multi-minute waits within a
// few iterations. Use `from_millis(2).factor(500)` for a 2x ramp scaled to
// 1s/2s/4s/8s..., capped by `MODIFY_RPC_MAX_DELAY` so total worst-case wait
// stays well under the server's session disconnect timeout.
const MODIFY_RPC_BACKOFF_BASE: u64 = 2;
const MODIFY_RPC_BACKOFF_FACTOR_MS: u64 = 500;
const MODIFY_RPC_MAX_DELAY: Duration = Duration::from_secs(10);

const BPF_LOADER_UPGRADEABLE: &str = "BPFLoaderUpgradeab1e11111111111111111111111";

/// Error returned by [`BacktestSession::modify_program`](crate::BacktestSession::modify_program).
#[derive(Debug, Error)]
pub enum ProgramModError {
    #[error("session has no rpc_endpoint (was the session created?)")]
    NoRpcEndpoint,

    #[error("invalid program id `{id}`")]
    InvalidProgramId { id: String },

    #[error("RPC error: {source}")]
    Rpc {
        #[source]
        source: Box<dyn std::error::Error + Send + Sync>,
    },
}

/// Build a BPF Loader Upgradeable `ProgramData` account modification from raw ELF bytes.
///
/// Returns a map of `{programdata_address: AccountData}` ready to pass to
/// [`Continue::builder().modify_accounts(...)`](crate::Continue).
///
/// The ELF is wrapped in the standard `ProgramData` header format:
/// ```text
/// [0..4]    variant = 3  (u32 LE, UpgradeableLoaderState::ProgramData)
/// [4..12]   deploy_slot  (u64 LE)
/// [12]      upgrade_authority discriminant (0 = None, 1 = Some)
/// [13..45]  upgrade_authority pubkey bytes (32 bytes, present only when Some)
/// [45..]    ELF bytecode
/// ```
///
/// `deploy_slot` should be set to `start_slot.saturating_sub(1)` so the program
/// appears deployed *before* the first executed slot. Deploying at `start_slot`
/// itself triggers the SVM's same-slot restriction, marking the program unloaded.
///
/// `lamports` must be at least rent-exempt for the resulting account size.
/// With `upgrade_authority = None` the data length is `13 + elf.len()`;
/// with `Some(authority)` it is `45 + elf.len()`.
/// Fetch the exact minimum with
/// `rpc.get_minimum_balance_for_rent_exemption(data_len).await?`.
///
/// ## Example
///
/// ```no_run
/// use simulator_client::build_program_injection;
/// use solana_address::Address;
///
/// let program_id: Address = "YourProgramId...".parse().unwrap();
/// // Compute the programdata PDA using solana_loader_v3_interface::get_program_data_address,
/// // then convert to Address.
/// let programdata_addr: Address = "ProgramDataAddr...".parse().unwrap();
/// let elf = std::fs::read("my_program.so").unwrap();
/// let deploy_slot = 399_834_991; // start_slot - 1
///
/// let mods = build_program_injection(programdata_addr, &elf, deploy_slot, None, 10_000_000_000);
/// // Pass mods to Continue::builder().modify_accounts(mods).build()
/// ```
pub fn build_program_injection(
    programdata_address: Address,
    elf: &[u8],
    deploy_slot: u64,
    upgrade_authority: Option<Address>,
    lamports: u64,
) -> BTreeMap<Address, AccountData> {
    let data = build_programdata_bytes(elf, deploy_slot, upgrade_authority.as_ref());

    let account = AccountData {
        space: data.len() as u64,
        data: EncodedBinary::from_bytes(&data, BinaryEncoding::Base64),
        executable: false,
        lamports,
        owner: BPF_LOADER_UPGRADEABLE
            .parse::<Address>()
            .expect("valid BPF loader address"),
    };

    let mut map = BTreeMap::new();
    map.insert(programdata_address, account);
    map
}

/// Serialize ELF bytes into a `ProgramData` account data blob.
///
/// Exposed as a building block if you need to construct the account yourself
/// (e.g. to set a custom lamport amount after calling
/// `rpc.get_minimum_balance_for_rent_exemption`).
pub fn build_programdata_bytes(
    elf: &[u8],
    deploy_slot: u64,
    upgrade_authority: Option<&Address>,
) -> Vec<u8> {
    let header_len = if upgrade_authority.is_some() { 45 } else { 13 };
    let mut data = Vec::with_capacity(header_len + elf.len());

    // variant = 3 (UpgradeableLoaderState::ProgramData)
    data.extend_from_slice(&3u32.to_le_bytes());
    // deployment slot
    data.extend_from_slice(&deploy_slot.to_le_bytes());

    match upgrade_authority {
        None => {
            data.push(0); // Option::None
        }
        Some(authority) => {
            data.push(1); // Option::Some
            data.extend_from_slice(authority.as_ref());
        }
    }

    data.extend_from_slice(elf);
    data
}

/// Build a program-data account modification, fetching the deploy slot and rent
/// exemption from the given RPC client.
///
/// Equivalent to [`BacktestSession::modify_program`](crate::BacktestSession::modify_program)
/// but usable without holding a session — pass any RPC client that can see the
/// current cluster state.
pub async fn modify_program_via_rpc(
    rpc: &RpcClient,
    program_id: &str,
    elf: &[u8],
) -> Result<BTreeMap<Address, AccountData>, ProgramModError> {
    let program_addr: Address =
        program_id
            .parse()
            .map_err(|_| ProgramModError::InvalidProgramId {
                id: program_id.to_string(),
            })?;
    let programdata_addr = solana_loader_v3_interface::get_program_data_address(&program_addr);

    let strategy = ExponentialBackoff::from_millis(MODIFY_RPC_BACKOFF_BASE)
        .factor(MODIFY_RPC_BACKOFF_FACTOR_MS)
        .max_delay(MODIFY_RPC_MAX_DELAY)
        .map(jitter)
        .take(MODIFY_RPC_RETRIES);

    Retry::spawn(strategy, || async {
        modify_program_via_rpc_once(rpc, programdata_addr, elf)
            .await
            .inspect_err(|e| {
                warn!(
                    program_id,
                    error = %err_chain(e),
                    "modify_program_via_rpc attempt failed"
                )
            })
    })
    .await
}

async fn modify_program_via_rpc_once(
    rpc: &RpcClient,
    programdata_addr: Address,
    elf: &[u8],
) -> Result<BTreeMap<Address, AccountData>, ProgramModError> {
    let slot = rpc.get_slot().await.map_err(|e| ProgramModError::Rpc {
        source: Box::new(e),
    })?;
    let deploy_slot = slot.saturating_sub(1);

    let existing = rpc
        .get_account(&programdata_addr)
        .await
        .map_err(|e| ProgramModError::Rpc {
            source: Box::new(e),
        })?;

    let upgrade_authority = if existing.data.get(12).copied() == Some(1) {
        existing.data.get(13..45).and_then(|b| {
            let bytes: [u8; 32] = b.try_into().ok()?;
            Some(Address::from(bytes))
        })
    } else {
        None
    };

    let data_len = upgrade_authority.map_or(13, |_| 45) + elf.len();
    let lamports = rpc
        .get_minimum_balance_for_rent_exemption(data_len)
        .await
        .map_err(|e| ProgramModError::Rpc {
            source: Box::new(e),
        })?;

    Ok(build_program_injection(
        programdata_addr,
        elf,
        deploy_slot,
        upgrade_authority,
        lamports,
    ))
}