//! Runtime data-flow helpers for the composition executor.
//!
//! At dispatch time, a node's args carry `${id.field}` references
//! that point at parent outputs. The validator already proved every
//! reference resolves to a declared parent in [`super::resolve`],
//! so substitution here cannot encounter an unknown id; an unknown
//! field on a real parent surfaces as a substitution error the
//! executor records on the node's trace, which downstream nodes then
//! skip.
//!
//! Reference forms supported in this module:
//!
//! - `${id.field}` and `${id.path.to.field}`: pull a field out of
//!   parent `id`'s output, walking dot-separated keys.
//!
//! Reserved namespaces (`${var.X}`, `${item}`, `${env.X}`) are
//! passed through unchanged because they belong to other epic
//! children (cases, fan-out, env interpolation).

use std::collections::HashMap;

use serde_json::Value;
use thiserror::Error;

/// Reasons a substitution pass fails. Surfaces on the node's trace
/// and marks every downstream node skipped.
#[derive(Debug, Clone, Error, PartialEq, Eq)]
pub enum SubstitutionError {
    /// A parent named in a `${id.field}` reference was not in the
    /// trace map at substitution time. Should not happen in practice
    /// because the load-time validator catches this, but reported
    /// rather than panicking so a regression surfaces clearly.
    #[error("reference `{template}` points at unknown parent `{referent}`")]
    UnknownParent {
        /// Full `${...}` template the runner saw.
        template: String,
        /// The parent id pulled out of the template.
        referent: String,
    },
    /// A `${id.field}` reference resolved to a parent that has no
    /// `.field` (or the path walked past a non-object). The runner
    /// records this on the node's trace and skips the node.
    #[error("reference `{template}` could not find `{path}` on parent `{referent}`")]
    MissingField {
        /// Full `${...}` template.
        template: String,
        /// The parent id.
        referent: String,
        /// The dotted path past the parent id.
        path: String,
    },
}

/// Walk every string scalar inside `value` and replace each `${...}`
/// template with the value resolved from `parent_outputs`. Returns
/// either the rewritten value or the first substitution error.
///
/// Strings that contain only one whole template are replaced with
/// the resolved value (so a number stays a number, an object stays
/// an object). Strings that mix one or more templates with literal
/// text are stringified: every template is rendered to its string
/// form and the result is glued back together.
pub fn substitute(
    value: &Value,
    parent_outputs: &HashMap<String, Value>,
) -> Result<Value, SubstitutionError> {
    substitute_with(value, parent_outputs, None, None)
}

/// Like [`substitute`] but also resolves the `${item}` reserved
/// template against the supplied current-element value. Used by the
/// fan-out runtime so a `for_each` iteration sees the element it is
/// processing without re-implementing the walk.
pub fn substitute_with_item(
    value: &Value,
    parent_outputs: &HashMap<String, Value>,
    item: Option<&Value>,
) -> Result<Value, SubstitutionError> {
    substitute_with(value, parent_outputs, item, None)
}

/// Substitute `${...}` templates resolving parents, an optional
/// item, and an optional `${var.X}` map. The variable map is
/// supplied by the cases dataset; when `None`, `${var.X}`
/// references pass through unchanged so a single-run composition
/// keeps the existing behavior.
pub fn substitute_with(
    value: &Value,
    parent_outputs: &HashMap<String, Value>,
    item: Option<&Value>,
    vars: Option<&HashMap<String, Value>>,
) -> Result<Value, SubstitutionError> {
    match value {
        Value::String(s) => substitute_string(s, parent_outputs, item, vars),
        Value::Array(arr) => {
            let mut out = Vec::with_capacity(arr.len());
            for entry in arr {
                out.push(substitute_with(entry, parent_outputs, item, vars)?);
            }
            Ok(Value::Array(out))
        }
        Value::Object(map) => {
            let mut out = serde_json::Map::with_capacity(map.len());
            for (k, v) in map {
                out.insert(k.clone(), substitute_with(v, parent_outputs, item, vars)?);
            }
            Ok(Value::Object(out))
        }
        other => Ok(other.clone()),
    }
}

fn substitute_string(
    s: &str,
    parent_outputs: &HashMap<String, Value>,
    item: Option<&Value>,
    vars: Option<&HashMap<String, Value>>,
) -> Result<Value, SubstitutionError> {
    let parts = split_templates(s);
    // No templates at all: return the string as-is.
    if parts.iter().all(|p| matches!(p, Segment::Text(_))) {
        return Ok(Value::String(s.to_string()));
    }
    // Single whole-string template: return the resolved value with
    // its original JSON type intact (so a number stays a number).
    if parts.len() == 1 {
        if let Some(Segment::Template(template)) = parts.first() {
            return resolve_template(template, parent_outputs, item, vars);
        }
    }
    // Mixed template + literal text: stringify each resolved value
    // and concatenate.
    let mut buf = String::new();
    for part in parts {
        match part {
            Segment::Text(t) => buf.push_str(&t),
            Segment::Template(template) => {
                let resolved = resolve_template(&template, parent_outputs, item, vars)?;
                buf.push_str(&stringify(&resolved));
            }
        }
    }
    Ok(Value::String(buf))
}

