//! Composition DAG model and load-time validator.
//!
//! A composition wires several tool calls together into one declared
//! graph that runs in topological order, producing a combined output
//! plus a per-node trace. Edges are explicit: each node names its
//! parents in `needs:`, and the validator rejects a `${id.field}`
//! reference whose `id` is not in that node's `needs:` list. The
//! intent is that a broken composition fails at load, not at 2am.
//!
//! This module owns the typed in-memory model
//! ([`ResolvedComposition`]), the topological-order computation, the
//! reference and cycle validators, and the trace shape the executor
//! emits. The CLI hands the loader's parsed [`CompositionSpec`] to
//! [`ResolvedComposition::resolve`], which either returns a runnable
//! composition or a [`CompositionError`] the loader surfaces with a
//! user-facing pointer.
//!
//! The actual node dispatch (calling `tools/call` against a
//! [`crate::connector::ServerPool`] and assembling the trace) lives
//! in `executor` so the model stays I/O-free and easy to unit-test.

use std::collections::{BTreeMap, HashMap, HashSet};

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

mod data_flow;
mod evaluate;
mod executor;
mod spec;

pub use data_flow::{substitute, substitute_with_item, SubstitutionError};
pub use evaluate::{evaluate_expect, CompositionOutcome};
pub use executor::{run, run_cases, CaseTrace};
pub use spec::{
    AssembleSpec, BudgetSpec, CaseSpec, CompositionSpec, MockSpec, NodeSpec, OutputSpec,
    TransformSpec,
};

/// Reasons a composition fails its load-time validation.
#[derive(Debug, Clone, Error, PartialEq, Eq)]
pub enum CompositionError {
    /// Two nodes share the same `id:`. Node ids must be unique within
    /// a composition because every reference and trace key is
    /// indexed by id.
    #[error("composition `{composition}` declares node `{id}` twice")]
    DuplicateNode {
        /// Composition name.
        composition: String,
        /// Offending node id.
        id: String,
    },
    /// A node lists a parent in `needs:` that does not exist in the
    /// composition. Usually a typo.
    #[error("composition `{composition}` node `{node}` needs `{missing}`, which is not declared")]
    UnknownNeed {
        /// Composition name.
        composition: String,
        /// Node that names the missing parent.
        node: String,
        /// The missing parent id.
        missing: String,
    },
    /// A node's args reference `${id.field}` where `id` is not in the
    /// node's `needs:` list. The validator rejects implicit edges so
    /// the DAG stays accurate.
    #[error(
        "composition `{composition}` node `{node}` references `{referent}` (via `{template}`) without listing it in `needs:`"
    )]
    UndeclaredReference {
        /// Composition name.
        composition: String,
        /// Node whose args carry the reference.
        node: String,
        /// The referenced node id.
        referent: String,
        /// The full `${...}` template that surfaced the reference.
        template: String,
    },
    /// The graph has a cycle. The executor cannot run a cyclic
    /// composition, so we reject it at load.
    #[error("composition `{composition}` has a cycle through node `{node}`")]
    Cycle {
        /// Composition name.
        composition: String,
        /// One node on the cycle (any is fine; the error points the
        /// user at the offending region).
        node: String,
    },
    /// `output:` names a node that does not exist.
    #[error("composition `{composition}` `output: {output}` does not match any declared node")]
    UnknownOutputNode {
        /// Composition name.
        composition: String,
        /// The offending `output:` value.
        output: String,
    },
    /// A per-node `transform.jq:` program failed to compile. Caught
    /// at load time so a typo points at the suite source rather than
    /// surfacing at run time.
    #[error("composition `{composition}` node `{node}` transform.jq did not compile: {detail}")]
    BadJq {
        /// Composition name.
        composition: String,
        /// Node id carrying the offending jq.
        node: String,
        /// One-line summary of the jq compile failure.
        detail: String,
    },
    /// A `for_each:` template did not parse to a `${parent_id}` or
    /// `${parent_id[*]}` form. The runner cannot tell which parent
    /// supplies the fan-out source if the template is freeform.
    #[error("composition `{composition}` node `{node}` for_each `{value}` must be `${{parent}}` or `${{parent[*]}}`")]
    BadForEach {
        /// Composition name.
        composition: String,
        /// Offending node id.
        node: String,
        /// Raw value the user wrote.
        value: String,
    },
    /// A `for_each:` template names a parent that is not in this
    /// node's `needs:` list. The DAG is wrong by construction; reject
    /// at load so the runner never tries to fan out from a node it
    /// did not declare.
    #[error("composition `{composition}` node `{node}` for_each references `{parent}` which is not in needs")]
    ForEachParentNotInNeeds {
        /// Composition name.
        composition: String,
        /// Offending node id.
        node: String,
        /// Parent the for_each tried to fan out from.
        parent: String,
    },
    /// A `max_iterations:` override outside 1..=100,000 (WOR-1322). The
    /// hard ceiling exists so a typo cannot authorize a six-figure
    /// fan-out against the user's server.
    #[error("composition `{composition}` node `{node}` sets max_iterations to {value}; the supported range is 1 to 100000")]
    MaxIterationsOutOfRange {
        /// Composition name.
        composition: String,
        /// Offending node id.
        node: String,
        /// The out-of-range value the user wrote.
        value: usize,
    },
    /// A `max_iterations:` on a node without `for_each:` would be
    /// silently inert, so the loader rejects it (WOR-1322).
    #[error("composition `{composition}` node `{node}` sets max_iterations but declares no for_each; the cap only applies to a fan-out")]
    MaxIterationsWithoutForEach {
        /// Composition name.
        composition: String,
        /// Offending node id.
        node: String,
    },
}

