//! Step-level budget gate for tool tests (WOR-1280).
//!
//! Why this exists: the long-form `expect:` mapping carries budgets
//! (`max_duration_ms`, `max_response_tokens`) and header checks alongside the
//! assertions, and before WOR-1280 nothing in the run path evaluated them, so
//! a budget-carrying test silently passed. This module is the single
//! evaluation point: it folds the parsed [`ExpectBudgets`] against the
//! measured call duration and the response envelope, reusing the shared
//! [`crate::matchers::budgets`] evaluators so no second copy of the budget
//! math can drift.

use serde_json::Value;

use crate::matchers::budgets::{
    evaluate_duration, evaluate_response_tokens, tokenizer_by_name, TokenBudget,
};
use crate::matchers::headers::HeaderMatch;
use crate::matchers::AssertionResult;
use crate::suite::{ExpectBudgets, ServerSpec};

/// Transport class of a test's server target, as far as budgets care.
///
/// Header checks need to know whether the target can have HTTP response
/// headers at all; the executor's [`crate::connector::LiveServer`] does not
/// carry its origin spec, so the plan builder derives this from the suite's
/// `servers:` entry and threads it onto the plan.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
pub enum TargetTransport {
    /// Subprocess server over stdio. No HTTP headers exist.
    #[default]
    Stdio,
    /// Streamable HTTP or legacy SSE server. Headers exist on the wire.
    Http,
    /// Cassette replay. No live wire, so no headers.
    Cassette,
}

impl TargetTransport {
    /// Classify a suite server spec. Lives here rather than at the CLI
    /// plan builder so the mapping is owned by the same module that
    /// consumes it.
    pub fn from_spec(spec: &ServerSpec) -> Self {
        match spec {
            ServerSpec::Stdio { .. } => TargetTransport::Stdio,
            ServerSpec::StreamableHttp { .. } | ServerSpec::LegacySse { .. } => {
                TargetTransport::Http
            }
            ServerSpec::Cassette { .. } => TargetTransport::Cassette,
        }
    }

    /// Stable lowercase label for the transport, used by the coverage
    /// recorder so observed and declared transports compare as plain strings
    /// (WOR-1386).
    pub fn label(self) -> &'static str {
        match self {
            TargetTransport::Stdio => "stdio",
            TargetTransport::Http => "http",
            TargetTransport::Cassette => "cassette",
        }
    }
}

/// Evaluate every declared budget against the finished step. Returns the
/// first breach as a rendered failure message, mirroring how assertion
/// failures surface.
///
/// `measured` is the wall-clock duration of the live call; `None` means no
/// live call was timed (fixture injection), which skips the duration gate.
/// `response_headers` carries the HTTP response headers the transport
/// captured for this call (WOR-1287); `None` means no headers were captured
/// (non-HTTP target, or no live call), which the header gate reports rather
/// than skips. Cache-aware semantics are preserved structurally: a cached
/// result never reaches the executor, so `evaluate_duration` is always
/// called with `cache_hit = false` here and the skip-on-cache-hit branch
/// stays with the shared evaluator for any future caller that does see
/// cache hits.
pub(super) fn evaluate_budget_gate(
    budgets: &ExpectBudgets,
    server: &str,
    transport: TargetTransport,
    measured: Option<std::time::Duration>,
    response: &Value,
    response_headers: Option<&std::collections::HashMap<String, String>>,
) -> Result<(), String> {
    if budgets.is_empty() {
        return Ok(());
    }
    if let (Some(budget_ms), Some(duration)) = (budgets.max_duration_ms, measured) {
        let actual_ms = u64::try_from(duration.as_millis()).unwrap_or(u64::MAX);
        let result = evaluate_duration(actual_ms, budget_ms, false);
        if !result.pass {
            return Err(render_failure("max_duration_ms", &result));
        }
    }
    if let Some(spec) = &budgets.max_response_tokens {
        check_token_budget(spec, response)?;
    }
    if budgets.has_header_checks() {
        check_headers(budgets, server, transport, response_headers)?;
    }
    Ok(())
}

