hamelin_lib 0.11.2

//! Substitute concrete values for `${name}` template parameters in the query AST.
//!
//! Substitution is a single AST traversal: each [`TemplateParameter`] expression is replaced with
//! a literal derived from its [`TemplateParameterValue`], and each [`TemplatedTablePath`] is
//! materialized by string-substitution into an [`Identifier`]. There is intentionally no separate
//! "collect used parameters" pre-pass: the substitution traversal itself surfaces missing or
//! wrong-typed values as it visits the holes, so the AST is walked exactly once.

use std::collections::HashMap;
use std::sync::Arc;

use vecmap::VecMap;

use crate::err::TemplateSubstitutionError;
use crate::tree::ast::clause::{
    Assignment, AssignmentClause, FromClause, SortExpression, TableAlias, TableReference,
    TableReferenceTarget, TableSegmentPart, TemplatedTablePath,
};
use crate::tree::ast::command::{
    AggCommand, AppendCommand, Command, CommandKind, DistinctCommand, ExplodeCommand, FromCommand,
    JoinCommand, LimitCommand, LookupCommand, MatchCommand, ParseCommand, RowsCommand,
    SelectCommand, SetCommand, SortCommand, SuppressCommand, TrimstringsCommand, UnionCommand,
    UnnestCommand, WhereCommand, WindowCommand, WithinCommand,
};
use crate::tree::ast::expression::{
    ArrayLiteral, BinaryOperator, BooleanLiteral, Cast, DoubleLiteral, Expression, ExpressionKind,
    FieldLookup, FunctionCall, IndexAccess, IntLiteral, Lambda, PairLiteral, StringLiteral,
    StructLiteral, TemplateParameter, TsTrunc, TupleLiteral, UnaryPostfixOperator,
    UnaryPrefixOperator,
};
use crate::tree::ast::identifier::{Identifier, SimpleIdentifier};
use crate::tree::ast::pattern::{NestedPattern, Pattern, PatternKind, QuantifiedPattern};
use crate::tree::ast::pipeline::Pipeline;
use crate::tree::ast::query::{DefBody, DefStatement, Query, QueryKind, ValidQuery};

use super::{materialize_templated_table_path, TemplateParameterValue};

fn value_to_expression_kind(value: &TemplateParameterValue) -> ExpressionKind {
    match value {
        TemplateParameterValue::Boolean(b) => BooleanLiteral { value: *b }.into(),
        TemplateParameterValue::Int(i) => IntLiteral { int: *i }.into(),
        TemplateParameterValue::Double(d) => DoubleLiteral { value: *d }.into(),
        TemplateParameterValue::String(s) => StringLiteral { value: s.clone() }.into(),
    }
}

fn path_fragment_from_value(
    name: &str,
    value: &TemplateParameterValue,
) -> Result<String, TemplateSubstitutionError> {
    match value {
        TemplateParameterValue::String(s) => Ok(s.clone()),
        TemplateParameterValue::Boolean(_)
        | TemplateParameterValue::Int(_)
        | TemplateParameterValue::Double(_) => {
            let got = match value {
                TemplateParameterValue::Boolean(_) => "boolean",
                TemplateParameterValue::Int(_) => "int",
                TemplateParameterValue::Double(_) => "double",
                TemplateParameterValue::String(_) => "string",
            };
            Err(TemplateSubstitutionError::WrongValueKind {
                name: name.to_string(),
                expected: "string (table path fragment)".to_string(),
                got: got.to_string(),
            })
        }
    }
}

fn substitute_table_reference(
    tr: &TableReference,
    values: &HashMap<SimpleIdentifier, TemplateParameterValue>,
) -> Result<TableReference, TemplateSubstitutionError> {
    match &tr.target {
        TableReferenceTarget::Static(_) => Ok(tr.clone()),
        TableReferenceTarget::Templated(path) => {
            let ident = materialize_templated_path_with_values(path, values)?;
            Ok(TableReference {
                span: tr.span,
                target: ident.into(),
            })
        }
    }
}

