//! Comparative matrix-level assertions for the live runner.
//!
//! A single agent test fans out across a model matrix, one cell per model.
//! The per-cell `threshold:` judges each cell alone; a `matrix_assert:`
//! compares cells against each other. This module owns the two constructs:
//!
//! - `max_score`: the GATING floor. The run fails unless the best cell's
//!   `combined_score` across the matrix is at or above it.
//! - `select_best`: reporting only. Names the highest-scoring cell.
//!
//! The aggregation runs after the report is built (so each cell carries its
//! `report.tests[].score`), is pulled out of `run_live.rs` to keep that
//! module under the production-size guard, and is a pure function over the
//! report plus the recorded cells so it is unit-testable offline.

use mcptest_core::report::Report;
use mcptest_core::suite::{MatrixAssert, ModelRef, ResolvedSuite};

use crate::cli::args::RunArgs;
use crate::cli::global::Reporter;

/// One matrix cell dispatched live, recorded so the post-run aggregation can
/// group cells by their owning agent test.
///
/// Mirrors the `recorded_cells` resume-journal pattern but tracks what the
/// matrix gate needs: the agent name (group key), the model label (for the
/// report line), and the reporter row names this cell produced so the best
/// per-cell score can be looked up from the report.
pub(super) struct MatrixCell {
    /// Agent test name, the group key. All cells of one agent test share it.
    pub(super) agent: String,
    /// Model display label, shown in the select-best / floor report line.
    pub(super) label: String,
    /// Reporter row names this cell contributed (one per `runs:` iteration).
    pub(super) rows: Vec<String>,
}

/// Outcome of evaluating one agent test's matrix-level assertion.
///
/// Separates the report line (always produced when cells exist) from the
/// gate verdict so the caller can print the former and gate on the latter.
pub(super) struct MatrixVerdict {
    /// Human-readable summary line for the run output.
    pub(super) line: String,
    /// `false` when a `max_score` floor was set and the best cell fell
    /// below it. `true` otherwise (no floor, or the floor was met).
    pub(super) gate_passed: bool,
}

/// Best cell found across one agent test's matrix.
struct BestCell {
    label: String,
    score: f64,
}

/// Find the highest-scoring cell across an agent test's matrix.
///
/// Reads `report.tests[].score` for every row each cell produced and keeps
/// the cell with the largest score. A cell whose rows carry no score (the
/// score model was not engaged) is skipped. Returns `None` when no cell
/// produced a comparable score, so the caller can report "no scored cell"
/// instead of inventing a zero.
fn best_cell(report: &Report, cells: &[&MatrixCell]) -> Option<BestCell> {
    let mut best: Option<BestCell> = None;
    for cell in cells {
        let cell_score = cell
            .rows
            .iter()
            .filter_map(|row| score_for_row(report, row))
            .fold(None, |acc: Option<f64>, s| {
                Some(acc.map_or(s, |a| a.max(s)))
            });
        let Some(score) = cell_score else {
            continue;
        };
        if best.as_ref().is_none_or(|b| score > b.score) {
            best = Some(BestCell {
                label: cell.label.clone(),
                score,
            });
        }
    }
    best
}

/// Look up the combined score the report recorded for one reporter row.
fn score_for_row(report: &Report, row: &str) -> Option<f64> {
    report
        .tests
        .iter()
        .find(|t| t.name == row)
        .and_then(|t| t.score)
}

/// Evaluate one agent test's `matrix_assert:` against the run report.
///
/// Picks the best cell, builds a report line covering both the select-best
/// and the floor result, and decides whether the floor gate passed. When no
/// cell produced a score, a `max_score` floor cannot be met (so it gates the
/// run), while a `select_best`-only assertion just reports the gap.
pub(super) fn evaluate_matrix_assert(
    report: &Report,
    agent: &str,
    assert: &MatrixAssert,
    cells: &[&MatrixCell],
) -> MatrixVerdict {
    let best = best_cell(report, cells);
    match (best, assert.max_score) {
        (Some(best), Some(floor)) => {
            let passed = best.score >= floor;
            let verdict = if passed { "PASS" } else { "FAIL" };
            MatrixVerdict {
                line: format!(
                    "matrix-assert [{verdict}] {agent}: best cell `{}` scored {:.3} \
                     (max-score floor {floor:.3})",
                    best.label, best.score
                ),
                gate_passed: passed,
            }
        }
        (Some(best), None) => MatrixVerdict {
            line: format!(
                "matrix-assert [BEST] {agent}: best cell `{}` scored {:.3}",
                best.label, best.score
            ),
            gate_passed: true,
        },
        (None, Some(floor)) => MatrixVerdict {
            // No cell produced a score, so a floor cannot be cleared.
            line: format!(
                "matrix-assert [FAIL] {agent}: no cell produced a score, \
                 cannot clear max-score floor {floor:.3}"
            ),
            gate_passed: false,
        },
        (None, None) => MatrixVerdict {
            line: format!("matrix-assert [BEST] {agent}: no cell produced a score"),
            gate_passed: true,
        },
    }
}

