//! Recovery scoring for calls against an unresponsive server.
//!
//! Why this module exists: mcptest tests servers that *respond*. It does not
//! score what an agent does when a server is reachable but never answers (the
//! "dead but connected" failure, distinct from a refused connection or a
//! malformed reply). A practitioner injecting that fault wants three numbers
//! out of one observation: how long it took the agent to *detect* the
//! unresponsive call, how long until it *recovered* (succeeded another way,
//! gave up cleanly, or fell back), and how much of that was *user-visible
//! hang*. This module is the pure, transport-free core of that scoring so it
//! is unit-testable and reusable by the runner, the SDK hosts, and the
//! enterprise audit chain.
//!
//! The metric vocabulary follows the offline failure-injection and recovery
//! work in ToolMisuseBench (arXiv:2604.01508): inject an offline/unresponsive
//! tool, then score task success and time-to-recovery rather than only the
//! happy path. The shapes here are deliberately I/O-free; a caller times the
//! call and hands us the millisecond marks.

use std::collections::HashMap;

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

use crate::suite::{FaultKind, FaultSpec, RecoveryGate};

/// One timed observation of a call against a faulty (unresponsive) server.
///
/// All marks are milliseconds on the same clock as `call_start_ms` (a
/// monotonic clock the caller chooses). The fields capture "when did we
/// start", "when did the agent notice the call was not coming back", and
/// "when did the agent recover", against the timeout budget the agent was
/// told to honor. `detected_ms` and `recovered_ms` are optional because a
/// run can hang forever (never detect) or detect-but-never-recover.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
pub struct RecoveryObservation {
    /// Clock value when the call was issued, milliseconds.
    pub call_start_ms: u64,
    /// Clock value when the agent detected the call would not return
    /// (a timeout fired, a cancellation landed), or `None` if it never
    /// noticed and is still hanging.
    pub detected_ms: Option<u64>,
    /// Clock value when the agent recovered (retried elsewhere, fell back,
    /// or surfaced a clean error to the user), or `None` if it never did.
    pub recovered_ms: Option<u64>,
    /// The timeout budget the agent was configured with, milliseconds. A
    /// detection at or before this mark counts as a clean timeout; later
    /// (or never) means the agent hung past its own budget.
    pub timeout_budget_ms: u64,
}

/// Scored outcome of one [`RecoveryObservation`].
///
/// Serializable so reporters and the enterprise audit chain can persist the
/// recovery picture alongside the per-call rows. The three durations are
/// relative to `call_start_ms`; `recovered` and `clean_timeout` are the two
/// pass/fail-shaped booleans a gate reads.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
pub struct RecoveryReport {
    /// Time from call start to the agent detecting the unresponsive call,
    /// milliseconds, or `None` if it never detected (still hanging).
    pub time_to_detection_ms: Option<u64>,
    /// Time from call start to the agent recovering, milliseconds, or
    /// `None` if it never recovered.
    pub time_to_recovery_ms: Option<u64>,
    /// How long the user saw the call hang. This is the time the agent
    /// spent stuck before it either detected the fault (the usual case) or,
    /// failing detection, recovered. When neither ever happens it is the
    /// time spent up to the timeout budget, because that is the longest a
    /// correctly-configured agent should ever expose to the user.
    pub user_visible_hang_ms: u64,
    /// True when the agent recovered (recovered mark present and not before
    /// the call started).
    pub recovered: bool,
    /// True when the agent detected the fault at or before its timeout
    /// budget: it timed out *cleanly* rather than hanging past its own
    /// deadline (or forever).
    pub clean_timeout: bool,
}

/// Score one observation into a [`RecoveryReport`].
///
/// Pure and total: out-of-order marks (a detection or recovery before the
/// call started) are treated as "did not happen" rather than producing a
/// negative duration, so the report never carries nonsense.
pub fn score_recovery(obs: &RecoveryObservation) -> RecoveryReport {
    let start = obs.call_start_ms;
    let detected_after = obs.detected_ms.and_then(|m| m.checked_sub(start));
    let recovered_after = obs.recovered_ms.and_then(|m| m.checked_sub(start));
    let clean_timeout = match obs.detected_ms {
        Some(detected) => {
            detected >= start && detected <= start.saturating_add(obs.timeout_budget_ms)
        }
        None => false,
    };
    let recovered = recovered_after.is_some();
    let user_visible_hang_ms =
        user_visible_hang(detected_after, recovered_after, obs.timeout_budget_ms);
    RecoveryReport {
        time_to_detection_ms: detected_after,
        time_to_recovery_ms: recovered_after,
        user_visible_hang_ms,
        recovered,
        clean_timeout,
    }
}

