hamelin_lib 0.11.2

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

use super::node::Span;

use derive_more::derive::From;

use crate::{
    antlr::{
        hamelinparser::{
            AggCommandContextAttrs, AppendCommandContextAttrs, CommandContextAll,
            DistinctCommandContextAttrs, DropCommandContextAttrs, ExplodeCommandContextAttrs,
            FromCommandContextAttrs, GroupClauseContextAttrs, JoinCommandContextAttrs,
            LimitCommandContextAttrs, MatchCommandContextAttrs, NestCommandContextAttrs,
            ParseCommandContextAttrs, SelectCommandContextAttrs, SetCommandContextAttrs,
            SortCommandContextAttrs, SuppressCommandContextAttrs, TrimstringsCommandContextAttrs,
            UnionCommandContextAttrs, UnnestCommandContextAttrs, WhereCommandContextAttrs,
            WindowCommandContextAttrs, WithinCommandContextAttrs,
        },
        interval,
    },
    err::TranslationError,
    tree::ast::{
        context::{FromCst, ParseContext, TryFromCst},
        display::write_comma_list,
    },
    write_utils::pad,
};

use super::{
    clause::{Assignment, AssignmentClause, FromClause, Selection, SortExpression, TableReference},
    expression::{ErrorExpression, Expression, ExpressionKind, StringLiteral},
    identifier::ParsedIdentifier,
    node::Spannable,
    pattern::Pattern,
};

// ============================================================================
// Command wrapper and enum
// ============================================================================

#[derive(Debug, Clone)]
pub struct Command {
    pub span: Span,
    pub kind: CommandKind,
}

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

impl From<CommandKind> for Command {
    fn from(kind: CommandKind) -> Self {
        Self {
            span: Span::NONE,
            kind,
        }
    }
}

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

#[derive(Debug, Clone, PartialEq, From)]
pub enum CommandKind {
    Set(SetCommand),
    Where(WhereCommand),
    Select(SelectCommand),
    Drop(DropCommand),
    Limit(LimitCommand),
    Trimstrings(TrimstringsCommand),
    Within(WithinCommand),
    Sort(SortCommand),
    From(FromCommand),
    Union(UnionCommand),
    Parse(ParseCommand),
    Agg(AggCommand),
    Distinct(DistinctCommand),
    Suppress(SuppressCommand),
    Window(WindowCommand),
    Append(AppendCommand),
    Join(JoinCommand),
    Lookup(LookupCommand),
    Explode(ExplodeCommand),
    Unnest(UnnestCommand),
    Rows(RowsCommand),
    Nest(NestCommand),
    Match(MatchCommand),
    Error(Arc<TranslationError>),
}

/// Classification of commands based on their behavior
///
/// This enum provides a centralized way to classify commands as either
/// query operations (that return data) or DML operations (that modify data).
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum CommandClass {
    /// Commands that return data (query operations)
    /// Examples: FROM, WHERE, SELECT, AGG, etc.
    Query,

    /// Commands that modify data (write operations)
    /// Examples: APPEND
    Dml,

    /// Commands that modify table structure
    /// Examples: CREATE
    Ddl,
}

impl CommandKind {
    /// Get the classification of this command
    ///
    /// Returns `CommandClass::Dml` for data modification commands (APPEND),
    /// and `CommandClass::Query` for all other commands.
    pub fn command_class(&self) -> CommandClass {
        match self {
            CommandKind::Append(_) => CommandClass::Dml,
            _ => CommandClass::Query,
        }
    }

    /// Whether this command introduces a new data source and must appear
    /// as the first command in a pipeline (FROM, UNION, ROWS, MATCH).
    pub fn is_source_command(&self) -> bool {
        matches!(
            self,
            CommandKind::From(_)
                | CommandKind::Union(_)
                | CommandKind::Rows(_)
                | CommandKind::Match(_)
        )
    }

    /// Whether this command contains any template parameter references (`${...}`)
    /// in its expressions or table paths.
    pub fn has_template_parameters(&self) -> bool {
        let expr = |e: &Expression| e.contains_template_parameter();
        let assignment_clauses = |clauses: &[Arc<AssignmentClause>]| {
            clauses.iter().any(|c| match c.as_ref() {
                AssignmentClause::Assignment(a) => expr(&a.expression),
                AssignmentClause::Error(_) => false,
            })
        };
        let sort_exprs = |sorts: &[Arc<SortExpression>]| sorts.iter().any(|s| expr(&s.expression));
        let from_clauses = |clauses: &[Arc<FromClause>]| {
            clauses.iter().any(|c| c.table_reference().is_templated())
        };

        match self {
            CommandKind::Set(c) => c.assignments.iter().any(|a| expr(&a.expression)),
            CommandKind::Where(c) => expr(&c.condition),
            CommandKind::Select(c) => assignment_clauses(&c.clauses),
            CommandKind::Drop(_) | CommandKind::Nest(_) | CommandKind::Error(_) => false,
            CommandKind::Limit(c) => expr(&c.count),
            CommandKind::Trimstrings(c) => expr(&c.byte_len),
            CommandKind::Within(c) => expr(&c.duration),
            CommandKind::Sort(c) => sort_exprs(&c.expressions),
            CommandKind::From(c) => from_clauses(&c.clauses),
            CommandKind::Union(c) => from_clauses(&c.clauses),
            CommandKind::Parse(c) => c.source.as_ref().is_some_and(|e| expr(e)) || expr(&c.pattern),
            CommandKind::Agg(c) => {
                assignment_clauses(&c.clauses)
                    || assignment_clauses(&c.group_by)
                    || sort_exprs(&c.sort)
            }
            CommandKind::Distinct(c) => assignment_clauses(&c.clauses),
            CommandKind::Window(c) => {
                assignment_clauses(&c.clauses)
                    || assignment_clauses(&c.group_by)
                    || sort_exprs(&c.sort)
                    || c.within.as_ref().is_some_and(|e| expr(e))
            }
            CommandKind::Append(c) => c.table.is_templated(),
            CommandKind::Join(c) => {
                c.other.table_reference().is_templated()
                    || c.on_condition.as_ref().is_some_and(|e| expr(e))
            }
            CommandKind::Lookup(c) => {
                c.other.table_reference().is_templated()
                    || c.on_condition.as_ref().is_some_and(|e| expr(e))
            }
            CommandKind::Explode(c) => assignment_clauses(&c.clauses),
            CommandKind::Unnest(c) => expr(&c.expression),
            CommandKind::Rows(c) => expr(&c.expression),
            CommandKind::Match(c) => {
                c.pattern.iter().any(|p| p.has_template_parameters())
                    || assignment_clauses(&c.agg)
                    || assignment_clauses(&c.group_by)
                    || sort_exprs(&c.sort)
                    || c.within.as_ref().is_some_and(|e| expr(e))
            }
            CommandKind::Suppress(c) => expr(&c.interval) || assignment_clauses(&c.group_by),
        }
    }

