hamelin_lib 0.23.2

use std::collections::{HashMap, HashSet};
use std::fmt::{self, Display, Formatter};
use std::ops::RangeInclusive;
use std::rc::Rc;
use std::sync::Arc;

use derive_more::From;

use ordermap::OrderSet;

use super::clause::{TableReference, TableReferenceTarget};
use super::command::CommandKind;
use super::dataset_identifier::QualifiedDatasetIdentifier;
use super::expression::{ErrorExpression, Expression};
use super::identifier::{Identifier, ParsedIdentifier, ParsedSimpleIdentifier, SimpleIdentifier};
use super::node::{Span, Spannable};
use super::pipeline::Pipeline;
use crate::antlr::hamelinparser::{
    DefRhsContextAll, DefRhsExpressionContextAttrs, DefRhsPipelineContextAttrs,
    DefStatementContextAll, DefStatementContextAttrs, ExpressionQueryMainContextAttrs,
    PipelineQueryMainContextAttrs, QueryContextAll, QueryEOFContextAttrs, QueryMainContextAll,
    QueryWithDefsContextAttrs, QueryWithoutDefsContextAttrs, SimpleIdentifierContextAll,
};
use crate::antlr::interval;
use crate::antlr::resilient_parse_query;
use crate::err::{Context, TemplateSubstitutionError, TranslationError, TranslationErrors};
use crate::tree::ast::context::{FromCst, ParseContext};
use crate::tree::ast::{ParseWithContext, ParseWithErrors, TypeCheck};
use crate::tree::expand_templated_table_path;
use crate::tree::options::TemplateParameterKind;
use crate::tree::typed_ast::context::StatementTranslationContext;
use crate::tree::typed_ast::query::TypedStatement;
use crate::write_utils::pad;

/// Walk a table reference and add any catalog identifiers it points at to `datasets`, skipping
/// names that match a CTE in the surrounding query. Templated paths are expanded against
/// `template_parameters` (which must be provided when any templated reference is present);
/// expansion failures are surfaced as `TranslationError`s anchored at the reference span.
///
/// Shared by both top-level pipelines (FROM/UNION/JOIN/LOOKUP/APPEND) and MATCH patterns.
pub fn push_datasets_for_table_reference(
    tr: &TableReference,
    cte_names: &HashSet<Identifier>,
    datasets: &mut OrderSet<QualifiedDatasetIdentifier>,
    default_space: Option<&SimpleIdentifier>,
    template_parameters: Option<&HashMap<String, TemplateParameterKind>>,
) -> Result<(), Arc<TranslationError>> {
    let span = tr.span.to_range().unwrap_or(0..=0);
    let translation_error =
        |msg: &str| Arc::new(TranslationError::new(Context::new(span.clone(), msg)));

    match &tr.target {
        TableReferenceTarget::Static(_) => {
            let dataset = tr.static_valid_dataset_ref()?;
            if !dataset.is_cte_candidate(cte_names) {
                let qualified = dataset
                    .clone()
                    .qualify(default_space)
                    .map_err(|e| translation_error(&e.to_string()))?;
                datasets.insert(qualified);
            }
            Ok(())
        }
        TableReferenceTarget::Templated(path) => {
            let specs = template_parameters.ok_or_else(|| {
                translation_error(
                    "cannot list datasets: query contains templated table references; pass template_parameters to expand them",
                )
            })?;
            let lookup =
                |p: &SimpleIdentifier| -> Result<Arc<[String]>, TemplateSubstitutionError> {
                    match specs.get(p.as_str()) {
                        Some(TemplateParameterKind::IdentifierFragment(values)) => {
                            Ok(Arc::from(values.clone().into_boxed_slice()))
                        }
                        Some(_) => Err(
                            TemplateSubstitutionError::TablePathParameterNeedsIdentifierFragment(
                                p.to_string(),
                            ),
                        ),
                        None => Err(TemplateSubstitutionError::UndeclaredTemplateParameter(
                            p.to_string(),
                        )),
                    }
                };
            for id in expand_templated_table_path(path, lookup)
                .map_err(|e| translation_error(e.to_string().as_str()))?
            {
                if !id.is_cte_candidate(cte_names) {
                    let qualified = id
                        .clone()
                        .qualify(default_space)
                        .map_err(|e| translation_error(&e.to_string()))?;
                    datasets.insert(qualified);
                }
            }
            Ok(())
        }
    }
}