/// Apply a `--models` model-matrix sweep to the suite and pick its reporter.
///
/// When `--models` is set, every agent test's `models:` is replaced with the
/// CLI list so each agent fans across all of them in one run; the existing
/// per-(agent, model) cell machinery then feeds the comparison-matrix grid.
/// When the operator did not pass an explicit `--reporter`, the run defaults to
/// the `matrix` reporter, since a model sweep's natural output is the grid.
/// Returns the reporter to use, unchanged when no `--models` were requested.
pub(super) fn apply_model_matrix(
    suite: &mut ResolvedSuite,
    args: &RunArgs,
    reporter: Option<Reporter>,
) -> Option<Reporter> {
    if args.models.is_empty() {
        return reporter;
    }
    let models: Vec<ModelRef> = args
        .models
        .iter()
        .map(|m| ModelRef::AutoDetect(m.clone()))
        .collect();
    for agent in &mut suite.agents {
        agent.models = models.clone();
    }
    reporter.or(Some(Reporter::Matrix))
}

/// Attach the recorded matrix cells to the report so the `matrix` reporters
/// (and the cloud viewer) render the test-by-model grid from the same artifact
/// every other reporter consumes. No-op when no matrix cells were recorded, so
/// an ordinary run keeps `report.matrix` as `None`.
pub(super) fn attach_matrix_section(report: &mut Report, cells: &[MatrixCell]) {
    if cells.is_empty() {
        return;
    }
    report.matrix = Some(mcptest_core::report::MatrixReport {
        cells: cells
            .iter()
            .map(|c| mcptest_core::report::MatrixCell {
                agent: c.agent.clone(),
                label: c.label.clone(),
                rows: c.rows.clone(),
            })
            .collect(),
    });
}

/// Group recorded matrix cells by their owning agent test, preserving the
/// declaration order agents were dispatched in so the report lines read
/// stably across runs.
pub(super) fn group_by_agent(cells: &[MatrixCell]) -> Vec<(String, Vec<&MatrixCell>)> {
    let mut order: Vec<String> = Vec::new();
    let mut groups: std::collections::HashMap<String, Vec<&MatrixCell>> =
        std::collections::HashMap::new();
    for cell in cells {
        if !groups.contains_key(&cell.agent) {
            order.push(cell.agent.clone());
        }
        groups.entry(cell.agent.clone()).or_default().push(cell);
    }
    order
        .into_iter()
        .map(|agent| {
            let cells = groups.remove(&agent).unwrap_or_default();
            (agent, cells)
        })
        .collect()
}

/// Gate on the comparative matrix-level assertions. Groups the
/// recorded matrix cells by their owning agent test, evaluates each agent's
/// `matrix_assert:` against the report, prints the select-best / floor line,
/// and returns `true` when any `max_score` floor was breached. A breach is a
/// failing test outcome, so it reuses exit code 1 like the other gates.
pub(super) fn gate_matrix_asserts(
    report: &Report,
    matrix_cells: &[MatrixCell],
    suite: &mcptest_core::suite::ResolvedSuite,
    quiet: bool,
) -> bool {
    let mut any_failed = false;
    for (agent_name, cells) in group_by_agent(matrix_cells) {
        let Some(assert) = suite
            .agents
            .iter()
            .find(|a| a.name == agent_name)
            .and_then(|a| a.matrix_assert.as_ref())
        else {
            continue;
        };
        let verdict = evaluate_matrix_assert(report, &agent_name, assert, &cells);
        if !verdict.gate_passed {
            any_failed = true;
        }
        if !quiet {
            println!("{}", verdict.line);
        }
    }
    any_failed
}