/// Build the per-parameter substitution map lazily from `values`. Each parameter appearing in a
/// path must be a string; non-string values are surfaced here (rather than via a pre-pass) since
/// substitution is the only reader of values.
fn materialize_templated_path_with_values(
    path: &TemplatedTablePath,
    values: &HashMap<SimpleIdentifier, TemplateParameterValue>,
) -> Result<Identifier, TemplateSubstitutionError> {
    let mut subst_map: HashMap<SimpleIdentifier, String> = HashMap::new();
    for seg in &path.segments {
        for part in &seg.parts {
            if let TableSegmentPart::Parameter(psi) = part {
                let name = psi
                    .clone()
                    .valid()
                    .map_err(TemplateSubstitutionError::ParameterIdentifierParseError)?;
                if subst_map.contains_key(&name) {
                    continue;
                }
                let value = values.get(&name).ok_or_else(|| {
                    TemplateSubstitutionError::MissingTemplateParameter(name.as_str().to_string())
                })?;
                let s = path_fragment_from_value(name.as_str(), value)?;
                subst_map.insert(name, s);
            }
        }
    }
    materialize_templated_table_path(path, &subst_map)
}

fn substitute_from_clause(
    clause: &FromClause,
    values: &HashMap<SimpleIdentifier, TemplateParameterValue>,
) -> Result<FromClause, TemplateSubstitutionError> {
    match clause {
        FromClause::TableReference(tr) => Ok(FromClause::TableReference(Arc::new(
            substitute_table_reference(tr.as_ref(), values)?,
        ))),
        FromClause::TableAlias(a) => Ok(FromClause::TableAlias(Arc::new(TableAlias {
            span: a.span,
            alias: a.alias.clone(),
            table: substitute_table_reference(&a.table, values)?,
        }))),
    }
}

fn map_arc_expr(
    e: &Arc<Expression>,
    values: &HashMap<SimpleIdentifier, TemplateParameterValue>,
) -> Result<Arc<Expression>, TemplateSubstitutionError> {
    Ok(Arc::new(substitute_expression(e, values)?))
}

fn map_assignment_clauses(
    clauses: &[Arc<AssignmentClause>],
    values: &HashMap<SimpleIdentifier, TemplateParameterValue>,
) -> Result<Vec<Arc<AssignmentClause>>, TemplateSubstitutionError> {
    let mut out = Vec::with_capacity(clauses.len());
    for cl in clauses {
        let new_cl = match cl.as_ref() {
            AssignmentClause::Assignment(assignment) => AssignmentClause::Assignment(Assignment {
                span: assignment.span,
                identifier: assignment.identifier.clone(),
                expression: map_arc_expr(&assignment.expression, values)?,
            }),
            AssignmentClause::Error(error) => AssignmentClause::Error(error.clone()),
        };
        out.push(Arc::new(new_cl));
    }
    Ok(out)
}

fn map_sort_expressions(
    expressions: &[Arc<SortExpression>],
    values: &HashMap<SimpleIdentifier, TemplateParameterValue>,
) -> Result<Vec<Arc<SortExpression>>, TemplateSubstitutionError> {
    let mut out = Vec::with_capacity(expressions.len());
    for expression in expressions {
        out.push(Arc::new(SortExpression {
            span: expression.span,
            expression: map_arc_expr(&expression.expression, values)?,
            order: expression.order,
        }));
    }
    Ok(out)
}

