sim-cli 0.5.1

use std::sync::Arc;

use dashmap::DashMap;
use simulator_client::AccountDiffNotification;
use solana_transaction_status::{EncodedConfirmedTransactionWithStatusMeta, EncodedTransaction};
use tokio::sync::mpsc::UnboundedSender;

use crate::{
    fetch::transaction_from_encoded,
    output::{AccountDiff, AccountDiffRow, OutputEvent, Transaction},
};

/// Handle a single account-diff notification from the subscription manager.
///
/// Emits a `diff` event on the NDJSON stream. Enrichment (`pre_lamports`,
/// `post_lamports`, `pre_tokens`, `post_tokens`) is read from the shared
/// records map, which is populated by the `transactionSubscribe` stream
/// running alongside this subscription in account-diff mode. If the matching
/// tx record hasn't arrived yet (small race window between two WS channels),
/// the row's enrichment fields are emitted as `None`.
pub(crate) fn on_account_diff_notification(
    records: Arc<DashMap<String, Transaction>>,
    output_tx: UnboundedSender<OutputEvent>,
    notification: AccountDiffNotification,
) {
    let Some(sig) = notification.signature else {
        return;
    };
    let Some(account) = notification.account else {
        return;
    };
    let slot = notification.context.slot;
    let diff = AccountDiff {
        account,
        pre_state: notification.pre,
        post_state: notification.post,
    };

    let row = match records.get(&sig) {
        Some(entry) => AccountDiffRow::from_diff(
            entry.slot,
            &sig,
            &diff,
            &entry.sol_changes,
            &entry.token_changes,
        ),
        None => AccountDiffRow::from_diff(slot, &sig, &diff, &[], &[]),
    };

    let _ = output_tx.send(OutputEvent::Diff(row));
}

/// Handle a single `transactionSubscribe` notification.
///
/// The server pushes a fully-assembled, `getTransaction`-shaped payload, so we
/// build the [`Transaction`] record directly from the notification and skip
/// any per-tx fetch. Logs come from the meta in the payload.
pub(crate) fn on_transaction_notification(
    records: Arc<DashMap<String, Transaction>>,
    output_tx: UnboundedSender<OutputEvent>,
    notification: EncodedConfirmedTransactionWithStatusMeta,
) {
    let slot = notification.slot;
    let Some(signature) = primary_signature(&notification.transaction.transaction) else {
        tracing::warn!(
            slot,
            "transactionSubscribe notification missing primary signature; dropping"
        );
        return;
    };

    let logs = notification
        .transaction
        .meta
        .as_ref()
        .and_then(|meta| match &meta.log_messages {
            solana_transaction_status::option_serializer::OptionSerializer::Some(v) => {
                Some(v.clone())
            }
            _ => None,
        })
        .unwrap_or_default();

    let Some(mut tx) = transaction_from_encoded(notification, &signature, slot) else {
        tracing::warn!(
            slot,
            signature,
            "transactionSubscribe payload could not be decoded into a Transaction; dropping"
        );
        return;
    };
    tx.logs = logs;
    records.insert(signature, tx.clone());
    let _ = output_tx.send(OutputEvent::Tx(tx));
}

/// Pull the first signature off an [`EncodedTransaction`]. Returns `None` for
/// encodings we don't decode (binary forms) or for a transaction with no
/// signatures, which we'd have no way to key the records map by.
fn primary_signature(encoded: &EncodedTransaction) -> Option<String> {
    match encoded {
        EncodedTransaction::Json(ui_tx) => ui_tx.signatures.first().cloned(),
        EncodedTransaction::LegacyBinary(_) | EncodedTransaction::Binary(_, _) => None,
        EncodedTransaction::Accounts(parsed) => parsed.signatures.first().cloned(),
    }
}