/// Warning emitted by the validator. Warnings do not block the load;
/// the loader surfaces them as `mcptest run --lint` lines so the user
/// can decide whether to act.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct CompositionWarning {
    /// Composition name the warning applies to.
    pub composition: String,
    /// The kind of warning. Each variant carries enough context for
    /// the reporter to print a single line.
    pub kind: WarningKind,
}

/// Concrete warning kinds emitted by [`ResolvedComposition::resolve`].
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum WarningKind {
    /// A `needs:` entry was declared but the node's args never
    /// referenced it. The runner treats the entry as a control-only
    /// ordering edge, which is sometimes intentional and sometimes a
    /// leftover, so the lint asks the author.
    UnusedNeed {
        /// Node that carries the unused entry.
        node: String,
        /// The parent id that is not referenced.
        parent: String,
    },
}

/// A composition with its DAG validated. Carries the typed nodes plus
/// a topological order the executor consumes.
#[derive(Debug, Clone, PartialEq)]
pub struct ResolvedComposition {
    /// Human-readable name surfaced to the reporter.
    pub name: String,
    /// Optional one-line description for `mcptest list`.
    pub description: Option<String>,
    /// Nodes indexed by id. The map is built from the parsed list and
    /// the typed view enforces id uniqueness.
    pub nodes: BTreeMap<String, ResolvedNode>,
    /// Topological execution order. Nodes at the same "rank" are
    /// independent and the runner is free to schedule them in
    /// parallel. The vector is flattened by Kahn's algorithm so the
    /// runner can iterate in order without re-walking the graph.
    pub topo_order: Vec<String>,
    /// How to assemble the composition's combined output. `None`
    /// means "use the last entry of `topo_order`"; the
    /// [`OutputSpec::Single`] form selects one node's output
    /// verbatim, and [`OutputSpec::Assemble`] runs a jq program
    /// against `{nodes: {<id>: <output>, ...}}`.
    pub output: Option<OutputSpec>,
    /// Deterministic assertions evaluated against the combined output
    /// plus the trace. Carried as `Value` because the assertion
    /// matchers live in `crate::suite` and we do not want this module
    /// to depend on the matcher tree directly.
    pub expect: Vec<Value>,
    /// Dataset of cases. The runner iterates the composition once
    /// per case, binding the case's `vars` for `${var.X}`
    /// substitution. Empty for a single-run composition.
    pub cases: Vec<CaseSpec>,
    /// Optional composition-level cost / token / duration budget.
    pub budget: Option<BudgetSpec>,
    /// Judged `eval:` metrics scored over the combined output and
    /// node trace after the run. Populated by the suite loader from
    /// the YAML `eval:` block via
    /// `suite::eval_parse::parse_eval_metrics`. Empty when the
    /// composition declares no eval block.
    pub eval: Vec<crate::suite::EvalMetric>,
}

/// One node in a validated composition.
#[derive(Debug, Clone, PartialEq)]
pub struct ResolvedNode {
    /// Stable identifier.
    pub id: String,
    /// Parents declared in `needs:`. Order is preserved so the
    /// reporter shows the user's declaration order.
    pub needs: Vec<String>,
    /// Server key resolved through the suite's `servers:` map.
    pub server: Option<String>,
    /// MCP tool name.
    pub tool: String,
    /// Raw args from YAML. The runner walks the value to substitute
    /// `${...}` references at dispatch time; the validator only walks
    /// it to extract reference ids for the contract check.
    pub args: Option<Value>,
    /// Optional jq program applied to the node's output before
    /// downstream nodes (and the trace) see it.
    pub transform: Option<TransformSpec>,
    /// Optional jq predicate evaluated before dispatch. Truthy runs
    /// the node, falsy marks it skipped. Compiled at load time so a
    /// typo fails the suite immediately.
    pub when: Option<String>,
    /// Optional `for_each` parent id. When set, the node fans out
    /// into one sub-call per element of `parent.output`. Validated
    /// at load: the parent id must be in `needs:`.
    pub for_each: Option<String>,
    /// Per-node concurrency cap for `for_each`. `None` means "use
    /// the runner default (1)".
    pub max_concurrency: Option<usize>,
    /// Per-node override of the `for_each` fan-out cap (WOR-1322).
    /// `None` means "use the runner default (10,000)". Validated to
    /// 1..=100,000 at load.
    pub max_iterations: Option<usize>,
    /// Inline mock fixture that bypasses the live dispatch. `None`
    /// keeps the live path; `Some` makes the node deterministic.
    pub mock: Option<MockSpec>,
    /// Per-node cost cap in USD per call. `None` disables the cap.
    pub max_cost_per_call_usd: Option<f64>,
    /// Per-node token cap per call. `None` disables the cap.
    pub max_tokens_per_call: Option<u64>,
}