#[derive(Debug, Clone, Default, PartialEq, Eq)]
#[non_exhaustive]
pub struct QueryDatasetsByCommand {
    pub r#from: Vec<QualifiedDatasetIdentifier>,
    pub union: Vec<QualifiedDatasetIdentifier>,
    pub join: Vec<QualifiedDatasetIdentifier>,
    pub lookup: Vec<QualifiedDatasetIdentifier>,
}

#[derive(Default)]
struct CollectedQueryDatasetsByCommand {
    all: OrderSet<QualifiedDatasetIdentifier>,
    r#from: OrderSet<QualifiedDatasetIdentifier>,
    union: OrderSet<QualifiedDatasetIdentifier>,
    join: OrderSet<QualifiedDatasetIdentifier>,
    lookup: OrderSet<QualifiedDatasetIdentifier>,
}

fn push_categorized_datasets_for_table_reference(
    tr: &TableReference,
    cte_names: &HashSet<Identifier>,
    all: &mut OrderSet<QualifiedDatasetIdentifier>,
    category: &mut OrderSet<QualifiedDatasetIdentifier>,
    default_space: Option<&SimpleIdentifier>,
    template_parameters: Option<&HashMap<String, TemplateParameterKind>>,
) -> Result<(), Arc<TranslationError>> {
    let mut references = OrderSet::new();
    push_datasets_for_table_reference(
        tr,
        cte_names,
        &mut references,
        default_space,
        template_parameters,
    )?;
    for reference in references {
        all.insert(reference.clone());
        category.insert(reference);
    }
    Ok(())
}

/// Top-level query node
#[derive(Debug, Clone)]
pub struct Query {
    pub span: Span,
    pub kind: QueryKind,
}

impl Query {
    pub fn valid_ref(&self) -> Result<&ValidQuery, Arc<TranslationError>> {
        match &self.kind {
            QueryKind::Valid(valid_query) => Ok(valid_query),
            QueryKind::Error(translation_error) => Err(translation_error.clone()),
        }
    }

    /// Whether this query contains any `${...}` template parameter references
    /// in its expressions or table paths.
    pub fn has_template_parameters(&self) -> bool {
        let valid = match &self.kind {
            QueryKind::Valid(v) => v,
            QueryKind::Error(_) => return false,
        };
        let pipeline_has = |pipeline: &Pipeline| {
            pipeline
                .commands
                .iter()
                .any(|c| c.kind.has_template_parameters())
        };
        for def in &valid.defs {
            match &def.body {
                DefBody::Pipeline(p) => {
                    if pipeline_has(p) {
                        return true;
                    }
                }
                DefBody::Expression(e) => {
                    if e.contains_template_parameter() {
                        return true;
                    }
                }
            }
        }
        pipeline_has(&valid.main_pipeline)
    }

    /// Returns all dataset (table) references in the query.
    ///
    /// This walks the AST to find all tables referenced in:
    /// - FROM clauses
    /// - JOIN/LOOKUP clauses
    /// - APPEND target tables
    /// - MATCH pattern tables
    ///
    /// Tabular `DEF` names (`DefBody::Pipeline`, user or compiler-inserted) are treated as local
    /// CTEs and excluded from this list. Scalar `DEF` names do not shadow table references and are
    /// not excluded.
    ///
    /// Unqualified table paths (without a `space:` prefix) require a default space via
    /// `default_space`.
    pub fn datasets(
        &self,
        default_space: Option<&SimpleIdentifier>,
    ) -> Result<Vec<QualifiedDatasetIdentifier>, Arc<TranslationError>> {
        self.datasets_with_options(default_space, None)
    }

    /// Like [`Self::datasets`], but applies a default space and/or expands templated table paths.
    pub fn datasets_with_options(
        &self,
        default_space: Option<&SimpleIdentifier>,
        template_parameters: Option<&HashMap<String, TemplateParameterKind>>,
    ) -> Result<Vec<QualifiedDatasetIdentifier>, Arc<TranslationError>> {
        self.datasets_with_optional_template_parameters(default_space, template_parameters)
    }

