//! Within-session Agent Stability Index sub-scores.
//!
//! Why this module exists: the cross-run model-compat drift report
//! ([`crate::model_compat`]) catches a model that behaves differently
//! between two captured baselines. It says nothing about a *single*
//! long session quietly degrading as it runs. The community asked for a
//! drift signal that does NOT depend on HTTP errors (the Claude Code
//! "work feels worse but no metric moved" regression). This module is
//! that signal: deterministic, transport-free sub-scores folded from one
//! multi-turn trace envelope, mirroring the model-compat "signal +
//! threshold + flagged" shape so the two reports complement each other.
//!
//! Background, Agent Stability Index (arXiv:2601.04170): within a single
//! multi-turn session an agent can progressively lose coherence, thrash
//! between tools, loop on a failing call, or burn tokens without making
//! progress. The paper scores twelve dimensions with a model in the loop.
//! We implement a tractable DETERMINISTIC subset over the serialized
//! trace, so a suite can assert the floors in CI with no model judgement.
//! Each sub-score is a heuristic in `0.0..=1.0` where HIGHER means MORE
//! STABLE. They are heuristics, not a model's opinion: treat a low score
//! as "look here," not as a proof of regression.

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

/// Deterministic within-session stability sub-scores folded from one
/// multi-turn agent trace envelope.
///
/// Every sub-score is a heuristic in `0.0..=1.0` where higher is more
/// stable. Serializable so reporters and the enterprise audit chain can
/// persist the within-session picture next to the cross-run model-compat
/// drift. A short or empty trace is reported as trivially stable (all
/// scores `1.0`): there is no within-session history to degrade.
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub struct StabilityReport {
    /// Number of assistant turns the scores were folded over. One or zero
    /// assistant turns is trivially stable (no progression to measure).
    pub assistant_turns: usize,
    /// Total tool calls observed in the trace.
    pub tool_calls: usize,
    /// Tool-usage-pattern stability: how concentrated the tool repertoire
    /// is. Computed as `1 - (distinct_tools - 1) / (total_calls - 1)` over
    /// the tool-call sequence, clamped to `0.0..=1.0`. A session that
    /// stays on one or two tools scores high; one that thrashes across a
    /// fresh tool on nearly every call scores low. `1.0` when there are
    /// fewer than two calls (nothing to thrash between).
    pub tool_usage_stability: f64,
    /// Response-consistency: low variance in assistant-turn length is
    /// stable. Computed as `1 - min(1, coefficient_of_variation)` of the
    /// assistant-turn character counts (std dev over mean). A run whose
    /// answers stay a similar size scores high; one whose answers swing
    /// from a sentence to a wall of text scores low. `1.0` with fewer
    /// than two assistant turns or an all-empty set.
    pub response_consistency: f64,
    /// Redundancy / looping: the fraction of tool calls that are NOT exact
    /// repeats of an earlier call. Computed as `distinct_calls /
    /// total_calls` keyed on `(name, server, args)`. An agent that loops
    /// the same failing call drives this toward zero; one that never
    /// repeats itself scores `1.0`. `1.0` when there are no tool calls.
    pub redundancy: f64,
    /// Cost-per-progress: total tokens spent relative to useful (distinct)
    /// tool calls, mapped to a stability score. "Progress" is the count of
    /// distinct tool calls (a repeated call is not progress). The score is
    /// `target_tokens_per_call / max(target, tokens_per_call)` so a run at
    /// or under the target budget per useful call scores `1.0` and a run
    /// that spends multiples of it decays toward zero. `1.0` when there
    /// are no tokens recorded (nothing was spent).
    pub cost_per_progress: f64,
}

impl StabilityReport {
    /// The lowest of the four sub-scores. The CLI prints this as the
    /// headline within-session stability number so an operator sees the
    /// weakest dimension first.
    pub fn weakest_score(&self) -> f64 {
        self.tool_usage_stability
            .min(self.response_consistency)
            .min(self.redundancy)
            .min(self.cost_per_progress)
    }

    /// One-line human-readable summary of the four sub-scores. The CLI
    /// prints this so the operator sees every dimension together without
    /// re-scanning the struct.
    pub fn summary(&self) -> String {
        format!(
            "tool-usage {:.2}, response-consistency {:.2}, redundancy {:.2}, cost-per-progress {:.2}",
            self.tool_usage_stability,
            self.response_consistency,
            self.redundancy,
            self.cost_per_progress,
        )
    }
}