impl ResolvedComposition {
    /// Validate a parsed [`CompositionSpec`] and return a runnable
    /// envelope plus any non-blocking warnings. Returns the first
    /// error encountered so the user sees a single clear pointer
    /// rather than a wall of cascading messages.
    pub fn resolve(
        spec: &CompositionSpec,
    ) -> Result<(Self, Vec<CompositionWarning>), CompositionError> {
        let nodes = build_node_map(spec)?;
        validate_needs(spec, &nodes)?;
        let references = extract_references(spec);
        validate_references(spec, &references)?;
        validate_output(spec, &nodes)?;
        validate_transforms(spec)?;
        validate_for_each(spec)?;
        let topo_order = topological_order(spec, &nodes)?;
        let warnings = collect_warnings(spec, &references);

        let resolved_nodes: BTreeMap<String, ResolvedNode> = spec
            .nodes
            .iter()
            .map(|n| {
                (
                    n.id.clone(),
                    ResolvedNode {
                        id: n.id.clone(),
                        needs: n.needs.clone(),
                        server: n.server.clone(),
                        tool: n.tool.clone(),
                        args: n.args.clone(),
                        transform: n.transform.clone(),
                        when: n.when.clone(),
                        for_each: parse_for_each_id(n.for_each.as_deref()),
                        max_concurrency: n.max_concurrency,
                        max_iterations: n.max_iterations,
                        mock: n.mock.clone(),
                        max_cost_per_call_usd: n.max_cost_per_call_usd,
                        max_tokens_per_call: n.max_tokens_per_call,
                    },
                )
            })
            .collect();

        Ok((
            Self {
                name: spec.name.clone(),
                description: spec.description.clone(),
                nodes: resolved_nodes,
                topo_order,
                output: spec.output.clone(),
                expect: spec.expect.clone(),
                cases: spec.cases.clone(),
                budget: spec.budget.clone(),
                eval: Vec::new(),
            },
            warnings,
        ))
    }

    /// Attach a typed eval-metric list to a resolved composition.
    /// The suite loader calls this after `parse_eval_metrics` has
    /// converted the YAML `eval:` block; passing through the
    /// resolver avoids leaking the suite-side parser into the
    /// composition crate.
    pub fn with_eval(mut self, eval: Vec<crate::suite::EvalMetric>) -> Self {
        self.eval = eval;
        self
    }
}

/// Per-node trace recorded by the executor. The runner emits one of
/// these per node so the assertion pipeline can address
/// `nodes.<id>.output`, `nodes.<id>.status`, etc.
#[derive(Debug, Clone, Default, Serialize, Deserialize, PartialEq)]
pub struct NodeTrace {
    /// Node id.
    pub id: String,
    /// Final status after dispatch.
    pub status: NodeStatus,
    /// The substituted args sent to the tool, or `Null` when the
    /// node was skipped before dispatch.
    pub inputs: Value,
    /// The raw tool output (the `{"result":...}` envelope's `result`
    /// field), or `Null` for skip/error.
    pub output: Value,
    /// Wall-clock duration in milliseconds. `None` for skipped nodes.
    pub duration_ms: Option<u64>,
    /// Cost in USD attributed to the node. Populated by LLM-backed
    /// adapters; defaults to `0.0` for tool calls with no
    /// attribution.
    #[serde(default, skip_serializing_if = "is_zero_f64")]
    pub cost_usd: f64,
    /// Input + output tokens attributed to the node. Same
    /// attribution rules as `cost_usd`.
    #[serde(default, skip_serializing_if = "is_zero_u64")]
    pub tokens: u64,
    /// Model id attributed to the node, when one is known.
    /// Surfaces in the per-model attribution table of the reporter.
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub model: Option<String>,
}

fn is_zero_f64(v: &f64) -> bool {
    *v == 0.0
}

fn is_zero_u64(v: &u64) -> bool {
    *v == 0
}

/// Terminal status for a node.
#[derive(Debug, Clone, Copy, Default, Serialize, Deserialize, PartialEq, Eq)]
#[serde(rename_all = "snake_case")]
pub enum NodeStatus {
    /// Node ran and the tool returned successfully.
    Ok,
    /// Node ran and the tool returned an error envelope.
    Err,
    /// Node was skipped (a parent failed, or a future `when:` is
    /// false).
    #[default]
    Skipped,
}