/// Compute the user-visible hang.
///
/// The user stops seeing a hang the moment the agent either detects the
/// fault (shows a timeout, a spinner-to-error) or recovers. We take the
/// earlier of those two marks. When neither happened, the agent is still
/// stuck, but a correctly-configured agent caps exposure at its timeout
/// budget, so we report the budget rather than an unbounded number.
fn user_visible_hang(
    detected_after: Option<u64>,
    recovered_after: Option<u64>,
    budget_ms: u64,
) -> u64 {
    match (detected_after, recovered_after) {
        (Some(d), Some(r)) => d.min(r),
        (Some(d), None) => d,
        (None, Some(r)) => r,
        (None, None) => budget_ms,
    }
}

/// True when the run recovered within `max_recovery_ms` of the call start.
///
/// This is the assertable gate: a suite (or a Rust caller) can require
/// "recover within 5000ms". A run that never recovered fails the budget, and
/// a run that recovered but took longer than the budget also fails. Use it
/// alongside `RecoveryReport::clean_timeout` when you also want to forbid the
/// hang-past-budget case.
pub fn meets_recovery_budget(report: &RecoveryReport, max_recovery_ms: u64) -> bool {
    match report.time_to_recovery_ms {
        Some(ms) => ms <= max_recovery_ms,
        None => false,
    }
}

/// Decide whether a tool call should be synthesized as an unresponsive timeout
/// rather than dispatched to the server.
///
/// This is the executor-side directive the recovery design rests on: a hang is
/// never realized as a real block. The driver calls this before each tool
/// dispatch; a `true` return means "do not call the server, return a timeout
/// error result". `counts` carries the per-fault call counter so
/// [`FaultKind::RecoverAfter`] can hang its first N calls and then proceed; pass
/// the same map across every call in one run. `detection_budget_ms` is the
/// per-call timeout budget: a `slow` fault reads as unresponsive only when its
/// delay reaches the budget. The first matching fault in `plan` wins.
pub fn injected_timeout(
    plan: &[FaultSpec],
    counts: &mut HashMap<String, u32>,
    tool: &str,
    server: &str,
    detection_budget_ms: u64,
) -> bool {
    for spec in plan {
        if !spec.target.matches(tool, server) {
            continue;
        }
        return match spec.kind {
            FaultKind::Hang | FaultKind::Wedged => true,
            FaultKind::Slow { delay_ms } => delay_ms >= detection_budget_ms,
            FaultKind::RecoverAfter { failures } => {
                let seen = counts.entry(spec.name.clone()).or_insert(0);
                let hang = *seen < failures;
                *seen += 1;
                hang
            }
        };
    }
    false
}

/// The outcome of evaluating an agent run's `recovery:` gate.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct RecoveryOutcome {
    /// The scored recovery picture (detection, recovery, hang).
    pub report: RecoveryReport,
    /// How many tool calls were synthesized as unresponsive timeouts.
    pub detections: u32,
    /// True when the gate passed.
    pub passed: bool,
    /// Human-readable one-line explanation for the report row.
    pub detail: String,
}

/// Evaluate an agent run's `recovery:` gate over its trace envelope.
///
/// The hang is scored in virtual time, never wall-clock: every synthesized
/// timeout consumes one `max_detection_ms` budget, so a run that hit the fault
/// `k` times and then replied cleanly recovered at `k * max_detection_ms`. A run
/// that ended on a budget cap (turn/tool/time) without a final reply never
/// recovered. The gate fails (a quality failure, exit 1) when the agent did not
/// recover, took longer than `max_recovery_ms`, or, under
/// `require_clean_timeout`, did not time out within its detection budget.
///
/// `plan` is the agent's resolved fault list (its `inject:` names looked up in
/// the suite's `faults:`). With an empty `plan` no fault fires, so a clean run
/// passes vacuously.
pub fn evaluate_recovery(
    envelope: &Value,
    plan: &[FaultSpec],
    gate: &RecoveryGate,
) -> RecoveryOutcome {
    let budget = gate.max_detection_ms;
    let detections = count_detections(envelope, plan, budget);
    let recovered_clean = ran_to_clean_completion(envelope);

    let recovered_ms = if recovered_clean {
        Some(u64::from(detections) * budget)
    } else {
        None
    };
    // The agent "detected" the fault only if it then moved on and replied; a run
    // that hung until a budget cap never produced a clean detection.
    let detected_ms = if detections > 0 && recovered_clean {
        Some(budget)
    } else {
        None
    };
    let obs = RecoveryObservation {
        call_start_ms: 0,
        detected_ms,
        recovered_ms,
        timeout_budget_ms: budget,
    };
    let report = score_recovery(&obs);

    let passed = recovery_passes(&report, gate, detections);
    let detail = recovery_detail(&report, detections, gate, recovered_clean);
    RecoveryOutcome {
        report,
        detections,
        passed,
        detail,
    }
}

