//! Context-rot knee probe (`context_rot_probe:` gate) for the live runner.
//!
//! Why this module exists: a single accuracy number at one catalog size hides
//! the cliff where a long candidate list starts to confuse the model. The probe
//! sweeps the served tool-catalog size across the gate's `levels:` (injecting
//! that many distractor tools per level), re-runs the agent at each level, scores
//! selection accuracy per level, and feeds the `(context_tokens, quality)` points
//! to [`mcptest_core::eval::context_rot::summarize`]. The CLI then gates the
//! resulting `context_rot.*` targets with no model in the scoring loop.
//!
//! Split into a pure, offline-testable evaluator
//! ([`eval_context_rot_probe_gate`]) and a live async sweep driver
//! ([`run_context_rot_sweep`]) that re-runs the agent per level and computes each
//! level's served-catalog token count. Lives next to
//! [`super::run_live_agent_metrics`], whose per-gate evaluators it mirrors.

use anyhow::{Context, Result};

use mcptest_core::checks::CheckOutcome;
use mcptest_core::eval::context_rot::{summarize, ContextPoint};
use mcptest_core::eval::distractors::{score_selection, DistractorSpec, BUNDLED_CATALOG};
use mcptest_core::eval::provider::ToolSpec;
use mcptest_core::eval::tool_selection_f1::observed_calls;
use mcptest_core::executor::eval_metric_assertion;
use mcptest_core::suite::ContextRotProbeGate;
use serde_json::{json, Value};

/// Evaluate one agent test's `context_rot_probe:` gate over a per-level sample.
///
/// `per_level` is one `(context_tokens, envelopes)` entry per swept level, in the
/// same declared order as `gate.levels`, where `context_tokens` is the
/// served-catalog token count at that level and `envelopes` are the agent run
/// envelopes recorded at it. For each level this pools the observed tool calls,
/// scores selection accuracy against the gate's `correct` set and that level's
/// resolved distractor ids (the noise it injected), and turns the
/// `(context_tokens, accuracy)` pair into a [`ContextPoint`]. It then folds the
/// points through [`summarize`], applies the gate's `expect:` over the
/// `context_rot.*` targets (an empty `expect:` reports the curve without gating),
/// and returns a single PASS/FAIL [`CheckOutcome`]. Pure: no async, no I/O, so it
/// is unit-testable with synthetic envelopes. Mirrors
/// [`super::run_live_agent_metrics::eval_distractor_gate`].
pub(crate) fn eval_context_rot_probe_gate(
    case_name: &str,
    gate: &ContextRotProbeGate,
    per_level: &[(u32, Vec<Value>)],
) -> CheckOutcome {
    let points: Vec<ContextPoint> = per_level
        .iter()
        .enumerate()
        .map(|(i, (context_tokens, envelopes))| {
            // The distractor id set at this level is sized by the level's
            // injected count (`gate.levels[i]`), not the token count, so a chosen
            // distractor scores against the noise that level actually added. The
            // sweep produces `per_level` in declared level order, so the count is
            // `gate.levels[i]`; a defensive fallback of 0 keeps the pure
            // evaluator total if the caller passes a shorter `levels`.
            let count = gate.levels.get(i).copied().unwrap_or(0);
            let distractor_ids = DistractorSpec {
                count,
                source: gate.source.clone(),
                complexity: None,
            }
            .resolve_ids();
            let chosen: Vec<String> = envelopes.iter().flat_map(observed_calls).collect();
            let score = score_selection(&chosen, &gate.correct, &distractor_ids);
            ContextPoint {
                context_tokens: *context_tokens,
                quality_percent: score.accuracy_percent,
            }
        })
        .collect();
    let summary = summarize(&points);
    let metrics = summary.metrics_object();
    // An empty `expect:` reports the curve without gating: a context-rot floor is
    // a deliberate choice (any real sweep shows some degradation), so the least
    // surprising default is a no-op pass that still records the knee diagnostics.
    let mut failure: Option<String> = None;
    for (i, assertion) in gate.expect.iter().enumerate() {
        if let Err(reason) = eval_metric_assertion(&metrics, assertion) {
            if failure.is_none() {
                failure = Some(format!(
                    "context-rot assertion #{i} (`{}`) failed: {reason}; {}",
                    assertion.target,
                    summary.line()
                ));
            }
        }
    }
    CheckOutcome {
        name: format!("{case_name} (context-rot-probe)"),
        passed: failure.is_none(),
        failure,
    }
}

