sim-cli 0.7.0

use std::{
    collections::{BTreeMap, HashMap},
    fs::File,
    io::BufReader,
    path::PathBuf,
};

use chrono::{DateTime, Utc};
use clap::Parser;
use console::style;
use eyre::{Context, Result};

use crate::output::{
    NdjsonScan, OutputEvent, SimulationMetadata, SimulationSummary, Transaction,
    load_simulation_output, missing_metadata_envelope, scan_ndjson_events,
};

#[derive(Parser, Debug)]
pub struct SummarizeArgs {
    /// Simulation output file to summarize
    pub file: PathBuf,
    /// Comma-separated list of wallet addresses to trace balances across all transactions
    #[arg(long, value_delimiter = ',')]
    pub accounts: Vec<String>,
}

/// Opening and closing balance for one token/SOL across the full simulation
/// run, keyed by slot. Slots can arrive out of order (parallel sessions
/// interleave them), but within a slot the producer emits in on-chain
/// execution order — so the earliest change at the lowest slot opens, the
/// latest at the highest closes.
struct BalanceTrace {
    opening: u64,
    opening_slot: u64,
    closing: u64,
    closing_slot: u64,
    decimals: u8,
}

/// Per wallet: "SOL" or mint -> [`BalanceTrace`].
type WalletTraces<'a> = HashMap<&'a str, BTreeMap<String, BalanceTrace>>;

pub fn summarize(args: SummarizeArgs) -> Result<()> {
    let mut traces: WalletTraces<'_> = args
        .accounts
        .iter()
        .map(|a| (a.as_str(), BTreeMap::new()))
        .collect();

    let file = File::open(&args.file)
        .with_context(|| format!("failed to read {}", args.file.display()))?;

    // Stream the NDJSON envelopes: memory scales with the number of tracked
    // wallets, not the file size.
    let mut metadata: Option<SimulationMetadata> = None;
    let mut summary: Option<SimulationSummary> = None;
    let mut tx_count = 0_usize;
    let mut successes = 0_usize;

    let scan = scan_ndjson_events(BufReader::new(file), &args.file, &mut |event| {
        match event {
            OutputEvent::Metadata(m) => metadata = Some(m),
            OutputEvent::Summary(s) => summary = Some(s),
            OutputEvent::Tx(tx) => {
                tx_count += 1;
                if tx.success {
                    successes += 1;
                }
                fold_transaction(&mut traces, &tx);
            }
            OutputEvent::SessionStarted(_) | OutputEvent::Diff(_) => {}
        }
        Ok(())
    })?;

    let (metadata, summary) = if let NdjsonScan::NotAnEnvelopeStream { .. } = scan {
        // Legacy single-blob fallback: load in memory and run the same fold.
        let output = load_simulation_output(&args.file)?;
        for tx in &output.transactions {
            fold_transaction(&mut traces, tx);
        }
        (output.metadata, output.summary)
    } else {
        let mut metadata = metadata.ok_or_else(|| missing_metadata_envelope(&args.file))?;
        let summary = summary.unwrap_or_else(|| SimulationSummary {
            total_transactions: tx_count,
            successes,
            failures: tx_count - successes,
            session_ids: Vec::new(),
        });
        // Lift session_ids onto metadata so the header shows one consistent count.
        if metadata.session_ids.is_empty() && !summary.session_ids.is_empty() {
            metadata.session_ids = summary.session_ids.clone();
        }
        (metadata, summary)
    };

    let m = &metadata;
    println!("File         : {}", args.file.display());
    println!("Slot range   : {} → {}", m.start_slot, m.end_slot);
    if !m.program_ids.is_empty() {
        println!("Program IDs  : {}", m.program_ids.join(", "));
    }
    if !m.program_so.is_empty() {
        println!("Program.so   : {}", m.program_so.join(", "));
    }
    println!("Session(s)   : {}", m.session_ids.len());
    let ran_at = DateTime::<Utc>::from_timestamp(m.ran_at_unix_secs as i64, 0)
        .map(|dt| dt.format("%Y-%m-%d %H:%M:%S UTC").to_string())
        .unwrap_or_else(|| format!("{} (unix)", m.ran_at_unix_secs));
    println!("Ran at       : {ran_at}");
    println!();

    println!(
        "{} Transactions : {}  ({} successes, {} failures)",
        style("i").blue(),
        summary.total_transactions,
        summary.successes,
        summary.failures,
    );

    if args.accounts.is_empty() {
        return Ok(());
    }

    println!();
    for account in &args.accounts {
        println!("Account: {account}");
        let wallet_traces = &traces[account.as_str()];
        if wallet_traces.is_empty() {
            println!("  (no activity found)");
        }
        for (label, t) in wallet_traces {
            let dec = t.decimals as usize;
            let div = 10u64.pow(t.decimals as u32) as f64;
            let open = t.opening as f64 / div;
            let close = t.closing as f64 / div;
            let delta = close - open;
            println!("  {label:<50}  {open:.dec$} → {close:.dec$}  (Δ {delta:+.dec$})");
        }
        println!();
    }

    Ok(())
}

