//! Label-free cross-run consistency metrics for a multi-run agent test
//! (WOR-1369).
//!
//! Why this module exists: a flaky agent can clear the per-run stability
//! sub-scores yet still take a different path on every run. These three
//! deterministic, model-free metrics measure how consistent the N run traces
//! are *with each other*, so a `runs: N` test gates on reproducibility, not just
//! on each run's internal stability. They read the recorded tool-call trace and
//! call no model, so a fixed set of envelopes always yields the same numbers.
//!
//! - **Tool-sequence similarity (TSS)**: how alike the per-run tool-name
//!   sequences are, as the mean pairwise longest-common-subsequence ratio.
//! - **Argument consistency (AC)**: among tool calls that line up by position
//!   and name across two runs, the fraction that were called with identical
//!   arguments.
//! - **Early divergence**: whether the runs tend to split apart in the first
//!   two steps (where most agent divergence concentrates) rather than late.

use serde_json::Value;

use super::tool_calls_from_envelope;

/// The tool-name sequence and the per-call argument values of one run, pulled
/// from its recorded envelope. Kept together so the pairwise passes walk one
/// allocation per run instead of re-parsing the envelope.
struct RunTrace {
    names: Vec<String>,
    args: Vec<Value>,
}

/// Project one run envelope onto its ordered tool-name sequence and the
/// argument object of each call. A call with no `name` is skipped; an absent
/// `args` (or the `arguments` spelling some providers use) becomes `null`.
fn run_trace(envelope: &Value) -> RunTrace {
    let calls = tool_calls_from_envelope(envelope);
    let mut names = Vec::with_capacity(calls.len());
    let mut args = Vec::with_capacity(calls.len());
    for call in &calls {
        let Some(name) = call.get("name").and_then(Value::as_str) else {
            continue;
        };
        names.push(name.to_string());
        let arg = call
            .get("args")
            .or_else(|| call.get("arguments"))
            .cloned()
            .unwrap_or(Value::Null);
        args.push(arg);
    }
    RunTrace { names, args }
}

/// Length of the longest common subsequence of two name sequences. Standard
/// dynamic-programming fill; the sequences are short (one per tool call), so the
/// `O(n*m)` table is cheap.
fn lcs_len(a: &[String], b: &[String]) -> usize {
    if a.is_empty() || b.is_empty() {
        return 0;
    }
    let mut prev = vec![0usize; b.len() + 1];
    let mut curr = vec![0usize; b.len() + 1];
    for ai in a {
        for (j, bj) in b.iter().enumerate() {
            curr[j + 1] = if ai == bj {
                prev[j] + 1
            } else {
                prev[j + 1].max(curr[j])
            };
        }
        std::mem::swap(&mut prev, &mut curr);
        curr.iter_mut().for_each(|c| *c = 0);
    }
    prev[b.len()]
}

/// Iterate every unordered pair of runs `(i < j)`. Returns an empty iterator
/// for fewer than two runs, where consistency is vacuous.
fn pairs(n: usize) -> impl Iterator<Item = (usize, usize)> {
    (0..n).flat_map(move |i| (i + 1..n).map(move |j| (i, j)))
}

/// Tool-sequence similarity across the run traces, in `0.0..=1.0`.
///
/// The mean over every run pair of `2 * lcs / (len_a + len_b)` (the
/// sequence Sorensen-Dice coefficient): `1.0` when every run called the same
/// tools in the same order, lower as the orders diverge. Two empty traces are
/// identical (`1.0`); fewer than two runs is trivially consistent (`1.0`).
pub fn tool_sequence_similarity(envelopes: &[Value]) -> f64 {
    let traces: Vec<RunTrace> = envelopes.iter().map(run_trace).collect();
    let mut total = 0.0;
    let mut count = 0u32;
    for (i, j) in pairs(traces.len()) {
        let (a, b) = (&traces[i].names, &traces[j].names);
        let denom = a.len() + b.len();
        let sim = if denom == 0 {
            1.0
        } else {
            2.0 * lcs_len(a, b) as f64 / denom as f64
        };
        total += sim;
        count += 1;
    }
    if count == 0 {
        1.0
    } else {
        total / count as f64
    }
}

/// Argument consistency across the run traces, in `0.0..=1.0`.
///
/// For each run pair, walk the positions both runs reached; at the positions
/// where the same tool was called, score the fraction whose argument objects
/// are byte-identical, then average that fraction across pairs. `1.0` when
/// aligned calls always used the same args; vacuous cases (no aligned same-name
/// positions, or fewer than two runs) score `1.0`.
pub fn argument_consistency(envelopes: &[Value]) -> f64 {
    let traces: Vec<RunTrace> = envelopes.iter().map(run_trace).collect();
    let mut total = 0.0;
    let mut count = 0u32;
    for (i, j) in pairs(traces.len()) {
        let (a, b) = (&traces[i], &traces[j]);
        let common = a.names.len().min(b.names.len());
        let mut aligned = 0u32;
        let mut matching = 0u32;
        for pos in 0..common {
            if a.names[pos] == b.names[pos] {
                aligned += 1;
                if a.args[pos] == b.args[pos] {
                    matching += 1;
                }
            }
        }
        let pair_score = if aligned == 0 {
            1.0
        } else {
            matching as f64 / aligned as f64
        };
        total += pair_score;
        count += 1;
    }
    if count == 0 {
        1.0
    } else {
        total / count as f64
    }
}