    /// Like [`Self::datasets`], but expands templated table paths using `IdentifierFragment` domains.
    pub fn datasets_with_template_parameters(
        &self,
        template_parameters: &HashMap<String, TemplateParameterKind>,
    ) -> Result<Vec<QualifiedDatasetIdentifier>, Arc<TranslationError>> {
        self.datasets_with_options(None, Some(template_parameters))
    }

    /// Returns external dataset references grouped by the command that references them.
    ///
    /// This walks `FROM`, `UNION`, `JOIN`, and `LOOKUP` commands in the main pipeline and
    /// tabular `DEF` bodies. Tabular `DEF` names are local CTE references and are excluded. A
    /// catalog dataset referenced by more than one command kind is present in each corresponding
    /// list.
    pub fn datasets_by_command(
        &self,
        default_space: Option<&SimpleIdentifier>,
    ) -> Result<QueryDatasetsByCommand, Arc<TranslationError>> {
        self.datasets_by_command_with_options(default_space, None)
    }

    /// Like [`Self::datasets_by_command`], but expands templated table paths.
    pub fn datasets_by_command_with_options(
        &self,
        default_space: Option<&SimpleIdentifier>,
        template_parameters: Option<&HashMap<String, TemplateParameterKind>>,
    ) -> Result<QueryDatasetsByCommand, Arc<TranslationError>> {
        let references = self.collect_dataset_references(default_space, template_parameters)?;

        Ok(QueryDatasetsByCommand {
            r#from: references.r#from.into_iter().collect(),
            union: references.union.into_iter().collect(),
            join: references.join.into_iter().collect(),
            lookup: references.lookup.into_iter().collect(),
        })
    }

    fn datasets_with_optional_template_parameters(
        &self,
        default_space: Option<&SimpleIdentifier>,
        template_parameters: Option<&HashMap<String, TemplateParameterKind>>,
    ) -> Result<Vec<QualifiedDatasetIdentifier>, Arc<TranslationError>> {
        Ok(self
            .collect_dataset_references(default_space, template_parameters)?
            .all
            .into_iter()
            .collect())
    }

    fn collect_dataset_references(
        &self,
        default_space: Option<&SimpleIdentifier>,
        template_parameters: Option<&HashMap<String, TemplateParameterKind>>,
    ) -> Result<CollectedQueryDatasetsByCommand, Arc<TranslationError>> {
        let valid = self.valid_ref()?;

        let mut cte_names: HashSet<Identifier> = HashSet::new();
        for d in &valid.defs {
            if let DefBody::Pipeline(_) = &d.body {
                if let Ok(id) = d.name.valid_ref() {
                    cte_names.insert(id.clone());
                }
            }
        }

        let mut references = CollectedQueryDatasetsByCommand::default();

        // Collect all pipelines to process (main + DEF clauses)
        let mut pipelines: Vec<&Pipeline> = vec![&valid.main_pipeline];
        for d in &valid.defs {
            if let DefBody::Pipeline(p) = &d.body {
                pipelines.push(&p);
            }
        }
        for pipeline in pipelines {
            for cmd in &pipeline.commands {
                match &cmd.kind {
                    CommandKind::From(from_cmd) => {
                        for clause in &from_cmd.clauses {
                            push_categorized_datasets_for_table_reference(
                                &clause.table_reference(),
                                &cte_names,
                                &mut references.all,
                                &mut references.r#from,
                                default_space,
                                template_parameters,
                            )?;
                        }
                    }
                    CommandKind::Union(union_cmd) => {
                        for clause in &union_cmd.clauses {
                            push_categorized_datasets_for_table_reference(
                                &clause.table_reference(),
                                &cte_names,
                                &mut references.all,
                                &mut references.union,
                                default_space,
                                template_parameters,
                            )?;
                        }
                    }
                    CommandKind::Join(join_cmd) => {
                        push_categorized_datasets_for_table_reference(
                            &join_cmd.other.table_reference(),
                            &cte_names,
                            &mut references.all,
                            &mut references.join,
                            default_space,
                            template_parameters,
                        )?;
                    }
                    CommandKind::Lookup(lookup_cmd) => {
                        push_categorized_datasets_for_table_reference(
                            &lookup_cmd.other.table_reference(),
                            &cte_names,
                            &mut references.all,
                            &mut references.lookup,
                            default_space,
                            template_parameters,
                        )?;
                    }
                    CommandKind::Append(append_cmd) => {
                        if append_cmd.table.is_templated() {
                            let span = append_cmd.table.span.to_range().unwrap_or(0..=0);
                            return Err(Arc::new(TranslationError::new(Context::new(
                                span,
                                "cannot list datasets: templated APPEND target is not allowed",
                            ))));
                        }
                        push_datasets_for_table_reference(
                            append_cmd.table.as_ref(),
                            &cte_names,
                            &mut references.all,
                            default_space,
                            template_parameters,
                        )?;
                    }
                    CommandKind::Match(match_cmd) => {
                        for pattern in &match_cmd.pattern {
                            pattern.extract_datasets(
                                &cte_names,
                                &mut references.all,
                                default_space,
                                template_parameters,
                            )?;
                        }
                    }
                    _ => {}
                }
            }
        }

        Ok(references)
    }
}