    /// Get the display name of this command (for error messages)
    pub fn command_name(&self) -> &'static str {
        match self {
            CommandKind::Set(_) => "SET",
            CommandKind::Where(_) => "WHERE",
            CommandKind::Select(_) => "SELECT",
            CommandKind::Drop(_) => "DROP",
            CommandKind::Limit(_) => "LIMIT",
            CommandKind::Trimstrings(_) => "TRIMSTRINGS",
            CommandKind::Within(_) => "WITHIN",
            CommandKind::Sort(_) => "SORT",
            CommandKind::From(_) => "FROM",
            CommandKind::Union(_) => "UNION",
            CommandKind::Parse(_) => "PARSE",
            CommandKind::Agg(_) => "AGG",
            CommandKind::Distinct(_) => "DISTINCT",
            CommandKind::Suppress(_) => "SUPPRESS",
            CommandKind::Window(_) => "WINDOW",
            CommandKind::Append(_) => "APPEND",
            CommandKind::Join(_) => "JOIN",
            CommandKind::Lookup(_) => "LOOKUP",
            CommandKind::Explode(_) => "EXPLODE",
            CommandKind::Unnest(_) => "UNNEST",
            CommandKind::Rows(_) => "ROWS",
            CommandKind::Nest(_) => "NEST",
            CommandKind::Match(_) => "MATCH",
            CommandKind::Error(_) => "ERROR",
        }
    }
}

// ============================================================================
// SET Command
// ============================================================================

#[derive(Debug, Clone)]
pub struct SetCommand {
    pub assignments: Vec<Arc<Assignment>>,
}

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

// ============================================================================
// WHERE Command
// ============================================================================

#[derive(Debug, Clone)]
pub struct WhereCommand {
    pub condition: Arc<Expression>,
}

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

// ============================================================================
// LIMIT Command
// ============================================================================

#[derive(Debug, Clone)]
pub struct LimitCommand {
    pub count: Arc<Expression>,
}

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

// ============================================================================
// TRIMSTRINGS Command
// ============================================================================

#[derive(Debug, Clone)]
pub struct TrimstringsCommand {
    pub byte_len: Arc<Expression>,
}

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

// ============================================================================
// WITHIN Command
// ============================================================================

#[derive(Debug, Clone)]
pub struct WithinCommand {
    pub duration: Arc<Expression>,
}

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

// ============================================================================
// SORT Command
// ============================================================================

#[derive(Debug, Clone)]
pub struct SortCommand {
    pub expressions: Vec<Arc<SortExpression>>,
}

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

// ============================================================================
// FROM Command
// ============================================================================

#[derive(Debug, Clone)]
pub struct FromCommand {
    pub clauses: Vec<Arc<FromClause>>,
}

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

// ============================================================================
// SELECT Command
// ============================================================================

#[derive(Debug, Clone)]
pub struct SelectCommand {
    pub clauses: Vec<Arc<AssignmentClause>>,
}

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

// ============================================================================
// DROP Command
// ============================================================================

#[derive(Debug, Clone)]
pub struct DropCommand {
    pub selections: Vec<Arc<Selection>>,
}

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

// ============================================================================
// UNION Command
// ============================================================================

#[derive(Debug, Clone)]
pub struct UnionCommand {
    pub clauses: Vec<Arc<FromClause>>,
}

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

// ============================================================================
// PARSE Command
// ============================================================================

#[derive(Debug, Clone)]
pub struct ParseCommand {
    pub source: Option<Arc<Expression>>,
    pub pattern: Arc<Expression>,
    pub identifiers: Vec<ParsedIdentifier>,
    pub nodrop: bool,
}

impl PartialEq for ParseCommand {
    fn eq(&self, other: &Self) -> bool {
        self.source == other.source
            && self.pattern == other.pattern
            && self.identifiers == other.identifiers
            && self.nodrop == other.nodrop
    }
}

// ============================================================================
// AGG Command
// ============================================================================

#[derive(Debug, Clone)]
pub struct AggCommand {
    pub clauses: Vec<Arc<AssignmentClause>>,
    pub group_by: Vec<Arc<AssignmentClause>>,
    pub sort: Vec<Arc<SortExpression>>,
}

impl PartialEq for AggCommand {
    fn eq(&self, other: &Self) -> bool {
        self.clauses == other.clauses && self.group_by == other.group_by && self.sort == other.sort
    }
}

// ============================================================================
// DISTINCT Command
// ============================================================================

#[derive(Debug, Clone)]
pub struct DistinctCommand {
    pub clauses: Vec<Arc<AssignmentClause>>,
}

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

