//! Cross-run tool-selection aggregation and the pass^k metric.
//!
//! Why this module exists: a single agent run can pick the right tool by
//! luck, so a multi-run agent test asserts a *rate* across N independent
//! runs instead of a single outcome. This module is the pure core of that
//! aggregation. It takes the N trace envelopes and a
//! [`crate::suite::ToolSelectionFloor`] and folds them into a
//! [`SelectionReport`]; the CLI prints the report and gates on
//! [`SelectionReport::meets_floor`]. Keeping the math here (no provider,
//! no transport, no I/O) means it is unit-testable and reusable by the
//! SDK hosts and the enterprise audit chain.
//!
//! The headline metric is pass^k: the fraction of runs that BOTH select
//! the expected tool AND stay within the optional token budget. It
//! captures the tokens-vs-correctness tradeoff in one number, because a
//! model that gets the tool right only by burning an unbounded budget is
//! not actually reliable.

use serde::{Deserialize, Serialize};

use crate::suite::ToolSelectionFloor;

/// Aggregated outcome of the N runs of one (agent, model) pair against a
/// [`ToolSelectionFloor`].
///
/// Serializable so reporters and the enterprise audit chain can persist
/// the cost-vs-accuracy picture alongside the per-run executor rows.
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub struct SelectionReport {
    /// Number of runs aggregated (the agent test's `runs:` value).
    pub runs: usize,
    /// Tool name the floor expected the model to select.
    pub expected_tool: String,
    /// How many runs selected the expected tool (ignoring budget).
    pub selected_count: usize,
    /// `selected_count / runs` as a fraction in `0.0..=1.0`. Zero when
    /// there were no runs.
    pub selection_rate: f64,
    /// The pass^k metric: the fraction of runs that BOTH selected the
    /// tool AND stayed within the token budget. Zero when there were no
    /// runs.
    pub passk: f64,
    /// Largest `conversation.tokens.total` observed across the runs.
    pub max_total_tokens_observed: u32,
    /// How many runs stayed within the token budget. Equal to `runs`
    /// when the floor sets no `max_total_tokens`.
    pub within_budget_count: usize,
}

impl SelectionReport {
    /// True when the runs clear the floor: the selection rate meets the
    /// minimum AND every run stayed within the token budget.
    ///
    /// The budget check is all-or-nothing on purpose: a single run that
    /// blows the cap means the model cannot be trusted to stay in budget,
    /// even if it picked the right tool every time.
    pub fn meets_floor(&self, floor: &ToolSelectionFloor) -> bool {
        self.selection_rate >= floor.min_selection_rate && self.within_budget_count == self.runs
    }

    /// One-line human-readable summary pairing tool-selection accuracy
    /// with the pass^k metric and the worst-case token spend. The CLI
    /// prints this so the operator sees cost and correctness together.
    pub fn cost_vs_accuracy(&self) -> String {
        let rate_pct = (self.selection_rate * 100.0).round() as i64;
        let passk_pct = (self.passk * 100.0).round() as i64;
        format!(
            "selection {}/{} ({}%), pass^k {}%, max tokens {}",
            self.selected_count, self.runs, rate_pct, passk_pct, self.max_total_tokens_observed
        )
    }
}

/// Fold the N trace envelopes of a multi-run agent test into a
/// [`SelectionReport`] against `floor`.
///
/// Each envelope is expected to carry `tool_calls: [{name, ...}]` and
/// `conversation.tokens.total` (the shape `ConversationTrace::to_envelope`
/// produces). A run "selected" when `floor.expected_tool` matches any
/// `tool_calls[].name`. A run is "within budget" when the floor sets no
/// `max_total_tokens` or the run's `tokens.total` is at or under it. The
/// pass^k count requires both conditions. Missing or malformed fields are
/// treated conservatively: a run with no readable tool calls did not
/// select, and an unreadable token total counts as zero (so it cannot
/// breach a budget on its own).
pub fn aggregate_selection(
    trace_envelopes: &[serde_json::Value],
    floor: &ToolSelectionFloor,
) -> SelectionReport {
    let runs = trace_envelopes.len();
    let mut selected_count = 0usize;
    let mut within_budget_count = 0usize;
    let mut passk_count = 0usize;
    let mut max_total_tokens_observed = 0u32;

    for envelope in trace_envelopes {
        let selected = envelope_selected_tool(envelope, &floor.expected_tool);
        let tokens = envelope_total_tokens(envelope);
        max_total_tokens_observed = max_total_tokens_observed.max(tokens);
        let within_budget = match floor.max_total_tokens {
            Some(cap) => tokens <= cap,
            None => true,
        };
        if selected {
            selected_count += 1;
        }
        if within_budget {
            within_budget_count += 1;
        }
        if selected && within_budget {
            passk_count += 1;
        }
    }

    let selection_rate = ratio(selected_count, runs);
    let passk = ratio(passk_count, runs);
    SelectionReport {
        runs,
        expected_tool: floor.expected_tool.clone(),
        selected_count,
        selection_rate,
        passk,
        max_total_tokens_observed,
        within_budget_count,
    }
}

