//! Execution safety policy for features that call real tools with
//! synthesized arguments.
//!
//! Why this exists: scaffold, propose-assertions, and the probe tier all
//! execute live tool calls against a server the user may care about. A
//! synthesized call to `delete_file` is an agent autonomously causing side
//! effects. MCP tool annotations (`readOnlyHint`, `destructiveHint`,
//! `idempotentHint`, `openWorldHint`) exist for exactly this situation, so
//! this module turns them, plus a name heuristic for unannotated tools, into
//! one shared classification and decision layer those callers consume instead
//! of each re-deriving "is this safe to run" locally.
//!
//! See `docs/execution-safety.md` for the operator-facing description.

use std::sync::atomic::{AtomicU32, Ordering};
use std::time::Duration;

use serde::Deserialize;
use serde_json::Value;
use thiserror::Error;

/// Words in an unannotated tool name that imply the tool mutates state.
///
/// Public so docs and reporters can show the operator exactly which words
/// triggered a classification, rather than describing the heuristic vaguely.
pub const MUTATING_NAME_WORDS: &[&str] = &[
    "create", "update", "set", "send", "write", "post", "put", "insert", "patch", "add", "upload",
];

/// Words in an unannotated tool name that imply the tool destroys data.
///
/// Public for the same reason as [`MUTATING_NAME_WORDS`]. A destructive word
/// always outranks a mutating word in the same name. This is the single
/// source of truth shared with the SEC-009 surface check (WOR-1278), so the
/// security scan and the execution policy can never give contradictory
/// verdicts on the same name. Reconciliation notes from the merge of the two
/// former lists: `overwrite` and `truncate` destroy existing data in place,
/// `kill` terminates running work, and `revoke` invalidates credentials, so
/// all four are destructive rather than merely mutating.
pub const DESTRUCTIVE_NAME_WORDS: &[&str] = &[
    "delete",
    "remove",
    "drop",
    "destroy",
    "overwrite",
    "wipe",
    "purge",
    "truncate",
    "erase",
    "kill",
    "revoke",
];

/// Marker text callers prepend to generated-but-not-executed tests.
///
/// Centralized so every generator emits the identical warning and tests can
/// assert on one string. The cleanup sentence is deliberate: mcptest never
/// cleans up data a mutating test created.
pub const REVIEW_MARKER: &str = "# review before first run\n\
     # mcptest did not auto-execute this tool. Data cleanup after mutating or\n\
     # destructive test runs is the developer's responsibility.";

/// The MCP tool annotation hints relevant to execution safety.
///
/// Typed mirror of the `annotations` object in a `tools/list` descriptor, so
/// callers that already parsed the catalog can classify without re-walking
/// JSON. All fields are optional because every hint is optional in the spec.
#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct ToolAnnotations {
    /// `readOnlyHint`: the tool does not modify its environment.
    pub read_only_hint: Option<bool>,
    /// `destructiveHint`: the tool may perform destructive updates.
    pub destructive_hint: Option<bool>,
    /// `idempotentHint`: repeated calls with the same args have no extra effect.
    pub idempotent_hint: Option<bool>,
    /// `openWorldHint`: the tool interacts with external entities.
    pub open_world_hint: Option<bool>,
}

/// Safety classification of a single tool.
///
/// Split into four classes (not three) so callers can distinguish a tool the
/// server declared read-only from one we merely presume is read-only because
/// its name looks harmless. Probe tiers may want to trust the former more.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum ToolClass {
    /// Annotated `readOnlyHint: true`. Safe to auto-execute freely.
    ReadOnly,
    /// Unannotated and the name looks harmless. Treated like read-only but
    /// callers can surface the lower confidence.
    ReadOnlyPresumed,
    /// Mutates state but is not destructive (annotated non-idempotent, or an
    /// unannotated mutating-looking name). Execute at most once, never
    /// double-call for stability, and mark generated tests serial.
    Mutating,
    /// Annotated `destructiveHint: true` or a destructive-looking name.
    /// Never auto-executed without an explicit override.
    Destructive,
}

/// Whether the caller can fall back to emitting a test instead of running it.
///
/// Exists because the right answer for a destructive tool differs by caller:
/// scaffold can write the test and let a human run it, while a probe that
/// must observe a live response can only refuse.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum CallContext {
    /// The caller can emit a test without executing it (scaffold, propose).
    CanGenerate,
    /// The caller needs a live call to make progress (probe).
    MustExecute,
}

/// What a caller is allowed to do with one tool call.
///
/// Returned by [`ExecutionPolicy::decide`] so every consumer applies the same
/// rules instead of re-implementing the override and marker logic.
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum Decision {
    /// Execute freely, including stability double-calls.
    Execute,
    /// Execute at most once. Never double-call: the second call of a
    /// non-idempotent tool is itself a side effect.
    ExecuteOnce,
    /// Do not execute. Emit the test with `marker` prepended so a human
    /// reviews it before its first run.
    GenerateOnly {
        /// Comment block to prepend to the generated test.
        marker: String,
    },
    /// Do not execute and nothing useful can be generated in this context.
    Refuse {
        /// Human-readable explanation for the report or log line.
        reason: String,
    },
}