// ============================================================================
// SUPPRESS Command
// ============================================================================

#[derive(Debug, Clone)]
pub struct SuppressCommand {
    pub interval: Arc<Expression>,
    pub group_by: Vec<Arc<AssignmentClause>>,
}

impl PartialEq for SuppressCommand {
    fn eq(&self, other: &Self) -> bool {
        self.interval == other.interval && self.group_by == other.group_by
    }
}

// ============================================================================
// WINDOW Command
// ============================================================================

#[derive(Debug, Clone)]
pub struct WindowCommand {
    pub clauses: Vec<Arc<AssignmentClause>>,
    pub group_by: Vec<Arc<AssignmentClause>>,
    pub sort: Vec<Arc<SortExpression>>,
    pub within: Option<Arc<Expression>>,
}

impl PartialEq for WindowCommand {
    fn eq(&self, other: &Self) -> bool {
        self.clauses == other.clauses
            && self.group_by == other.group_by
            && self.sort == other.sort
            && self.within == other.within
    }
}

// ============================================================================
// APPEND Command
// ============================================================================

#[derive(Debug, Clone)]
pub struct AppendCommand {
    pub table: Arc<TableReference>,
    pub distinct_by: Vec<Arc<Selection>>,
}

impl PartialEq for AppendCommand {
    fn eq(&self, other: &Self) -> bool {
        self.table == other.table && self.distinct_by == other.distinct_by
    }
}

// ============================================================================
// JOIN Command
// ============================================================================

#[derive(Debug, Clone)]
pub struct JoinCommand {
    pub other: Arc<FromClause>,
    pub on_condition: Option<Arc<Expression>>,
}

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

// ============================================================================
// LOOKUP Command
// ============================================================================

#[derive(Debug, Clone)]
pub struct LookupCommand {
    pub other: Arc<FromClause>,
    pub on_condition: Option<Arc<Expression>>,
}

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

// ============================================================================
// EXPLODE Command
// ============================================================================

/// EXPLODE command - expands array fields into rows.
///
/// Supports multiple clauses for parallel expansion (zipping arrays together).
/// When multiple arrays are exploded together, elements at the same index are
/// expanded together, with shorter arrays padded with NULL.
///
/// Note: Multi-clause EXPLODE is only available via the builder API for internal use.
/// The grammar only supports single-clause EXPLODE.
#[derive(Debug, Clone)]
pub struct ExplodeCommand {
    /// The clauses to explode. Usually a single clause from parsing,
    /// but can be multiple for internal lowering passes.
    pub clauses: Vec<Arc<AssignmentClause>>,
}

impl ExplodeCommand {
    /// Create with a single clause (backward compatibility).
    pub fn single(clause: Arc<AssignmentClause>) -> Self {
        Self {
            clauses: vec![clause],
        }
    }

    /// Create with multiple clauses (for internal lowering passes).
    pub fn multiple(clauses: Vec<Arc<AssignmentClause>>) -> Self {
        Self { clauses }
    }
}

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

// ============================================================================
// UNNEST Command
// ============================================================================

#[derive(Debug, Clone)]
pub struct UnnestCommand {
    pub expression: Arc<Expression>,
}

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

// ============================================================================
// ROWS Command
// ============================================================================

#[derive(Debug, Clone)]
pub struct RowsCommand {
    pub expression: Arc<Expression>,
}

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

// ============================================================================
// NEST Command
// ============================================================================

#[derive(Debug, Clone)]
pub struct NestCommand {
    pub identifier: Arc<ParsedIdentifier>,
}

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

// ============================================================================
// MATCH Command
// ============================================================================

#[derive(Debug, Clone)]
pub struct MatchCommand {
    pub pattern: Vec<Arc<Pattern>>,
    pub agg: Vec<Arc<AssignmentClause>>,
    pub group_by: Vec<Arc<AssignmentClause>>,
    pub sort: Vec<Arc<SortExpression>>,
    pub within: Option<Arc<Expression>>,
}

impl PartialEq for MatchCommand {
    fn eq(&self, other: &Self) -> bool {
        self.pattern == other.pattern
            && self.agg == other.agg
            && self.group_by == other.group_by
            && self.sort == other.sort
            && self.within == other.within
    }
}

// ============================================================================
// CST Conversions
// ============================================================================