impl PartialEq for Query {
    fn eq(&self, other: &Self) -> bool {
        // Compare only the semantic content (kind), ignoring CST
        self.kind == other.kind
    }
}

impl ParseWithContext for Query {
    fn parse_with_context(input: String, mut ctx: ParseContext) -> (Self, TranslationErrors) {
        match resilient_parse_query(input) {
            Ok((csteof, parse_errors)) => match csteof.query() {
                Some(cst) => {
                    let ast = Self::from_cst_with_context(cst, &mut ctx);
                    let mut errors = ctx.take_errors();
                    errors.extend(parse_errors);
                    (ast, errors)
                }
                None => {
                    let err = ctx.error("parse did not resolve a query tree").emit();
                    let mut errors = ctx.take_errors();
                    errors.extend(parse_errors);
                    (
                        Self {
                            span: Span::NONE,
                            kind: QueryKind::Error(err),
                        },
                        errors,
                    )
                }
            },
            Err(e) => {
                let error = ctx
                    .error("fatal error initializing parser")
                    .with_source_boxed(e.into())
                    .emit();
                (
                    Self {
                        span: Span::NONE,
                        kind: QueryKind::Error(error.clone()),
                    },
                    (*error).clone().single(),
                )
            }
        }
    }
}

impl ParseWithErrors for Query {
    fn parse_with_errors(input: impl Into<String>) -> (Self, TranslationErrors) {
        Self::parse_with_context(input.into(), ParseContext::new())
    }
}

impl Spannable for Query {
    fn span(&self) -> Option<RangeInclusive<usize>> {
        self.span.to_range()
    }
}

#[derive(Debug, Clone, PartialEq, From)]
pub enum QueryKind {
    Valid(ValidQuery),
    Error(Arc<TranslationError>),
}

/// A valid query: optional `DEF` prefix (user and/or compiler-generated tabular or scalar
/// bindings), plus a main pipeline (or expression lowered to a pipeline).
#[derive(Debug, Clone)]
pub struct ValidQuery {
    pub span: Span,
    /// `DEF` bindings in source order (all precede the main body). Tabular `DEF`s are CTE-shaped
    /// (`DefBody::Pipeline`); scalars use `DefBody::Expression`.
    pub defs: Vec<DefStatement>,
    pub main_pipeline: Arc<Pipeline>,
}

/// `DEF name = … ;` — tabular (pipeline) or scalar (`expression`) binding.
#[derive(Debug, Clone, PartialEq)]
pub struct DefStatement {
    pub span: Span,
    pub name: ParsedIdentifier,
    pub body: DefBody,
}

#[derive(Debug, Clone, PartialEq)]
pub enum DefBody {
    Pipeline(Arc<Pipeline>),
    Expression(Arc<Expression>),
}

