//! Parser for the budget half of the long-form `expect:` mapping (WOR-1280).
//!
//! Why a separate module: `test_parse` already carries the whole matcher
//! grammar and sits over the module-size ceiling, and budgets are a
//! self-contained vocabulary (`max_duration_ms`, `max_response_tokens`,
//! `response_headers`, `response_headers_absent`). The assertion half of the
//! long form is handled by `test_parse::parse_expectations`; this module only
//! lifts the budget fields onto the typed test so the executor can gate on
//! them instead of silently dropping them, which was the WOR-1280 bug.

use serde_yaml::Value as YamlValue;

use crate::matchers::budgets::{tokenizer_by_name, BudgetMode};
use crate::matchers::headers::HeaderMatch;

use super::test_parse::yaml_to_json;
use super::types::{ExpectBudgets, TokenBudgetSpec};
use super::ParseError;

/// Default tokenizer when the long-form token budget omits `tokenizer:`.
/// Matches the schema's documented default.
const DEFAULT_TOKENIZER: &str = "cl100k_base";

/// Default per-image token cost, mirroring
/// [`crate::matchers::budgets::TokenBudget::default`].
const DEFAULT_IMAGE_COST: usize = 1000;

/// Lift the budget fields off a test's long-form `expect:` mapping.
///
/// The short-form array and a missing `expect:` both return the empty
/// [`ExpectBudgets`], so legacy suites parse unchanged. Every field is
/// validated here so a bad budget fails the suite load with a path-bearing
/// error instead of surfacing mid-run.
pub(super) fn parse_expect_budgets(
    entry: &YamlValue,
    path: &str,
) -> Result<ExpectBudgets, ParseError> {
    let Some(node) = entry.get("expect") else {
        return Ok(ExpectBudgets::default());
    };
    if !node.is_mapping() {
        return Ok(ExpectBudgets::default());
    }
    let block_path = format!("{path}/expect");
    Ok(ExpectBudgets {
        max_duration_ms: parse_max_duration(node, &block_path)?,
        max_response_tokens: parse_token_budget(node, &block_path)?,
        response_headers: parse_response_headers(node, &block_path)?,
        response_headers_absent: parse_headers_absent(node, &block_path)?,
    })
}

/// Parse `max_duration_ms` as a positive integer. Zero is rejected (the
/// schema pins `minimum: 1`) because a zero budget can never pass and
/// always indicates a typo.
fn parse_max_duration(node: &YamlValue, path: &str) -> Result<Option<u64>, ParseError> {
    let Some(value) = node.get("max_duration_ms") else {
        return Ok(None);
    };
    match value.as_u64() {
        Some(ms) if ms >= 1 => Ok(Some(ms)),
        _ => Err(ParseError::BadShape {
            path: format!("{path}/max_duration_ms"),
            field: "max_duration_ms",
            expected: "a positive integer number of milliseconds",
        }),
    }
}

/// Parse `max_response_tokens`, both shapes: the bare positive integer and
/// the long object form (`budget`, `tokenizer`, `mode`, `image_cost`). The
/// tokenizer name is resolved here so a typo is a load error, not a
/// first-run failure.
fn parse_token_budget(node: &YamlValue, path: &str) -> Result<Option<TokenBudgetSpec>, ParseError> {
    let Some(value) = node.get("max_response_tokens") else {
        return Ok(None);
    };
    let field_path = format!("{path}/max_response_tokens");
    if let Some(budget) = value.as_u64() {
        if budget == 0 {
            return Err(ParseError::BadShape {
                path: field_path,
                field: "max_response_tokens",
                expected: "a positive integer token budget",
            });
        }
        return Ok(Some(TokenBudgetSpec {
            budget: budget as usize,
            tokenizer: DEFAULT_TOKENIZER.to_string(),
            mode: BudgetMode::Fail,
            image_cost: DEFAULT_IMAGE_COST,
        }));
    }
    if !value.is_mapping() {
        return Err(ParseError::BadShape {
            path: field_path,
            field: "max_response_tokens",
            expected: "a positive integer or an object with a `budget` field",
        });
    }
    let budget = match value.get("budget").and_then(|v| v.as_u64()) {
        Some(b) if b >= 1 => b as usize,
        _ => {
            return Err(ParseError::BadShape {
                path: format!("{field_path}/budget"),
                field: "budget",
                expected: "a positive integer token budget",
            })
        }
    };
    let tokenizer = match value.get("tokenizer") {
        None => DEFAULT_TOKENIZER.to_string(),
        Some(t) => validate_tokenizer(t, &field_path)?,
    };
    let mode = match value.get("mode").and_then(|v| v.as_str()) {
        None | Some("fail") => BudgetMode::Fail,
        Some("warn") => BudgetMode::Warn,
        Some(_) => {
            return Err(ParseError::BadShape {
                path: format!("{field_path}/mode"),
                field: "mode",
                expected: "`fail` or `warn`",
            })
        }
    };
    let image_cost = match value.get("image_cost") {
        None => DEFAULT_IMAGE_COST,
        Some(c) => c.as_u64().ok_or(ParseError::BadShape {
            path: format!("{field_path}/image_cost"),
            field: "image_cost",
            expected: "a non-negative integer token cost",
        })? as usize,
    };
    Ok(Some(TokenBudgetSpec {
        budget,
        tokenizer,
        mode,
        image_cost,
    }))
}