impl FromCst<Rc<CommandContextAll<'static>>> for Command {
    fn from_cst_with_context(cst: Rc<CommandContextAll<'static>>, ctx: &mut ParseContext) -> Self {
        use crate::antlr::hamelinparser::CommandContextAll::*;

        let kind = match cst.as_ref() {
            SetCommandContext(lctx) => {
                let assignments = lctx
                    .assignment_all()
                    .into_iter()
                    .map(|a| Arc::new(Assignment::from_cst_with_context(a, ctx)))
                    .collect();

                SetCommand { assignments }.into()
            }

            WhereCommandContext(wctx) => {
                let condition = if let Some(expr_ctx) = wctx.expression() {
                    Expression::from_cst_with_context(expr_ctx, ctx)
                } else {
                    let err = ctx
                        .error("missing condition in WHERE command")
                        .at(wctx)
                        .emit();
                    Expression::from_kind(ErrorExpression { error: err })
                };

                WhereCommand {
                    condition: Arc::new(condition),
                }
                .into()
            }

            LimitCommandContext(lctx) => {
                let count = if let Some(expr_ctx) = lctx.expression() {
                    Expression::from_cst_with_context(expr_ctx, ctx)
                } else {
                    let err = ctx.error("missing count in LIMIT command").at(lctx).emit();
                    Expression::from_kind(ErrorExpression { error: err })
                };

                LimitCommand {
                    count: Arc::new(count),
                }
                .into()
            }

            TrimstringsCommandContext(tctx) => {
                let byte_len = if let Some(expr_ctx) = tctx.expression() {
                    Expression::from_cst_with_context(expr_ctx, ctx)
                } else {
                    let err = ctx
                        .error("missing byte length in TRIMSTRINGS command")
                        .at(tctx)
                        .emit();
                    Expression::from_kind(ErrorExpression { error: err })
                };

                TrimstringsCommand {
                    byte_len: Arc::new(byte_len),
                }
                .into()
            }

            WithinCommandContext(wictx) => {
                let duration = if let Some(expr_ctx) = wictx.expression() {
                    Expression::from_cst_with_context(expr_ctx, ctx)
                } else {
                    let err = ctx
                        .error("missing duration in WITHIN command")
                        .at(wictx)
                        .emit();
                    Expression::from_kind(ErrorExpression { error: err })
                };

                WithinCommand {
                    duration: Arc::new(duration),
                }
                .into()
            }

            SortCommandContext(sctx) => {
                let expressions = sctx
                    .sortExpression_all()
                    .into_iter()
                    .map(|e| Arc::new(SortExpression::from_cst_with_context(e, ctx)))
                    .collect();

                SortCommand { expressions }.into()
            }

            FromCommandContext(fctx) => {
                let clauses = fctx
                    .fromClause_all()
                    .into_iter()
                    .map(|c| Arc::new(FromClause::from_cst_with_context(c, ctx)))
                    .collect();

                FromCommand { clauses }.into()
            }

            SelectCommandContext(sctx) => {
                let clauses = sctx
                    .assignmentClause_all()
                    .into_iter()
                    .map(|c| Arc::new(AssignmentClause::from_cst_with_context(c, ctx)))
                    .collect();

                SelectCommand { clauses }.into()
            }

            DropCommandContext(dctx) => {
                let selections = dctx
                    .selection_all()
                    .into_iter()
                    .map(|s| Arc::new(Selection::from_cst_with_context(s, ctx)))
                    .collect();

                DropCommand { selections }.into()
            }

            UnionCommandContext(uctx) => {
                let clauses = uctx
                    .fromClause_all()
                    .into_iter()
                    .map(|c| Arc::new(FromClause::from_cst_with_context(c, ctx)))
                    .collect();

                UnionCommand { clauses }.into()
            }

            ParseCommandContext(pctx) => {
                let source = pctx
                    .src
                    .clone()
                    .map(|expr_ctx| Arc::new(Expression::from_cst_with_context(expr_ctx, ctx)));

                let pattern = pctx
                    .string()
                    .map(|string_ctx| {
                        match StringLiteral::try_from_cst_with_context(&string_ctx, ctx) {
                            Ok(string_lit) => ExpressionKind::from(string_lit).into(),
                            Err(error) => Expression::from_kind(ErrorExpression { error }),
                        }
                    })
                    .unwrap_or_else(|| {
                        let err = ctx
                            .error("missing pattern in PARSE command")
                            .at(pctx)
                            .emit();
                        Expression::from_kind(ErrorExpression { error: err })
                    });

                let identifiers = pctx
                    .identifier_all()
                    .into_iter()
                    .map(|id| ParsedIdentifier::from_cst_with_context(id, ctx))
                    .collect();

                let nodrop = pctx.NODROP().is_some();

                ParseCommand {
                    source,
                    pattern: Arc::new(pattern),
                    identifiers,
                    nodrop,
                }
                .into()
            }

            AggCommandContext(actx) => {
                let clauses = actx
                    .assignmentClause_all()
                    .into_iter()
                    .map(|c| Arc::new(AssignmentClause::from_cst_with_context(c, ctx)))
                    .collect();

                let group_by = actx
                    .groupClause_all()
                    .into_iter()
                    .filter_map(|gc| gc.assignmentClause())
                    .map(|c| Arc::new(AssignmentClause::from_cst_with_context(c, ctx)))
                    .collect();

                let sort = actx
                    .sortExpression_all()
                    .into_iter()
                    .map(|s| Arc::new(SortExpression::from_cst_with_context(s, ctx)))
                    .collect();

                AggCommand {
                    clauses,
                    group_by,
                    sort,
                }
                .into()
            }

            DistinctCommandContext(dctx) => {
                let clauses = dctx
                    .groupClause_all()
                    .into_iter()
                    .filter_map(|gc| gc.assignmentClause())
                    .map(|c| Arc::new(AssignmentClause::from_cst_with_context(c, ctx)))
                    .collect();

                DistinctCommand { clauses }.into()
            }

            SuppressCommandContext(sctx) => {
                let interval = if let Some(expr_ctx) = sctx.expression() {
                    Arc::new(Expression::from_cst_with_context(expr_ctx, ctx))
                } else {
                    let err = ctx
                        .error("missing interval expression in SUPPRESS command")
                        .at(sctx)
                        .emit();
                    Arc::new(Expression::from_kind(ErrorExpression { error: err }))
                };

                let group_by = sctx
                    .groupClause_all()
                    .into_iter()
                    .filter_map(|gc| gc.assignmentClause())
                    .map(|c| Arc::new(AssignmentClause::from_cst_with_context(c, ctx)))
                    .collect();

                SuppressCommand { interval, group_by }.into()
            }

            WindowCommandContext(wctx) => {
                let clauses = wctx
                    .assignmentClause_all()
                    .into_iter()
                    .map(|c| Arc::new(AssignmentClause::from_cst_with_context(c, ctx)))
                    .collect();

                let group_by = wctx
                    .groupClause_all()
                    .into_iter()
                    .filter_map(|gc| gc.assignmentClause())
                    .map(|c| Arc::new(AssignmentClause::from_cst_with_context(c, ctx)))
                    .collect();

                let sort = wctx
                    .sortExpression_all()
                    .into_iter()
                    .map(|s| Arc::new(SortExpression::from_cst_with_context(s, ctx)))
                    .collect();

                let within = wctx
                    .within
                    .clone()
                    .map(|expr_ctx| Arc::new(Expression::from_cst_with_context(expr_ctx, ctx)));

                WindowCommand {
                    clauses,
                    group_by,
                    sort,
                    within,
                }
                .into()
            }

            AppendCommandContext(apctx) => {
                let table = apctx
                    .tableReference()
                    .map(|tr| TableReference::from_cst_with_context(tr, ctx))
                    .unwrap_or_else(|| TableReference {
                        span: Span::NONE,
                        target: ctx
                            .error("missing table reference in APPEND command")
                            .at(apctx)
                            .emit()
                            .into(),
                    });

                let distinct_by = apctx
                    .selection_all()
                    .into_iter()
                    .map(|s| Arc::new(Selection::from_cst_with_context(s, ctx)))
                    .collect();

                AppendCommand {
                    table: Arc::new(table),
                    distinct_by,
                }
                .into()
            }

            JoinCommandContext(jctx) => {
                let is_lookup = jctx.LOOKUP_COMMAND().is_some();

                let other = Arc::new(
                    jctx.fromClause()
                        .map(|fc| FromClause::from_cst_with_context(fc, ctx))
                        .unwrap_or_else(|| {
                            Arc::new(TableReference {
                                span: Span::NONE,
                                target: ctx
                                    .error("missing from clause in JOIN/LOOKUP command")
                                    .at(jctx)
                                    .emit()
                                    .into(),
                            })
                            .into()
                        }),
                );

                let on_condition = jctx
                    .on
                    .clone()
                    .map(|expr_ctx| Arc::new(Expression::from_cst_with_context(expr_ctx, ctx)));

                if is_lookup {
                    LookupCommand {
                        other,
                        on_condition,
                    }
                    .into()
                } else {
                    JoinCommand {
                        other,
                        on_condition,
                    }
                    .into()
                }
            }

            ExplodeCommandContext(exctx) => {
                let clause = Arc::new(
                    exctx
                        .assignmentClause()
                        .map(|ac| AssignmentClause::from_cst_with_context(ac, ctx))
                        .unwrap_or_else(|| {
                            AssignmentClause::Error(
                                ctx.error("missing clause in EXPLODE command")
                                    .at(exctx)
                                    .emit(),
                            )
                        }),
                );

                ExplodeCommand::single(clause).into()
            }

            UnnestCommandContext(unctx) => {
                let is_rows = unctx.ROWS_COMMAND().is_some();

                let expression = if let Some(expr_ctx) = unctx.expression() {
                    Expression::from_cst_with_context(expr_ctx, ctx)
                } else {
                    let err = ctx
                        .error("missing expression in UNNEST/ROWS command")
                        .at(unctx)
                        .emit();
                    Expression::from_kind(ErrorExpression { error: err })
                };

                let rc_expression = Arc::new(expression);

                if is_rows {
                    RowsCommand {
                        expression: rc_expression.clone(),
                    }
                    .into()
                } else {
                    UnnestCommand {
                        expression: rc_expression,
                    }
                    .into()
                }
            }

            NestCommandContext(nctx) => {
                let identifier = Arc::new(
                    nctx.identifier()
                        .map(|id| ParsedIdentifier::from_cst_with_context(id, ctx))
                        .unwrap_or_else(|| {
                            ctx.error("missing identifier in NEST command")
                                .at(nctx)
                                .emit()
                                .into()
                        }),
                );

                NestCommand { identifier }.into()
            }

            MatchCommandContext(mctx) => {
                let pattern = mctx
                    .pattern_all()
                    .into_iter()
                    .map(|p| Arc::new(Pattern::from_cst_with_context(p, ctx)))
                    .collect();

                // AGG clauses (assignmentClause directly under MATCH, not via groupClause)
                let agg = mctx
                    .assignmentClause_all()
                    .into_iter()
                    .map(|c| Arc::new(AssignmentClause::from_cst_with_context(c, ctx)))
                    .collect();

                let group_by = mctx
                    .groupClause_all()
                    .into_iter()
                    .filter_map(|gc| {
                        gc.assignmentClause()
                            .map(|c| Arc::new(AssignmentClause::from_cst_with_context(c, ctx)))
                    })
                    .collect();

                let sort = mctx
                    .sortExpression_all()
                    .into_iter()
                    .map(|sort_ctx| Arc::new(SortExpression::from_cst_with_context(sort_ctx, ctx)))
                    .collect();

                let within = mctx.within.as_ref().map(|expr_ctx| {
                    Arc::new(Expression::from_cst_with_context(expr_ctx.clone(), ctx))
                });

                MatchCommand {
                    pattern,
                    agg,
                    group_by,
                    sort,
                    within,
                }
                .into()
            }

            _ => CommandKind::Error(
                ctx.error("command not yet implemented in AST")
                    .at(cst.as_ref())
                    .emit(),
            ),
        };

        Command {
            span: interval(cst.as_ref()).into(),
            kind,
        }
    }
}

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