/// Top-level composition trace. The executor builds this after every
/// node has either run or been marked skipped.
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
pub struct CompositionTrace {
    /// Composition name (echoed for reporter clarity).
    pub name: String,
    /// Executed nodes in the order the runner dispatched them.
    pub order: Vec<String>,
    /// True when every node either ran successfully or was skipped
    /// via the documented skip rule. False when any node errored.
    pub ran: bool,
    /// Per-node trace indexed by id.
    pub nodes: BTreeMap<String, NodeTrace>,
    /// The combined output: the trace of the `output:` node, or the
    /// trace of the last node in topological order. `Null` when no
    /// node has run yet (the executor never emits this case for a
    /// real run; it is the [`Default`] value used in tests).
    pub output: Value,
    /// Total cost in USD across every node. Sum of
    /// `nodes.<id>.cost_usd`. Addressable in assertions as
    /// `cost.total`.
    #[serde(default, skip_serializing_if = "is_zero_f64")]
    pub cost_total: f64,
    /// Total tokens across every node. Sum of `nodes.<id>.tokens`.
    /// Addressable in assertions as `tokens.total`.
    #[serde(default, skip_serializing_if = "is_zero_u64")]
    pub tokens_total: u64,
    /// Name of the budget cap that aborted the run, if any. `None`
    /// means the run finished within every cap. The runner sets this
    /// when an [`BudgetSpec`] field trips; the assertion pipeline
    /// can address `composition.budget_tripped` to fail the verdict
    /// with the cap that fired.
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub budget_tripped: Option<String>,
}

impl Default for CompositionTrace {
    fn default() -> Self {
        Self {
            name: String::new(),
            order: Vec::new(),
            ran: false,
            nodes: BTreeMap::new(),
            output: Value::Null,
            cost_total: 0.0,
            tokens_total: 0,
            budget_tripped: None,
        }
    }
}

// Internal helpers below. These are deliberately small and pure so
// the load-time validation tests can drive each one in isolation.

fn build_node_map(spec: &CompositionSpec) -> Result<HashMap<&str, &NodeSpec>, CompositionError> {
    let mut map = HashMap::with_capacity(spec.nodes.len());
    for node in &spec.nodes {
        if map.insert(node.id.as_str(), node).is_some() {
            return Err(CompositionError::DuplicateNode {
                composition: spec.name.clone(),
                id: node.id.clone(),
            });
        }
    }
    Ok(map)
}

fn validate_needs(
    spec: &CompositionSpec,
    nodes: &HashMap<&str, &NodeSpec>,
) -> Result<(), CompositionError> {
    for node in &spec.nodes {
        for parent in &node.needs {
            if !nodes.contains_key(parent.as_str()) {
                return Err(CompositionError::UnknownNeed {
                    composition: spec.name.clone(),
                    node: node.id.clone(),
                    missing: parent.clone(),
                });
            }
        }
    }
    Ok(())
}

/// Per-node reference index built once and walked twice (validator +
/// warning collector). Keyed by node id, value is the set of referenced
/// parents plus the template the reference came from. The template is
/// retained so an error can point at the user's literal source string.
type RefIndex = HashMap<String, Vec<RefHit>>;

#[derive(Debug, Clone)]
struct RefHit {
    referent: String,
    template: String,
}

fn extract_references(spec: &CompositionSpec) -> RefIndex {
    let mut index: RefIndex = HashMap::new();
    for node in &spec.nodes {
        let Some(args) = &node.args else { continue };
        let mut hits = Vec::new();
        walk_value_for_refs(args, &mut hits);
        if !hits.is_empty() {
            index.insert(node.id.clone(), hits);
        }
    }
    index
}

fn walk_value_for_refs(value: &Value, hits: &mut Vec<RefHit>) {
    match value {
        Value::String(s) => {
            for (template, referent) in scan_template(s) {
                hits.push(RefHit { referent, template });
            }
        }
        Value::Array(arr) => {
            for v in arr {
                walk_value_for_refs(v, hits);
            }
        }
        Value::Object(map) => {
            for (_, v) in map {
                walk_value_for_refs(v, hits);
            }
        }
        _ => {}
    }
}

/// Scan a string for `${...}` templates and pull the leading
/// identifier out of each one. Only references shaped like
/// `${ident.<rest>}` or `${ident[<rest>]}` count as DAG edges; bare
/// `${var.X}` and `${item}` are runner-side concepts the validator
/// passes through.
fn scan_template(s: &str) -> Vec<(String, String)> {
    const RESERVED: &[&str] = &["var", "item", "env"];
    let mut out = Vec::new();
    let mut rest = s;
    while let Some(start) = rest.find("${") {
        let after = &rest[start + 2..];
        let Some(end) = after.find('}') else { break };
        let inner = &after[..end];
        let template = format!("${{{inner}}}");
        let ident: String = inner
            .chars()
            .take_while(|c| c.is_ascii_alphanumeric() || *c == '_')
            .collect();
        if !ident.is_empty() && !RESERVED.contains(&ident.as_str()) {
            out.push((template, ident));
        }
        rest = &after[end + 1..];
    }
    out
}

