//! Objective multi-server tool-selection F1 via equal-function sets.
//!
//! Why this module exists: judging whether an agent picked the right tool in a
//! multi-server scenario usually needs a model in the loop, which is slow,
//! costs money, and is not byte-stable. MSC-Bench (arXiv:2510.19423) removes
//! the model: tools that accomplish the same job are grouped into an
//! "equal-function set" (an equivalence class), and a tool choice counts as
//! correct when it lands in the right class. That makes precision, recall, and
//! F1 computable deterministically from a recorded trace alone, with zero LLM
//! calls. The same trace always yields the same score.
//!
//! This module is the pure core of that scoring. It takes the observed tool
//! calls from one run's trace envelope (the shape
//! [`crate::eval::trace_validation`] and the agent driver produce) plus the
//! declared classes, and folds them into a [`F1Score`] with per-class detail.
//! The cross-run aggregate ([`aggregate_f1`]) micro-averages over the N run
//! envelopes of a multi-run agent test, mirroring how
//! [`crate::eval::selection`] folds the pass^k floor and how
//! [`crate::eval::stability_aggregate`] folds the stability targets. 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.
//!
//! # Counting rule
//!
//! Scoring runs over the multiset of observed tool calls against the declared
//! classes, each class consumed at most once (set-level recall, not call-level):
//!
//! - A class is *matched* (a true positive) the first time an observed call
//!   names a `server.tool` belonging to that class. Each class matches at most
//!   once, so calling two interchangeable members of one class does not inflate
//!   the score: the first satisfies the class, and any later member of an
//!   already-matched class is ignored (neither a true positive nor an extra).
//! - A declared class with no observed member is *missed* (a false negative).
//! - An observed call whose `server.tool` belongs to no declared class, and is
//!   not a repeat crediting an already-matched class, is *unexpected* (a false
//!   positive).
//!
//! Then, with tp = matched classes, fn = missed classes, fp = unexpected calls:
//!
//! - precision = tp / (tp + fp)
//! - recall    = tp / (tp + fn)
//! - f1        = 2 * precision * recall / (precision + recall)
//!
//! Edge cases are pinned so the output is byte-stable: with no declared classes
//! and no observed calls every metric is 1.0 (a vacuously perfect selection); a
//! zero denominator yields 0.0 for that metric rather than NaN.

use serde::{Deserialize, Serialize};
use serde_json::{json, Value};

use crate::suite::{Assertion, Matcher};

/// A named group of interchangeable `server.tool` references (one MSC-Bench
/// equal-function set).
///
/// Any member accomplishes the task, so selecting any one satisfies the class.
/// Modeled as a named set so the per-class breakdown can name exactly which
/// capability was missed rather than report an opaque index.
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct EqualFunctionSet {
    /// Class name, unique within a scenario. Surfaced in the breakdown so
    /// authors read which capability was missed, not a positional index.
    pub name: String,
    /// Interchangeable tool ids; selecting any one satisfies the class.
    ///
    /// Each id is matched against an observed call's qualified id (see
    /// [`qualified_id`]) so multi-server scenarios disambiguate identically
    /// named tools on different servers. A bare `tool` id (no `server.`
    /// prefix) matches a call's tool name on any server.
    pub members: Vec<String>,
}

impl EqualFunctionSet {
    /// Build a class from its name and member ids. Kept so call sites and tests
    /// share one construction path and stay source-compatible if private fields
    /// are added later.
    pub fn new(name: impl Into<String>, members: impl IntoIterator<Item = String>) -> Self {
        Self {
            name: name.into(),
            members: members.into_iter().collect(),
        }
    }

    /// True when `observed` (a `server.tool` qualified id, or a bare tool name)
    /// is a member of this class. Centralizes the membership test so the
    /// counting rule never duplicates it.
    pub fn contains(&self, observed: &str) -> bool {
        self.members.iter().any(|m| member_matches(m, observed))
    }
}

