//! YAML-level model for `compositions:`.
//!
//! Parsing lives here so the model and validator are in the same
//! crate as the typed runner contract. `mcptest_config` re-exports
//! these types so the loader can pull them off a YAML document.

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

/// One composition declared in `compositions:`.
#[derive(Debug, Clone, Default, Serialize, Deserialize, PartialEq)]
#[serde(deny_unknown_fields)]
pub struct CompositionSpec {
    /// Human-readable name surfaced to the reporter.
    pub name: String,
    /// Optional one-line description for `mcptest list` output.
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub description: Option<String>,
    /// Composition nodes in declaration order. Execution order is
    /// computed from `needs:` edges in the runner, not from this
    /// list order; declaration order is preserved only for reporter
    /// readability and per-suite stability.
    #[serde(default, skip_serializing_if = "Vec::is_empty")]
    pub nodes: Vec<NodeSpec>,
    /// Which node's output is the composition's combined output, or
    /// a jq assembly that fuses several nodes into one shape. When
    /// omitted, the executor falls back to the last node in
    /// topological order. See [`OutputSpec`].
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub output: Option<OutputSpec>,
    /// Deterministic assertions evaluated against the combined
    /// output + node trace after the run. The shape mirrors
    /// `tools[].expect:` so existing matchers work unchanged.
    #[serde(default, skip_serializing_if = "Vec::is_empty")]
    pub expect: Vec<Value>,
    /// Optional dataset of cases. When present, the composition
    /// runs once per row; each row's `vars` populate the `${var.X}`
    /// templates inside node args, and per-case `golden` and
    /// `expect` overrides ride alongside the trace for the reporter.
    /// Empty when omitted, in which case the runner takes the
    /// single-run path unchanged.
    #[serde(default, skip_serializing_if = "Vec::is_empty")]
    pub cases: Vec<CaseSpec>,
    /// Optional cost / token / duration budget. When set, the
    /// runner aborts a node mid-graph the first time the
    /// accumulated total crosses a cap. The trace's verdict carries
    /// which cap tripped so a CI gate can fail with a clear pointer
    /// at spend.
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub budget: Option<BudgetSpec>,
    /// Optional `eval:` block carrying the WOR-871 judged-metric
    /// mapping (`rubric`, `mcp_task_completion`, `mcp_use`,
    /// `argument_correctness`, `plan_quality`). Stored as a raw
    /// `serde_json::Value` so the existing
    /// `suite::eval_parse::parse_eval_metrics` helper does the
    /// typed conversion; that helper lives in `mcptest-core::suite`
    /// and is the same one agent tests use.
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub eval: Option<Value>,
}

/// Cost / tokens / duration cap for a composition or one of its
/// nodes. Every field is optional; an unset field disables that
/// dimension's gate.
#[derive(Debug, Clone, Default, Serialize, Deserialize, PartialEq)]
#[serde(deny_unknown_fields)]
pub struct BudgetSpec {
    /// Maximum total cost in USD across every node in the
    /// composition. Reported by the LLM-backed adapter for each
    /// node; nodes without an attribution contribute zero.
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub max_cost_usd: Option<f64>,
    /// Maximum total input + output tokens across every node in the
    /// composition.
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub max_tokens: Option<u64>,
    /// Maximum total wall-clock duration of the composition run, in
    /// milliseconds.
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub max_duration_ms: Option<u64>,
}

/// One row of the `cases:` dataset. Each row produces one execution
/// of the composition with its `vars` bound for `${var.X}`
/// substitution, plus optional per-case `golden` and `expect` that
/// override the composition-level shape.
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
#[serde(deny_unknown_fields)]
pub struct CaseSpec {
    /// Human-readable case name surfaced to the reporter. Used by
    /// the `mcptest run --case <name>` filter.
    pub name: String,
    /// Variables bound for `${var.X}` substitution. Each key in this
    /// map can be referenced from any node's `args:` as
    /// `${var.<key>}`. Optional; omit when the case relies only on
    /// per-node mocks or live calls.
    #[serde(default, skip_serializing_if = "std::collections::BTreeMap::is_empty")]
    pub vars: std::collections::BTreeMap<String, Value>,
    /// Optional golden value the per-case assertion pipeline can
    /// compare the combined output against. The shape and matcher
    /// language are deferred to the reporter; the runner records it
    /// verbatim on the trace.
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub golden: Option<Value>,
    /// Optional per-case assertions. Each entry mirrors a top-level
    /// `expect:` row; when present, these are appended to the
    /// composition's own `expect:` for this case only.
    #[serde(default, skip_serializing_if = "Vec::is_empty")]
    pub expect: Vec<Value>,
}

/// Either the id of a single node whose output is the composition's
/// combined output, or a jq assembly program that fuses several
/// nodes into one shape. The untagged shape lets a YAML author write
/// the simple form (`output: readme`) without any nesting.
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
#[serde(untagged)]
pub enum OutputSpec {
    /// `output: <node-id>` selects one node's output verbatim.
    Single(String),
    /// `output: { assemble: <jq> }` runs a jq program over
    /// `{nodes: {<id>: <output>, ...}}` and uses the result as the
    /// combined output.
    Assemble(AssembleSpec),
}

/// jq assembly expression that runs against `{nodes: {<id>: <output>, ...}}`.
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
#[serde(deny_unknown_fields)]
pub struct AssembleSpec {
    /// jq program. Compiled at load so a typo fails the suite.
    pub assemble: String,
}