/// Per-sub-score floors below which a within-session drift flag fires
///. A score must be strictly below its floor to flag, so a
/// floor of `0.0` never fires and a floor of `1.0` flags anything short
/// of perfect. Mirrors the model-compat threshold style so suites
/// configure both reports the same way.
#[derive(Debug, Clone, Copy, PartialEq, Serialize, Deserialize)]
pub struct StabilityFloors {
    /// Minimum acceptable [`StabilityReport::tool_usage_stability`].
    pub tool_usage_stability: f64,
    /// Minimum acceptable [`StabilityReport::response_consistency`].
    pub response_consistency: f64,
    /// Minimum acceptable [`StabilityReport::redundancy`].
    pub redundancy: f64,
    /// Minimum acceptable [`StabilityReport::cost_per_progress`].
    pub cost_per_progress: f64,
}

impl Default for StabilityFloors {
    /// Conservative defaults that fire only on clear within-session
    /// degradation: a `0.5` floor on every sub-score. A score under half
    /// means the dimension is more degraded than stable.
    fn default() -> Self {
        Self {
            tool_usage_stability: 0.5,
            response_consistency: 0.5,
            redundancy: 0.5,
            cost_per_progress: 0.5,
        }
    }
}

/// Target token budget per useful (distinct) tool call. A run at or under
/// this spend per useful call scores a perfect cost-per-progress; above
/// it the score decays. Chosen as a round mid-size budget so the score is
/// a relative signal, not an absolute pricing claim.
const TARGET_TOKENS_PER_CALL: f64 = 2000.0;

/// Fold one multi-turn agent trace envelope into a [`StabilityReport`].
///
/// The envelope is expected to carry the shape
/// `ConversationTrace::to_envelope` produces: `tool_calls: [{name,
/// server, args}]` and `conversation: { tokens: { total }, turns:
/// [{role, content}] }`. Missing or malformed fields are tolerated: a
/// trace with no readable turns or calls reports trivially stable, since
/// there is no within-session progression to degrade.
pub fn agent_stability(trace_envelope: &Value) -> StabilityReport {
    let calls = read_tool_calls(trace_envelope);
    let assistant_lengths = read_assistant_turn_lengths(trace_envelope);
    let total_tokens = read_total_tokens(trace_envelope);

    let distinct_calls = distinct_call_count(&calls);
    StabilityReport {
        assistant_turns: assistant_lengths.len(),
        tool_calls: calls.len(),
        tool_usage_stability: tool_usage_stability(&calls),
        response_consistency: response_consistency(&assistant_lengths),
        redundancy: redundancy(calls.len(), distinct_calls),
        cost_per_progress: cost_per_progress(total_tokens, distinct_calls),
    }
}

/// Return the names of the sub-scores that fell strictly below their
/// configured floor. This is the assertable within-session
/// drift signal: an empty vector means no dimension drifted, and the CLI
/// gates on it the same way it gates on the cross-run model-compat diff.
/// Names are stable snake_case identifiers in a fixed order so callers
/// and tests can match on them.
pub fn drift_flags(report: &StabilityReport, floors: &StabilityFloors) -> Vec<String> {
    let mut flags = Vec::new();
    if report.tool_usage_stability < floors.tool_usage_stability {
        flags.push("tool_usage_stability".to_string());
    }
    if report.response_consistency < floors.response_consistency {
        flags.push("response_consistency".to_string());
    }
    if report.redundancy < floors.redundancy {
        flags.push("redundancy".to_string());
    }
    if report.cost_per_progress < floors.cost_per_progress {
        flags.push("cost_per_progress".to_string());
    }
    flags
}

/// A normalized fingerprint of one tool call used to detect exact repeats
/// and count the distinct repertoire. Keyed on name, server, and the
/// canonical JSON of the arguments so a re-issued identical call collides.
type CallKey = (String, String, String);

/// Pull the tool-call fingerprints from the envelope in order. A missing
/// or non-array `tool_calls` yields an empty vector (no calls to score).
fn read_tool_calls(envelope: &Value) -> Vec<CallKey> {
    let Some(calls) = envelope.get("tool_calls").and_then(|v| v.as_array()) else {
        return Vec::new();
    };
    calls
        .iter()
        .map(|call| {
            let name = call
                .get("name")
                .and_then(|n| n.as_str())
                .unwrap_or_default()
                .to_string();
            let server = call
                .get("server")
                .and_then(|s| s.as_str())
                .unwrap_or_default()
                .to_string();
            // Canonical JSON string of the args so two equal arg bags
            // collide regardless of in-memory key order.
            let args = call
                .get("args")
                .map(canonical_json)
                .unwrap_or_else(|| "null".to_string());
            (name, server, args)
        })
        .collect()
}

