sim-cli 0.2.0

use std::{
    collections::{BTreeSet, HashMap},
    time::Duration,
};

use solana_client::nonblocking::rpc_client::RpcClient;
use solana_commitment_config::CommitmentConfig;
use solana_rpc_client_api::config::RpcTransactionConfig;
use solana_transaction_status::{
    EncodedTransaction, UiMessage, UiTransactionEncoding, option_serializer::OptionSerializer,
};
use tracing::warn;

use crate::output::{SolChange, TokenChange, Transaction};

const FETCH_MAX_ATTEMPTS: u32 = 5;
const FETCH_BASE_DELAY_MS: u64 = 200;

/// Fetch a transaction via `getTransaction` and return a partially populated
/// `Transaction`. The `logs` and `account_diffs` fields are left empty for
/// callers to fill in.
///
/// Retries up to `FETCH_MAX_ATTEMPTS` times with exponential backoff.
/// Returns `None` if all attempts fail or the transaction has no meta.
pub async fn fetch_transaction(rpc: &RpcClient, signature: &str, slot: u64) -> Option<Transaction> {
    let sig: solana_signature::Signature = signature.parse().ok()?;
    let config = RpcTransactionConfig {
        encoding: Some(UiTransactionEncoding::Json),
        commitment: Some(CommitmentConfig::confirmed()),
        max_supported_transaction_version: Some(0),
    };

    let mut attempt = 0u32;
    let confirmed = loop {
        match rpc.get_transaction_with_config(&sig, config).await {
            Ok(tx) => break tx,
            Err(e) => {
                attempt += 1;
                if attempt >= FETCH_MAX_ATTEMPTS {
                    warn!(
                        signature,
                        slot,
                        attempts = attempt,
                        error = %e,
                        "fetch_transaction failed after all retries, dropping event"
                    );
                    return None;
                }
                let delay = Duration::from_millis(FETCH_BASE_DELAY_MS * (1 << attempt));
                warn!(
                    signature,
                    slot,
                    attempt,
                    error = %e,
                    delay_ms = delay.as_millis(),
                    "fetch_transaction failed, retrying"
                );
                tokio::time::sleep(delay).await;
            }
        }
    };
    let block_time = confirmed.block_time.map(|t| t as u64);
    let inner = confirmed.transaction;
    let meta = inner.meta.as_ref()?;
    let error = meta.err.as_ref().map(|e| e.to_string());

    let static_keys: Vec<String> = match &inner.transaction {
        EncodedTransaction::Json(ui_tx) => match &ui_tx.message {
            UiMessage::Raw(raw) => raw.account_keys.clone(),
            UiMessage::Parsed(parsed) => parsed
                .account_keys
                .iter()
                .map(|k| k.pubkey.clone())
                .collect(),
        },
        _ => return None,
    };

    let mut keys = static_keys;
    if let OptionSerializer::Some(addresses) = &meta.loaded_addresses {
        keys.extend(addresses.writable.iter().cloned());
        keys.extend(addresses.readonly.iter().cloned());
    }

    let sol_changes: Vec<SolChange> = keys
        .iter()
        .enumerate()
        .filter_map(|(i, pubkey)| {
            let pre = *meta.pre_balances.get(i)?;
            let post = *meta.post_balances.get(i)?;
            if pre == post {
                return None;
            }
            Some(SolChange {
                pubkey: pubkey.clone(),
                pre_lamports: pre,
                post_lamports: post,
            })
        })
        .collect();

    let empty = vec![];
    let pre_tokens = match &meta.pre_token_balances {
        OptionSerializer::Some(v) => v,
        _ => &empty,
    };
    let post_tokens = match &meta.post_token_balances {
        OptionSerializer::Some(v) => v,
        _ => &empty,
    };

    let mut pre_map: HashMap<u8, (u64, String, String, u8)> = HashMap::new();
    for t in pre_tokens {
        let amount = t.ui_token_amount.amount.parse::<u64>().unwrap_or(0);
        let owner = match &t.owner {
            OptionSerializer::Some(o) => o.clone(),
            _ => String::new(),
        };
        pre_map.insert(
            t.account_index,
            (amount, t.mint.clone(), owner, t.ui_token_amount.decimals),
        );
    }

    let mut post_map: HashMap<u8, (u64, String, String, u8)> = HashMap::new();
    for t in post_tokens {
        let amount = t.ui_token_amount.amount.parse::<u64>().unwrap_or(0);
        let owner = match &t.owner {
            OptionSerializer::Some(o) => o.clone(),
            _ => String::new(),
        };
        post_map.insert(
            t.account_index,
            (amount, t.mint.clone(), owner, t.ui_token_amount.decimals),
        );
    }

    let mut all_indices: BTreeSet<u8> = BTreeSet::new();
    all_indices.extend(pre_map.keys().copied());
    all_indices.extend(post_map.keys().copied());

    let token_changes: Vec<TokenChange> = all_indices
        .into_iter()
        .filter_map(|idx| {
            let pubkey = keys.get(idx as usize)?.clone();
            let (pre_amount, mint, owner, decimals) = if let Some(info) = pre_map.get(&idx) {
                (info.0, info.1.clone(), info.2.clone(), info.3)
            } else {
                let info = post_map.get(&idx)?;
                (0, info.1.clone(), info.2.clone(), info.3)
            };
            let post_amount = post_map.get(&idx).map(|i| i.0).unwrap_or(0);
            if pre_amount == post_amount {
                return None;
            }
            Some(TokenChange {
                pubkey,
                mint,
                owner,
                pre_amount,
                post_amount,
                decimals,
            })
        })
        .collect();

    Some(Transaction {
        slot,
        timestamp: block_time,
        signature: signature.to_string(),
        success: error.is_none(),
        error,
        sol_changes,
        token_changes,
        logs: Vec::new(),
        account_diffs: Vec::new(),
    })
}