/// Knobs that govern auto-execution of synthesized tool calls.
///
/// One struct rather than loose parameters so scaffold, propose, and probe
/// share a single value built once from CLI flags in a later ticket. Defaults
/// are deliberately timid: no destructive execution, two concurrent calls,
/// a polite pause between HTTP calls.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct ExecutionPolicy {
    /// Allow executing tools classified [`ToolClass::Destructive`]. Off by
    /// default; the operator must opt in explicitly.
    pub execute_destructive: bool,
    /// Upper bound on total tool calls for the run, `None` for unlimited.
    /// Enforced via [`CallBudget`].
    pub max_calls: Option<u32>,
    /// Maximum tool calls in flight at once. Small by default so a probe
    /// never looks like a load test.
    pub concurrency: usize,
    /// Pause between consecutive HTTP calls. Callers decide whether the
    /// target transport is HTTP; stdio callers may ignore it.
    pub call_delay: Duration,
}

impl Default for ExecutionPolicy {
    fn default() -> Self {
        Self {
            execute_destructive: false,
            max_calls: None,
            concurrency: 2,
            call_delay: Duration::from_millis(100),
        }
    }
}

impl ExecutionPolicy {
    /// Decide what to do with one call to a tool of the given class.
    ///
    /// Centralizes the override logic so a caller cannot accidentally
    /// execute a destructive tool by forgetting a check: destructive tools
    /// yield [`Decision::GenerateOnly`] or [`Decision::Refuse`] unless
    /// [`ExecutionPolicy::execute_destructive`] is set.
    pub fn decide(&self, class: ToolClass, context: CallContext) -> Decision {
        match class {
            ToolClass::ReadOnly | ToolClass::ReadOnlyPresumed => Decision::Execute,
            ToolClass::Mutating => Decision::ExecuteOnce,
            // Even with the override, a destructive call is non-idempotent,
            // so it gets at most one call.
            ToolClass::Destructive if self.execute_destructive => Decision::ExecuteOnce,
            ToolClass::Destructive => match context {
                CallContext::CanGenerate => Decision::GenerateOnly {
                    marker: REVIEW_MARKER.to_string(),
                },
                CallContext::MustExecute => Decision::Refuse {
                    reason: "tool is destructive and execute_destructive is not set".to_string(),
                },
            },
        }
    }

    /// Build the [`CallBudget`] for one run from this policy.
    ///
    /// Exists so the budget limit and the policy cannot drift apart when a
    /// caller wires them up separately.
    pub fn budget(&self) -> CallBudget {
        CallBudget::new(self.max_calls)
    }
}

/// Classify a tool from its name and optional parsed annotations.
///
/// Explicit annotations always win over the name heuristic: a server that
/// declares `readOnlyHint: true` on `delete_draft` knows its own semantics
/// better than a word list does. Among hints, read-only is checked first
/// (the spec says the other hints are only meaningful when it is false),
/// then destructive, then non-idempotent.
pub fn classify(name: &str, annotations: Option<&ToolAnnotations>) -> ToolClass {
    if let Some(hints) = annotations {
        if hints.read_only_hint == Some(true) {
            return ToolClass::ReadOnly;
        }
        if hints.destructive_hint == Some(true) {
            return ToolClass::Destructive;
        }
        if hints.idempotent_hint == Some(false) {
            return ToolClass::Mutating;
        }
    }
    classify_by_name(name)
}

/// Classify a raw `tools/list` tool descriptor.
///
/// The codebase passes tool descriptors around as `serde_json::Value`
/// (see `stub_gen`, `lint`, `security`), so this entry point saves every
/// caller the same name-and-annotations extraction dance.
pub fn classify_descriptor(tool: &Value) -> ToolClass {
    let name = tool.get("name").and_then(Value::as_str).unwrap_or("");
    let annotations = tool
        .get("annotations")
        // A malformed annotations object (e.g. `"destructiveHint": "yes"`)
        // is a description-quality problem that the lint module flags; for
        // safety classification we deliberately treat it as unannotated and
        // let the name heuristic decide.
        .and_then(|a| serde_json::from_value::<ToolAnnotations>(a.clone()).ok());
    classify(name, annotations.as_ref())
}

/// Heuristic classification for tools with no usable annotations.
fn classify_by_name(name: &str) -> ToolClass {
    let tokens = name_tokens(name);
    let has = |words: &[&str]| tokens.iter().any(|t| words.contains(&t.as_str()));
    if has(DESTRUCTIVE_NAME_WORDS) {
        ToolClass::Destructive
    } else if has(MUTATING_NAME_WORDS) {
        ToolClass::Mutating
    } else {
        ToolClass::ReadOnlyPresumed
    }
}