/// The first `count` bundled catalog entries mapped to [`ToolSpec`]s the driver
/// can advertise as extra tools.
///
/// Each gets a trivial `{"type":"object"}` input schema: the distractors are
/// padding for the candidate list, never dispatched, so a richer schema would
/// only inflate the token count without changing what they probe. Capped at the
/// catalog size by `take`, so an over-large level injects every bundled tool.
fn distractor_tool_specs(count: u32) -> Vec<ToolSpec> {
    BUNDLED_CATALOG
        .iter()
        .take(count as usize)
        .map(|entry| ToolSpec {
            name: entry.name.to_string(),
            description: entry.description.to_string(),
            input_schema: json!({ "type": "object" }),
        })
        .collect()
}

/// Count the tokens of a served tool catalog (real tools plus the level's
/// injected distractors), the `context_tokens` x-axis of the context-rot curve.
///
/// Mirrors `super::run_live_agent_metrics::agent_tool_surface_tokens`: it counts
/// name + description + `inputSchema` per tool with the `cl100k_base` BPE
/// tokenizer and the doctor catalog-token report at maxed-out thresholds (so it
/// returns the catalog total, not the doctor warnings). Returns 0 when the
/// tokenizer cannot load, so a level still contributes a point rather than
/// stranding the whole sweep.
fn catalog_tokens(tools: &[ToolSpec]) -> u32 {
    let tools_list = json!({
        "tools": tools
            .iter()
            .map(|t| json!({
                "name": t.name,
                "description": t.description,
                "inputSchema": t.input_schema,
            }))
            .collect::<Vec<_>>()
    });
    let Ok(tokenizer) = mcptest_core::matchers::budgets::BpeTokenizer::cl100k_base() else {
        return 0;
    };
    let report =
        crate::doctor::compute_catalog_tokens(&tools_list, &tokenizer, usize::MAX, usize::MAX);
    report.total.min(u32::MAX as usize) as u32
}

/// Read-only inputs for the live context-rot sweep. Bundled so
/// [`run_context_rot_sweep`] takes a single context parameter.
pub(crate) struct SweepCtx<'a> {
    /// The agent test whose `context_rot_probe:` gate drives the sweep.
    pub(crate) agent: &'a mcptest_core::suite::AgentTest,
    /// The base per-cell driver config; cloned per level with `extra_tools` set.
    pub(crate) cfg: &'a mcptest_agent::AgentConfig,
    /// The connected MCP server handles the agent runs against.
    pub(crate) handles: &'a [mcptest_agent::ServerHandle<'a>],
    /// The shared dollar/token budget the per-level runs charge against.
    pub(crate) budget: &'a std::sync::Arc<mcptest_agent::Budget>,
    /// The resolved provider for this cell (the same one the normal passes used).
    pub(crate) provider: &'a dyn mcptest_core::eval::provider::LlmProvider,
    /// Real merged server tools the model sees before any distractor padding, so
    /// each level's token count includes the real surface, not just the noise.
    pub(crate) base_tools: &'a [ToolSpec],
    /// Report row name the outcome is attributed to.
    pub(crate) case_name: &'a str,
}

/// Drive the live context-rot sweep and evaluate the gate.
///
/// For each level in `gate.levels`, this clones the per-cell config, sets
/// `extra_tools` to the first `level` bundled distractors, re-runs the agent
/// `runs:` times to collect that level's envelopes, and computes the served
/// catalog's token count (real tools plus that level's distractors). It then
/// hands the `(context_tokens, envelopes)` pairs to
/// [`eval_context_rot_probe_gate`] and returns the single PASS/FAIL outcome. Each
/// level re-runs the agent, so this is live-only (a cassette replay carries a
/// single catalog and does not call it).
pub(crate) async fn run_context_rot_sweep(ctx: SweepCtx<'_>) -> Result<CheckOutcome> {
    let Some(gate) = &ctx.agent.context_rot_probe else {
        // The caller only invokes this when the gate is present; an absent gate
        // is a trivially-passing no-op so the call site stays unconditional-safe.
        return Ok(CheckOutcome {
            name: format!("{} (context-rot-probe)", ctx.case_name),
            passed: true,
            failure: None,
        });
    };
    let run_count = ctx.agent.runs.max(1) as usize;
    let mut per_level: Vec<(u32, Vec<Value>)> = Vec::with_capacity(gate.levels.len());
    for level in &gate.levels {
        let extra_tools = distractor_tool_specs(*level);
        // The served catalog at this level: the real merged tools plus the
        // injected distractors, the same surface the driver advertises below.
        let mut served = ctx.base_tools.to_vec();
        served.extend_from_slice(&extra_tools);
        let tokens = catalog_tokens(&served);
        let mut level_cfg = ctx.cfg.clone();
        level_cfg.extra_tools = extra_tools;
        let mut envelopes: Vec<Value> = Vec::with_capacity(run_count);
        for _ in 0..run_count {
            let trace = super::run_live_agents::drive_agent(
                ctx.agent,
                &level_cfg,
                ctx.provider,
                ctx.handles,
                std::sync::Arc::clone(ctx.budget),
            )
            .await
            .map_err(|e| {
                anyhow::anyhow!(
                    "agent `{}` context-rot sweep (`{}`) level {level} run failed: {}",
                    ctx.agent.name,
                    ctx.case_name,
                    mcptest_agent::redact_secrets(&e.to_string())
                )
            })?;
            envelopes.push(trace.to_envelope());
        }
        per_level.push((tokens, envelopes));
    }
    Ok(eval_context_rot_probe_gate(ctx.case_name, gate, &per_level))
}