/// Apply the gate's budgets to a scored report. Recovery is mandatory; the
/// recovery-time cap and the clean-timeout requirement are optional.
fn recovery_passes(report: &RecoveryReport, gate: &RecoveryGate, detections: u32) -> bool {
    if !report.recovered {
        return false;
    }
    if let Some(max) = gate.max_recovery_ms {
        if !meets_recovery_budget(report, max) {
            return false;
        }
    }
    // A clean timeout only makes sense when the agent actually engaged the
    // fault; an agent that avoided the faulted tool entirely has nothing to time
    // out on, so the requirement is vacuous there.
    if gate.require_clean_timeout && detections > 0 && !report.clean_timeout {
        return false;
    }
    true
}

/// One-line explanation for the report row.
fn recovery_detail(
    report: &RecoveryReport,
    detections: u32,
    gate: &RecoveryGate,
    recovered_clean: bool,
) -> String {
    if !recovered_clean {
        return format!(
            "did not recover: the run hit the injected fault {detections} time(s) and never \
             returned a clean final reply within {} detection-budget ms",
            gate.max_detection_ms
        );
    }
    match report.time_to_recovery_ms {
        Some(ms) => format!(
            "recovered in {ms} ms (virtual) after {detections} detection(s); detection budget {} ms",
            gate.max_detection_ms
        ),
        None => "recovered without engaging the injected fault".to_string(),
    }
}

/// Count how many tool calls in the trace would have been synthesized as
/// unresponsive timeouts, replaying the same per-fault decision the driver made.
fn count_detections(envelope: &Value, plan: &[FaultSpec], budget: u64) -> u32 {
    let Some(calls) = envelope.get("tool_calls").and_then(|v| v.as_array()) else {
        return 0;
    };
    let mut counts: HashMap<String, u32> = HashMap::new();
    let mut detections = 0u32;
    for call in calls {
        let tool = call.get("name").and_then(|v| v.as_str()).unwrap_or("");
        let server = call.get("server").and_then(|v| v.as_str()).unwrap_or("");
        if injected_timeout(plan, &mut counts, tool, server, budget) {
            detections += 1;
        }
    }
    detections
}