/// True when a declared member id matches an observed qualified id.
///
/// A declared id with a `server.` prefix must match the full qualified id; a
/// bare declared id matches the observed call's tool name regardless of server,
/// so a single-server scenario can declare `members: [search]` without naming
/// the server. Free function so both the membership test and its unit tests
/// share one rule.
fn member_matches(declared: &str, observed: &str) -> bool {
    if declared == observed {
        return true;
    }
    // Bare declared id matches the tool half of `server.tool`.
    !declared.contains('.') && observed.rsplit('.').next() == Some(declared)
}

/// Per-class outcome within a single run's F1 scoring.
///
/// Carries enough detail for the report to explain why a run scored as it did:
/// which class was satisfied by which observed tool, and which class was missed.
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct ClassOutcome {
    /// Class name this outcome describes.
    pub class: String,
    /// The observed `server.tool` that satisfied the class, when matched.
    /// `None` means the class was missed (the false negative the author needs
    /// to see in the breakdown).
    pub matched_by: Option<String>,
}

/// Deterministic F1 outcome of scoring one run's observed tool calls against
/// the declared classes.
///
/// Serializable so the report can carry the full per-class detail. Percentages
/// are integers in `0..=100` to match the existing floor convention (the
/// pass^k floor and the stability targets), to keep the output byte-stable (no
/// float-formatting drift across platforms), and so a suite gates on a plain
/// percentage.
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct F1Score {
    /// Per-class breakdown: which class matched (and by what) or was missed.
    pub class_outcomes: Vec<ClassOutcome>,
    /// Observed `server.tool` calls that matched no declared class (the
    /// "extra" false positives the breakdown reports).
    pub unexpected: Vec<String>,
    /// Matched classes (true positives).
    pub true_positives: u32,
    /// Missed classes (false negatives).
    pub false_negatives: u32,
    /// Unexpected calls (false positives).
    pub false_positives: u32,
    /// Precision as an integer percent in `0..=100`.
    pub precision_percent: u32,
    /// Recall as an integer percent in `0..=100`.
    pub recall_percent: u32,
    /// F1 as an integer percent in `0..=100`.
    pub f1_percent: u32,
}

/// Compute integer-percent precision, recall, and F1 from the raw counts.
///
/// Returns 100 for every metric only in the vacuous case (no classes and no
/// calls), and 0 for a zero denominator otherwise, so callers never see NaN and
/// the result is byte-stable. Free function so the per-run and aggregate paths
/// share one arithmetic.
fn metrics(tp: u32, fp: u32, fn_: u32) -> (u32, u32, u32) {
    // Vacuous perfection: nothing was expected and nothing was wrongly chosen,
    // so precision, recall, and F1 are all perfect. Special-cased up front so
    // every metric reads 100 (a zero-denominator ratio would otherwise report
    // 0 for precision and recall).
    if tp == 0 && fp == 0 && fn_ == 0 {
        return (100, 100, 100);
    }
    let precision = ratio_percent(tp, tp + fp);
    let recall = ratio_percent(tp, tp + fn_);
    // A zero denominator here means one side had no signal; report 0 for F1
    // rather than NaN. The vacuous all-zero case is handled above.
    let f1 = (2 * precision * recall)
        .checked_div(precision + recall)
        .unwrap_or(0);
    (precision, recall, f1)
}

/// Compute `numerator / denominator` as an integer percent in `0..=100`.
///
/// A zero denominator yields 0 so the arithmetic never divides by zero and the
/// output is identical on every platform.
fn ratio_percent(numerator: u32, denominator: u32) -> u32 {
    (numerator * 100).checked_div(denominator).unwrap_or(0)
}