fn validate_references(spec: &CompositionSpec, refs: &RefIndex) -> Result<(), CompositionError> {
    for node in &spec.nodes {
        let Some(hits) = refs.get(node.id.as_str()) else {
            continue;
        };
        let declared: HashSet<&str> = node.needs.iter().map(String::as_str).collect();
        for hit in hits {
            if !declared.contains(hit.referent.as_str()) {
                return Err(CompositionError::UndeclaredReference {
                    composition: spec.name.clone(),
                    node: node.id.clone(),
                    referent: hit.referent.clone(),
                    template: hit.template.clone(),
                });
            }
        }
    }
    Ok(())
}

fn validate_output(
    spec: &CompositionSpec,
    nodes: &HashMap<&str, &NodeSpec>,
) -> Result<(), CompositionError> {
    match &spec.output {
        None => Ok(()),
        Some(OutputSpec::Single(out)) => {
            if !nodes.contains_key(out.as_str()) {
                return Err(CompositionError::UnknownOutputNode {
                    composition: spec.name.clone(),
                    output: out.clone(),
                });
            }
            Ok(())
        }
        Some(OutputSpec::Assemble(spec_assemble)) => {
            crate::transform_jq::compile(&spec_assemble.assemble).map_err(|e| {
                CompositionError::BadJq {
                    composition: spec.name.clone(),
                    node: "<output.assemble>".into(),
                    detail: format!("output.assemble: {e}"),
                }
            })
        }
    }
}

/// Parse a `for_each:` template ("${id}" or "${id[*]}") into the
/// bare parent id. Returns `None` when the input is `None` (no
/// for_each declared) or when the shape does not match; the
/// validator surfaces the latter as [`CompositionError::BadForEach`].
pub(super) fn parse_for_each_id(template: Option<&str>) -> Option<String> {
    let raw = template?.trim();
    let inner = raw.strip_prefix("${")?.strip_suffix('}')?;
    // Accept `id` and `id[*]` shapes; reject anything past the
    // optional `[*]` so a typo like `${id.field}` does not silently
    // become an edge.
    let ident = inner.split_once('[').map_or(inner, |(head, _)| head);
    let tail = inner.split_once('[').map(|(_, t)| t);
    if let Some(t) = tail {
        if t != "*]" {
            return None;
        }
    }
    let id = ident.trim();
    if id.is_empty() || !id.chars().all(|c| c.is_ascii_alphanumeric() || c == '_') {
        return None;
    }
    Some(id.to_string())
}

/// Validate the `for_each:` template on every node: the shape must
/// parse, and the referenced parent must be in `needs:`. The
/// per-node concurrency cap is a hint; the runner clamps it later.
fn validate_for_each(spec: &CompositionSpec) -> Result<(), CompositionError> {
    for node in &spec.nodes {
        // WOR-1322: bound the `max_iterations:` override and reject it on
        // a node it could never apply to, where it would be silently inert.
        if let Some(value) = node.max_iterations {
            if node.for_each.is_none() {
                return Err(CompositionError::MaxIterationsWithoutForEach {
                    composition: spec.name.clone(),
                    node: node.id.clone(),
                });
            }
            if value == 0 || value > executor::FOR_EACH_ITERATIONS_CEILING {
                return Err(CompositionError::MaxIterationsOutOfRange {
                    composition: spec.name.clone(),
                    node: node.id.clone(),
                    value,
                });
            }
        }
        let Some(raw) = &node.for_each else {
            continue;
        };
        let Some(parent) = parse_for_each_id(Some(raw.as_str())) else {
            return Err(CompositionError::BadForEach {
                composition: spec.name.clone(),
                node: node.id.clone(),
                value: raw.clone(),
            });
        };
        if !node.needs.iter().any(|p| p == &parent) {
            return Err(CompositionError::ForEachParentNotInNeeds {
                composition: spec.name.clone(),
                node: node.id.clone(),
                parent,
            });
        }
    }
    Ok(())
}

/// Compile every node's `transform.jq:` program and every node's
/// `when:` predicate at load so a typo fails the suite rather than
/// the run. The runner re-runs the same jq engine at dispatch and
/// trusts that what compiled at load still compiles.
fn validate_transforms(spec: &CompositionSpec) -> Result<(), CompositionError> {
    for node in &spec.nodes {
        if let Some(transform) = &node.transform {
            if let Err(err) = crate::transform_jq::compile(&transform.jq) {
                return Err(CompositionError::BadJq {
                    composition: spec.name.clone(),
                    node: node.id.clone(),
                    detail: err.to_string(),
                });
            }
        }
        if let Some(predicate) = &node.when {
            if let Err(err) = crate::transform_jq::compile(predicate) {
                return Err(CompositionError::BadJq {
                    composition: spec.name.clone(),
                    node: node.id.clone(),
                    detail: format!("when predicate: {err}"),
                });
            }
        }
    }
    Ok(())
}