#[cfg(test)]
mod tests {
    use super::*;
    use mcptest_core::report::{Report, RunSummary, TestResult, TestStatus};

    fn scored_row(name: &str, score: Option<f64>, status: TestStatus) -> TestResult {
        TestResult {
            name: name.into(),
            status,
            duration_ms: 1,
            cache_hit: false,
            started_at: None,
            ended_at: None,
            failure: None,
            file: None,
            line: None,
            steps: Vec::new(),
            consensus: None,
            rule_id: None,
            severity: None,
            score,
            derived_metrics: Vec::new(),
            natural_language: false,
        }
    }

    fn report_with(rows: Vec<TestResult>) -> Report {
        Report {
            schema_version: "0.1".into(),
            mcptest_version: "test".into(),
            generated_at: "2026-05-24T00:00:00Z".into(),
            run_id: "test-run".into(),
            // Derived (WOR-1317) so the fixture cannot drift from its rows.
            summary: RunSummary::from_tests(&rows),
            tests: rows,
            cost: None,
            metadata: None,
            matrix: None,
            protocol_violations: Vec::new(),
        }
    }

    fn cell(agent: &str, label: &str, rows: &[&str]) -> MatrixCell {
        MatrixCell {
            agent: agent.into(),
            label: label.into(),
            rows: rows.iter().map(|r| r.to_string()).collect(),
        }
    }

    #[test]
    fn best_cell_picks_the_highest_score_across_the_matrix() {
        let report = report_with(vec![
            scored_row("a [m1]", Some(0.4), TestStatus::Fail),
            scored_row("a [m2]", Some(0.9), TestStatus::Pass),
            scored_row("a [m3]", Some(0.6), TestStatus::Fail),
        ]);
        let cells = [
            cell("a", "m1", &["a [m1]"]),
            cell("a", "m2", &["a [m2]"]),
            cell("a", "m3", &["a [m3]"]),
        ];
        let refs: Vec<&MatrixCell> = cells.iter().collect();
        let best = best_cell(&report, &refs).expect("a scored cell");
        assert_eq!(best.label, "m2");
        assert!((best.score - 0.9).abs() < f64::EPSILON);
    }

    #[test]
    fn max_score_floor_passes_when_the_best_cell_clears_it() {
        let report = report_with(vec![
            scored_row("a [m1]", Some(0.4), TestStatus::Fail),
            scored_row("a [m2]", Some(0.85), TestStatus::Pass),
        ]);
        let cells = [cell("a", "m1", &["a [m1]"]), cell("a", "m2", &["a [m2]"])];
        let refs: Vec<&MatrixCell> = cells.iter().collect();
        let assert = MatrixAssert {
            max_score: Some(0.8),
            select_best: false,
        };
        let verdict = evaluate_matrix_assert(&report, "a", &assert, &refs);
        assert!(verdict.gate_passed, "best 0.85 clears floor 0.8");
        assert!(verdict.line.contains("PASS"), "{}", verdict.line);
        assert!(verdict.line.contains("m2"), "{}", verdict.line);
    }

    #[test]
    fn max_score_floor_fails_when_no_cell_clears_it() {
        let report = report_with(vec![
            scored_row("a [m1]", Some(0.4), TestStatus::Fail),
            scored_row("a [m2]", Some(0.55), TestStatus::Fail),
        ]);
        let cells = [cell("a", "m1", &["a [m1]"]), cell("a", "m2", &["a [m2]"])];
        let refs: Vec<&MatrixCell> = cells.iter().collect();
        let assert = MatrixAssert {
            max_score: Some(0.8),
            select_best: false,
        };
        let verdict = evaluate_matrix_assert(&report, "a", &assert, &refs);
        assert!(!verdict.gate_passed, "best 0.55 is below floor 0.8");
        assert!(verdict.line.contains("FAIL"), "{}", verdict.line);
    }

    #[test]
    fn select_best_reports_the_top_cell_and_never_gates() {
        let report = report_with(vec![
            scored_row("a [m1]", Some(0.2), TestStatus::Fail),
            scored_row("a [m2]", Some(0.7), TestStatus::Fail),
        ]);
        let cells = [cell("a", "m1", &["a [m1]"]), cell("a", "m2", &["a [m2]"])];
        let refs: Vec<&MatrixCell> = cells.iter().collect();
        let assert = MatrixAssert {
            max_score: None,
            select_best: true,
        };
        let verdict = evaluate_matrix_assert(&report, "a", &assert, &refs);
        // Reporting only: even with both cells failing their own verdict,
        // the matrix assertion never fails the run.
        assert!(verdict.gate_passed, "select_best never gates");
        assert!(verdict.line.contains("BEST"), "{}", verdict.line);
        assert!(verdict.line.contains("m2"), "names the top cell");
    }