/// Check the declared tokenizer name resolves. Goes through the same
/// factory the executor uses so the load-time check and the run-time
/// lookup cannot drift apart.
fn validate_tokenizer(value: &YamlValue, field_path: &str) -> Result<String, ParseError> {
    let err = || ParseError::BadShape {
        path: format!("{field_path}/tokenizer"),
        field: "tokenizer",
        expected: "one of `whitespace`, `cl100k_base`, `o200k_base`, `gpt2`, `claude`",
    };
    let name = value.as_str().ok_or_else(err)?;
    match tokenizer_by_name(name) {
        Ok(Some(_)) => Ok(name.to_string()),
        _ => Err(err()),
    }
}

/// Parse the `response_headers:` map into `(name, expectation)` pairs in
/// declaration order. Each value is a bare string (literal equality) or a
/// one-key object naming `regex`, `schema`, `exists`, or `contains`.
fn parse_response_headers(
    node: &YamlValue,
    path: &str,
) -> Result<Vec<(String, HeaderMatch)>, ParseError> {
    let Some(value) = node.get("response_headers") else {
        return Ok(Vec::new());
    };
    let field_path = format!("{path}/response_headers");
    let map = value.as_mapping().ok_or(ParseError::BadShape {
        path: field_path.clone(),
        field: "response_headers",
        expected: "a map of header name to expected value",
    })?;
    let mut out = Vec::with_capacity(map.len());
    for (key, expectation) in map {
        let name = key.as_str().ok_or(ParseError::BadShape {
            path: field_path.clone(),
            field: "response_headers",
            expected: "string header names",
        })?;
        let entry_path = format!("{field_path}/{name}");
        out.push((
            name.to_string(),
            parse_header_match(expectation, &entry_path)?,
        ));
    }
    Ok(out)
}

/// Parse one header expectation value into a [`HeaderMatch`].
fn parse_header_match(value: &YamlValue, path: &str) -> Result<HeaderMatch, ParseError> {
    if let Some(literal) = value.as_str() {
        return Ok(HeaderMatch::Literal(literal.to_string()));
    }
    let bad = |expected: &'static str| ParseError::BadShape {
        path: path.to_string(),
        field: "response_headers",
        expected,
    };
    let map = value
        .as_mapping()
        .ok_or_else(|| bad("a string or a one-key object (regex/schema/exists/contains)"))?;
    if map.len() != 1 {
        return Err(bad(
            "exactly one of `regex`, `schema`, `exists`, or `contains`",
        ));
    }
    if let Some(pattern) = map.get(YamlValue::String("regex".into())) {
        let p = pattern.as_str().ok_or_else(|| bad("a regex string"))?;
        return Ok(HeaderMatch::Regex(p.to_string()));
    }
    if let Some(schema) = map.get(YamlValue::String("schema".into())) {
        return Ok(HeaderMatch::Schema(yaml_to_json(schema)?));
    }
    if let Some(exists) = map.get(YamlValue::String("exists".into())) {
        // `exists: false` is rejected rather than aliased to absence:
        // `response_headers_absent` is the documented spelling for that.
        if exists.as_bool() == Some(true) {
            return Ok(HeaderMatch::Exists);
        }
        return Err(bad(
            "`exists: true` (use `response_headers_absent` for absence)",
        ));
    }
    if let Some(needle) = map.get(YamlValue::String("contains".into())) {
        let n = needle.as_str().ok_or_else(|| bad("a substring"))?;
        return Ok(HeaderMatch::Contains(n.to_string()));
    }
    Err(bad(
        "exactly one of `regex`, `schema`, `exists`, or `contains`",
    ))
}