enum Segment {
    Text(String),
    Template(String),
}

fn split_templates(s: &str) -> Vec<Segment> {
    let mut out = Vec::new();
    let mut rest = s;
    while let Some(start) = rest.find("${") {
        if start > 0 {
            out.push(Segment::Text(rest[..start].to_string()));
        }
        let after = &rest[start + 2..];
        let Some(end) = after.find('}') else {
            // No closing brace: treat the remainder as literal text.
            out.push(Segment::Text(rest[start..].to_string()));
            rest = "";
            break;
        };
        out.push(Segment::Template(after[..end].to_string()));
        rest = &after[end + 1..];
    }
    if !rest.is_empty() {
        out.push(Segment::Text(rest.to_string()));
    }
    out
}

fn resolve_template(
    inner: &str,
    parent_outputs: &HashMap<String, Value>,
    item: Option<&Value>,
    vars: Option<&HashMap<String, Value>>,
) -> Result<Value, SubstitutionError> {
    let template = format!("${{{inner}}}");
    // `${item}` and `${item.field}` resolve against the current
    // fan-out element when one is in scope; outside a fan-out the
    // template passes through unchanged so a non-iterating node does
    // not break.
    let prefix = inner.split(['.', '[']).next().unwrap_or("");
    if prefix == "item" {
        let Some(current) = item else {
            return Ok(Value::String(template));
        };
        if inner == "item" {
            return Ok(current.clone());
        }
        let path = inner.trim_start_matches("item.");
        return walk_path(current, path)
            .map(Ok)
            .unwrap_or(Err(SubstitutionError::MissingField {
                template,
                referent: "item".into(),
                path: path.to_string(),
            }));
    }
    // `${var.X}` resolves against the case's vars map when one was
    // supplied (cases dataset path); outside a case the template
    // passes through unchanged so a single-run composition stays
    // backwards-compatible.
    if prefix == "var" {
        let Some(map) = vars else {
            return Ok(Value::String(template));
        };
        let path = inner.strip_prefix("var.").unwrap_or("");
        if path.is_empty() {
            return Ok(Value::String(template));
        }
        let (head, rest) = path.split_once('.').map_or((path, ""), |(h, r)| (h, r));
        let Some(root) = map.get(head) else {
            return Err(SubstitutionError::UnknownParent {
                template,
                referent: format!("var.{head}"),
            });
        };
        if rest.is_empty() {
            return Ok(root.clone());
        }
        return walk_path(root, rest)
            .map(Ok)
            .unwrap_or(Err(SubstitutionError::MissingField {
                template,
                referent: format!("var.{head}"),
                path: rest.to_string(),
            }));
    }
    if prefix == "env" {
        return Ok(Value::String(template));
    }
    let (referent, path) = match inner.split_once('.') {
        Some((ident, rest)) => (ident.to_string(), rest.to_string()),
        None => (inner.to_string(), String::new()),
    };
    let parent = parent_outputs
        .get(&referent)
        .ok_or_else(|| SubstitutionError::UnknownParent {
            template: template.clone(),
            referent: referent.clone(),
        })?;
    let resolved = if path.is_empty() {
        parent.clone()
    } else {
        walk_path(parent, &path).ok_or_else(|| SubstitutionError::MissingField {
            template,
            referent: referent.clone(),
            path: path.clone(),
        })?
    };
    Ok(resolved)
}

fn walk_path(value: &Value, dotted: &str) -> Option<Value> {
    let mut cursor = value;
    for segment in dotted.split('.') {
        cursor = cursor.get(segment)?;
    }
    Some(cursor.clone())
}