impl FromCst<Rc<DefRhsContextAll<'static>>> for DefBody {
    fn from_cst_with_context(rhs: Rc<DefRhsContextAll<'static>>, ctx: &mut ParseContext) -> Self {
        match rhs.as_ref() {
            DefRhsContextAll::DefRhsPipelineContext(pctx) => {
                let pipe = pctx
                    .pipeline()
                    .map(|p| Arc::new(Pipeline::from_cst_with_context(p, ctx)));
                pipe.map(DefBody::Pipeline).unwrap_or_else(|| {
                    DefBody::Expression(Arc::new(Expression::from_kind(ErrorExpression {
                        error: ctx.error("DEF pipeline missing").at(rhs.as_ref()).emit(),
                    })))
                })
            }
            DefRhsContextAll::DefRhsExpressionContext(ectx) => {
                let expr = ectx
                    .expression()
                    .map(|e| Arc::new(Expression::from_cst_with_context(e, ctx)));
                expr.map(DefBody::Expression).unwrap_or_else(|| {
                    DefBody::Expression(Arc::new(Expression::from_kind(ErrorExpression {
                        error: ctx.error("DEF expression missing").at(rhs.as_ref()).emit(),
                    })))
                })
            }
            DefRhsContextAll::Error(_) => {
                DefBody::Expression(Arc::new(Expression::from_kind(ErrorExpression {
                    error: ctx
                        .error("DEF right-hand side parse error")
                        .at(rhs.as_ref())
                        .emit(),
                })))
            }
        }
    }
}

impl DefStatement {
    fn from_name_and_rhs(
        name_cst: Rc<SimpleIdentifierContextAll<'static>>,
        rhs: Rc<DefRhsContextAll<'static>>,
        ctx: &mut ParseContext,
    ) -> Self {
        let name_iv = interval(name_cst.as_ref());
        let rhs_iv = interval(rhs.as_ref());
        let span = Span::from_range(
            *name_iv.start().min(rhs_iv.start())..=*name_iv.end().max(rhs_iv.end()),
        );
        let name = ParsedSimpleIdentifier::from_cst_with_context(name_cst.as_ref(), ctx).into();
        let body = DefBody::from_cst_with_context(rhs, ctx);
        DefStatement { span, name, body }
    }

    fn from_cst_with_context(
        def_cst: Rc<DefStatementContextAll<'static>>,
        ctx: &mut ParseContext,
    ) -> Option<Self> {
        match (def_cst.simpleIdentifier(), def_cst.defRhs()) {
            (Some(name), Some(rhs)) => {
                let mut def = Self::from_name_and_rhs(name, rhs, ctx);
                def.span = interval(def_cst.as_ref()).into();
                Some(def)
            }
            _ => None,
        }
    }

    pub fn fmt_indented(&self, f: &mut Formatter<'_>, indentation: usize) -> fmt::Result {
        pad(f, indentation)?;
        write!(f, "DEF {} = ", self.name)?;
        match &self.body {
            DefBody::Pipeline(p) => p.fmt_indented(f, indentation)?,
            DefBody::Expression(e) => write!(f, "{e}")?,
        }
        write!(f, ";")
    }
}

impl Spannable for DefStatement {
    fn span(&self) -> Option<RangeInclusive<usize>> {
        self.span.to_range()
    }
}

impl PartialEq for ValidQuery {
    fn eq(&self, other: &Self) -> bool {
        self.defs == other.defs && self.main_pipeline == other.main_pipeline
    }
}

impl Spannable for ValidQuery {
    fn span(&self) -> Option<RangeInclusive<usize>> {
        self.span.to_range()
    }
}

// ============================================================================
// Display / fmt_indented implementations
// ============================================================================

impl Query {
    pub fn fmt_indented(&self, f: &mut Formatter<'_>, indentation: usize) -> fmt::Result {
        self.kind.fmt_indented(f, indentation)
    }
}

impl Display for Query {
    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
        self.fmt_indented(f, 0)
    }
}