/// Evaluate the token budget over the response's `result.content` blocks
/// through the shared evaluator. The tokenizer name was validated at suite
/// load, so a lookup miss here is a defensive error, not a user typo.
fn check_token_budget(
    spec: &crate::suite::TokenBudgetSpec,
    response: &Value,
) -> Result<(), String> {
    let tokenizer = tokenizer_by_name(&spec.tokenizer)
        .map_err(|e| format!("max_response_tokens: tokenizer `{}`: {e}", spec.tokenizer))?
        .ok_or_else(|| {
            format!(
                "max_response_tokens: unknown tokenizer `{}`",
                spec.tokenizer
            )
        })?;
    let budget = TokenBudget {
        budget: spec.budget,
        mode: spec.mode,
        image_cost: spec.image_cost,
    };
    let empty = Vec::new();
    let blocks = response
        .get("result")
        .and_then(|r| r.get("content"))
        .and_then(Value::as_array)
        .unwrap_or(&empty);
    let result = evaluate_response_tokens(blocks, &budget, tokenizer.as_ref())
        .map_err(|e| format!("max_response_tokens: {e}"))?;
    if result.pass {
        // Warn mode lands here on overage; surface the note in the log so
        // the warning is visible without failing the run.
        if let Some(note) = result.diffs.first().and_then(|d| d.note.as_deref()) {
            tracing::warn!(
                target: "mcptest_core::executor",
                note,
                "max_response_tokens budget exceeded in warn mode"
            );
        }
        return Ok(());
    }
    Err(render_failure("max_response_tokens", &result))
}

/// Gate the header checks on what the target transport can actually offer.
///
/// Stdio and cassette targets have no HTTP headers, so a header check there
/// is a suite bug and fails with a pointer at the fix. HTTP targets evaluate
/// the per-call headers the transport captured (WOR-1287) through
/// [`crate::matchers::headers::evaluate_response_headers`]; when nothing was
/// captured for the call (fixture injection, or a response that never
/// arrived) the check fails honestly instead of silently passing (the
/// WOR-1280 failure mode this module exists to remove).
fn check_headers(
    budgets: &ExpectBudgets,
    server: &str,
    transport: TargetTransport,
    captured: Option<&std::collections::HashMap<String, String>>,
) -> Result<(), String> {
    match transport {
        TargetTransport::Stdio => Err(format!(
            "`response_headers` checks are declared but server `{server}` is a stdio \
             target and stdio transports have no HTTP response headers; point the test \
             at a `url:` server or remove the header checks"
        )),
        TargetTransport::Cassette => Err(format!(
            "`response_headers` checks are declared but server `{server}` replays a \
             cassette, which carries no live HTTP response headers; run against the \
             live `url:` server or remove the header checks"
        )),
        TargetTransport::Http => {
            let Some(captured) = captured else {
                return Err(format!(
                    "`response_headers` checks are declared for HTTP server `{server}`, \
                     but no response headers were captured for this call, so the check \
                     cannot be evaluated"
                ));
            };
            let expected: std::collections::HashMap<String, HeaderMatch> =
                budgets.response_headers.iter().cloned().collect();
            let result = crate::matchers::headers::evaluate_response_headers(
                crate::matchers::headers::Transport::Http,
                &expected,
                &budgets.response_headers_absent,
                captured,
            )
            .map_err(|e| format!("response_headers: {e}"))?;
            if result.pass {
                return Ok(());
            }
            Err(render_failure("response_headers", &result))
        }
    }
}