fn substitute_expression(
    expr: &Expression,
    values: &HashMap<SimpleIdentifier, TemplateParameterValue>,
) -> Result<Expression, TemplateSubstitutionError> {
    let span = expr.span;
    let kind = match &expr.kind {
        ExpressionKind::TemplateParameter(TemplateParameter { name }) => {
            let si = name
                .clone()
                .valid()
                .map_err(TemplateSubstitutionError::ParameterIdentifierParseError)?;
            let value = values.get(&si).ok_or_else(|| {
                TemplateSubstitutionError::MissingTemplateParameter(si.as_str().to_string())
            })?;
            value_to_expression_kind(value)
        }
        ExpressionKind::IntLiteral(_)
        | ExpressionKind::DecimalLiteral(_)
        | ExpressionKind::DoubleLiteral(_)
        | ExpressionKind::ScientificLiteral(_)
        | ExpressionKind::BooleanLiteral(_)
        | ExpressionKind::StringLiteral(_)
        | ExpressionKind::BinaryLiteral(_)
        | ExpressionKind::NullLiteral(_)
        | ExpressionKind::FieldReference(_)
        | ExpressionKind::UnboundRangeLiteral(_)
        | ExpressionKind::RowsLiteral(_)
        | ExpressionKind::IntervalLiteral(_)
        | ExpressionKind::Error(_) => expr.kind.clone(),
        ExpressionKind::ArrayLiteral(a) => {
            let mut elements = Vec::with_capacity(a.elements.len());
            for e in &a.elements {
                elements.push(map_arc_expr(e, values)?);
            }
            ArrayLiteral { elements }.into()
        }
        ExpressionKind::TupleLiteral(t) => {
            let mut elements = Vec::with_capacity(t.elements.len());
            for e in &t.elements {
                elements.push(map_arc_expr(e, values)?);
            }
            TupleLiteral { elements }.into()
        }
        ExpressionKind::PairLiteral(p) => PairLiteral {
            left: map_arc_expr(&p.left, values)?,
            right: map_arc_expr(&p.right, values)?,
        }
        .into(),
        ExpressionKind::StructLiteral(s) => {
            let mut fields = Vec::with_capacity(s.fields.len());
            for (id, e) in &s.fields {
                fields.push((id.clone(), map_arc_expr(e, values)?));
            }
            StructLiteral { fields }.into()
        }
        ExpressionKind::UnaryPrefixOperator(u) => UnaryPrefixOperator {
            operator: u.operator.clone(),
            operand: map_arc_expr(&u.operand, values)?,
        }
        .into(),
        ExpressionKind::UnaryPostfixOperator(u) => UnaryPostfixOperator {
            operator: u.operator.clone(),
            operand: map_arc_expr(&u.operand, values)?,
        }
        .into(),
        ExpressionKind::BinaryOperator(b) => BinaryOperator {
            left: map_arc_expr(&b.left, values)?,
            operator: b.operator.clone(),
            right: map_arc_expr(&b.right, values)?,
        }
        .into(),
        ExpressionKind::FunctionCall(f) => {
            let mut positional_args = Vec::with_capacity(f.positional_args.len());
            for e in &f.positional_args {
                positional_args.push(map_arc_expr(e, values)?);
            }
            let mut named_args = VecMap::new();
            for (k, e) in f.named_args.iter() {
                named_args.insert(k.clone(), map_arc_expr(e, values)?);
            }
            FunctionCall {
                name: f.name.clone(),
                positional_args,
                named_args,
            }
            .into()
        }
        ExpressionKind::IndexAccess(i) => IndexAccess {
            value: map_arc_expr(&i.value, values)?,
            index: map_arc_expr(&i.index, values)?,
        }
        .into(),
        ExpressionKind::FieldLookup(f) => FieldLookup {
            value: map_arc_expr(&f.value, values)?,
            field_identifier: f.field_identifier.clone(),
        }
        .into(),
        ExpressionKind::Cast(c) => Cast {
            expression: map_arc_expr(&c.expression, values)?,
            target_type: c.target_type.clone(),
        }
        .into(),
        ExpressionKind::TsTrunc(t) => TsTrunc {
            expression: map_arc_expr(&t.expression, values)?,
            unit: t.unit.clone(),
            multiplier: t.multiplier,
        }
        .into(),
        ExpressionKind::Lambda(l) => Lambda {
            parameters: l.parameters.clone(),
            body: map_arc_expr(&l.body, values)?,
        }
        .into(),
    };
    Ok(Expression { span, kind })
}