/// Render a JSON value as a canonical string. `serde_json::Value` already
/// stores object keys in a `BTreeMap`, so serialization is deterministic;
/// the fallback string keeps the function total without an unwrap.
fn canonical_json(v: &Value) -> String {
    serde_json::to_string(v).unwrap_or_else(|_| "null".to_string())
}

/// Pull the character length of each assistant turn from
/// `conversation.turns`, skipping non-assistant roles. A missing turns
/// array yields an empty vector.
fn read_assistant_turn_lengths(envelope: &Value) -> Vec<usize> {
    let Some(turns) = envelope
        .get("conversation")
        .and_then(|c| c.get("turns"))
        .and_then(|t| t.as_array())
    else {
        return Vec::new();
    };
    turns
        .iter()
        .filter(|turn| turn.get("role").and_then(|r| r.as_str()) == Some("assistant"))
        .map(|turn| {
            turn.get("content")
                .and_then(|c| c.as_str())
                .map(|s| s.chars().count())
                .unwrap_or(0)
        })
        .collect()
}

/// Read `conversation.tokens.total`, defaulting to zero when the path is
/// absent or not an unsigned integer.
fn read_total_tokens(envelope: &Value) -> u64 {
    envelope
        .get("conversation")
        .and_then(|c| c.get("tokens"))
        .and_then(|t| t.get("total"))
        .and_then(|v| v.as_u64())
        .unwrap_or(0)
}

/// Count the distinct tool-call fingerprints (the useful repertoire).
fn distinct_call_count(calls: &[CallKey]) -> usize {
    let mut seen: Vec<&CallKey> = Vec::new();
    for call in calls {
        if !seen.contains(&call) {
            seen.push(call);
        }
    }
    seen.len()
}

/// Tool-usage-pattern stability. See [`StabilityReport::tool_usage_stability`].
fn tool_usage_stability(calls: &[CallKey]) -> f64 {
    if calls.len() < 2 {
        return 1.0;
    }
    let distinct_tools = distinct_tool_names(calls);
    // Spread of the repertoire across the call sequence. One tool over N
    // calls is fully concentrated (0 spread); a fresh tool every call is
    // maximally spread (1.0). Stability is one minus that spread.
    let spread = (distinct_tools.saturating_sub(1)) as f64 / (calls.len() - 1) as f64;
    (1.0 - spread).clamp(0.0, 1.0)
}

/// Count distinct tool NAMES (ignoring args), the size of the repertoire.
fn distinct_tool_names(calls: &[CallKey]) -> usize {
    let mut names: Vec<&str> = Vec::new();
    for (name, _, _) in calls {
        if !names.contains(&name.as_str()) {
            names.push(name);
        }
    }
    names.len()
}

/// Response-consistency. See [`StabilityReport::response_consistency`].
fn response_consistency(lengths: &[usize]) -> f64 {
    if lengths.len() < 2 {
        return 1.0;
    }
    let n = lengths.len() as f64;
    let mean = lengths.iter().map(|&l| l as f64).sum::<f64>() / n;
    if mean <= 0.0 {
        // All turns empty: no shape to vary, treat as consistent.
        return 1.0;
    }
    let variance = lengths
        .iter()
        .map(|&l| {
            let d = l as f64 - mean;
            d * d
        })
        .sum::<f64>()
        / n;
    let coefficient_of_variation = variance.sqrt() / mean;
    (1.0 - coefficient_of_variation.min(1.0)).clamp(0.0, 1.0)
}

/// Redundancy / looping. See [`StabilityReport::redundancy`].
fn redundancy(total_calls: usize, distinct_calls: usize) -> f64 {
    if total_calls == 0 {
        return 1.0;
    }
    (distinct_calls as f64 / total_calls as f64).clamp(0.0, 1.0)
}