/// Score one run's observed tool-call ids against the declared classes.
///
/// `observed` is the ordered list of `server.tool` qualified ids the run made
/// (extract them from a trace envelope with [`observed_calls`]). Deterministic
/// and free of model calls, so a fixed observation yields a byte-stable result.
/// See the module doc comment for the exact counting rule.
pub fn score_f1(classes: &[EqualFunctionSet], observed: &[String]) -> F1Score {
    let mut matched = vec![false; classes.len()];
    let mut class_outcomes: Vec<ClassOutcome> = Vec::with_capacity(classes.len());
    let mut unexpected: Vec<String> = Vec::new();

    for id in observed {
        match classes.iter().position(|c| c.contains(id)) {
            Some(idx) => {
                if !matched[idx] {
                    matched[idx] = true;
                    class_outcomes.push(ClassOutcome {
                        class: classes[idx].name.clone(),
                        matched_by: Some(id.clone()),
                    });
                }
            }
            None => unexpected.push(id.clone()),
        }
    }

    for (idx, class) in classes.iter().enumerate() {
        if !matched[idx] {
            class_outcomes.push(ClassOutcome {
                class: class.name.clone(),
                matched_by: None,
            });
        }
    }

    let true_positives = matched.iter().filter(|m| **m).count() as u32;
    let false_negatives = matched.iter().filter(|m| !**m).count() as u32;
    let false_positives = unexpected.len() as u32;
    let (precision_percent, recall_percent, f1_percent) =
        metrics(true_positives, false_positives, false_negatives);

    F1Score {
        class_outcomes,
        unexpected,
        true_positives,
        false_negatives,
        false_positives,
        precision_percent,
        recall_percent,
        f1_percent,
    }
}

/// Extract the ordered `server.tool` qualified ids from one run's trace
/// envelope.
///
/// Reads the same `tool_calls: [{name, server}]` shape the agent driver and
/// the composition executor produce (see [`crate::eval::selection`] and the
/// agent `ConversationTrace`). A call with no readable `name` is skipped, and a
/// missing or empty `server` falls back to the bare tool name so a
/// single-server trace still scores. Pure: no I/O.
pub fn observed_calls(envelope: &Value) -> Vec<String> {
    let Some(calls) = envelope.get("tool_calls").and_then(|v| v.as_array()) else {
        return Vec::new();
    };
    calls
        .iter()
        .filter_map(|call| {
            let name = call.get("name").and_then(|n| n.as_str())?;
            let server = call.get("server").and_then(|s| s.as_str()).unwrap_or("");
            Some(qualified_id(server, name))
        })
        .collect()
}

/// Build a `server.tool` qualified id from a server key and tool name.
///
/// An empty server yields the bare tool name so single-server traces (and
/// classes that declare bare member ids) match without a server prefix.
pub fn qualified_id(server: &str, tool: &str) -> String {
    if server.is_empty() {
        tool.to_string()
    } else {
        format!("{server}.{tool}")
    }
}

/// Cross-run F1 aggregate over the N run envelopes of a multi-run agent test.
///
/// Aggregates by micro-averaging: it sums true/false positives and negatives
/// across every run, then derives precision/recall/F1 from the totals. The sums
/// are commutative, so run order never changes the aggregate and a fixed set of
/// envelopes is byte-stable. Carries the per-run scores too so the report shows
/// exactly which class each run missed. Serializable so the audit chain can
/// persist the picture next to the per-run rows.
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct F1Aggregate {
    /// Number of runs aggregated (the agent test's `runs:` value).
    pub runs: usize,
    /// Per-run F1 scores in run order (the byte-stable per-class breakdowns).
    pub per_run: Vec<F1Score>,
    /// Micro-averaged precision percent across all runs.
    pub precision_percent: u32,
    /// Micro-averaged recall percent across all runs.
    pub recall_percent: u32,
    /// Micro-averaged F1 percent across all runs.
    pub f1_percent: u32,
}

impl F1Aggregate {
    /// One-line human-readable summary the CLI prints so the operator sees the
    /// three numbers together. Mirrors
    /// [`crate::eval::stability_aggregate::StabilitySummary::line`].
    pub fn line(&self) -> String {
        format!(
            "F1 {}%, precision {}%, recall {}% over {} runs",
            self.f1_percent, self.precision_percent, self.recall_percent, self.runs
        )
    }