/// One node in a composition.
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
#[serde(deny_unknown_fields)]
pub struct NodeSpec {
    /// Stable identifier used in `needs:` lists and `${id.field}`
    /// references. Must be unique within a composition.
    pub id: String,
    /// Parent nodes whose outputs this node consumes. The loader
    /// rejects a `${x.field}` reference in `args:` whose `x` is not
    /// in this list. An unreferenced `needs:` entry survives as a
    /// control-only ordering edge but the validator emits a lint
    /// warning.
    #[serde(default, skip_serializing_if = "Vec::is_empty")]
    pub needs: Vec<String>,
    /// Server key resolved through the same `servers:` map every
    /// tool / resource / prompt test uses. The `server:` key is
    /// mandatory when the suite declares more than one server; the
    /// multi-server loader's "single default" rule applies when only
    /// one server is declared.
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub server: Option<String>,
    /// MCP tool name the node calls.
    pub tool: String,
    /// Arguments forwarded to the tool. Supports `${id.field}` and
    /// `${var.X}` substitution; the substitution happens in the
    /// runner, not in the loader.
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub args: Option<Value>,
    /// Optional jq expression that reshapes the node's output
    /// before downstream nodes (and the trace) see it. Downstream
    /// `${id.field}` references resolve against the transformed
    /// shape. Omit to keep the raw tool result.
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub transform: Option<TransformSpec>,
    /// Optional jq predicate evaluated before dispatch. The program
    /// runs against `{parents: {<id>: <output>, ...}}`; a truthy
    /// result runs the node, a falsy result marks it `skipped` in
    /// the trace and propagates the skip to descendants. Omit to
    /// always run the node when its parents succeeded.
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub when: Option<String>,
    /// Optional `for_each: "${parent_id}"` template that fans the
    /// node into one sub-call per element of `parent_id`'s output
    /// array. Inside the sub-call's args, `${item}` substitutes the
    /// current element. Outputs are collected back into an array
    /// (input order) on the node's trace.
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub for_each: Option<String>,
    /// Per-node concurrency cap for `for_each`. Bounds how many
    /// sub-calls may dispatch in parallel. Defaults to 1 (sequential)
    /// when omitted so the runner never silently multiplies wall-clock
    /// load on a user's test server.
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub max_concurrency: Option<usize>,
    /// Per-node override of the `for_each` fan-out cap (WOR-1322). The
    /// runner refuses a fan-out whose parent array is larger than this.
    /// Defaults to 10,000 when omitted; the loader rejects values outside
    /// 1..=100,000 and an override on a node without `for_each`, where it
    /// would be silently inert.
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub max_iterations: Option<usize>,
    /// Optional inline mock that skips the live dispatch. When set,
    /// the node records `mock.result` as its output without touching
    /// the transport. The trace's `status` stays `ok` and the per-node
    /// duration is reported as `0`; assertions that read
    /// `nodes.<id>.output` resolve against the mock value just as
    /// they would against a real call. Used for the deterministic end
    /// of the spectrum: freeze an expensive node (an LLM, a paid API)
    /// with a fixed output so the assertion under test is the only
    /// signal that can flip the verdict.
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub mock: Option<MockSpec>,
    /// Per-node cost cap in USD per call. `None` disables it.
    /// Applied to each sub-call when the node is a `for_each` fan-out.
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub max_cost_per_call_usd: Option<f64>,
    /// Per-node token cap per call. `None` disables it. Applied per
    /// sub-call for `for_each` nodes.
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub max_tokens_per_call: Option<u64>,
}

/// Inline mock fixture used in place of a live dispatch.
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
#[serde(deny_unknown_fields)]
pub struct MockSpec {
    /// The value recorded as the node's output. Substituted into
    /// `nodes.<id>.output` verbatim, so downstream `${id.field}`
    /// references resolve against it exactly as they would against a
    /// real tool response.
    pub result: Value,
}

/// Per-node response transform. Today the runner accepts a `jq:`
/// program string. Extensions (a subprocess command, a CEL
/// expression) can land as new variants without breaking existing
/// YAML.
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
#[serde(deny_unknown_fields)]
pub struct TransformSpec {
    /// The jq program applied to the node's raw output. Compiled and
    /// validated at load time so a typo fails the suite rather than a
    /// 2am surprise.
    pub jq: String,
}

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

    #[test]
    fn parses_a_composition_with_one_node() {
        let yaml = r#"
name: render readme
nodes:
  - id: search
    tool: search_repos
    args:
      q: "mcptest"
"#;
        let spec: CompositionSpec = serde_yaml::from_str(yaml).unwrap();
        assert_eq!(spec.name, "render readme");
        assert_eq!(spec.nodes.len(), 1);
        let node = &spec.nodes[0];
        assert_eq!(node.id, "search");
        assert_eq!(node.tool, "search_repos");
        assert!(node.needs.is_empty());
    }

    #[test]
    fn parses_explicit_needs_edges() {
        let yaml = r#"
name: pipeline
nodes:
  - id: search
    tool: search_repos
    args: { q: "mcptest" }
  - id: readme
    needs: [search]
    tool: get_readme
    args:
      repo: "${search.top_hit}"
"#;
        let spec: CompositionSpec = serde_yaml::from_str(yaml).unwrap();
        assert_eq!(spec.nodes[1].needs, vec!["search".to_string()]);
    }

    #[test]
    fn rejects_unknown_fields_at_the_node_level() {
        let yaml = r#"
name: bad
nodes:
  - id: a
    tool: t
    unexpected_field: 42
"#;
        assert!(serde_yaml::from_str::<CompositionSpec>(yaml).is_err());
    }

    #[test]
    fn defaults_when_optional_fields_are_omitted() {
        let yaml = r#"
name: bare
nodes:
  - id: only
    tool: noop
"#;
        let spec: CompositionSpec = serde_yaml::from_str(yaml).unwrap();
        assert_eq!(spec.output, None);
        assert!(spec.expect.is_empty());
        assert_eq!(spec.nodes[0].args, None);
        assert_eq!(spec.nodes[0].server, None);
    }
}