/// Compute a topological order using Kahn's algorithm. Node order in
/// the declaration list is used as the tiebreaker so the resulting
/// vector is stable across runs (important for reporter output).
fn topological_order(
    spec: &CompositionSpec,
    _nodes: &HashMap<&str, &NodeSpec>,
) -> Result<Vec<String>, CompositionError> {
    // indegree[id] = how many parents must run before this node
    let mut indegree: BTreeMap<String, usize> = BTreeMap::new();
    let mut children: BTreeMap<String, Vec<String>> = BTreeMap::new();
    // Preserve declaration order in `ready` so the reporter is stable.
    let order_pos: HashMap<String, usize> = spec
        .nodes
        .iter()
        .enumerate()
        .map(|(i, n)| (n.id.clone(), i))
        .collect();
    for node in &spec.nodes {
        indegree.entry(node.id.clone()).or_insert(0);
        for parent in &node.needs {
            *indegree.entry(node.id.clone()).or_insert(0) += 1;
            children
                .entry(parent.clone())
                .or_default()
                .push(node.id.clone());
        }
    }
    let mut ready: Vec<String> = indegree
        .iter()
        .filter(|(_, n)| **n == 0)
        .map(|(id, _)| id.clone())
        .collect();
    ready.sort_by_key(|id| order_pos.get(id).copied().unwrap_or(usize::MAX));
    let mut order = Vec::with_capacity(spec.nodes.len());
    while let Some(id) = ready.first().cloned() {
        ready.remove(0);
        order.push(id.clone());
        if let Some(kids) = children.get(&id) {
            for kid in kids {
                if let Some(deg) = indegree.get_mut(kid) {
                    *deg = deg.saturating_sub(1);
                    if *deg == 0 {
                        ready.push(kid.clone());
                    }
                }
            }
            ready.sort_by_key(|id| order_pos.get(id).copied().unwrap_or(usize::MAX));
        }
    }
    if order.len() != spec.nodes.len() {
        // Any node not in `order` is part of a cycle; point the user
        // at the first one we declared so the error is stable.
        let in_order: HashSet<&str> = order.iter().map(String::as_str).collect();
        let offender = spec
            .nodes
            .iter()
            .map(|n| n.id.as_str())
            .find(|id| !in_order.contains(*id))
            .unwrap_or("(unknown)")
            .to_string();
        return Err(CompositionError::Cycle {
            composition: spec.name.clone(),
            node: offender,
        });
    }
    Ok(order)
}