/// Run the `context_rot_probe:` sweep for one cell when it declares the gate,
/// pushing the gate outcome onto `acc`.
///
/// Live-only: the sweep re-runs the agent at each catalog level, so a
/// stub-provider cell (no API key resolved) is skipped with a SKIP note rather
/// than sweeping canned responses (which would show no rot). Called from
/// [`super::run_live_agents::dispatch_agent_cell`]; it lives here, not there, so
/// that module stays under the production-module-size guard.
#[allow(clippy::too_many_arguments)]
pub(super) async fn run_context_rot_sweep_cell(
    agent: &mcptest_core::suite::AgentTest,
    cfg: &mcptest_agent::AgentConfig,
    handles: &[mcptest_agent::ServerHandle<'_>],
    budget: &std::sync::Arc<mcptest_agent::Budget>,
    lookup: &mcptest_agent::ProviderLookup,
    case_name: &str,
    quiet: bool,
    acc: &mut super::run_live_agents::CellAcc<'_>,
) -> Result<()> {
    use mcptest_agent::ProviderLookup;
    if agent.context_rot_probe.is_none() {
        return Ok(());
    }
    // The sweep needs a real provider to be meaningful; a stub returns canned
    // text regardless of catalog size, so a sweep over it would show no rot.
    let ProviderLookup::Resolved { provider, .. } = lookup else {
        if !quiet {
            println!(
                "context-rot probe [SKIP] {case_name}: no real provider (set the model's API key)"
            );
        }
        return Ok(());
    };
    // List the real merged server tools once so each level's token count includes
    // the real surface, not just the injected distractors.
    let base_tools = mcptest_agent::list_merged_tools(handles)
        .await
        .with_context(|| {
            format!("context-rot probe for `{case_name}`: listing the server tools failed")
        })?;
    let ctx = SweepCtx {
        agent,
        cfg,
        handles,
        budget,
        provider: provider.as_ref(),
        base_tools: &base_tools,
        case_name,
    };
    let outcome = run_context_rot_sweep(ctx).await?;
    if !quiet {
        let verdict = if outcome.passed { "PASS" } else { "FAIL" };
        let note = outcome.failure.as_deref().unwrap_or("curve recorded");
        println!("context-rot probe [{verdict}] {case_name}: {note}");
    }
    acc.f1_outcomes.push(outcome);
    Ok(())
}

#[cfg(test)]
mod tests {
    use super::*;
    use mcptest_core::eval::distractors::DistractorSource;
    use mcptest_core::suite::{Assertion, Matcher};

    /// Build a run envelope whose `tool_calls` name the given tools with no
    /// server key, so [`observed_calls`] yields the bare ids the catalog
    /// distractors and bare correct ids compare against. Mirrors
    /// `run_live_agent_metrics::envelope_with_calls`.
    fn envelope_with_calls(names: &[&str]) -> Value {
        let calls: Vec<Value> = names.iter().map(|n| json!({ "name": n })).collect();
        json!({ "tool_calls": calls })
    }

    fn gate(levels: Vec<u32>, correct: &[&str], expect: Vec<Assertion>) -> ContextRotProbeGate {
        ContextRotProbeGate {
            levels,
            source: DistractorSource::Catalog,
            correct: correct.iter().map(|s| s.to_string()).collect(),
            expect,
        }
    }

    /// A schema-matcher assertion over a `context_rot.*` target.
    fn assertion(target: &str, schema: Value) -> Assertion {
        Assertion {
            target: target.to_string(),
            matcher: Matcher::Schema(schema),
            weight: 1.0,
            name: None,
            transform: None,
        }
    }

    #[test]
    fn distractor_tool_specs_caps_at_catalog_size() {
        assert_eq!(distractor_tool_specs(2).len(), 2);
        assert_eq!(distractor_tool_specs(9999).len(), BUNDLED_CATALOG.len());
        // Each carries the bundled name and a trivial object schema.
        let specs = distractor_tool_specs(1);
        assert_eq!(specs[0].name, BUNDLED_CATALOG[0].name);
        assert_eq!(specs[0].input_schema, json!({ "type": "object" }));
    }

    #[test]
    fn an_empty_expect_reports_the_curve_without_gating() {
        // A real knee curve: quality holds, then a cliff at the largest catalog.
        let per_level = vec![
            (1000, vec![envelope_with_calls(&["search"])]),
            (4000, vec![envelope_with_calls(&["search"])]),
            // At the fattest catalog the model fell for the first bundled
            // distractor (weather_forecast), so accuracy at this level is 0.
            (16000, vec![envelope_with_calls(&["weather_forecast"])]),
        ];
        let g = gate(vec![0, 4, 8], &["search"], vec![]);
        let outcome = eval_context_rot_probe_gate("picker", &g, &per_level);
        // No expect: means report-only, so even a steep cliff passes.
        assert!(outcome.passed, "{:?}", outcome.failure);
        assert_eq!(outcome.name, "picker (context-rot-probe)");
    }

    #[test]
    fn a_curve_with_a_cliff_fails_a_max_degradation_gate() {
        // Quality 100 at the two small catalogs, then 0 at the largest: the
        // agent picks `weather_forecast` (the first bundled distractor, injected
        // once 4+ distractors are present) instead of the correct `search`.
        let per_level = vec![
            (1000, vec![envelope_with_calls(&["search"])]),
            (4000, vec![envelope_with_calls(&["search"])]),
            (16000, vec![envelope_with_calls(&["weather_forecast"])]),
        ];
        // `max_degradation` is the numeric gate target (knee_detected is a
        // boolean diagnostic). The drop from 100 to 0 is 100 points, so a
        // `maximum: 0` floor must fail.
        let g = gate(
            vec![0, 4, 8],
            &["search"],
            vec![assertion(
                "context_rot.max_degradation",
                json!({ "maximum": 0 }),
            )],
        );
        let outcome = eval_context_rot_probe_gate("picker", &g, &per_level);
        assert!(!outcome.passed);
        assert!(outcome
            .failure
            .unwrap()
            .contains("context_rot.max_degradation"));
    }

    #[test]
    fn a_flat_curve_passes_a_max_degradation_floor() {
        // Quality holds near-perfect across every catalog size: no real rot.
        let per_level = vec![
            (1000, vec![envelope_with_calls(&["search"])]),
            (4000, vec![envelope_with_calls(&["search"])]),
            (16000, vec![envelope_with_calls(&["search"])]),
        ];
        // Allow up to 10pt of degradation; a flat curve degrades 0.
        let g = gate(
            vec![0, 4, 8],
            &["search"],
            vec![assertion(
                "context_rot.max_degradation",
                json!({ "maximum": 10 }),
            )],
        );
        let outcome = eval_context_rot_probe_gate("picker", &g, &per_level);
        assert!(outcome.passed, "{:?}", outcome.failure);
    }

    #[test]
    fn a_steep_curve_fails_a_max_degradation_floor() {
        // Perfect at the small catalog, all-distractor at the large one: a 100pt
        // drop, which breaches a 10pt degradation ceiling.
        let per_level = vec![
            (1000, vec![envelope_with_calls(&["search"])]),
            (16000, vec![envelope_with_calls(&["weather_forecast"])]),
        ];
        let g = gate(
            vec![0, 8],
            &["search"],
            vec![assertion(
                "context_rot.max_degradation",
                json!({ "maximum": 10 }),
            )],
        );
        let outcome = eval_context_rot_probe_gate("picker", &g, &per_level);
        assert!(!outcome.passed);
        assert!(outcome
            .failure
            .unwrap()
            .contains("context_rot.max_degradation"));
    }

    #[test]
    fn an_empty_sample_is_trivially_robust() {
        // No levels collected: the summary is the empty-curve baseline (no rot),
        // so even a knee-forbidding gate passes (there is no knee).
        let g = gate(
            vec![],
            &["search"],
            vec![assertion(
                "context_rot.knee_detected",
                json!({ "maximum": 0 }),
            )],
        );
        let outcome = eval_context_rot_probe_gate("picker", &g, &[]);
        assert!(outcome.passed, "{:?}", outcome.failure);
    }
}
