//! Integration tests for the composition cost / token budget.
//!
//! Drives the trace-level cost_total and tokens_total accumulation
//! plus the budget cap that flips composition.budget_tripped. Cost
//! and token attribution comes from the live LLM-backed adapters in
//! the real path; for the test we synthesize the per-node values by
//! constructing the trace shape directly.

use mcptest_core::composition::{BudgetSpec, CompositionTrace, NodeStatus, NodeTrace};
use serde_json::Value;
use std::collections::BTreeMap;

fn node(id: &str, cost: f64, tokens: u64) -> NodeTrace {
    NodeTrace {
        id: id.into(),
        status: NodeStatus::Ok,
        inputs: Value::Null,
        output: Value::Null,
        duration_ms: Some(0),
        cost_usd: cost,
        tokens,
        model: None,
    }
}

#[test]
fn trace_default_cost_and_tokens_are_zero_and_serialize_compactly() {
    let trace = CompositionTrace::default();
    assert_eq!(trace.cost_total, 0.0);
    assert_eq!(trace.tokens_total, 0);
    assert_eq!(trace.budget_tripped, None);
    let json = serde_json::to_value(&trace).unwrap();
    // Default zeros + budget_tripped None are skipped at serialize
    // time so the existing report JSON shape is unchanged for
    // compositions that do not attribute spend.
    assert!(json.get("cost_total").is_none());
    assert!(json.get("tokens_total").is_none());
    assert!(json.get("budget_tripped").is_none());
}

#[test]
fn budget_spec_with_no_caps_does_not_trip() {
    let budget = BudgetSpec::default();
    // None all the way down: no cap means no trip even when totals are nonzero.
    assert!(budget.max_cost_usd.is_none());
    assert!(budget.max_tokens.is_none());
}

#[test]
fn node_trace_round_trips_through_serde_when_attribution_is_set() {
    // Lock the JSON wire shape so a future report consumer can rely
    // on cost_usd / tokens / model fields being present whenever the
    // adapter populated them.
    let trace = node("llm", 0.42, 4200);
    let value = serde_json::to_value(&trace).unwrap();
    assert_eq!(value["cost_usd"], 0.42);
    assert_eq!(value["tokens"], 4200);
    let back: NodeTrace = serde_json::from_value(value).unwrap();
    assert_eq!(back, trace);
}

#[test]
fn composition_trace_totals_accumulate_across_nodes() {
    // The executor sums every node's cost_usd / tokens to populate
    // CompositionTrace.cost_total / tokens_total. Locking the
    // semantics here means a future change that introduces
    // per-iteration tracking for for_each must keep the rolling
    // totals on the top-level trace stable.
    let mut nodes: BTreeMap<String, NodeTrace> = BTreeMap::new();
    nodes.insert("a".into(), node("a", 0.10, 100));
    nodes.insert("b".into(), node("b", 0.25, 250));
    let trace = CompositionTrace {
        nodes: nodes.clone(),
        cost_total: nodes.values().map(|t| t.cost_usd).sum(),
        tokens_total: nodes.values().map(|t| t.tokens).sum(),
        ..Default::default()
    };
    assert!((trace.cost_total - 0.35).abs() < 1e-9);
    assert_eq!(trace.tokens_total, 350);
}