    /// Names of classes missed by at least one run, in first-seen order.
    ///
    /// The actionable half of the breakdown: it tells the author which
    /// capability the agent never selected. Surfaced in the gate failure note.
    pub fn missed_classes(&self) -> Vec<String> {
        let mut missed: Vec<String> = Vec::new();
        for run in &self.per_run {
            for outcome in &run.class_outcomes {
                if outcome.matched_by.is_none() && !missed.contains(&outcome.class) {
                    missed.push(outcome.class.clone());
                }
            }
        }
        missed
    }

    /// Tool ids chosen by at least one run that matched no declared class, in
    /// first-seen order. The other half of the breakdown: which extra tool the
    /// agent wrongly chose.
    pub fn unexpected_tools(&self) -> Vec<String> {
        let mut extra: Vec<String> = Vec::new();
        for run in &self.per_run {
            for id in &run.unexpected {
                if !extra.contains(id) {
                    extra.push(id.clone());
                }
            }
        }
        extra
    }
}

/// Fold the N trace envelopes of a multi-run agent test into an
/// [`F1Aggregate`] against the declared classes.
///
/// Each envelope is scored with [`score_f1`] over its [`observed_calls`], then
/// the per-run counts are micro-averaged. An empty slice scores as a vacuously
/// perfect aggregate (every metric 100) so a gate never fires on "no runs",
/// matching the stability aggregate's empty-sample behavior.
pub fn aggregate_f1(classes: &[EqualFunctionSet], trace_envelopes: &[Value]) -> F1Aggregate {
    let per_run: Vec<F1Score> = trace_envelopes
        .iter()
        .map(|envelope| score_f1(classes, &observed_calls(envelope)))
        .collect();
    let mut tp = 0u32;
    let mut fp = 0u32;
    let mut fn_ = 0u32;
    for run in &per_run {
        tp += run.true_positives;
        fp += run.false_positives;
        fn_ += run.false_negatives;
    }
    let (precision_percent, recall_percent, f1_percent) = metrics(tp, fp, fn_);
    F1Aggregate {
        runs: per_run.len(),
        per_run,
        precision_percent,
        recall_percent,
        f1_percent,
    }
}

/// Default F1 floor when an `equal_function_sets:` block declares no `expect:`.
///
/// Mirrors the stability gate's default of `0.5`: a selection F1 under half is
/// more wrong than right. Expressed as the percent the assertable target uses.
const DEFAULT_MIN_F1_PERCENT: i64 = 50;

/// Build the targets object the `equal_function_sets:` assertions resolve
/// against.
///
/// Shape: `{ "tool_selection": { "f1": <num>, "precision": <num>,
/// "recall": <num> } }`, all integer percents in `0..=100`, so the existing
/// dot-path resolver walks `tool_selection.f1` and friends directly, the same
/// way the `stability.*` and `jury.*` targets resolve against their metrics
/// object. Reusing the resolver means no new matcher and no new reporter.
pub fn metrics_object(aggregate: &F1Aggregate) -> Value {
    json!({
        "tool_selection": {
            "f1": aggregate.f1_percent,
            "precision": aggregate.precision_percent,
            "recall": aggregate.recall_percent,
        }
    })
}