fn substitute_command(
    cmd: &Arc<Command>,
    values: &HashMap<SimpleIdentifier, TemplateParameterValue>,
) -> Result<Arc<Command>, TemplateSubstitutionError> {
    let span = cmd.span;
    let kind = match &cmd.kind {
        CommandKind::Set(c) => {
            let mut assignments = Vec::with_capacity(c.assignments.len());
            for a in &c.assignments {
                assignments.push(Arc::new(Assignment {
                    span: a.span,
                    identifier: a.identifier.clone(),
                    expression: map_arc_expr(&a.expression, values)?,
                }));
            }
            SetCommand { assignments }.into()
        }
        CommandKind::Where(c) => WhereCommand {
            condition: map_arc_expr(&c.condition, values)?,
        }
        .into(),
        CommandKind::Select(c) => SelectCommand {
            clauses: map_assignment_clauses(&c.clauses, values)?,
        }
        .into(),
        CommandKind::Limit(c) => LimitCommand {
            count: map_arc_expr(&c.count, values)?,
        }
        .into(),
        CommandKind::Trimstrings(c) => TrimstringsCommand {
            byte_len: map_arc_expr(&c.byte_len, values)?,
        }
        .into(),
        CommandKind::Within(c) => WithinCommand {
            duration: map_arc_expr(&c.duration, values)?,
        }
        .into(),
        CommandKind::From(c) => {
            let mut clauses = Vec::with_capacity(c.clauses.len());
            for cl in &c.clauses {
                clauses.push(Arc::new(substitute_from_clause(cl, values)?));
            }
            FromCommand { clauses }.into()
        }
        CommandKind::Union(c) => {
            let mut clauses = Vec::with_capacity(c.clauses.len());
            for cl in &c.clauses {
                clauses.push(Arc::new(substitute_from_clause(cl, values)?));
            }
            UnionCommand { clauses }.into()
        }
        CommandKind::Parse(c) => ParseCommand {
            source: c
                .source
                .as_ref()
                .map(|s| map_arc_expr(s, values))
                .transpose()?,
            pattern: map_arc_expr(&c.pattern, values)?,
            identifiers: c.identifiers.clone(),
            nodrop: c.nodrop,
        }
        .into(),
        CommandKind::Agg(c) => AggCommand {
            clauses: map_assignment_clauses(&c.clauses, values)?,
            group_by: map_assignment_clauses(&c.group_by, values)?,
            sort: map_sort_expressions(&c.sort, values)?,
        }
        .into(),
        CommandKind::Window(c) => WindowCommand {
            clauses: map_assignment_clauses(&c.clauses, values)?,
            group_by: map_assignment_clauses(&c.group_by, values)?,
            sort: map_sort_expressions(&c.sort, values)?,
            within: c
                .within
                .as_ref()
                .map(|w| map_arc_expr(w, values))
                .transpose()?,
        }
        .into(),
        // APPEND is mutating: template parameters are not allowed. The grammar already keeps
        // them out of `distinct_by` (only plain `Selection`/`identifier` is accepted there), so
        // the only place a template can sneak in is the target table reference; reject that up
        // front instead of silently materializing it.
        CommandKind::Append(c) => {
            if matches!(c.table.target, TableReferenceTarget::Templated(_)) {
                return Err(TemplateSubstitutionError::TemplatedAppendTarget);
            }
            AppendCommand {
                table: c.table.clone(),
                distinct_by: c.distinct_by.clone(),
            }
            .into()
        }
        CommandKind::Join(c) => JoinCommand {
            other: Arc::new(substitute_from_clause(c.other.as_ref(), values)?),
            on_condition: c
                .on_condition
                .as_ref()
                .map(|e| map_arc_expr(e, values))
                .transpose()?,
        }
        .into(),
        CommandKind::Lookup(c) => LookupCommand {
            other: Arc::new(substitute_from_clause(c.other.as_ref(), values)?),
            on_condition: c
                .on_condition
                .as_ref()
                .map(|e| map_arc_expr(e, values))
                .transpose()?,
        }
        .into(),
        CommandKind::Explode(c) => ExplodeCommand {
            clauses: map_assignment_clauses(&c.clauses, values)?,
        }
        .into(),
        CommandKind::Distinct(c) => DistinctCommand {
            clauses: map_assignment_clauses(&c.clauses, values)?,
        }
        .into(),
        CommandKind::Suppress(c) => SuppressCommand {
            interval: map_arc_expr(&c.interval, values)?,
            group_by: map_assignment_clauses(&c.group_by, values)?,
        }
        .into(),
        CommandKind::Unnest(c) => UnnestCommand {
            expression: map_arc_expr(&c.expression, values)?,
        }
        .into(),
        CommandKind::Rows(c) => RowsCommand {
            expression: map_arc_expr(&c.expression, values)?,
        }
        .into(),
        CommandKind::Nest(c) => c.clone().into(),
        CommandKind::Match(c) => {
            let mut pattern = Vec::with_capacity(c.pattern.len());
            for p in &c.pattern {
                pattern.push(Arc::new(substitute_pattern(p, values)?));
            }
            MatchCommand {
                pattern,
                agg: map_assignment_clauses(&c.agg, values)?,
                group_by: map_assignment_clauses(&c.group_by, values)?,
                sort: map_sort_expressions(&c.sort, values)?,
                within: c
                    .within
                    .as_ref()
                    .map(|w| map_arc_expr(w, values))
                    .transpose()?,
            }
            .into()
        }
        CommandKind::Drop(c) => c.clone().into(),
        CommandKind::Sort(c) => SortCommand {
            expressions: map_sort_expressions(&c.expressions, values)?,
        }
        .into(),
        CommandKind::Error(e) => e.clone().into(),
    };
    Ok(Arc::new(Command { span, kind }))
}