/// Whether the runs tend to diverge early (in steps 1-2).
///
/// For each run pair, find the first position where the tool-name sequences
/// differ (a shorter sequence that is a prefix of a longer one diverges at its
/// end). `true` when a strict majority of the pairs that diverge at all do so at
/// step index 0 or 1, the band where agent divergence concentrates. When no
/// pair diverges, the runs are reproducible and this is `false`.
pub fn early_divergence(envelopes: &[Value]) -> bool {
    let traces: Vec<RunTrace> = envelopes.iter().map(run_trace).collect();
    let mut early = 0u32;
    let mut diverged = 0u32;
    for (i, j) in pairs(traces.len()) {
        let (a, b) = (&traces[i].names, &traces[j].names);
        let common = a.len().min(b.len());
        let first_diff = (0..common).find(|&p| a[p] != b[p]);
        let divergence_at = match first_diff {
            Some(pos) => Some(pos),
            // No mismatch in the shared prefix: a length difference diverges at
            // the shorter length, identical sequences do not diverge at all.
            None if a.len() != b.len() => Some(common),
            None => None,
        };
        if let Some(pos) = divergence_at {
            diverged += 1;
            if pos <= 1 {
                early += 1;
            }
        }
    }
    diverged > 0 && early * 2 > diverged
}

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

    /// Envelope whose recorded `tool_calls` are the given (name, args) pairs.
    fn envelope(calls: &[(&str, Value)]) -> Value {
        let tool_calls: Vec<Value> = calls
            .iter()
            .map(|(name, args)| json!({ "name": name, "args": args }))
            .collect();
        json!({ "tool_calls": tool_calls })
    }

    fn names_only(names: &[&str]) -> Value {
        let calls: Vec<(&str, Value)> = names.iter().map(|n| (*n, json!({}))).collect();
        envelope(&calls)
    }

    #[test]
    fn identical_runs_are_fully_consistent() {
        let runs = vec![names_only(&["a", "b", "c"]), names_only(&["a", "b", "c"])];
        assert_eq!(tool_sequence_similarity(&runs), 1.0);
        assert_eq!(argument_consistency(&runs), 1.0);
        assert!(!early_divergence(&runs));
    }

    #[test]
    fn a_single_run_is_trivially_consistent() {
        let runs = vec![names_only(&["a", "b"])];
        assert_eq!(tool_sequence_similarity(&runs), 1.0);
        assert_eq!(argument_consistency(&runs), 1.0);
        assert!(!early_divergence(&runs));
    }

    #[test]
    fn tss_drops_when_the_sequences_share_only_part() {
        // LCS of [a,b,c] and [a,x,c] is [a,c] = 2; 2*2/(3+3) = 0.666...
        let runs = vec![names_only(&["a", "b", "c"]), names_only(&["a", "x", "c"])];
        let tss = tool_sequence_similarity(&runs);
        assert!((tss - 2.0 / 3.0).abs() < 1e-9, "tss was {tss}");
    }

    #[test]
    fn ac_drops_when_aligned_calls_use_different_args() {
        // Same tool at both positions, but the second call's args differ in one
        // of the two runs: 1 of 2 aligned positions match -> 0.5.
        let a = envelope(&[("search", json!({"q": "x"})), ("fetch", json!({"id": 1}))]);
        let b = envelope(&[("search", json!({"q": "x"})), ("fetch", json!({"id": 2}))]);
        assert_eq!(argument_consistency(&[a, b]), 0.5);
    }

    #[test]
    fn early_divergence_flags_a_step_one_split() {
        // The two runs differ at the very first call.
        let runs = vec![names_only(&["a", "b", "c"]), names_only(&["z", "b", "c"])];
        assert!(early_divergence(&runs));
    }

    #[test]
    fn late_divergence_is_not_flagged_early() {
        // The runs agree through steps 1-4 and only differ at step 5.
        let runs = vec![
            names_only(&["a", "b", "c", "d", "e"]),
            names_only(&["a", "b", "c", "d", "z"]),
        ];
        assert!(!early_divergence(&runs));
    }

    #[test]
    fn a_prefix_run_diverges_at_its_end() {
        // [a,b] vs [a,b,c]: identical prefix, divergence at index 2 (late).
        let runs = vec![names_only(&["a", "b"]), names_only(&["a", "b", "c"])];
        assert!(!early_divergence(&runs));
        // A two-call prefix vs a longer run still shares its whole length.
        let tss = tool_sequence_similarity(&runs);
        assert!((tss - 2.0 * 2.0 / 5.0).abs() < 1e-9, "tss was {tss}");
    }
}