/// The default assertion set applied when an `equal_function_sets:` block
/// declares no `expect:`: `tool_selection.f1 >= 50`.
///
/// The smallest useful form of the block gates F1; a user who wants to gate
/// precision or recall instead (or in addition) writes an explicit `expect:`.
pub fn default_expect() -> Vec<Assertion> {
    vec![Assertion {
        target: "tool_selection.f1".to_string(),
        matcher: Matcher::Schema(json!({ "minimum": DEFAULT_MIN_F1_PERCENT })),
        weight: 1.0,
        name: None,
        transform: None,
    }]
}

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

    fn classes() -> Vec<EqualFunctionSet> {
        vec![
            EqualFunctionSet::new(
                "search",
                ["brave.web_search".to_string(), "google.search".to_string()],
            ),
            EqualFunctionSet::new("fetch", ["http.get".to_string()]),
        ]
    }

    fn ids(items: &[&str]) -> Vec<String> {
        items.iter().map(|s| s.to_string()).collect()
    }

    #[test]
    fn perfect_score_when_each_class_matched_once() {
        let score = score_f1(&classes(), &ids(&["google.search", "http.get"]));
        assert_eq!(score.true_positives, 2);
        assert_eq!(score.false_negatives, 0);
        assert_eq!(score.false_positives, 0);
        assert_eq!(score.f1_percent, 100);
        assert_eq!(score.precision_percent, 100);
        assert_eq!(score.recall_percent, 100);
    }

    #[test]
    fn equivalent_member_satisfies_class() {
        // brave.web_search is interchangeable with google.search in "search".
        let score = score_f1(&classes(), &ids(&["brave.web_search", "http.get"]));
        assert_eq!(score.f1_percent, 100);
        let search = score
            .class_outcomes
            .iter()
            .find(|c| c.class == "search")
            .expect("search class present");
        assert_eq!(search.matched_by.as_deref(), Some("brave.web_search"));
    }

    #[test]
    fn missed_class_is_a_false_negative_in_breakdown() {
        let score = score_f1(&classes(), &ids(&["google.search"]));
        assert_eq!(score.true_positives, 1);
        assert_eq!(score.false_negatives, 1);
        let fetch = score
            .class_outcomes
            .iter()
            .find(|c| c.class == "fetch")
            .expect("fetch class present");
        assert_eq!(fetch.matched_by, None);
        // recall = 1/2 = 50, precision = 1/1 = 100, f1 = 2*100*50/150 = 66.
        assert_eq!(score.recall_percent, 50);
        assert_eq!(score.precision_percent, 100);
        assert_eq!(score.f1_percent, 66);
    }

    #[test]
    fn extra_tool_is_a_false_positive() {
        let score = score_f1(
            &classes(),
            &ids(&["google.search", "http.get", "shell.exec"]),
        );
        assert_eq!(score.false_positives, 1);
        assert_eq!(score.unexpected, vec!["shell.exec".to_string()]);
        // precision = 2/3 = 66, recall = 2/2 = 100, f1 = 2*66*100/166 = 79.
        assert_eq!(score.precision_percent, 66);
        assert_eq!(score.recall_percent, 100);
        assert_eq!(score.f1_percent, 79);
    }

    #[test]
    fn repeating_a_matched_class_member_does_not_inflate_or_penalize() {
        let score = score_f1(
            &classes(),
            &ids(&["google.search", "brave.web_search", "http.get"]),
        );
        // The second "search" member is neither a new true positive nor extra.
        assert_eq!(score.true_positives, 2);
        assert_eq!(score.false_positives, 0);
        assert_eq!(score.f1_percent, 100);
    }

    #[test]
    fn empty_observation_misses_every_class() {
        let score = score_f1(&classes(), &[]);
        assert_eq!(score.true_positives, 0);
        assert_eq!(score.false_negatives, 2);
        assert_eq!(score.f1_percent, 0);
        assert_eq!(score.recall_percent, 0);
    }

    #[test]
    fn no_classes_and_no_calls_is_vacuously_perfect() {
        let score = score_f1(&[], &[]);
        assert_eq!(score.f1_percent, 100);
        assert_eq!(score.precision_percent, 100);
        assert_eq!(score.recall_percent, 100);
    }

    #[test]
    fn bare_member_id_matches_tool_on_any_server() {
        let classes = vec![EqualFunctionSet::new("search", ["search".to_string()])];
        let score = score_f1(&classes, &ids(&["google.search"]));
        assert_eq!(score.f1_percent, 100);
    }

    #[test]
    fn observed_calls_reads_server_and_name_from_envelope() {
        let envelope = json!({
            "tool_calls": [
                { "name": "search", "server": "google", "args": {} },
                { "name": "get", "server": "http", "args": {} }
            ]
        });
        assert_eq!(
            observed_calls(&envelope),
            ids(&["google.search", "http.get"])
        );
    }

    #[test]
    fn observed_calls_falls_back_to_bare_name_without_server() {
        let envelope = json!({ "tool_calls": [{ "name": "search", "args": {} }] });
        assert_eq!(observed_calls(&envelope), ids(&["search"]));
    }

    #[test]
    fn aggregate_micro_averages_across_runs() {
        // Run one: perfect (tp 2). Run two: misses fetch (tp 1, fn 1).
        let perfect = json!({
            "tool_calls": [
                { "name": "search", "server": "google" },
                { "name": "get", "server": "http" }
            ]
        });
        let partial = json!({ "tool_calls": [{ "name": "search", "server": "google" }] });
        let agg = aggregate_f1(&classes(), &[perfect, partial]);
        assert_eq!(agg.runs, 2);
        // tp = 3, fp = 0, fn = 1 -> recall 75, precision 100, f1 = 85.
        assert_eq!(agg.recall_percent, 75);
        assert_eq!(agg.precision_percent, 100);
        assert_eq!(agg.f1_percent, 85);
        assert_eq!(agg.missed_classes(), vec!["fetch".to_string()]);
    }

    #[test]
    fn aggregate_is_order_independent() {
        let a = json!({ "tool_calls": [{ "name": "search", "server": "google" }] });
        let b = json!({ "tool_calls": [{ "name": "get", "server": "http" }] });
        let forward = aggregate_f1(&classes(), &[a.clone(), b.clone()]);
        let reversed = aggregate_f1(&classes(), &[b, a]);
        assert_eq!(forward.f1_percent, reversed.f1_percent);
        assert_eq!(forward.precision_percent, reversed.precision_percent);
        assert_eq!(forward.recall_percent, reversed.recall_percent);
    }

    #[test]
    fn empty_sample_is_vacuously_perfect() {
        let agg = aggregate_f1(&classes(), &[]);
        assert_eq!(agg.runs, 0);
        assert_eq!(agg.f1_percent, 100);
        assert_eq!(agg.precision_percent, 100);
        assert_eq!(agg.recall_percent, 100);
    }

    #[test]
    fn aggregate_collects_unexpected_tools_across_runs() {
        let a = json!({ "tool_calls": [{ "name": "exec", "server": "shell" }] });
        let b = json!({ "tool_calls": [{ "name": "rm", "server": "shell" }] });
        let agg = aggregate_f1(&classes(), &[a, b]);
        assert_eq!(
            agg.unexpected_tools(),
            vec!["shell.exec".to_string(), "shell.rm".to_string()]
        );
    }

    #[test]
    fn metrics_object_exposes_the_three_targets() {
        use crate::executor::resolve_target;
        let agg = aggregate_f1(&classes(), &[json!({ "tool_calls": [] })]);
        let obj = metrics_object(&agg);
        assert!(resolve_target(&obj, "tool_selection.f1").is_number());
        assert!(resolve_target(&obj, "tool_selection.precision").is_number());
        assert!(resolve_target(&obj, "tool_selection.recall").is_number());
    }

    #[test]
    fn default_expect_gates_only_f1() {
        let defaults = default_expect();
        let targets: Vec<&str> = defaults.iter().map(|a| a.target.as_str()).collect();
        assert_eq!(targets, vec!["tool_selection.f1"]);
    }

    #[test]
    fn line_pairs_all_three_numbers() {
        let agg = F1Aggregate {
            runs: 3,
            per_run: vec![],
            precision_percent: 90,
            recall_percent: 80,
            f1_percent: 84,
        };
        assert_eq!(agg.line(), "F1 84%, precision 90%, recall 80% over 3 runs");
    }
}