fn substitute_pattern(
    pattern: &Pattern,
    values: &HashMap<SimpleIdentifier, TemplateParameterValue>,
) -> Result<Pattern, TemplateSubstitutionError> {
    let span = pattern.span;
    let kind = match &pattern.kind {
        PatternKind::Quantified(q) => QuantifiedPattern {
            span: q.span,
            from_clause: Arc::new(substitute_from_clause(q.from_clause.as_ref(), values)?),
            quantifier: q.quantifier.clone(),
        }
        .into(),
        PatternKind::Nested(n) => {
            let mut patterns = Vec::with_capacity(n.patterns.len());
            for p in &n.patterns {
                patterns.push(Arc::new(substitute_pattern(p, values)?));
            }
            NestedPattern {
                span: n.span,
                patterns,
                quantifier: n.quantifier.clone(),
            }
            .into()
        }
        PatternKind::Error(e) => e.clone().into(),
    };
    Ok(Pattern { span, kind })
}

fn substitute_pipeline(
    pipeline: &Pipeline,
    values: &HashMap<SimpleIdentifier, TemplateParameterValue>,
) -> Result<Pipeline, TemplateSubstitutionError> {
    let mut commands = Vec::with_capacity(pipeline.commands.len());
    for cmd in &pipeline.commands {
        commands.push(substitute_command(cmd, values)?);
    }
    Ok(Pipeline {
        span: pipeline.span,
        commands,
    })
}