impl QueryKind {
    fn from_query_main_with_defs(
        main_cst: Rc<QueryMainContextAll<'static>>,
        query_span: Span,
        defs: Vec<DefStatement>,
        ctx: &mut ParseContext,
    ) -> Self {
        match main_cst.as_ref() {
            QueryMainContextAll::PipelineQueryMainContext(pctx) => match pctx.pipeline() {
                Some(pipeline_cst) => QueryKind::Valid(ValidQuery {
                    span: query_span,
                    defs,
                    main_pipeline: Arc::new(Pipeline::from_cst_with_context(pipeline_cst, ctx)),
                }),
                None => QueryKind::Error(ctx.error("Query missing its main pipeline").emit()),
            },
            QueryMainContextAll::ExpressionQueryMainContext(ectx) => match ectx.expression() {
                Some(expr_cst) => {
                    let expr = Expression::from_cst_with_context(expr_cst, ctx);
                    let pipeline = Pipeline::from(expr);
                    QueryKind::Valid(ValidQuery {
                        span: query_span,
                        defs,
                        main_pipeline: Arc::new(pipeline),
                    })
                }
                None => QueryKind::Error(ctx.error("Expression did not parse").emit()),
            },
            QueryMainContextAll::Error(_) => {
                let error = ctx
                    .error("query main body parse error")
                    .at(main_cst.as_ref())
                    .emit();
                if defs.is_empty() {
                    QueryKind::Error(error)
                } else {
                    let fallback_main = Pipeline::from(Expression::from_kind(ErrorExpression {
                        error: error.clone(),
                    }));
                    QueryKind::Valid(ValidQuery {
                        span: query_span,
                        defs,
                        main_pipeline: Arc::new(fallback_main),
                    })
                }
            }
        }
    }

    pub fn fmt_indented(&self, f: &mut Formatter<'_>, indentation: usize) -> fmt::Result {
        match self {
            QueryKind::Valid(v) => v.fmt_indented(f, indentation),
            QueryKind::Error(_) => write!(f, "?!"),
        }
    }
}

impl ValidQuery {
    pub fn fmt_indented(&self, f: &mut Formatter<'_>, indentation: usize) -> fmt::Result {
        for d in &self.defs {
            d.fmt_indented(f, indentation)?;
            writeln!(f)?;
            writeln!(f)?;
            pad(f, indentation)?;
        }
        self.main_pipeline.fmt_indented(f, indentation)
    }
}

impl Display for DefStatement {
    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
        write!(f, "DEF {} = ", self.name)?;
        match &self.body {
            DefBody::Pipeline(p) => write!(f, "{p}")?,
            DefBody::Expression(e) => write!(f, "{e}")?,
        }
        write!(f, ";")
    }
}

impl Display for ValidQuery {
    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
        self.fmt_indented(f, 0)
    }
}

impl FromCst<Rc<QueryContextAll<'static>>> for Query {
    fn from_cst_with_context(cst: Rc<QueryContextAll<'static>>, ctx: &mut ParseContext) -> Self {
        let query_span = Some(interval(cst.as_ref()));

        let kind = match cst.as_ref() {
            QueryContextAll::QueryWithDefsContext(sctx) => {
                let defs: Vec<DefStatement> = sctx
                    .defStatement_all()
                    .into_iter()
                    .filter_map(|def_cst| DefStatement::from_cst_with_context(def_cst, ctx))
                    .collect();

                match sctx.queryMain() {
                    Some(main_cst) => QueryKind::from_query_main_with_defs(
                        main_cst,
                        query_span.clone().into(),
                        defs,
                        ctx,
                    ),
                    None => {
                        let error = ctx
                            .error("Query with DEFs missing its main body")
                            .at(cst.as_ref())
                            .emit();
                        let fallback_main =
                            Pipeline::from(Expression::from_kind(ErrorExpression {
                                error: error.clone(),
                            }));
                        QueryKind::Valid(ValidQuery {
                            span: query_span.clone().into(),
                            defs,
                            main_pipeline: Arc::new(fallback_main),
                        })
                    }
                }
            }
            QueryContextAll::QueryWithoutDefsContext(sctx) => match sctx.queryMain() {
                Some(main_cst) => QueryKind::from_query_main_with_defs(
                    main_cst,
                    query_span.clone().into(),
                    vec![],
                    ctx,
                ),
                None => QueryKind::Error(
                    ctx.error("Query missing its main body")
                        .at(cst.as_ref())
                        .emit(),
                ),
            },
            QueryContextAll::Error(_) => ctx
                .error("query parse error")
                .at(cst.as_ref())
                .emit()
                .into(),
        };

        Query {
            span: query_span.into(),
            kind,
        }
    }
}

impl TypeCheck for Query {
    type Output = TypedStatement;

    fn type_check_with_context(
        ast: Arc<Self>,
        ctx: &mut StatementTranslationContext,
    ) -> Self::Output {
        TypedStatement::from_ast_with_context(ast, ctx)
    }
}