/// Split a tool name into lowercase word tokens at separators and camelCase
/// boundaries, so `deleteFile`, `delete-file`, and `delete_file` all yield
/// the token `delete`. Shared with the SEC-009 surface check (WOR-1278) so
/// both consumers see the same words in the same name.
pub(crate) fn name_tokens(name: &str) -> Vec<String> {
    let mut tokens = Vec::new();
    let mut current = String::new();
    for ch in name.chars() {
        if ch.is_alphanumeric() {
            if ch.is_uppercase() && !current.is_empty() {
                tokens.push(std::mem::take(&mut current));
            }
            current.extend(ch.to_lowercase());
        } else if !current.is_empty() {
            tokens.push(std::mem::take(&mut current));
        }
    }
    if !current.is_empty() {
        tokens.push(current);
    }
    tokens
}

/// Thread-safe counter capping the total tool calls in one run.
///
/// Exists because the future probe tier fans calls out across tasks; a
/// plain integer would either need a mutex or undercount. Acquire before
/// every call; a denied acquire means the run hit its `--max-calls` limit.
#[derive(Debug)]
pub struct CallBudget {
    limit: Option<u32>,
    used: AtomicU32,
}

impl CallBudget {
    /// Create a budget with the given limit, `None` for unlimited.
    pub fn new(limit: Option<u32>) -> Self {
        Self {
            limit,
            used: AtomicU32::new(0),
        }
    }

    /// Reserve one call, or fail with [`ExecPolicyError::BudgetExhausted`].
    ///
    /// Compare-and-swap rather than a blind increment so that under
    /// concurrency exactly `limit` acquisitions succeed, never more.
    pub fn try_acquire(&self) -> Result<(), ExecPolicyError> {
        let Some(limit) = self.limit else {
            self.used.fetch_add(1, Ordering::Relaxed);
            return Ok(());
        };
        self.used
            .fetch_update(Ordering::SeqCst, Ordering::SeqCst, |n| {
                (n < limit).then_some(n + 1)
            })
            .map(|_| ())
            .map_err(|_| ExecPolicyError::BudgetExhausted { limit })
    }

    /// Calls acquired so far. For reporting how much of the budget a run used.
    pub fn used(&self) -> u32 {
        self.used.load(Ordering::SeqCst)
    }
}

/// Errors this module can produce.
///
/// Typed (per the crate convention) so callers can match on budget
/// exhaustion and stop scheduling calls instead of string-sniffing.
#[derive(Debug, Error, PartialEq, Eq)]
pub enum ExecPolicyError {
    /// The run's call budget is spent; no further tool calls are allowed.
    #[error("call budget exhausted: limit of {limit} tool calls reached")]
    BudgetExhausted {
        /// The configured maximum that was reached.
        limit: u32,
    },
}

#[cfg(test)]
mod tests {
    use std::sync::Arc;

    use serde_json::json;

    use super::*;

    fn annotations(value: Value) -> ToolAnnotations {
        serde_json::from_value(value).unwrap()
    }

    #[test]
    fn read_only_annotation_classifies_as_read_only() {
        let hints = annotations(json!({ "readOnlyHint": true }));
        assert_eq!(classify("fetch_page", Some(&hints)), ToolClass::ReadOnly);
    }

    #[test]
    fn destructive_annotation_classifies_as_destructive() {
        let hints = annotations(json!({ "destructiveHint": true }));
        assert_eq!(classify("apply", Some(&hints)), ToolClass::Destructive);
    }

    #[test]
    fn non_idempotent_annotation_classifies_as_mutating() {
        let hints = annotations(json!({ "idempotentHint": false }));
        assert_eq!(classify("notify", Some(&hints)), ToolClass::Mutating);
    }

    #[test]
    fn annotations_outrank_the_name_heuristic() {
        let hints = annotations(json!({ "readOnlyHint": true }));
        assert_eq!(classify("delete_draft", Some(&hints)), ToolClass::ReadOnly);
    }

    #[test]
    fn every_mutating_word_classifies_unannotated_name_as_mutating() {
        for word in MUTATING_NAME_WORDS {
            let name = format!("{word}_thing");
            assert_eq!(classify(&name, None), ToolClass::Mutating, "name {name}");
        }
    }

    #[test]
    fn every_destructive_word_classifies_unannotated_name_as_destructive() {
        for word in DESTRUCTIVE_NAME_WORDS {
            let name = format!("{word}_thing");
            assert_eq!(classify(&name, None), ToolClass::Destructive, "name {name}");
        }
    }