/// `numerator / denominator` as an `f64`, returning `0.0` when the
/// denominator is zero so a zero-run aggregation never divides by zero.
fn ratio(numerator: usize, denominator: usize) -> f64 {
    if denominator == 0 {
        0.0
    } else {
        numerator as f64 / denominator as f64
    }
}

/// True when `expected_tool` appears in any `tool_calls[].name` of the
/// envelope. A missing or non-array `tool_calls` means no selection.
fn envelope_selected_tool(envelope: &serde_json::Value, expected_tool: &str) -> bool {
    let Some(calls) = envelope.get("tool_calls").and_then(|v| v.as_array()) else {
        return false;
    };
    calls
        .iter()
        .filter_map(|call| call.get("name").and_then(|n| n.as_str()))
        .any(|name| name == expected_tool)
}

/// Read `conversation.tokens.total` as a `u32`, defaulting to zero when
/// the path is absent or not an unsigned integer.
fn envelope_total_tokens(envelope: &serde_json::Value) -> u32 {
    envelope
        .get("conversation")
        .and_then(|c| c.get("tokens"))
        .and_then(|t| t.get("total"))
        .and_then(|v| v.as_u64())
        .map(|n| n as u32)
        .unwrap_or(0)
}

#[cfg(test)]
mod tests {
    use super::*;
    use serde_json::json;

    fn floor(expected: &str, min_rate: f64, cap: Option<u32>) -> ToolSelectionFloor {
        ToolSelectionFloor {
            expected_tool: expected.to_string(),
            min_selection_rate: min_rate,
            max_total_tokens: cap,
        }
    }

    fn envelope(tool: Option<&str>, total: u32) -> serde_json::Value {
        let calls = match tool {
            Some(name) => json!([{ "name": name }]),
            None => json!([]),
        };
        json!({
            "tool_calls": calls,
            "conversation": { "tokens": { "total": total } }
        })
    }

    #[test]
    fn all_runs_select_within_budget_meets_floor_with_passk_one() {
        let floor = floor("get_weather", 0.8, Some(1000));
        let runs = vec![
            envelope(Some("get_weather"), 100),
            envelope(Some("get_weather"), 200),
            envelope(Some("get_weather"), 300),
        ];
        let report = aggregate_selection(&runs, &floor);
        assert_eq!(report.runs, 3);
        assert_eq!(report.selected_count, 3);
        assert!((report.selection_rate - 1.0).abs() < f64::EPSILON);
        assert!((report.passk - 1.0).abs() < f64::EPSILON);
        assert_eq!(report.within_budget_count, 3);
        assert_eq!(report.max_total_tokens_observed, 300);
        assert!(report.meets_floor(&floor));
    }

    #[test]
    fn missing_the_tool_drops_selection_rate_below_min_and_fails_floor() {
        let floor = floor("get_weather", 0.8, None);
        let runs = vec![
            envelope(Some("get_weather"), 100),
            envelope(None, 100),
            envelope(Some("other_tool"), 100),
        ];
        let report = aggregate_selection(&runs, &floor);
        assert_eq!(report.selected_count, 1);
        assert!((report.selection_rate - (1.0 / 3.0)).abs() < 1e-9);
        // No token cap means every run is within budget.
        assert_eq!(report.within_budget_count, 3);
        assert!(!report.meets_floor(&floor));
    }

    #[test]
    fn one_run_over_budget_fails_floor_even_at_full_selection() {
        let floor = floor("get_weather", 1.0, Some(500));
        let runs = vec![
            envelope(Some("get_weather"), 400),
            envelope(Some("get_weather"), 600), // over cap
            envelope(Some("get_weather"), 200),
        ];
        let report = aggregate_selection(&runs, &floor);
        // Every run picked the tool, so selection rate is 100%.
        assert!((report.selection_rate - 1.0).abs() < f64::EPSILON);
        // But one run blew the budget, so within_budget_count < runs.
        assert_eq!(report.within_budget_count, 2);
        // pass^k folds in the budget breach: only 2 of 3 runs both
        // selected and stayed in budget.
        assert!((report.passk - (2.0 / 3.0)).abs() < 1e-9);
        assert!(!report.meets_floor(&floor));
    }

    #[test]
    fn cost_vs_accuracy_string_pairs_selection_passk_and_max_tokens() {
        let floor = floor("get_weather", 0.5, Some(1000));
        let runs = vec![envelope(Some("get_weather"), 100), envelope(None, 900)];
        let report = aggregate_selection(&runs, &floor);
        assert_eq!(
            report.cost_vs_accuracy(),
            "selection 1/2 (50%), pass^k 50%, max tokens 900"
        );
    }

    #[test]
    fn zero_runs_yields_zero_rates_and_meets_an_open_floor() {
        let floor = floor("get_weather", 0.0, None);
        let report = aggregate_selection(&[], &floor);
        assert_eq!(report.runs, 0);
        assert!((report.selection_rate - 0.0).abs() < f64::EPSILON);
        assert!((report.passk - 0.0).abs() < f64::EPSILON);
        // within_budget_count == runs (both zero) and min_rate is 0.0,
        // so a vacuous floor is met. The CLI only aggregates when runs
        // are present, so this is the degenerate-but-safe case.
        assert!(report.meets_floor(&floor));
    }
}