/// Helper: write a comma-separated clause list (e.g. AGG expressions, BY group_by).
///
/// `keyword_len` is the length of the keyword prefix (e.g. "AGG " = 4) used to
/// compute the hang indent for wrapping items within this clause.
fn write_clause_items<T>(
    f: &mut Formatter<'_>,
    items: &[Arc<T>],
    base_indent: usize,
    keyword_len: usize,
    total_subelements: usize,
    fmt_item: impl Fn(&mut Formatter<'_>, &T, usize) -> fmt::Result,
) -> fmt::Result
where
    T: ?Sized,
{
    let hang_indent = base_indent + keyword_len;
    write_comma_list(
        f,
        items,
        hang_indent,
        total_subelements,
        4,
        |f, item, ind| fmt_item(f, item, ind),
    )
}

/// Count total subelements across a slice of items that have a `subelements()` method.
fn sum_subelements_assignment(items: &[Arc<super::clause::AssignmentClause>]) -> usize {
    items.iter().map(|i| i.subelements()).sum()
}

fn sum_subelements_sort(items: &[Arc<SortExpression>]) -> usize {
    items.iter().map(|i| i.subelements()).sum()
}

impl Command {
    pub fn subelements(&self) -> usize {
        self.kind.subelements()
    }

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

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

impl CommandKind {
    pub fn subelements(&self) -> usize {
        match self {
            CommandKind::Set(c) => c.assignments.iter().map(|a| a.subelements()).sum(),
            CommandKind::Where(c) => c.condition.subelements(),
            CommandKind::Select(c) => sum_subelements_assignment(&c.clauses),
            CommandKind::Drop(c) => c.selections.len(),
            CommandKind::Limit(c) => c.count.subelements(),
            CommandKind::Trimstrings(c) => c.byte_len.subelements(),
            CommandKind::Within(c) => c.duration.subelements(),
            CommandKind::Sort(c) => sum_subelements_sort(&c.expressions),
            CommandKind::From(c) => c.clauses.len(),
            CommandKind::Union(c) => c.clauses.len(),
            CommandKind::Parse(c) => {
                c.source.as_ref().map_or(0, |s| s.subelements() + 1)
                    + c.pattern.subelements()
                    + c.identifiers.len()
                    + 1
            }
            CommandKind::Agg(c) => {
                sum_subelements_assignment(&c.clauses)
                    + sum_subelements_assignment(&c.group_by)
                    + sum_subelements_sort(&c.sort)
            }
            CommandKind::Distinct(c) => sum_subelements_assignment(&c.clauses),
            CommandKind::Suppress(c) => {
                c.interval.subelements() + sum_subelements_assignment(&c.group_by)
            }
            CommandKind::Window(c) => {
                sum_subelements_assignment(&c.clauses)
                    + sum_subelements_assignment(&c.group_by)
                    + sum_subelements_sort(&c.sort)
                    + c.within.as_ref().map_or(0, |w| w.subelements() + 1)
            }
            CommandKind::Append(c) => 1 + c.distinct_by.len(),
            CommandKind::Join(c) => 1 + c.on_condition.as_ref().map_or(0, |e| e.subelements() + 1),
            CommandKind::Lookup(c) => {
                1 + c.on_condition.as_ref().map_or(0, |e| e.subelements() + 1)
            }
            CommandKind::Explode(c) => sum_subelements_assignment(&c.clauses),
            CommandKind::Unnest(c) => c.expression.subelements(),
            CommandKind::Rows(c) => c.expression.subelements(),
            CommandKind::Nest(_) => 1,
            CommandKind::Match(c) => {
                c.pattern.len()
                    + sum_subelements_assignment(&c.agg)
                    + sum_subelements_assignment(&c.group_by)
                    + sum_subelements_sort(&c.sort)
                    + c.within.as_ref().map_or(0, |w| w.subelements() + 1)
            }
            CommandKind::Error(_) => 0,
        }
    }

    pub fn fmt_indented(&self, f: &mut Formatter<'_>, indentation: usize) -> fmt::Result {
        match self {
            CommandKind::Set(c) => c.fmt_indented(f, indentation),
            CommandKind::Where(c) => c.fmt_indented(f, indentation),
            CommandKind::Select(c) => c.fmt_indented(f, indentation),
            CommandKind::Drop(c) => c.fmt_indented(f, indentation),
            CommandKind::Limit(c) => c.fmt_indented(f, indentation),
            CommandKind::Trimstrings(c) => c.fmt_indented(f, indentation),
            CommandKind::Within(c) => c.fmt_indented(f, indentation),
            CommandKind::Sort(c) => c.fmt_indented(f, indentation),
            CommandKind::From(c) => c.fmt_indented(f, indentation),
            CommandKind::Union(c) => c.fmt_indented(f, indentation),
            CommandKind::Parse(c) => c.fmt_indented(f, indentation),
            CommandKind::Agg(c) => c.fmt_indented(f, indentation),
            CommandKind::Distinct(c) => c.fmt_indented(f, indentation),
            CommandKind::Suppress(c) => c.fmt_indented(f, indentation),
            CommandKind::Window(c) => c.fmt_indented(f, indentation),
            CommandKind::Append(c) => c.fmt_indented(f, indentation),
            CommandKind::Join(c) => c.fmt_indented(f, indentation),
            CommandKind::Lookup(c) => c.fmt_indented(f, indentation),
            CommandKind::Explode(c) => c.fmt_indented(f, indentation),
            CommandKind::Unnest(c) => c.fmt_indented(f, indentation),
            CommandKind::Rows(c) => c.fmt_indented(f, indentation),
            CommandKind::Nest(c) => c.fmt_indented(f, indentation),
            CommandKind::Match(c) => c.fmt_indented(f, indentation),
            CommandKind::Error(_) => write!(f, "?!"),
        }
    }
}

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

impl SetCommand {
    pub fn fmt_indented(&self, f: &mut Formatter<'_>, indentation: usize) -> fmt::Result {
        write!(f, "SET ")?;
        // "SET " = 4 chars
        write_clause_items(
            f,
            &self.assignments,
            indentation,
            4,
            self.assignments.iter().map(|a| a.subelements()).sum(),
            |f, a, ind| a.fmt_indented(f, ind),
        )
    }
}

impl WhereCommand {
    pub fn fmt_indented(&self, f: &mut Formatter<'_>, indentation: usize) -> fmt::Result {
        write!(f, "WHERE ")?;
        // "WHERE " = 6 chars
        self.condition.fmt_indented(f, indentation + 6)
    }
}

impl SelectCommand {
    pub fn fmt_indented(&self, f: &mut Formatter<'_>, indentation: usize) -> fmt::Result {
        write!(f, "SELECT ")?;
        // "SELECT " = 7 chars
        write_clause_items(
            f,
            &self.clauses,
            indentation,
            7,
            sum_subelements_assignment(&self.clauses),
            |f, c, ind| c.fmt_indented(f, ind),
        )
    }
}

impl DropCommand {
    pub fn fmt_indented(&self, f: &mut Formatter<'_>, indentation: usize) -> fmt::Result {
        write!(f, "DROP ")?;
        // "DROP " = 5 chars
        write_clause_items(
            f,
            &self.selections,
            indentation,
            5,
            self.selections.len(),
            |f, s, _ind| write!(f, "{}", s),
        )
    }
}

impl LimitCommand {
    pub fn fmt_indented(&self, f: &mut Formatter<'_>, indentation: usize) -> fmt::Result {
        write!(f, "LIMIT ")?;
        self.count.fmt_indented(f, indentation + 6)
    }
}

impl TrimstringsCommand {
    pub fn fmt_indented(&self, f: &mut Formatter<'_>, indentation: usize) -> fmt::Result {
        write!(f, "TRIMSTRINGS ")?;
        self.byte_len.fmt_indented(f, indentation + 12)
    }
}

impl WithinCommand {
    pub fn fmt_indented(&self, f: &mut Formatter<'_>, indentation: usize) -> fmt::Result {
        write!(f, "WITHIN ")?;
        self.duration.fmt_indented(f, indentation + 7)
    }
}

impl SortCommand {
    pub fn fmt_indented(&self, f: &mut Formatter<'_>, indentation: usize) -> fmt::Result {
        write!(f, "SORT ")?;
        // "SORT " = 5 chars
        write_clause_items(
            f,
            &self.expressions,
            indentation,
            5,
            sum_subelements_sort(&self.expressions),
            |f, e, ind| e.fmt_indented(f, ind),
        )
    }
}

impl FromCommand {
    pub fn fmt_indented(&self, f: &mut Formatter<'_>, indentation: usize) -> fmt::Result {
        write!(f, "FROM ")?;
        // "FROM " = 5 chars
        write_clause_items(
            f,
            &self.clauses,
            indentation,
            5,
            self.clauses.len(),
            |f, c, ind| c.fmt_indented(f, ind),
        )
    }
}

impl UnionCommand {
    pub fn fmt_indented(&self, f: &mut Formatter<'_>, indentation: usize) -> fmt::Result {
        write!(f, "UNION ")?;
        // "UNION " = 6 chars
        write_clause_items(
            f,
            &self.clauses,
            indentation,
            6,
            self.clauses.len(),
            |f, c, ind| c.fmt_indented(f, ind),
        )
    }
}

impl ParseCommand {
    pub fn fmt_indented(&self, f: &mut Formatter<'_>, indentation: usize) -> fmt::Result {
        write!(f, "PARSE ")?;
        let inner_indent = indentation + 6;
        if let Some(source) = &self.source {
            source.fmt_indented(f, inner_indent)?;
            write!(f, " ")?;
        }
        self.pattern.fmt_indented(f, inner_indent)?;
        if !self.identifiers.is_empty() {
            write!(f, " AS ")?;
            for (i, ident) in self.identifiers.iter().enumerate() {
                if i > 0 {
                    write!(f, ", ")?;
                }
                write!(f, "{}", ident)?;
            }
        }
        if self.nodrop {
            write!(f, " NODROP")?;
        }
        Ok(())
    }
}

/// Helper to write AGG-style sub-clauses (BY, SORT, WITHIN) on their own lines.
///
/// `base_indent` is the indentation where sub-clause keywords should start
/// (typically `indentation + 2` to align past `| `).
fn write_agg_subclauses(
    f: &mut Formatter<'_>,
    base_indent: usize,
    multiline: bool,
    group_by: &[Arc<super::clause::AssignmentClause>],
    sort: &[Arc<SortExpression>],
    within: Option<&Arc<Expression>>,
) -> fmt::Result {
    if !group_by.is_empty() {
        if multiline {
            writeln!(f)?;
            pad(f, base_indent)?;
        } else {
            write!(f, " ")?;
        }
        write!(f, "BY ")?;
        write_clause_items(
            f,
            group_by,
            base_indent,
            3,
            sum_subelements_assignment(group_by),
            |f, c, ind| c.fmt_indented(f, ind),
        )?;
    }
    if !sort.is_empty() {
        if multiline {
            writeln!(f)?;
            pad(f, base_indent)?;
        } else {
            write!(f, " ")?;
        }
        write!(f, "SORT ")?;
        write_clause_items(
            f,
            sort,
            base_indent,
            5,
            sum_subelements_sort(sort),
            |f, e, ind| e.fmt_indented(f, ind),
        )?;
    }
    if let Some(within_expr) = within {
        if multiline {
            writeln!(f)?;
            pad(f, base_indent)?;
        } else {
            write!(f, " ")?;
        }
        write!(f, "WITHIN ")?;
        within_expr.fmt_indented(f, base_indent + 7)?;
    }
    Ok(())
}

impl DistinctCommand {
    pub fn fmt_indented(&self, f: &mut Formatter<'_>, indentation: usize) -> fmt::Result {
        write!(f, "DISTINCT ")?;
        // "DISTINCT " = 9 chars
        write_clause_items(
            f,
            &self.clauses,
            indentation,
            9,
            sum_subelements_assignment(&self.clauses),
            |f, c, ind| c.fmt_indented(f, ind),
        )
    }
}

impl SuppressCommand {
    pub fn fmt_indented(&self, f: &mut Formatter<'_>, indentation: usize) -> fmt::Result {
        write!(f, "SUPPRESS ")?;
        let multiline =
            self.interval.subelements() + sum_subelements_assignment(&self.group_by) > 4;
        self.interval.fmt_indented(f, indentation + 9)?;
        if !self.group_by.is_empty() {
            write_agg_subclauses(f, indentation, multiline, &self.group_by, &[], None)?;
        }
        Ok(())
    }
}

impl AggCommand {
    pub fn fmt_indented(&self, f: &mut Formatter<'_>, indentation: usize) -> fmt::Result {
        let total = sum_subelements_assignment(&self.clauses)
            + sum_subelements_assignment(&self.group_by)
            + sum_subelements_sort(&self.sort);
        let multiline = total > 4;

        write!(f, "AGG ")?;
        // "AGG " = 4 chars
        write_clause_items(
            f,
            &self.clauses,
            indentation,
            4,
            sum_subelements_assignment(&self.clauses),
            |f, c, ind| c.fmt_indented(f, ind),
        )?;

        write_agg_subclauses(f, indentation, multiline, &self.group_by, &self.sort, None)
    }
}

impl WindowCommand {
    pub fn fmt_indented(&self, f: &mut Formatter<'_>, indentation: usize) -> fmt::Result {
        let total = sum_subelements_assignment(&self.clauses)
            + sum_subelements_assignment(&self.group_by)
            + sum_subelements_sort(&self.sort)
            + self.within.as_ref().map_or(0, |w| w.subelements() + 1);
        let multiline = total > 4;

        write!(f, "WINDOW ")?;
        // "WINDOW " = 7 chars
        write_clause_items(
            f,
            &self.clauses,
            indentation,
            7,
            sum_subelements_assignment(&self.clauses),
            |f, c, ind| c.fmt_indented(f, ind),
        )?;

        write_agg_subclauses(
            f,
            indentation,
            multiline,
            &self.group_by,
            &self.sort,
            self.within.as_ref(),
        )
    }
}

impl AppendCommand {
    pub fn fmt_indented(&self, f: &mut Formatter<'_>, _indentation: usize) -> fmt::Result {
        write!(f, "APPEND {}", self.table)?;
        if !self.distinct_by.is_empty() {
            write!(f, " DISTINCT ")?;
            for (i, sel) in self.distinct_by.iter().enumerate() {
                if i > 0 {
                    write!(f, ", ")?;
                }
                write!(f, "{}", sel)?;
            }
        }
        Ok(())
    }
}

impl JoinCommand {
    pub fn fmt_indented(&self, f: &mut Formatter<'_>, indentation: usize) -> fmt::Result {
        write!(f, "JOIN ")?;
        self.other.fmt_indented(f, indentation + 5)?;
        if let Some(on) = &self.on_condition {
            write!(f, " ON ")?;
            on.fmt_indented(f, indentation + 5)?;
        }
        Ok(())
    }
}

impl LookupCommand {
    pub fn fmt_indented(&self, f: &mut Formatter<'_>, indentation: usize) -> fmt::Result {
        write!(f, "LOOKUP ")?;
        self.other.fmt_indented(f, indentation + 7)?;
        if let Some(on) = &self.on_condition {
            write!(f, " ON ")?;
            on.fmt_indented(f, indentation + 7)?;
        }
        Ok(())
    }
}

impl ExplodeCommand {
    pub fn fmt_indented(&self, f: &mut Formatter<'_>, indentation: usize) -> fmt::Result {
        write!(f, "EXPLODE ")?;
        // "EXPLODE " = 8 chars
        write_clause_items(
            f,
            &self.clauses,
            indentation,
            8,
            sum_subelements_assignment(&self.clauses),
            |f, c, ind| c.fmt_indented(f, ind),
        )
    }
}

impl UnnestCommand {
    pub fn fmt_indented(&self, f: &mut Formatter<'_>, indentation: usize) -> fmt::Result {
        write!(f, "UNNEST ")?;
        self.expression.fmt_indented(f, indentation + 7)
    }
}

impl RowsCommand {
    pub fn fmt_indented(&self, f: &mut Formatter<'_>, indentation: usize) -> fmt::Result {
        write!(f, "ROWS ")?;
        self.expression.fmt_indented(f, indentation + 5)
    }
}

impl NestCommand {
    pub fn fmt_indented(&self, f: &mut Formatter<'_>, _indentation: usize) -> fmt::Result {
        write!(f, "NEST {}", self.identifier)
    }
}

impl MatchCommand {
    pub fn fmt_indented(&self, f: &mut Formatter<'_>, indentation: usize) -> fmt::Result {
        let total = self.pattern.len()
            + sum_subelements_assignment(&self.agg)
            + sum_subelements_assignment(&self.group_by)
            + sum_subelements_sort(&self.sort)
            + self.within.as_ref().map_or(0, |w| w.subelements() + 1);
        let multiline = total > 4;

        write!(f, "MATCH ")?;
        // Patterns separated by commas
        for (i, pat) in self.pattern.iter().enumerate() {
            if i > 0 {
                write!(f, ", ")?;
            }
            pat.fmt_indented(f, indentation + 6)?;
        }

        // AGG sub-clause
        if !self.agg.is_empty() {
            if multiline {
                writeln!(f)?;
                pad(f, indentation)?;
            } else {
                write!(f, " ")?;
            }
            write!(f, "AGG ")?;
            write_clause_items(
                f,
                &self.agg,
                indentation,
                4,
                sum_subelements_assignment(&self.agg),
                |f, c, ind| c.fmt_indented(f, ind),
            )?;
        }

        write_agg_subclauses(
            f,
            indentation,
            multiline,
            &self.group_by,
            &self.sort,
            self.within.as_ref(),
        )
    }
}