/// Replace all template holes using concrete values only (no parameter specs). Expression holes use
/// the value's shape; table path holes require a string value.
pub fn substitute_query(
    query: &Query,
    values: HashMap<String, TemplateParameterValue>,
) -> Result<Query, TemplateSubstitutionError> {
    let valid = match &query.kind {
        QueryKind::Valid(v) => v,
        QueryKind::Error(_) => return Err(TemplateSubstitutionError::QueryHasParseErrors),
    };

    // Convert the JSON-shaped key map to SimpleIdentifiers once. Every visited template hole
    // (expression or table path) looks values up in this map; missing/wrong-typed values are
    // surfaced lazily by the visitor that hits the hole, so no separate "collect used"
    // pre-pass is needed.
    let values_si: HashMap<SimpleIdentifier, TemplateParameterValue> = values
        .into_iter()
        .map(|(k, v)| (SimpleIdentifier::new(k.as_str()), v))
        .collect();

    let mut defs = Vec::with_capacity(valid.defs.len());
    for d in &valid.defs {
        let body = match &d.body {
            DefBody::Pipeline(p) => {
                DefBody::Pipeline(Arc::new(substitute_pipeline(p.as_ref(), &values_si)?))
            }
            DefBody::Expression(e) => DefBody::Expression(map_arc_expr(e, &values_si)?),
        };
        defs.push(DefStatement {
            span: d.span,
            name: d.name.clone(),
            body,
        });
    }

    let main_pipeline = Arc::new(substitute_pipeline(
        valid.main_pipeline.as_ref(),
        &values_si,
    )?);

    Ok(Query {
        span: query.span,
        kind: ValidQuery {
            span: valid.span,
            defs,
            main_pipeline,
        }
        .into(),
    })
}

#[cfg(test)]
mod tests {
    use std::collections::HashMap;

    use super::substitute_query;
    use crate::tree::ast::query::Query;
    use crate::tree::ast::ParseWithErrors;
    use crate::tree::template::TemplateParameterValue;

    #[test]
    fn substitute_replaces_expression_template() {
        let (q, errs) = Query::parse_with_errors("FROM x | WHERE ${p}");
        assert!(errs.is_empty(), "{errs:?}");
        let mut vals = HashMap::new();
        vals.insert("p".to_string(), TemplateParameterValue::Boolean(true));
        let out = substitute_query(&q, vals).unwrap();
        let rendered = out.to_string();
        assert!(!rendered.contains("${"), "{rendered}");
    }

    #[test]
    fn substitute_table_path_reports_missing_value() {
        let (q, errs) = Query::parse_with_errors("FROM db.${tenant}");
        assert!(errs.is_empty(), "{errs:?}");

        let values = HashMap::new();

        let err = substitute_query(&q, values).unwrap_err();
        assert_eq!(
            err.to_string(),
            "missing value for template parameter 'tenant'"
        );
    }

    #[test]
    fn substitute_table_path_with_string_value() {
        let (q, errs) = Query::parse_with_errors("FROM db.${tenant}_events");
        assert!(errs.is_empty(), "{errs:?}");

        let mut values = HashMap::new();
        values.insert(
            "tenant".to_string(),
            TemplateParameterValue::String("acme".to_string()),
        );

        let out = substitute_query(&q, values).unwrap();
        assert_eq!(out.to_string(), "FROM db.acme_events");
    }

    #[test]
    fn table_path_accepts_non_standard_chars_after_substitution() {
        let (q, errs) = Query::parse_with_errors("FROM db.${t}");
        assert!(errs.is_empty(), "{errs:?}");
        let mut values = HashMap::new();
        values.insert(
            "t".to_string(),
            TemplateParameterValue::String("events-main".to_string()),
        );
        let out = substitute_query(&q, values).unwrap();
        assert_eq!(out.to_string(), "FROM db.`events-main`");
    }

    #[test]
    fn backquoted_literal_in_templated_table_path() {
        let (q, errs) = Query::parse_with_errors("FROM db.`events-`${tenant}");
        assert!(errs.is_empty(), "{errs:?}");
        let mut values = HashMap::new();
        values.insert(
            "tenant".to_string(),
            TemplateParameterValue::String("acme".to_string()),
        );
        let out = substitute_query(&q, values).unwrap();
        assert_eq!(out.to_string(), "FROM db.`events-acme`");
    }
}