/// Parse `response_headers_absent:` as a list of header-name strings.
fn parse_headers_absent(node: &YamlValue, path: &str) -> Result<Vec<String>, ParseError> {
    let Some(value) = node.get("response_headers_absent") else {
        return Ok(Vec::new());
    };
    let field_path = format!("{path}/response_headers_absent");
    let seq = value.as_sequence().ok_or(ParseError::BadShape {
        path: field_path.clone(),
        field: "response_headers_absent",
        expected: "a list of header name strings",
    })?;
    let mut out = Vec::with_capacity(seq.len());
    for item in seq {
        let name = item.as_str().ok_or(ParseError::BadShape {
            path: field_path.clone(),
            field: "response_headers_absent",
            expected: "a list of header name strings",
        })?;
        out.push(name.to_string());
    }
    Ok(out)
}

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

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

    #[test]
    fn short_form_expect_array_yields_empty_budgets() {
        let e = entry("expect:\n  - target: result\n    matcher: { exact: 1 }\n");
        let b = parse_expect_budgets(&e, "/tools/0").unwrap();
        assert!(b.is_empty());
    }

    #[test]
    fn missing_expect_yields_empty_budgets() {
        let e = entry("name: t\n");
        let b = parse_expect_budgets(&e, "/tools/0").unwrap();
        assert!(b.is_empty());
    }

    #[test]
    fn mapping_form_lifts_max_duration_ms() {
        let e = entry("expect:\n  assertions: []\n  max_duration_ms: 500\n");
        let b = parse_expect_budgets(&e, "/tools/0").unwrap();
        assert_eq!(b.max_duration_ms, Some(500));
    }

    #[test]
    fn zero_max_duration_ms_is_a_load_error() {
        let e = entry("expect:\n  max_duration_ms: 0\n");
        let err = parse_expect_budgets(&e, "/tools/0").unwrap_err();
        assert!(err.to_string().contains("max_duration_ms"), "{err}");
    }

    #[test]
    fn short_token_budget_defaults_tokenizer_mode_and_image_cost() {
        let e = entry("expect:\n  max_response_tokens: 200\n");
        let b = parse_expect_budgets(&e, "/tools/0").unwrap();
        let spec = b.max_response_tokens.expect("budget parsed");
        assert_eq!(spec.budget, 200);
        assert_eq!(spec.tokenizer, "cl100k_base");
        assert_eq!(spec.mode, BudgetMode::Fail);
        assert_eq!(spec.image_cost, 1000);
    }

    #[test]
    fn long_token_budget_reads_every_field() {
        let e = entry(
            "expect:\n  max_response_tokens:\n    budget: 64\n    tokenizer: whitespace\n    mode: warn\n    image_cost: 10\n",
        );
        let b = parse_expect_budgets(&e, "/tools/0").unwrap();
        let spec = b.max_response_tokens.expect("budget parsed");
        assert_eq!(spec.budget, 64);
        assert_eq!(spec.tokenizer, "whitespace");
        assert_eq!(spec.mode, BudgetMode::Warn);
        assert_eq!(spec.image_cost, 10);
    }

    #[test]
    fn unknown_tokenizer_name_fails_the_load() {
        let e = entry("expect:\n  max_response_tokens:\n    budget: 64\n    tokenizer: bogus\n");
        let err = parse_expect_budgets(&e, "/tools/0").unwrap_err();
        assert!(err.to_string().contains("tokenizer"), "{err}");
    }

    #[test]
    fn unknown_mode_fails_the_load() {
        let e = entry("expect:\n  max_response_tokens:\n    budget: 64\n    mode: yell\n");
        let err = parse_expect_budgets(&e, "/tools/0").unwrap_err();
        assert!(err.to_string().contains("mode"), "{err}");
    }

    #[test]
    fn response_headers_parse_every_value_variant() {
        let e = entry(
            "expect:\n  response_headers:\n    content-type: application/json\n    x-trace: { regex: \"^t-\\\\d+$\" }\n    x-count: { schema: { type: string } }\n    x-any: { exists: true }\n    server: { contains: nginx }\n  response_headers_absent: [set-cookie]\n",
        );
        let b = parse_expect_budgets(&e, "/tools/0").unwrap();
        assert_eq!(b.response_headers.len(), 5);
        assert_eq!(
            b.response_headers[0],
            (
                "content-type".to_string(),
                HeaderMatch::Literal("application/json".to_string())
            )
        );
        assert!(matches!(b.response_headers[1].1, HeaderMatch::Regex(_)));
        assert!(matches!(b.response_headers[2].1, HeaderMatch::Schema(_)));
        assert!(matches!(b.response_headers[3].1, HeaderMatch::Exists));
        assert!(matches!(b.response_headers[4].1, HeaderMatch::Contains(_)));
        assert_eq!(b.response_headers_absent, vec!["set-cookie".to_string()]);
        assert!(b.has_header_checks());
    }

    #[test]
    fn exists_false_is_rejected_with_a_pointer_to_absent() {
        let e = entry("expect:\n  response_headers:\n    x-y: { exists: false }\n");
        let err = parse_expect_budgets(&e, "/tools/0").unwrap_err();
        assert!(err.to_string().contains("response_headers_absent"), "{err}");
    }

    #[test]
    fn header_value_with_two_variant_keys_is_rejected() {
        let e = entry("expect:\n  response_headers:\n    x-y: { regex: a, contains: b }\n");
        assert!(parse_expect_budgets(&e, "/tools/0").is_err());
    }
}