    #[test]
    fn camel_case_and_kebab_names_tokenize_for_the_heuristic() {
        assert_eq!(classify("deleteFile", None), ToolClass::Destructive);
        assert_eq!(classify("send-message", None), ToolClass::Mutating);
    }

    #[test]
    fn destructive_word_outranks_mutating_word_in_one_name() {
        assert_eq!(classify("create_or_delete", None), ToolClass::Destructive);
    }

    #[test]
    fn neutral_unannotated_name_is_read_only_presumed() {
        assert_eq!(classify("search", None), ToolClass::ReadOnlyPresumed);
        assert_eq!(classify("get_weather", None), ToolClass::ReadOnlyPresumed);
    }

    #[test]
    fn descriptor_with_annotations_classifies_from_annotations() {
        let tool = json!({
            "name": "search",
            "annotations": { "destructiveHint": true }
        });
        assert_eq!(classify_descriptor(&tool), ToolClass::Destructive);
    }

    #[test]
    fn descriptor_with_malformed_annotations_falls_back_to_name() {
        let tool = json!({
            "name": "purge_cache",
            "annotations": { "destructiveHint": "yes" }
        });
        assert_eq!(classify_descriptor(&tool), ToolClass::Destructive);
    }

    #[test]
    fn read_only_tool_executes_freely() {
        let policy = ExecutionPolicy::default();
        let decision = policy.decide(ToolClass::ReadOnly, CallContext::MustExecute);
        assert_eq!(decision, Decision::Execute);
        let presumed = policy.decide(ToolClass::ReadOnlyPresumed, CallContext::CanGenerate);
        assert_eq!(presumed, Decision::Execute);
    }

    #[test]
    fn mutating_tool_yields_execute_once() {
        let policy = ExecutionPolicy::default();
        let decision = policy.decide(ToolClass::Mutating, CallContext::MustExecute);
        assert_eq!(decision, Decision::ExecuteOnce);
    }

    #[test]
    fn destructive_tool_is_generated_but_not_executed() {
        let policy = ExecutionPolicy::default();
        match policy.decide(ToolClass::Destructive, CallContext::CanGenerate) {
            Decision::GenerateOnly { marker } => {
                assert!(marker.contains("# review before first run"));
                assert!(marker.contains("developer's responsibility"));
            }
            other => panic!("expected GenerateOnly, got {other:?}"),
        }
    }

    #[test]
    fn destructive_tool_refuses_without_override_when_execution_required() {
        let policy = ExecutionPolicy::default();
        match policy.decide(ToolClass::Destructive, CallContext::MustExecute) {
            Decision::Refuse { reason } => assert!(reason.contains("execute_destructive")),
            other => panic!("expected Refuse, got {other:?}"),
        }
    }

    #[test]
    fn destructive_tool_executes_once_with_override() {
        let policy = ExecutionPolicy {
            execute_destructive: true,
            ..ExecutionPolicy::default()
        };
        let decision = policy.decide(ToolClass::Destructive, CallContext::MustExecute);
        assert_eq!(decision, Decision::ExecuteOnce);
    }

    #[test]
    fn default_policy_is_timid() {
        let policy = ExecutionPolicy::default();
        assert!(!policy.execute_destructive);
        assert_eq!(policy.max_calls, None);
        assert_eq!(policy.concurrency, 2);
        assert_eq!(policy.call_delay, Duration::from_millis(100));
    }

    #[test]
    fn budget_exhausts_at_exactly_max_calls_under_concurrency() {
        let budget = Arc::new(CallBudget::new(Some(25)));
        let mut handles = Vec::new();
        for _ in 0..8 {
            let budget = Arc::clone(&budget);
            handles.push(std::thread::spawn(move || {
                (0..10).filter(|_| budget.try_acquire().is_ok()).count()
            }));
        }
        let granted: usize = handles.into_iter().map(|h| h.join().unwrap()).sum();
        assert_eq!(granted, 25);
        assert_eq!(budget.used(), 25);
    }

    #[test]
    fn exhausted_budget_error_carries_the_limit() {
        let budget = CallBudget::new(Some(1));
        budget.try_acquire().unwrap();
        let err = budget.try_acquire().unwrap_err();
        assert_eq!(err, ExecPolicyError::BudgetExhausted { limit: 1 });
    }

    #[test]
    fn unlimited_budget_never_refuses() {
        let budget = CallBudget::new(None);
        for _ in 0..1000 {
            budget.try_acquire().unwrap();
        }
        assert_eq!(budget.used(), 1000);
    }

    #[test]
    fn policy_budget_uses_the_policy_limit() {
        let policy = ExecutionPolicy {
            max_calls: Some(2),
            ..ExecutionPolicy::default()
        };
        let budget = policy.budget();
        budget.try_acquire().unwrap();
        budget.try_acquire().unwrap();
        assert!(budget.try_acquire().is_err());
    }
}