    #[test]
    fn a_multi_run_cell_uses_its_highest_row_score() {
        // One cell, three `runs:` rows. The cell's score is the best row.
        let report = report_with(vec![
            scored_row("a #1", Some(0.3), TestStatus::Fail),
            scored_row("a #2", Some(0.8), TestStatus::Pass),
            scored_row("a #3", Some(0.5), TestStatus::Fail),
        ]);
        let cells = [cell("a", "m1", &["a #1", "a #2", "a #3"])];
        let refs: Vec<&MatrixCell> = cells.iter().collect();
        let best = best_cell(&report, &refs).expect("scored");
        assert!((best.score - 0.8).abs() < f64::EPSILON);
    }

    #[test]
    fn no_scored_cell_fails_a_floor_but_not_select_best() {
        let report = report_with(vec![scored_row("a [m1]", None, TestStatus::Pass)]);
        let cells = [cell("a", "m1", &["a [m1]"])];
        let refs: Vec<&MatrixCell> = cells.iter().collect();
        let gated = MatrixAssert {
            max_score: Some(0.5),
            select_best: false,
        };
        assert!(!evaluate_matrix_assert(&report, "a", &gated, &refs).gate_passed);
        let reporting = MatrixAssert {
            max_score: None,
            select_best: true,
        };
        assert!(evaluate_matrix_assert(&report, "a", &reporting, &refs).gate_passed);
    }

    fn suite_with_one_agent() -> ResolvedSuite {
        let yaml: serde_yaml::Value = serde_yaml::from_str(
            "servers:\n  s: { command: [\"echo\"] }\nagents:\n  - name: a\n    servers: [s]\n    prompt: hi\n    model: gpt-4o\n",
        )
        .unwrap();
        mcptest_core::suite::parse(&yaml).unwrap()
    }

    #[test]
    fn model_matrix_overrides_agent_models_and_defaults_to_matrix_reporter() {
        let mut suite = suite_with_one_agent();
        let args = RunArgs {
            models: vec!["gpt-4o".into(), "claude-3-5-sonnet".into()],
            ..Default::default()
        };
        let reporter = apply_model_matrix(&mut suite, &args, None);
        assert_eq!(suite.agents[0].models.len(), 2, "both models injected");
        assert!(
            matches!(reporter, Some(Reporter::Matrix)),
            "a model sweep defaults to the matrix reporter"
        );
    }

    #[test]
    fn no_models_flag_leaves_suite_and_reporter_untouched() {
        let mut suite = suite_with_one_agent();
        let reporter = apply_model_matrix(&mut suite, &RunArgs::default(), None);
        assert_eq!(suite.agents[0].models.len(), 1, "models untouched");
        assert!(reporter.is_none(), "no reporter forced without --models");
    }

    #[test]
    fn explicit_reporter_is_kept_under_a_model_sweep() {
        let mut suite = suite_with_one_agent();
        let args = RunArgs {
            models: vec!["a".into(), "b".into()],
            ..Default::default()
        };
        let reporter = apply_model_matrix(&mut suite, &args, Some(Reporter::Json));
        assert!(
            matches!(reporter, Some(Reporter::Json)),
            "an explicit --reporter wins over the matrix default"
        );
    }

    #[test]
    fn group_by_agent_keeps_cells_under_their_owner() {
        let cells = [
            cell("alpha", "m1", &["alpha [m1]"]),
            cell("beta", "m1", &["beta [m1]"]),
            cell("alpha", "m2", &["alpha [m2]"]),
        ];
        let groups = group_by_agent(&cells);
        assert_eq!(groups.len(), 2, "two distinct agents");
        assert_eq!(groups[0].0, "alpha", "declaration order preserved");
        assert_eq!(groups[0].1.len(), 2, "both alpha cells grouped");
        assert_eq!(groups[1].0, "beta");
        assert_eq!(groups[1].1.len(), 1);
    }
}
