//! The mutation-backed tool-description benchmark.
//!
//! [`run`] scores a catalog three ways and ties them together: the
//! deterministic TDQS description quality, the named smell counts, and the
//! lexical selection accuracy (with a Wilson confidence interval) over a set
//! of intents. It then applies each requested description mutation one at a
//! time and reports how selection moved, so a reviewer can see whether a
//! description edit improves or degrades selection, not just its style score.
//! Everything here is deterministic and model-free.

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

use super::mutation::{self, Mutation};
use super::selection::{self, Intent};
use super::smells::{self, Smell};
use super::stats::{wilson, Proportion, Z_95};
use crate::lint::tdqs;
use crate::lint::tool_descriptions::{default_rules, lint_catalog, Rule};
use std::collections::BTreeMap;

/// Quality scores for one catalog (baseline or mutated).
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub struct CatalogScore {
    /// Selection accuracy over the intent set, with a 95% Wilson interval.
    pub selection: Proportion,
    /// Mean per-tool TDQS score, `None` for an empty catalog.
    pub tdqs_mean: Option<f64>,
    /// Minimum per-tool TDQS score, `None` for an empty catalog.
    pub tdqs_min: Option<f64>,
    /// Smell counts by category.
    pub smells: BTreeMap<Smell, usize>,
    /// Total smell findings across all categories.
    pub smell_total: usize,
    /// Whitespace-token estimate of the catalog's name + description text.
    pub catalog_tokens: u64,
}

/// How one mutation moved selection and smells relative to the baseline.
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub struct MutationOutcome {
    /// Tool whose description was mutated.
    pub tool: String,
    /// The mutation applied.
    pub mutation: Mutation,
    /// Short label, e.g. `drop_keyword(weather)`.
    pub label: String,
    /// Selection accuracy after the mutation, with its interval.
    pub selection: Proportion,
    /// `baseline_accuracy - mutated_accuracy`. Positive means the mutation
    /// degraded selection; negative means it improved it.
    pub selection_drop: f64,
    /// `mutated_smell_total - baseline_smell_total`.
    pub smell_delta: i64,
    /// `mutated_tokens - baseline_tokens` (negative when text was removed).
    pub token_delta: i64,
}

/// The full benchmark result.
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub struct ToolBenchReport {
    /// The unmutated catalog's scores.
    pub baseline: CatalogScore,
    /// One entry per applied mutation, in order.
    pub mutations: Vec<MutationOutcome>,
}

impl ToolBenchReport {
    /// The largest selection drop any single mutation caused. A positive value
    /// is the headline finding: degrading a description measurably hurt
    /// selection. Zero when no mutation moved selection.
    pub fn max_selection_drop(&self) -> f64 {
        self.mutations
            .iter()
            .map(|m| m.selection_drop)
            .fold(0.0, f64::max)
    }
}

/// Score one catalog against the intent set.
pub fn score_catalog(catalog: &Value, intents: &[Intent]) -> CatalogScore {
    let (correct, total) = selection::score(intents, catalog);
    let tdqs = tdqs::score_catalog(catalog);
    let rules = default_rules();
    let refs: Vec<&dyn Rule> = rules.iter().map(AsRef::as_ref).collect();
    let smell_counts = smells::summarize(&lint_catalog(catalog, &refs));
    let smell_total = smell_counts.values().sum();
    CatalogScore {
        selection: wilson(correct, total, Z_95),
        tdqs_mean: tdqs.mean_tool_score(),
        tdqs_min: tdqs.min_tool_score(),
        smells: smell_counts,
        smell_total,
        catalog_tokens: catalog_tokens(catalog),
    }
}

/// Run the benchmark: score the baseline, then score each mutation in turn.
pub fn run(
    catalog: &Value,
    intents: &[Intent],
    mutations: &[(String, Mutation)],
) -> ToolBenchReport {
    let baseline = score_catalog(catalog, intents);
    let outcomes = mutations
        .iter()
        .map(|(tool, mutation)| {
            let mutated = mutation::apply(catalog, tool, mutation);
            let score = score_catalog(&mutated, intents);
            MutationOutcome {
                tool: tool.clone(),
                mutation: mutation.clone(),
                label: mutation.label(),
                selection_drop: baseline.selection.estimate - score.selection.estimate,
                smell_delta: score.smell_total as i64 - baseline.smell_total as i64,
                token_delta: score.catalog_tokens as i64 - baseline.catalog_tokens as i64,
                selection: score.selection,
            }
        })
        .collect();
    ToolBenchReport {
        baseline,
        mutations: outcomes,
    }
}