fn collect_warnings(spec: &CompositionSpec, refs: &RefIndex) -> Vec<CompositionWarning> {
    let mut out = Vec::new();
    for node in &spec.nodes {
        let referenced: HashSet<&str> = refs
            .get(node.id.as_str())
            .into_iter()
            .flatten()
            .map(|h| h.referent.as_str())
            .collect();
        for parent in &node.needs {
            if !referenced.contains(parent.as_str()) {
                out.push(CompositionWarning {
                    composition: spec.name.clone(),
                    kind: WarningKind::UnusedNeed {
                        node: node.id.clone(),
                        parent: parent.clone(),
                    },
                });
            }
        }
    }
    out
}

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

    fn parse(yaml: &str) -> CompositionSpec {
        serde_yaml::from_str(yaml).unwrap()
    }

    #[test]
    fn linear_chain_resolves_in_declaration_order() {
        let spec = parse(
            r#"
name: linear
nodes:
  - id: a
    tool: t
    args: { q: "q" }
  - id: b
    needs: [a]
    tool: u
    args: { x: "${a.id}" }
"#,
        );
        let (resolved, warnings) = ResolvedComposition::resolve(&spec).unwrap();
        assert_eq!(resolved.topo_order, vec!["a", "b"]);
        assert!(warnings.is_empty());
    }

    #[test]
    fn duplicate_node_id_is_rejected_at_load() {
        let spec = parse(
            r#"
name: dup
nodes:
  - id: a
    tool: t
  - id: a
    tool: u
"#,
        );
        match ResolvedComposition::resolve(&spec).unwrap_err() {
            CompositionError::DuplicateNode { id, .. } => assert_eq!(id, "a"),
            other => panic!("expected duplicate node, got {other:?}"),
        }
    }

    #[test]
    fn needs_pointing_at_a_missing_id_is_rejected() {
        let spec = parse(
            r#"
name: missing
nodes:
  - id: child
    needs: [parent]
    tool: t
"#,
        );
        match ResolvedComposition::resolve(&spec).unwrap_err() {
            CompositionError::UnknownNeed { node, missing, .. } => {
                assert_eq!(node, "child");
                assert_eq!(missing, "parent");
            }
            other => panic!("expected unknown need, got {other:?}"),
        }
    }

    #[test]
    fn template_reference_without_needs_is_rejected() {
        let spec = parse(
            r#"
name: lazy
nodes:
  - id: a
    tool: t
    args: { q: "q" }
  - id: b
    tool: u
    args: { x: "${a.id}" }
"#,
        );
        match ResolvedComposition::resolve(&spec).unwrap_err() {
            CompositionError::UndeclaredReference {
                node,
                referent,
                template,
                ..
            } => {
                assert_eq!(node, "b");
                assert_eq!(referent, "a");
                assert_eq!(template, "${a.id}");
            }
            other => panic!("expected undeclared reference, got {other:?}"),
        }
    }

    #[test]
    fn unreferenced_need_surfaces_as_warning_not_error() {
        let spec = parse(
            r#"
name: lint
nodes:
  - id: a
    tool: t
    args: { q: "q" }
  - id: b
    needs: [a]
    tool: u
    args: { x: "constant" }
"#,
        );
        let (_resolved, warnings) = ResolvedComposition::resolve(&spec).unwrap();
        assert_eq!(warnings.len(), 1);
        match &warnings[0].kind {
            WarningKind::UnusedNeed { node, parent } => {
                assert_eq!(node, "b");
                assert_eq!(parent, "a");
            }
        }
    }

    #[test]
    fn cycle_is_rejected() {
        // Cycle: a -> b -> a. We declare each node so the
        // duplicate-node check does not fire first.
        let spec = parse(
            r#"
name: cyc
nodes:
  - id: a
    needs: [b]
    tool: t
  - id: b
    needs: [a]
    tool: u
"#,
        );
        match ResolvedComposition::resolve(&spec).unwrap_err() {
            CompositionError::Cycle { node, .. } => {
                assert!(node == "a" || node == "b");
            }
            other => panic!("expected cycle, got {other:?}"),
        }
    }

    #[test]
    fn output_must_name_a_declared_node() {
        let spec = parse(
            r#"
name: bad output
output: "nope"
nodes:
  - id: a
    tool: t
"#,
        );
        match ResolvedComposition::resolve(&spec).unwrap_err() {
            CompositionError::UnknownOutputNode { output, .. } => {
                assert_eq!(output, "nope");
            }
            other => panic!("expected unknown output node, got {other:?}"),
        }
    }

    #[test]
    fn reserved_template_namespaces_are_not_treated_as_edges() {
        let spec = parse(
            r#"
name: reserved
nodes:
  - id: a
    tool: t
    args:
      q: "${var.search_term}"
      idx: "${item}"
      e: "${env.HOME}"
"#,
        );
        // No needs declared, no error expected: var/item/env are
        // runtime concepts, not DAG edges.
        let (resolved, warnings) = ResolvedComposition::resolve(&spec).unwrap();
        assert_eq!(resolved.topo_order, vec!["a"]);
        assert!(warnings.is_empty());
    }

    #[test]
    fn references_inside_arrays_and_nested_objects_are_walked() {
        let spec = parse(
            r#"
name: nested
nodes:
  - id: a
    tool: t
  - id: b
    needs: [a]
    tool: u
    args:
      pairs:
        - key: "${a.id}"
        - key: "literal"
      nested:
        inner: ["${a.score}"]
"#,
        );
        // Both references resolve through `needs:[a]`; no error.
        let (resolved, _) = ResolvedComposition::resolve(&spec).unwrap();
        assert!(resolved.topo_order.contains(&"b".to_string()));
    }

    #[test]
    fn default_composition_trace_is_empty() {
        // Used by tests + the no-op path before the executor records
        // anything; documenting the contract via a test means a
        // regression that adds a non-empty default will fail here.
        let trace = CompositionTrace::default();
        assert_eq!(trace.output, json!(null));
        assert!(trace.nodes.is_empty());
        assert!(!trace.ran);
    }

    #[test]
    fn transform_with_a_valid_jq_program_passes_validation() {
        let spec = parse(
            r#"
name: shape
nodes:
  - id: a
    tool: t
    transform:
      jq: ".result | { name }"
"#,
        );
        let (resolved, _) = ResolvedComposition::resolve(&spec).unwrap();
        assert!(resolved.nodes.get("a").unwrap().transform.is_some());
    }

    #[test]
    fn transform_with_a_broken_jq_program_is_rejected_at_load() {
        let spec = parse(
            r#"
name: bad
nodes:
  - id: a
    tool: t
    transform:
      jq: ".result | ("
"#,
        );
        match ResolvedComposition::resolve(&spec).unwrap_err() {
            CompositionError::BadJq { node, .. } => assert_eq!(node, "a"),
            other => panic!("expected BadJq, got {other:?}"),
        }
    }

    #[test]
    fn when_with_a_valid_jq_predicate_passes_validation() {
        let spec = parse(
            r#"
name: gated
nodes:
  - id: a
    tool: t
  - id: b
    needs: [a]
    tool: u
    when: ".parents.a.count > 0"
"#,
        );
        let (resolved, _) = ResolvedComposition::resolve(&spec).unwrap();
        assert_eq!(
            resolved.nodes.get("b").unwrap().when.as_deref(),
            Some(".parents.a.count > 0")
        );
    }

    #[test]
    fn when_with_a_broken_jq_predicate_is_rejected_at_load() {
        let spec = parse(
            r#"
name: gated bad
nodes:
  - id: a
    tool: t
    when: "(((("
"#,
        );
        match ResolvedComposition::resolve(&spec).unwrap_err() {
            CompositionError::BadJq { node, detail, .. } => {
                assert_eq!(node, "a");
                assert!(detail.contains("when predicate"), "{detail}");
            }
            other => panic!("expected BadJq, got {other:?}"),
        }
    }

    #[test]
    fn for_each_template_parses_simple_and_star_forms() {
        assert_eq!(parse_for_each_id(Some("${data}")).as_deref(), Some("data"));
        assert_eq!(
            parse_for_each_id(Some("${hits[*]}")).as_deref(),
            Some("hits")
        );
        assert_eq!(parse_for_each_id(None), None);
        assert_eq!(parse_for_each_id(Some("data")), None);
        assert_eq!(parse_for_each_id(Some("${id.field}")), None);
        assert_eq!(parse_for_each_id(Some("${id[0]}")), None);
    }

    #[test]
    fn for_each_with_valid_parent_in_needs_passes() {
        let spec = parse(
            r#"
name: fan
nodes:
  - id: search
    tool: t
  - id: enrich
    needs: [search]
    for_each: "${search}"
    tool: u
"#,
        );
        let (resolved, _) = ResolvedComposition::resolve(&spec).unwrap();
        assert_eq!(
            resolved.nodes.get("enrich").unwrap().for_each.as_deref(),
            Some("search")
        );
    }

    #[test]
    fn max_iterations_override_is_carried_onto_the_resolved_node() {
        let spec = parse(
            r#"
name: fan
nodes:
  - id: search
    tool: t
  - id: enrich
    needs: [search]
    for_each: "${search}"
    max_iterations: 25000
    tool: u
"#,
        );
        let (resolved, _) = ResolvedComposition::resolve(&spec).unwrap();
        assert_eq!(
            resolved.nodes.get("enrich").unwrap().max_iterations,
            Some(25_000)
        );
    }

    #[test]
    fn max_iterations_above_the_ceiling_is_rejected_at_load() {
        let spec = parse(
            r#"
name: fan
nodes:
  - id: search
    tool: t
  - id: enrich
    needs: [search]
    for_each: "${search}"
    max_iterations: 100001
    tool: u
"#,
        );
        match ResolvedComposition::resolve(&spec).unwrap_err() {
            CompositionError::MaxIterationsOutOfRange { node, value, .. } => {
                assert_eq!(node, "enrich");
                assert_eq!(value, 100_001);
            }
            other => panic!("expected MaxIterationsOutOfRange, got {other:?}"),
        }
    }

    #[test]
    fn max_iterations_without_for_each_is_rejected_as_inert() {
        let spec = parse(
            r#"
name: fan
nodes:
  - id: search
    max_iterations: 10
    tool: t
"#,
        );
        match ResolvedComposition::resolve(&spec).unwrap_err() {
            CompositionError::MaxIterationsWithoutForEach { node, .. } => {
                assert_eq!(node, "search");
            }
            other => panic!("expected MaxIterationsWithoutForEach, got {other:?}"),
        }
    }

    #[test]
    fn for_each_with_malformed_template_is_rejected() {
        let spec = parse(
            r#"
name: bad
nodes:
  - id: a
    tool: t
  - id: b
    needs: [a]
    for_each: "not a template"
    tool: u
"#,
        );
        match ResolvedComposition::resolve(&spec).unwrap_err() {
            CompositionError::BadForEach { node, value, .. } => {
                assert_eq!(node, "b");
                assert_eq!(value, "not a template");
            }
            other => panic!("expected BadForEach, got {other:?}"),
        }
    }

    #[test]
    fn for_each_parent_not_in_needs_is_rejected() {
        let spec = parse(
            r#"
name: bad
nodes:
  - id: a
    tool: t
  - id: b
    tool: u
    for_each: "${a}"
"#,
        );
        match ResolvedComposition::resolve(&spec).unwrap_err() {
            CompositionError::ForEachParentNotInNeeds { node, parent, .. } => {
                assert_eq!(node, "b");
                assert_eq!(parent, "a");
            }
            other => panic!("expected ForEachParentNotInNeeds, got {other:?}"),
        }
    }
}