/// True when the run ended with a clean final reply: stop reason `completed`
/// and a non-empty `final_response`. A run that tripped a turn/tool/time cap, or
/// produced no text, did not recover.
fn ran_to_clean_completion(envelope: &Value) -> bool {
    let stop_reason = envelope
        .get("conversation")
        .and_then(|c| c.get("stop_reason"))
        .and_then(|v| v.as_str())
        // A trace recorded before stop_reason existed reads back as completed.
        .unwrap_or("completed");
    let has_final = envelope
        .get("final_response")
        .and_then(|v| v.as_str())
        .is_some_and(|s| !s.is_empty());
    stop_reason == "completed" && has_final
}

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

    #[test]
    fn clean_timeout_detects_at_budget_and_recovers() {
        // Agent timed out exactly at its 3000ms budget, then fell back at
        // 3200ms. Clean timeout, recovered, hang stops at detection.
        let obs = RecoveryObservation {
            call_start_ms: 1_000,
            detected_ms: Some(4_000),
            recovered_ms: Some(4_200),
            timeout_budget_ms: 3_000,
        };
        let report = score_recovery(&obs);
        assert_eq!(report.time_to_detection_ms, Some(3_000));
        assert_eq!(report.time_to_recovery_ms, Some(3_200));
        assert_eq!(report.user_visible_hang_ms, 3_000);
        assert!(report.recovered);
        assert!(report.clean_timeout);
        assert!(meets_recovery_budget(&report, 5_000));
    }

    #[test]
    fn hang_past_budget_is_not_a_clean_timeout() {
        // Agent kept hanging until 9000ms, well past its 3000ms budget,
        // before it noticed. Detection happened, but late: not clean.
        let obs = RecoveryObservation {
            call_start_ms: 0,
            detected_ms: Some(9_000),
            recovered_ms: Some(9_100),
            timeout_budget_ms: 3_000,
        };
        let report = score_recovery(&obs);
        assert_eq!(report.time_to_detection_ms, Some(9_000));
        assert!(!report.clean_timeout, "9s > 3s budget must not be clean");
        assert!(report.recovered);
        // The user stared at a hang for the full 9s before it broke.
        assert_eq!(report.user_visible_hang_ms, 9_000);
    }

    #[test]
    fn recover_after_n_within_budget_passes_the_gate() {
        let obs = RecoveryObservation {
            call_start_ms: 500,
            detected_ms: Some(1_500),
            recovered_ms: Some(2_000),
            timeout_budget_ms: 2_000,
        };
        let report = score_recovery(&obs);
        assert_eq!(report.time_to_recovery_ms, Some(1_500));
        assert!(meets_recovery_budget(&report, 5_000));
        assert!(
            !meets_recovery_budget(&report, 1_000),
            "1500ms exceeds a 1000ms budget"
        );
    }

    #[test]
    fn never_recovers_and_never_detects_reports_budget_hang() {
        // The pathological case: the call hangs forever, the agent never
        // notices, never recovers. Hang is capped at the budget the agent
        // should have honored.
        let obs = RecoveryObservation {
            call_start_ms: 100,
            detected_ms: None,
            recovered_ms: None,
            timeout_budget_ms: 4_000,
        };
        let report = score_recovery(&obs);
        assert_eq!(report.time_to_detection_ms, None);
        assert_eq!(report.time_to_recovery_ms, None);
        assert!(!report.recovered);
        assert!(!report.clean_timeout);
        assert_eq!(report.user_visible_hang_ms, 4_000);
        assert!(!meets_recovery_budget(&report, 10_000));
    }

    #[test]
    fn detects_cleanly_but_never_recovers() {
        // Agent timed out cleanly at 2000ms but then gave up without a
        // fallback. Clean timeout true, recovered false, gate fails.
        let obs = RecoveryObservation {
            call_start_ms: 0,
            detected_ms: Some(2_000),
            recovered_ms: None,
            timeout_budget_ms: 2_000,
        };
        let report = score_recovery(&obs);
        assert!(report.clean_timeout);
        assert!(!report.recovered);
        assert_eq!(report.user_visible_hang_ms, 2_000);
        assert!(!meets_recovery_budget(&report, 5_000));
    }

    #[test]
    fn out_of_order_marks_are_treated_as_not_happened() {
        // A recovery mark before the call start is nonsense; treat it as
        // "no recovery" rather than emit a negative or wrapped duration.
        let obs = RecoveryObservation {
            call_start_ms: 5_000,
            detected_ms: Some(4_000),
            recovered_ms: Some(3_000),
            timeout_budget_ms: 1_000,
        };
        let report = score_recovery(&obs);
        assert_eq!(report.time_to_detection_ms, None);
        assert_eq!(report.time_to_recovery_ms, None);
        assert!(!report.recovered);
        assert!(!report.clean_timeout);
        assert_eq!(report.user_visible_hang_ms, 1_000);
    }

    #[test]
    fn recovery_before_detection_uses_earlier_mark_for_hang() {
        // If somehow recovery (a fallback path) lands before the explicit
        // detection mark, the user stopped seeing a hang at the earlier one.
        let obs = RecoveryObservation {
            call_start_ms: 0,
            detected_ms: Some(2_500),
            recovered_ms: Some(1_800),
            timeout_budget_ms: 3_000,
        };
        let report = score_recovery(&obs);
        assert_eq!(report.user_visible_hang_ms, 1_800);
        assert!(report.recovered);
        assert!(report.clean_timeout);
    }

    use crate::suite::{FaultKind, FaultSpec, FaultTarget, RecoveryGate};
    use serde_json::json;

    fn hang_fault(tool: &str) -> FaultSpec {
        FaultSpec {
            name: format!("hung-{tool}"),
            target: FaultTarget {
                tool: Some(tool.to_string()),
                server: None,
            },
            kind: FaultKind::Hang,
        }
    }

    #[test]
    fn injected_timeout_hangs_matching_tool_only() {
        let plan = vec![hang_fault("search")];
        let mut counts = HashMap::new();
        assert!(injected_timeout(
            &plan,
            &mut counts,
            "search",
            "backend",
            3_000
        ));
        assert!(!injected_timeout(
            &plan,
            &mut counts,
            "fetch",
            "backend",
            3_000
        ));
    }

    #[test]
    fn injected_timeout_recover_after_hangs_n_then_proceeds() {
        let plan = vec![FaultSpec {
            name: "flaky".into(),
            target: FaultTarget {
                tool: Some("search".into()),
                server: None,
            },
            kind: FaultKind::RecoverAfter { failures: 2 },
        }];
        let mut counts = HashMap::new();
        assert!(injected_timeout(&plan, &mut counts, "search", "s", 3_000));
        assert!(injected_timeout(&plan, &mut counts, "search", "s", 3_000));
        assert!(!injected_timeout(&plan, &mut counts, "search", "s", 3_000));
    }

    #[test]
    fn injected_timeout_slow_reads_as_unresponsive_at_or_past_budget() {
        let plan = vec![FaultSpec {
            name: "slow".into(),
            target: FaultTarget {
                tool: None,
                server: Some("backend".into()),
            },
            kind: FaultKind::Slow { delay_ms: 4_000 },
        }];
        let mut counts = HashMap::new();
        // 4000ms delay >= 3000ms budget: unresponsive.
        assert!(injected_timeout(&plan, &mut counts, "x", "backend", 3_000));
        // Same delay under a 5000ms budget: answers in time, no timeout.
        assert!(!injected_timeout(&plan, &mut counts, "x", "backend", 5_000));
    }

    fn gate(max_recovery_ms: Option<u64>, require_clean_timeout: bool) -> RecoveryGate {
        RecoveryGate {
            max_detection_ms: 3_000,
            max_recovery_ms,
            require_clean_timeout,
        }
    }

    #[test]
    fn evaluate_recovery_passes_a_clean_fallback() {
        // The agent called the hung tool once, got the synthesized timeout, then
        // replied cleanly: recovered at one detection budget (3000ms).
        let envelope = json!({
            "tool_calls": [{ "name": "search", "server": "backend" }],
            "final_response": "The search tool is unavailable right now.",
            "conversation": { "stop_reason": "completed" }
        });
        let plan = vec![hang_fault("search")];
        let outcome = evaluate_recovery(&envelope, &plan, &gate(Some(5_000), true));
        assert_eq!(outcome.detections, 1);
        assert!(outcome.passed, "{}", outcome.detail);
        assert_eq!(outcome.report.time_to_recovery_ms, Some(3_000));
        assert!(outcome.report.clean_timeout);
    }

    #[test]
    fn evaluate_recovery_fails_a_hung_agent_that_never_replies() {
        // The agent kept calling the hung tool until the turn budget tripped and
        // never produced a final reply: it did not recover (exit 1 quality).
        let envelope = json!({
            "tool_calls": [
                { "name": "search", "server": "backend" },
                { "name": "search", "server": "backend" }
            ],
            "final_response": "",
            "conversation": { "stop_reason": "turn_budget" }
        });
        let plan = vec![hang_fault("search")];
        let outcome = evaluate_recovery(&envelope, &plan, &gate(Some(5_000), true));
        assert!(!outcome.passed);
        assert!(!outcome.report.recovered);
        assert!(outcome.detail.contains("did not recover"));
    }

    #[test]
    fn evaluate_recovery_fails_when_recovery_exceeds_budget() {
        // Two hung detections cost 6000ms of virtual recovery, over the 5000ms
        // cap, so the agent recovered but too slowly.
        let envelope = json!({
            "tool_calls": [
                { "name": "search", "server": "backend" },
                { "name": "search", "server": "backend" }
            ],
            "final_response": "Giving up on search.",
            "conversation": { "stop_reason": "completed" }
        });
        let plan = vec![hang_fault("search")];
        let outcome = evaluate_recovery(&envelope, &plan, &gate(Some(5_000), false));
        assert_eq!(outcome.report.time_to_recovery_ms, Some(6_000));
        assert!(!outcome.passed);
    }

    #[test]
    fn evaluate_recovery_passes_when_agent_avoids_the_faulted_tool() {
        // The agent never called the faulted tool and replied cleanly: nothing to
        // time out on, so a require_clean_timeout gate still passes.
        let envelope = json!({
            "tool_calls": [{ "name": "lookup", "server": "backend" }],
            "final_response": "Here is the answer.",
            "conversation": { "stop_reason": "completed" }
        });
        let plan = vec![hang_fault("search")];
        let outcome = evaluate_recovery(&envelope, &plan, &gate(None, true));
        assert_eq!(outcome.detections, 0);
        assert!(outcome.passed, "{}", outcome.detail);
    }
}