/// Whitespace-token estimate over every tool's name and description. A crude
/// but deterministic proxy for the catalog token cost a model would pay.
fn catalog_tokens(catalog: &Value) -> u64 {
    let Some(tools) = catalog.get("tools").and_then(Value::as_array) else {
        return 0;
    };
    let mut count = 0u64;
    for tool in tools {
        for field in ["name", "description"] {
            if let Some(text) = tool.get(field).and_then(Value::as_str) {
                count += text.split_whitespace().count() as u64;
            }
        }
    }
    count
}

/// Flatten the report into the assertion target object the `expect:` grammar
/// resolves against.
pub fn report_envelope(report: &ToolBenchReport) -> Value {
    json!({
        "selection": {
            "accuracy": report.baseline.selection.estimate,
            "low": report.baseline.selection.low,
            "high": report.baseline.selection.high,
        },
        "tdqs": {
            "mean": report.baseline.tdqs_mean,
            "min": report.baseline.tdqs_min,
        },
        "smells": smells::metrics_object(&report.baseline.smells),
        "smell_total": report.baseline.smell_total,
        "catalog_tokens": report.baseline.catalog_tokens,
        "max_selection_drop": report.max_selection_drop(),
    })
}

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

    // Opaque tool names so the description, not the name, carries the
    // selection signal. This is the case the benchmark is built for: when a
    // name is unhelpful, description quality drives whether the right tool is
    // chosen.
    fn catalog() -> Value {
        json!({
            "tools": [
                {
                    "name": "alpha_tool",
                    "description": "Return the current weather forecast and temperature for a city."
                },
                {
                    "name": "beta_tool",
                    "description": "Return the latest news headlines for a topic or region."
                }
            ]
        })
    }

    fn intents() -> Vec<Intent> {
        vec![
            Intent {
                query: "what is the weather forecast in Paris".into(),
                gold: vec!["alpha_tool".into()],
            },
            Intent {
                query: "show me the latest news headlines".into(),
                gold: vec!["beta_tool".into()],
            },
        ]
    }

    #[test]
    fn a_good_catalog_selects_correctly() {
        let score = score_catalog(&catalog(), &intents());
        assert_eq!(score.selection.estimate, 1.0);
        // The descriptions are specific, so neither tool is underspecified.
        assert_eq!(
            score
                .smells
                .get(&Smell::Underspecified)
                .copied()
                .unwrap_or(0),
            0
        );
        assert!(score.tdqs_mean.unwrap() > 0.4);
    }

    #[test]
    fn blanking_a_description_degrades_selection() {
        let mutations = vec![("alpha_tool".to_string(), Mutation::Blank)];
        let report = run(&catalog(), &intents(), &mutations);
        assert_eq!(report.baseline.selection.estimate, 1.0);
        // With no description and an opaque name, the weather intent no longer
        // matches its tool.
        let outcome = &report.mutations[0];
        assert!(outcome.selection_drop > 0.0, "{outcome:?}");
        assert!(outcome.smell_delta > 0, "blanking adds smells: {outcome:?}");
        assert!(outcome.token_delta < 0, "blanking removes tokens");
        assert!(report.max_selection_drop() > 0.0);
    }

    #[test]
    fn genericizing_a_description_degrades_selection() {
        let mutations = vec![("beta_tool".to_string(), Mutation::Genericize)];
        let report = run(&catalog(), &intents(), &mutations);
        assert!(
            report.mutations[0].selection_drop > 0.0,
            "{:?}",
            report.mutations[0]
        );
    }

    #[test]
    fn dropping_the_sole_discriminating_keyword_loses_the_selection() {
        // The intent's only content overlap with the correct tool is the one
        // keyword. Dropping it leaves nothing to match on, so the selector
        // abstains and the pick is lost.
        let catalog = json!({
            "tools": [
                { "name": "alpha_tool", "description": "Translate text between human languages." },
                { "name": "beta_tool", "description": "Fetch the weather forecast for a city." }
            ]
        });
        let intents = vec![Intent {
            query: "what is the weather".into(),
            gold: vec!["beta_tool".into()],
        }];
        let mutations = vec![(
            "beta_tool".to_string(),
            Mutation::DropKeyword {
                word: "weather".into(),
            },
        )];
        let report = run(&catalog, &intents, &mutations);
        assert_eq!(report.baseline.selection.estimate, 1.0);
        assert!(
            report.mutations[0].selection_drop > 0.0,
            "{:?}",
            report.mutations[0]
        );
    }

    #[test]
    fn envelope_exposes_selection_smells_and_drop() {
        let mutations = vec![("alpha_tool".to_string(), Mutation::Blank)];
        let env = report_envelope(&run(&catalog(), &intents(), &mutations));
        assert_eq!(env["selection"]["accuracy"], json!(1.0));
        assert!(env["max_selection_drop"].as_f64().unwrap() > 0.0);
        assert!(env["catalog_tokens"].as_u64().unwrap() > 0);
    }
}