/// Fold one transaction's balance changes into the traces, in their vector
/// order so multiple changes to the same balance within one transaction keep
/// the first as pre and the last as post — matching a sorted walk.
fn fold_transaction(traces: &mut WalletTraces<'_>, tx: &Transaction) {
    for change in &tx.sol_changes {
        if let Some(wallet) = traces.get_mut(change.pubkey.as_str()) {
            apply_change(
                wallet,
                "SOL",
                tx.slot,
                change.pre_lamports,
                change.post_lamports,
                9,
            );
        }
    }
    for change in &tx.token_changes {
        if let Some(wallet) = traces.get_mut(change.owner.as_str()) {
            apply_change(
                wallet,
                &change.mint,
                tx.slot,
                change.pre_amount,
                change.post_amount,
                change.decimals,
            );
        }
    }
}

/// `opening` follows the lowest slot, `closing` the highest; within a slot,
/// arrival order breaks ties (earliest opens, latest closes).
fn apply_change(
    traces: &mut BTreeMap<String, BalanceTrace>,
    label: &str,
    slot: u64,
    pre: u64,
    post: u64,
    decimals: u8,
) {
    if let Some(trace) = traces.get_mut(label) {
        if slot < trace.opening_slot {
            trace.opening = pre;
            trace.opening_slot = slot;
            trace.decimals = decimals;
        }
        if slot >= trace.closing_slot {
            trace.closing = post;
            trace.closing_slot = slot;
        }
    } else {
        traces.insert(
            label.to_string(),
            BalanceTrace {
                opening: pre,
                opening_slot: slot,
                closing: post,
                closing_slot: slot,
                decimals,
            },
        );
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::output::{SolChange, TokenChange};

    fn tx(slot: u64, signature: &str, sol: &[(&str, u64, u64)]) -> Transaction {
        Transaction {
            slot,
            timestamp: None,
            signature: signature.to_string(),
            success: true,
            error: None,
            logs: Vec::new(),
            sol_changes: sol
                .iter()
                .map(|(pubkey, pre, post)| SolChange {
                    pubkey: pubkey.to_string(),
                    pre_lamports: *pre,
                    post_lamports: *post,
                })
                .collect(),
            token_changes: Vec::new(),
            account_diffs: Vec::new(),
        }
    }

    fn traces_for(wallet: &str) -> WalletTraces<'_> {
        [(wallet, BTreeMap::new())].into_iter().collect()
    }

    #[test]
    fn out_of_order_stream_folds_like_a_sorted_walk() {
        let mut traces = traces_for("w");
        fold_transaction(&mut traces, &tx(5, "b", &[("w", 50, 40)]));
        fold_transaction(&mut traces, &tx(1, "a", &[("w", 100, 90)]));
        fold_transaction(&mut traces, &tx(9, "z", &[("w", 40, 30)]));

        let trace = &traces["w"]["SOL"];
        assert_eq!((trace.opening, trace.closing), (100, 30));
    }

    #[test]
    fn within_slot_folds_in_arrival_order() {
        let mut traces = traces_for("w");
        fold_transaction(&mut traces, &tx(3, "first", &[("w", 100, 80)]));
        fold_transaction(&mut traces, &tx(3, "second", &[("w", 80, 70)]));

        let trace = &traces["w"]["SOL"];
        assert_eq!((trace.opening, trace.closing), (100, 70));
    }

    #[test]
    fn within_tx_changes_apply_in_vec_order() {
        let mut traces = traces_for("w");
        fold_transaction(&mut traces, &tx(2, "a", &[("w", 100, 90), ("w", 90, 85)]));

        let trace = &traces["w"]["SOL"];
        assert_eq!((trace.opening, trace.closing), (100, 85));
    }

    #[test]
    fn token_decimals_follow_the_opening_change() {
        let mut traces = traces_for("w");
        let token = |slot, sig: &str, decimals| {
            let mut t = tx(slot, sig, &[]);
            t.token_changes = vec![TokenChange {
                pubkey: "acct".to_string(),
                mint: "mint-a".to_string(),
                owner: "w".to_string(),
                pre_amount: 10,
                post_amount: 20,
                decimals,
            }];
            t
        };
        // The later-slot change arrives first; decimals must end up from the
        // earlier (opening) change, as a sorted walk would produce.
        fold_transaction(&mut traces, &token(7, "b", 9));
        fold_transaction(&mut traces, &token(2, "a", 6));

        assert_eq!(traces["w"]["mint-a"].decimals, 6);
    }

    #[test]
    fn untracked_wallets_are_ignored() {
        let mut traces = traces_for("w");
        fold_transaction(&mut traces, &tx(1, "a", &[("other", 100, 90)]));
        assert!(traces["w"].is_empty());
    }
}