/// Render a failing budget result as a one-line failure message, preferring
/// the evaluator's note (it carries the actual value and the overage).
fn render_failure(field: &str, result: &AssertionResult) -> String {
    let detail = result
        .diffs
        .first()
        .and_then(|d| d.note.as_deref())
        .unwrap_or("budget exceeded");
    format!("{field} budget exceeded: {detail}")
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::matchers::budgets::BudgetMode;
    use crate::matchers::headers::HeaderMatch;
    use crate::suite::TokenBudgetSpec;
    use serde_json::json;
    use std::time::Duration;

    fn duration_budget(ms: u64) -> ExpectBudgets {
        ExpectBudgets {
            max_duration_ms: Some(ms),
            ..Default::default()
        }
    }

    fn token_budget(budget: usize, mode: BudgetMode) -> ExpectBudgets {
        ExpectBudgets {
            max_response_tokens: Some(TokenBudgetSpec {
                budget,
                tokenizer: "whitespace".to_string(),
                mode,
                image_cost: 1000,
            }),
            ..Default::default()
        }
    }

    fn text_response(text: &str) -> Value {
        json!({ "result": { "content": [{ "type": "text", "text": text }] } })
    }

    #[test]
    fn empty_budgets_pass_without_a_measurement() {
        let r = evaluate_budget_gate(
            &ExpectBudgets::default(),
            "fs",
            TargetTransport::Stdio,
            None,
            &json!({}),
            None,
        );
        assert!(r.is_ok());
    }

    #[test]
    fn duration_over_budget_fails_with_the_overage_in_the_message() {
        let err = evaluate_budget_gate(
            &duration_budget(100),
            "fs",
            TargetTransport::Stdio,
            Some(Duration::from_millis(250)),
            &json!({}),
            None,
        )
        .expect_err("250ms breaches a 100ms budget");
        assert!(err.contains("max_duration_ms"), "{err}");
        assert!(err.contains("over budget"), "{err}");
    }

    #[test]
    fn duration_under_budget_passes() {
        let r = evaluate_budget_gate(
            &duration_budget(1000),
            "fs",
            TargetTransport::Stdio,
            Some(Duration::from_millis(5)),
            &json!({}),
            None,
        );
        assert!(r.is_ok());
    }

    #[test]
    fn unmeasured_step_skips_the_duration_gate() {
        // Fixture injection makes no live call, so there is nothing to time
        // and the duration budget cannot honestly fail.
        let r = evaluate_budget_gate(
            &duration_budget(1),
            "fs",
            TargetTransport::Stdio,
            None,
            &json!({}),
            None,
        );
        assert!(r.is_ok());
    }

    #[test]
    fn token_budget_over_in_fail_mode_fails() {
        let err = evaluate_budget_gate(
            &token_budget(2, BudgetMode::Fail),
            "fs",
            TargetTransport::Stdio,
            None,
            &text_response("one two three four"),
            None,
        )
        .expect_err("4 tokens breach a 2-token budget");
        assert!(err.contains("max_response_tokens"), "{err}");
        assert!(err.contains("over budget"), "{err}");
    }

    #[test]
    fn token_budget_over_in_warn_mode_passes() {
        let r = evaluate_budget_gate(
            &token_budget(2, BudgetMode::Warn),
            "fs",
            TargetTransport::Stdio,
            None,
            &text_response("one two three four"),
            None,
        );
        assert!(r.is_ok());
    }

    #[test]
    fn token_budget_under_passes() {
        let r = evaluate_budget_gate(
            &token_budget(10, BudgetMode::Fail),
            "fs",
            TargetTransport::Stdio,
            None,
            &text_response("one two"),
            None,
        );
        assert!(r.is_ok());
    }

    #[test]
    fn header_checks_on_a_stdio_target_fail_with_an_actionable_message() {
        let budgets = ExpectBudgets {
            response_headers: vec![(
                "content-type".to_string(),
                HeaderMatch::Literal("application/json".to_string()),
            )],
            ..Default::default()
        };
        let err = evaluate_budget_gate(
            &budgets,
            "fs",
            TargetTransport::Stdio,
            None,
            &json!({}),
            None,
        )
        .expect_err("stdio targets have no headers");
        assert!(err.contains("stdio"), "{err}");
        assert!(err.contains("`fs`"), "{err}");
        assert!(err.contains("url:"), "{err}");
    }

    #[test]
    fn header_checks_on_a_cassette_target_fail_with_an_actionable_message() {
        let budgets = ExpectBudgets {
            response_headers_absent: vec!["set-cookie".to_string()],
            ..Default::default()
        };
        let err = evaluate_budget_gate(
            &budgets,
            "rec",
            TargetTransport::Cassette,
            None,
            &json!({}),
            None,
        )
        .expect_err("cassette targets have no live headers");
        assert!(err.contains("cassette"), "{err}");
    }

    /// Build a captured-headers map the way the transport records it:
    /// lowercased names, string values.
    fn captured(pairs: &[(&str, &str)]) -> std::collections::HashMap<String, String> {
        pairs
            .iter()
            .map(|(k, v)| ((*k).to_string(), (*v).to_string()))
            .collect()
    }

    #[test]
    fn header_checks_on_an_http_target_without_captured_headers_fail_honestly() {
        // The inject_error path (and any call whose response never landed)
        // has no captured headers; the gate must say so, not pass silently.
        let budgets = ExpectBudgets {
            response_headers: vec![("x-trace".to_string(), HeaderMatch::Exists)],
            ..Default::default()
        };
        let err = evaluate_budget_gate(
            &budgets,
            "api",
            TargetTransport::Http,
            None,
            &json!({}),
            None,
        )
        .expect_err("uncaptured headers cannot silently pass");
        assert!(err.contains("no response headers were captured"), "{err}");
        assert!(err.contains("`api`"), "{err}");
    }

    #[test]
    fn header_checks_on_an_http_target_pass_against_matching_captured_headers() {
        let budgets = ExpectBudgets {
            response_headers: vec![
                (
                    "content-type".to_string(),
                    HeaderMatch::Literal("application/json".to_string()),
                ),
                ("x-trace-id".to_string(), HeaderMatch::Exists),
            ],
            response_headers_absent: vec!["set-cookie".to_string()],
            ..Default::default()
        };
        let headers = captured(&[
            ("content-type", "application/json"),
            ("x-trace-id", "trace-1"),
        ]);
        let r = evaluate_budget_gate(
            &budgets,
            "api",
            TargetTransport::Http,
            None,
            &json!({}),
            Some(&headers),
        );
        assert!(r.is_ok(), "{r:?}");
    }

    #[test]
    fn header_checks_on_an_http_target_fail_on_a_mismatch_naming_the_field() {
        // Wrong content-type value plus a forbidden header that is present:
        // either alone fails; the message must name `response_headers`.
        let budgets = ExpectBudgets {
            response_headers: vec![(
                "content-type".to_string(),
                HeaderMatch::Literal("application/json".to_string()),
            )],
            response_headers_absent: vec!["set-cookie".to_string()],
            ..Default::default()
        };
        let headers = captured(&[("content-type", "text/html"), ("set-cookie", "session=abc")]);
        let err = evaluate_budget_gate(
            &budgets,
            "api",
            TargetTransport::Http,
            None,
            &json!({}),
            Some(&headers),
        )
        .expect_err("mismatch and forbidden header must fail");
        assert!(err.contains("response_headers"), "{err}");
    }

    #[test]
    fn target_transport_classifies_every_server_spec() {
        let stdio = ServerSpec::Stdio {
            command: "srv".into(),
            args: vec![],
            env: Default::default(),
        };
        let http = ServerSpec::StreamableHttp {
            url: "http://localhost/mcp".into(),
            headers: Default::default(),
            oauth: None,
            stateless: false,
        };
        let sse = ServerSpec::LegacySse {
            url: "http://localhost/sse".into(),
            headers: Default::default(),
            oauth: None,
        };
        let cassette = ServerSpec::Cassette {
            path: "rec.json".into(),
            refresh_command: None,
            refresh_url: None,
        };
        assert_eq!(TargetTransport::from_spec(&stdio), TargetTransport::Stdio);
        assert_eq!(TargetTransport::from_spec(&http), TargetTransport::Http);
        assert_eq!(TargetTransport::from_spec(&sse), TargetTransport::Http);
        assert_eq!(
            TargetTransport::from_spec(&cassette),
            TargetTransport::Cassette
        );
    }
}
