//! Parser for the `context_rot_probe:` agent-test gate (WOR-1402).
//!
//! Split out of [`super::agent_gates_parse`] so that module stays under the
//! production-module-size guard. Reuses the distractor `source:` and `correct:`
//! parsers from there, and the shared expectation parser from
//! [`super::test_parse`].

use serde_yaml::Value as YamlValue;

use crate::eval::distractors::DistractorSource;

use super::agent_gates_parse::{parse_correct_tools, parse_distractor_source};
use super::test_parse::parse_expectations;
use super::{ContextRotProbeGate, ParseError};

/// Parse the optional `context_rot_probe:` gate on an agent test.
///
/// The block is an object with a required non-empty `levels:` list (distractor
/// counts to sweep, each an unsigned integer), a `source:` (the same
/// `catalog` / `near_duplicate` choice the `distractors:` block parses;
/// defaults to `catalog` when omitted so the simplest case stays terse), a
/// required non-empty `correct:` list of tool ids, and an optional `expect:`
/// list over the `context_rot.*` targets. An empty or omitted `expect:` reports
/// the curve without gating (a degradation floor is a deliberate choice).
pub(super) fn parse_context_rot_probe_gate(
    entry: &YamlValue,
    path: &str,
) -> Result<Option<ContextRotProbeGate>, ParseError> {
    let Some(node) = entry.get("context_rot_probe") else {
        return Ok(None);
    };
    let block_path = format!("{path}/context_rot_probe");
    if node.as_mapping().is_none() {
        return Err(ParseError::BadShape {
            path: block_path,
            field: "context_rot_probe",
            expected: "an object with levels:, source:, and correct: keys",
        });
    }
    let levels = parse_levels(node, &block_path)?;
    // The source defaults to the bundled catalog so the simplest sweep needs
    // only levels: and correct:. An explicit source: parses the same shape the
    // distractors: block uses.
    let source = match node.get("source") {
        Some(_) => parse_distractor_source(node, &block_path)?,
        None => DistractorSource::Catalog,
    };
    let correct = parse_correct_tools(node, &block_path)?;
    let expect = parse_expectations(node, &block_path)?;
    Ok(Some(ContextRotProbeGate {
        levels,
        source,
        correct,
        expect,
    }))
}

/// Parse the required non-empty `levels:` list into distractor counts. Each
/// entry is an unsigned integer; the list must be present and non-empty so the
/// sweep always has at least one point.
fn parse_levels(node: &YamlValue, block_path: &str) -> Result<Vec<u32>, ParseError> {
    let seq = node
        .get("levels")
        .and_then(|v| v.as_sequence())
        .ok_or(ParseError::Missing {
            path: block_path.to_string(),
            field: "levels",
        })?;
    if seq.is_empty() {
        return Err(ParseError::BadShape {
            path: block_path.to_string(),
            field: "levels",
            expected: "at least one distractor-count level",
        });
    }
    let mut out = Vec::with_capacity(seq.len());
    for item in seq {
        let level = item.as_u64().ok_or(ParseError::BadShape {
            path: block_path.to_string(),
            field: "levels",
            expected: "a list of non-negative integers",
        })?;
        out.push(level as u32);
    }
    Ok(out)
}

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

    fn entry(yaml: &str) -> YamlValue {
        serde_yaml::from_str(yaml).expect("test yaml parses")
    }

    #[test]
    fn absent_context_rot_probe_block_parses_to_none() {
        let node = entry("name: x\nprompt: hi\n");
        assert!(parse_context_rot_probe_gate(&node, "/agents/0")
            .expect("parse ok")
            .is_none());
    }

    #[test]
    fn context_rot_probe_parses_levels_correct_and_defaults_source_and_expect() {
        let node = entry("context_rot_probe:\n  levels: [0, 8, 32]\n  correct: [catalog.search]\n");
        let gate = parse_context_rot_probe_gate(&node, "/agents/0")
            .expect("parse ok")
            .expect("gate present");
        assert_eq!(gate.levels, vec![0, 8, 32]);
        assert_eq!(gate.correct, vec!["catalog.search".to_string()]);
        // An omitted source defaults to the bundled catalog.
        assert_eq!(gate.source, DistractorSource::Catalog);
        // An absent expect: list reports the curve without gating.
        assert!(gate.expect.is_empty());
    }

    #[test]
    fn context_rot_probe_parses_an_explicit_near_duplicate_source() {
        let node = entry(
            "context_rot_probe:\n  levels: [4]\n  source:\n    from: near_duplicate\n    of: [search]\n  correct: [catalog.search]\n",
        );
        let gate = parse_context_rot_probe_gate(&node, "/agents/0")
            .expect("parse ok")
            .expect("gate present");
        assert_eq!(
            gate.source,
            DistractorSource::NearDuplicate {
                of: vec!["search".to_string()],
            }
        );
    }

    #[test]
    fn context_rot_probe_without_levels_is_a_load_error() {
        let node = entry("context_rot_probe:\n  correct: [catalog.search]\n");
        assert!(parse_context_rot_probe_gate(&node, "/agents/0").is_err());
    }

    #[test]
    fn context_rot_probe_with_empty_levels_is_a_load_error() {
        let node = entry("context_rot_probe:\n  levels: []\n  correct: [catalog.search]\n");
        assert!(parse_context_rot_probe_gate(&node, "/agents/0").is_err());
    }

    #[test]
    fn context_rot_probe_without_correct_is_a_load_error() {
        let node = entry("context_rot_probe:\n  levels: [0, 8]\n");
        assert!(parse_context_rot_probe_gate(&node, "/agents/0").is_err());
    }
}