fn stringify(value: &Value) -> String {
    match value {
        Value::String(s) => s.clone(),
        Value::Null => "null".to_string(),
        Value::Bool(b) => b.to_string(),
        Value::Number(n) => n.to_string(),
        other => other.to_string(),
    }
}

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

    fn parents() -> HashMap<String, Value> {
        HashMap::from([
            (
                "search".to_string(),
                json!({"top_hit": {"name": "mcptest", "score": 0.9}}),
            ),
            (
                "auth".to_string(),
                json!({"token": "abc123", "expires_in": 3600}),
            ),
        ])
    }

    #[test]
    fn whole_string_template_preserves_the_resolved_value_type() {
        let args = json!({"score": "${search.top_hit.score}"});
        let out = substitute(&args, &parents()).unwrap();
        assert_eq!(out, json!({"score": 0.9}));
    }

    #[test]
    fn whole_string_template_can_return_an_object() {
        let args = json!({"hit": "${search.top_hit}"});
        let out = substitute(&args, &parents()).unwrap();
        assert_eq!(out, json!({"hit": {"name": "mcptest", "score": 0.9}}));
    }

    #[test]
    fn mixed_template_and_literal_text_is_stringified() {
        let args = json!({"bearer": "Bearer ${auth.token}"});
        let out = substitute(&args, &parents()).unwrap();
        assert_eq!(out, json!({"bearer": "Bearer abc123"}));
    }

    #[test]
    fn nested_objects_and_arrays_are_walked() {
        let args = json!({
            "items": [
                {"name": "${search.top_hit.name}"},
                {"name": "literal"}
            ],
            "wrap": {"deep": "${search.top_hit.name}"}
        });
        let out = substitute(&args, &parents()).unwrap();
        assert_eq!(
            out,
            json!({
                "items": [
                    {"name": "mcptest"},
                    {"name": "literal"}
                ],
                "wrap": {"deep": "mcptest"}
            })
        );
    }

    #[test]
    fn unknown_parent_returns_an_error() {
        let args = json!({"x": "${ghost.field}"});
        match substitute(&args, &parents()).unwrap_err() {
            SubstitutionError::UnknownParent { referent, .. } => assert_eq!(referent, "ghost"),
            other => panic!("expected UnknownParent, got {other:?}"),
        }
    }

    #[test]
    fn missing_field_returns_an_error() {
        let args = json!({"x": "${search.does_not_exist}"});
        match substitute(&args, &parents()).unwrap_err() {
            SubstitutionError::MissingField { referent, path, .. } => {
                assert_eq!(referent, "search");
                assert_eq!(path, "does_not_exist");
            }
            other => panic!("expected MissingField, got {other:?}"),
        }
    }

    #[test]
    fn reserved_namespaces_are_passed_through_unchanged() {
        let args = json!({
            "v": "${var.search_term}",
            "i": "${item}",
            "e": "${env.HOME}"
        });
        let out = substitute(&args, &parents()).unwrap();
        // The runner stages owning var/item/env see the raw templates
        // and handle them after this pass.
        assert_eq!(
            out,
            json!({
                "v": "${var.search_term}",
                "i": "${item}",
                "e": "${env.HOME}"
            })
        );
    }

    #[test]
    fn whole_string_template_with_only_an_id_resolves_to_the_whole_parent_output() {
        let args = json!({"all": "${auth}"});
        let out = substitute(&args, &parents()).unwrap();
        assert_eq!(out, json!({"all": {"token": "abc123", "expires_in": 3600}}));
    }

    #[test]
    fn non_template_value_is_returned_unchanged() {
        let args = json!({"q": "search query", "limit": 5});
        let out = substitute(&args, &parents()).unwrap();
        assert_eq!(out, args);
    }

    #[test]
    fn malformed_template_with_no_closing_brace_is_treated_as_literal_text() {
        let args = json!({"q": "${not_a_template"});
        let out = substitute(&args, &parents()).unwrap();
        assert_eq!(out, json!({"q": "${not_a_template"}));
    }

    #[test]
    fn multiple_templates_in_one_string_are_concatenated() {
        let args = json!({
            "label": "${search.top_hit.name}/${auth.token}"
        });
        let out = substitute(&args, &parents()).unwrap();
        assert_eq!(out, json!({"label": "mcptest/abc123"}));
    }

    #[test]
    fn item_substitution_returns_the_current_element_when_supplied() {
        let args = json!({"x": "${item}"});
        let out = substitute_with_item(&args, &HashMap::new(), Some(&json!(42))).unwrap();
        assert_eq!(out, json!({"x": 42}));
    }

    #[test]
    fn item_substitution_walks_dotted_fields() {
        let args = json!({"name": "${item.label}"});
        let item = json!({"label": "first", "score": 0.9});
        let out = substitute_with_item(&args, &HashMap::new(), Some(&item)).unwrap();
        assert_eq!(out, json!({"name": "first"}));
    }

    #[test]
    fn item_substitution_without_item_in_scope_passes_through_unchanged() {
        let args = json!({"x": "${item}", "y": "${item.field}"});
        let out = substitute(&args, &HashMap::new()).unwrap();
        assert_eq!(out, args);
    }

    #[test]
    fn item_substitution_with_missing_field_errors() {
        let args = json!({"x": "${item.missing}"});
        let item = json!({"present": 1});
        let err = substitute_with_item(&args, &HashMap::new(), Some(&item)).unwrap_err();
        assert!(matches!(err, SubstitutionError::MissingField { .. }));
    }
}