/// Cost-per-progress. See [`StabilityReport::cost_per_progress`].
fn cost_per_progress(total_tokens: u64, distinct_calls: usize) -> f64 {
    if total_tokens == 0 {
        return 1.0;
    }
    // A run that spent tokens but made no useful progress (no distinct
    // calls) is the worst case: floor the score at zero.
    if distinct_calls == 0 {
        return 0.0;
    }
    let tokens_per_call = total_tokens as f64 / distinct_calls as f64;
    let ratio = TARGET_TOKENS_PER_CALL / tokens_per_call.max(TARGET_TOKENS_PER_CALL);
    ratio.clamp(0.0, 1.0)
}

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

    /// Build an envelope with the given tool calls, assistant-turn
    /// lengths (rendered as strings of that many `x`), and token total.
    fn envelope(calls: Value, assistant_lengths: &[usize], total_tokens: u64) -> Value {
        let turns: Vec<Value> = assistant_lengths
            .iter()
            .map(|&len| json!({ "role": "assistant", "content": "x".repeat(len) }))
            .collect();
        json!({
            "tool_calls": calls,
            "conversation": {
                "tokens": { "total": total_tokens },
                "turns": turns,
            }
        })
    }

    #[test]
    fn clean_focused_trace_scores_high_on_every_subscore() {
        // One tool used consistently, distinct calls, similar-length
        // answers, modest token spend: stable on all four dimensions.
        let calls = json!([
            { "name": "search", "server": "docs", "args": { "q": "a" } },
            { "name": "search", "server": "docs", "args": { "q": "b" } },
            { "name": "fetch", "server": "docs", "args": { "id": 1 } },
        ]);
        let report = agent_stability(&envelope(calls, &[100, 110, 95], 3000));
        assert!(report.tool_usage_stability >= 0.5);
        assert!(report.response_consistency >= 0.5);
        assert!(report.redundancy >= 0.5);
        assert!(report.cost_per_progress >= 0.5);
        assert!(drift_flags(&report, &StabilityFloors::default()).is_empty());
    }

    #[test]
    fn looping_the_same_tool_call_scores_low_on_redundancy() {
        // The same call re-issued five times: classic degradation loop.
        let one = json!({ "name": "retry", "server": "api", "args": { "x": 1 } });
        let calls = json!([one, one, one, one, one]);
        let report = agent_stability(&envelope(calls, &[100, 100], 1000));
        // Only one distinct call out of five: redundancy is 0.2.
        assert!((report.redundancy - 0.2).abs() < 1e-9);
        let flags = drift_flags(&report, &StabilityFloors::default());
        assert!(flags.contains(&"redundancy".to_string()));
    }

    #[test]
    fn escalating_cost_with_no_progress_flags_cost_per_progress() {
        // Heavy token spend against a single useful call: the run burned
        // budget without making progress.
        let calls = json!([{ "name": "think", "server": "api", "args": {} }]);
        let report = agent_stability(&envelope(calls, &[100], 50_000));
        // 50k tokens / 1 distinct call is 25x the target budget.
        assert!(report.cost_per_progress < 0.5);
        let flags = drift_flags(&report, &StabilityFloors::default());
        assert!(flags.contains(&"cost_per_progress".to_string()));
    }

    #[test]
    fn thrashing_across_fresh_tools_scores_low_on_tool_usage() {
        // A different tool on every call: maximal repertoire spread.
        let calls = json!([
            { "name": "a", "server": "s", "args": {} },
            { "name": "b", "server": "s", "args": {} },
            { "name": "c", "server": "s", "args": {} },
            { "name": "d", "server": "s", "args": {} },
        ]);
        let report = agent_stability(&envelope(calls, &[100, 100], 1000));
        // 4 distinct tools over 4 calls: spread (4-1)/(4-1) = 1.0, so
        // tool-usage stability is 0.0.
        assert!(report.tool_usage_stability < 0.5);
        let flags = drift_flags(&report, &StabilityFloors::default());
        assert!(flags.contains(&"tool_usage_stability".to_string()));
    }

    #[test]
    fn swinging_response_length_scores_low_on_consistency() {
        // Answers swing from one char to a wall of text: high variance.
        let calls = json!([]);
        let report = agent_stability(&envelope(calls, &[1, 1, 1, 2000], 0));
        assert!(report.response_consistency < 0.5);
        let flags = drift_flags(&report, &StabilityFloors::default());
        assert!(flags.contains(&"response_consistency".to_string()));
    }

    #[test]
    fn single_turn_trace_is_trivially_stable() {
        let calls = json!([{ "name": "go", "server": "s", "args": {} }]);
        let report = agent_stability(&envelope(calls, &[120], 500));
        assert_eq!(report.assistant_turns, 1);
        assert!((report.tool_usage_stability - 1.0).abs() < f64::EPSILON);
        assert!((report.response_consistency - 1.0).abs() < f64::EPSILON);
        assert!((report.redundancy - 1.0).abs() < f64::EPSILON);
        assert!((report.cost_per_progress - 1.0).abs() < f64::EPSILON);
        assert!(drift_flags(&report, &StabilityFloors::default()).is_empty());
    }

    #[test]
    fn empty_trace_is_trivially_stable() {
        let report = agent_stability(&json!({}));
        assert_eq!(report.assistant_turns, 0);
        assert_eq!(report.tool_calls, 0);
        assert!((report.weakest_score() - 1.0).abs() < f64::EPSILON);
        assert!(drift_flags(&report, &StabilityFloors::default()).is_empty());
    }

    #[test]
    fn drift_flags_returns_only_subscores_below_their_floor() {
        // Construct a report by hand so the floor logic is isolated from
        // the folding math.
        let report = StabilityReport {
            assistant_turns: 3,
            tool_calls: 4,
            tool_usage_stability: 0.9,
            response_consistency: 0.3,
            redundancy: 0.4,
            cost_per_progress: 0.8,
        };
        let floors = StabilityFloors {
            tool_usage_stability: 0.5,
            response_consistency: 0.5,
            redundancy: 0.5,
            cost_per_progress: 0.5,
        };
        let flags = drift_flags(&report, &floors);
        assert_eq!(
            flags,
            vec!["response_consistency".to_string(), "redundancy".to_string()]
        );
    }

    #[test]
    fn a_zero_floor_never_fires_and_a_one_floor_flags_anything_imperfect() {
        let report = StabilityReport {
            assistant_turns: 2,
            tool_calls: 2,
            tool_usage_stability: 0.99,
            response_consistency: 1.0,
            redundancy: 1.0,
            cost_per_progress: 1.0,
        };
        let open = StabilityFloors {
            tool_usage_stability: 0.0,
            response_consistency: 0.0,
            redundancy: 0.0,
            cost_per_progress: 0.0,
        };
        assert!(drift_flags(&report, &open).is_empty());
        let strict = StabilityFloors {
            tool_usage_stability: 1.0,
            response_consistency: 1.0,
            redundancy: 1.0,
            cost_per_progress: 1.0,
        };
        // Only the sub-score short of 1.0 flags under a strict floor.
        assert_eq!(
            drift_flags(&report, &strict),
            vec!["tool_usage_stability".to_string()]
        );
    }

    #[test]
    fn repeated_call_is_not_counted_as_progress_for_cost() {
        // Three calls but only one distinct: tokens divide by 1 useful
        // call, so the loop is penalized on cost as well as redundancy.
        let one = json!({ "name": "loop", "server": "s", "args": { "n": 1 } });
        let calls = json!([one, one, one]);
        let report = agent_stability(&envelope(calls, &[100, 100], 10_000));
        // 10k / 1 distinct = 10k tokens per useful call, 5x the target.
        assert!(report.cost_per_progress < 0.5);
    }

    #[test]
    fn args_differing_only_in_key_order_collide_as_one_distinct_call() {
        let calls = json!([
            { "name": "f", "server": "s", "args": { "a": 1, "b": 2 } },
            { "name": "f", "server": "s", "args": { "b": 2, "a": 1 } },
        ]);
        let report = agent_stability(&envelope(calls, &[100, 100], 1000));
        // The two calls are the same after canonical JSON, so redundancy
        // sees one distinct call out of two: 0.5.
        assert!((report.redundancy - 0.5).abs() < 1e-9);
    }

    #[test]
    fn report_round_trips_through_serde() {
        // Use exactly representable sub-scores so the round-trip is bit
        // for bit; this isolates serde from float printing, which can
        // round the last ULP of a non-terminating fraction.
        let report = StabilityReport {
            assistant_turns: 2,
            tool_calls: 1,
            tool_usage_stability: 1.0,
            response_consistency: 0.5,
            redundancy: 1.0,
            cost_per_progress: 0.25,
        };
        let text = serde_json::to_string(&report).expect("serialize");
        let back: StabilityReport = serde_json::from_str(&text).expect("deserialize");
        assert_eq!(report, back);
    }

    #[test]
    fn weakest_score_is_the_minimum_subscore() {
        let report = StabilityReport {
            assistant_turns: 3,
            tool_calls: 3,
            tool_usage_stability: 0.7,
            response_consistency: 0.2,
            redundancy: 0.9,
            cost_per_progress: 0.6,
        };
        assert!((report.weakest_score() - 0.2).abs() < f64::EPSILON);
    }
}
