hamelin_lib 0.13.0

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

use super::node::Span;

use derive_more::derive::{From, TryUnwrap};

use crate::{
    antlr::{
        hamelinparser::{
            AssignmentClauseContextAll, AssignmentClauseContextAttrs, AssignmentContextAll,
            AssignmentContextAttrs, FromClauseContextAll, FromClauseContextAttrs,
            SelectionContextAll, SelectionContextAttrs, SortExpressionContextAll,
            SortExpressionContextAttrs, TableAliasContextAttrs, TablePathContextAll,
            TablePathContextAttrs, TablePathSegmentContextAll, TablePathSegmentContextAttrs,
            TableReferenceContextAll, TableReferenceContextAttrs, TemplatePairContextAttrs,
        },
        interval,
    },
    err::TranslationError,
    tree::ast::{
        context::{FromCst, ParseContext},
        identifier::{CompoundIdentifier, Identifier},
    },
};

use super::{
    expression::{ErrorExpression, Expression},
    identifier::{ParsedIdentifier, ParsedSimpleIdentifier, SimpleIdentifier},
    node::Spannable,
};
use antlr_rust::tree::ParseTree;

// ============================================================================
// Assignment
// ============================================================================

#[derive(Debug, Clone)]
pub struct Assignment {
    pub span: Span,
    pub identifier: ParsedIdentifier,
    pub expression: Arc<Expression>,
}

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

impl Assignment {
    pub fn from_parts(identifier: ParsedIdentifier, expression: Expression) -> Self {
        Self {
            span: Span::NONE,
            identifier,
            expression: Arc::new(expression),
        }
    }
}

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

impl FromCst<Rc<AssignmentContextAll<'static>>> for Assignment {
    fn from_cst_with_context(
        cst: Rc<AssignmentContextAll<'static>>,
        ctx: &mut ParseContext,
    ) -> Self {
        let identifier = cst
            .identifier()
            .map(|id_cst| ParsedIdentifier::from_cst_with_context(id_cst, ctx))
            .unwrap_or_else(|| {
                ctx.error("missing identifier in assignment")
                    .at(cst.as_ref())
                    .emit()
                    .into()
            });

        let expression = if let Some(expr_ctx) = cst.expression() {
            Expression::from_cst_with_context(expr_ctx, ctx)
        } else {
            let err = ctx
                .error("missing expression in assignment")
                .at(cst.as_ref())
                .emit();
            Expression::from_kind(ErrorExpression { error: err })
        };

        Self {
            span: interval(cst.as_ref()).into(),
            identifier,
            expression: Arc::new(expression),
        }
    }
}

// ============================================================================
// AssignmentClause
// ============================================================================

/// An assignment clause is either an explicit assignment (identifier = expression)
/// or an error. Bare expressions are desugared to assignments at parse time
/// by inferring the identifier from the expression structure.
#[derive(Debug, Clone, From, PartialEq, TryUnwrap)]
pub enum AssignmentClause {
    Assignment(Assignment),
    Error(Arc<TranslationError>),
}

impl Spannable for AssignmentClause {
    fn span(&self) -> Option<RangeInclusive<usize>> {
        match self {
            AssignmentClause::Assignment(a) => a.span(),
            AssignmentClause::Error(_) => None,
        }
    }
}

impl FromCst<Rc<AssignmentClauseContextAll<'static>>> for AssignmentClause {
    fn from_cst_with_context(
        cst: Rc<AssignmentClauseContextAll<'static>>,
        ctx: &mut ParseContext,
    ) -> Self {
        if let Some(assignment_ctx) = cst.assignment() {
            Assignment::from_cst_with_context(assignment_ctx, ctx).into()
        } else if let Some(expr_ctx) = cst.expression() {
            // Bare expression: desugar to assignment by inferring identifier
            let expression = Expression::from_cst_with_context(expr_ctx.clone(), ctx);

            // Try to infer identifier from AST structure, fall back to source text
            // Strip backticks (Hamelin quoting syntax) so they don't become literal chars in the name
            let identifier = Identifier::infer_from_expression(&expression).unwrap_or_else(|| {
                SimpleIdentifier::new(expr_ctx.get_text().replace('`', "")).into()
            });

            Assignment::from_parts(identifier, expression).into()
        } else {
            ctx.error("invalid assignment clause")
                .at(cst.as_ref())
                .emit()
                .into()
        }
    }
}

// ============================================================================
// FromClause
// ============================================================================

#[derive(Debug, Clone)]
pub struct TableAlias {
    pub span: Span,
    pub alias: ParsedSimpleIdentifier,
    pub table: TableReference,
}

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

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

/// One segment of a table path (`catalog.schema.table` or `foo.${p}_logs`).
#[derive(Debug, Clone, PartialEq)]
pub struct TablePathSegment {
    pub span: Span,
    pub parts: Vec<TableSegmentPart>,
}

/// A piece of a table path segment: literal identifier text or `${param}`.
#[derive(Debug, Clone, PartialEq)]
pub enum TableSegmentPart {
    Text(ParsedSimpleIdentifier),
    Parameter(ParsedSimpleIdentifier),
}

/// Table path containing at least one template parameter or multi-part segment.
#[derive(Debug, Clone, PartialEq)]
pub struct TemplatedTablePath {
    pub span: Span,
    pub segments: Vec<TablePathSegment>,
}

#[derive(Debug, Clone, PartialEq, From)]
pub enum TableReferenceTarget {
    /// Statically known identifier that does not contain any `${...}` parameters.
    /// Resolvable against the catalog without substitution.
    Static(ParsedIdentifier),
    /// Path that contains at least one `${...}` parameter; must be substituted before
    /// it can be resolved to a single concrete identifier.
    Templated(TemplatedTablePath),
}

#[derive(Debug, Clone)]
pub struct TableReference {
    pub span: Span,
    pub target: TableReferenceTarget,
}

impl TableReference {
    /// Resolve to a statically-known catalog identifier, or fail for templated references.
    ///
    /// Use this when the caller fundamentally cannot operate on a templated reference (post-
    /// substitution lowering, normalization, single-table reflection). Code paths that *can* deal
    /// with templated references should match on `target` directly and call
    /// `expand_templated_table_path` for the templated arm.
    pub fn static_valid_ref(&self) -> Result<&Identifier, Arc<TranslationError>> {
        match &self.target {
            TableReferenceTarget::Static(id) => id.valid_ref(),
            TableReferenceTarget::Templated(_) => Err(Arc::new(TranslationError::msg(
                self,
                "templated table reference cannot be used here; substitute template parameters first",
            ))),
        }
    }

    /// Whether this reference contains `${...}` in the path.
    pub fn is_templated(&self) -> bool {
        matches!(self.target, TableReferenceTarget::Templated(_))
    }
}

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

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

// Shortcut From impls so callers can `.into()` directly to a target without
// hand-writing the nested `Static(ParsedIdentifier::Valid/Error(...))` chain.
impl From<Arc<TranslationError>> for TableReferenceTarget {
    fn from(err: Arc<TranslationError>) -> Self {
        TableReferenceTarget::Static(ParsedIdentifier::Error(err))
    }
}

impl From<Identifier> for TableReferenceTarget {
    fn from(id: Identifier) -> Self {
        TableReferenceTarget::Static(ParsedIdentifier::Valid(id))
    }
}

impl FromCst<Rc<TablePathSegmentContextAll<'static>>> for TablePathSegment {
    fn from_cst_with_context(
        cst: Rc<TablePathSegmentContextAll<'static>>,
        ctx: &mut ParseContext,
    ) -> Self {
        let mut parts: Vec<TableSegmentPart> = Vec::new();

        if let Some(tok) = cst.IDENTIFIER() {
            let mut si = SimpleIdentifier::with_arc_name(ctx.intern(&tok.get_text()));
            si.span = interval(tok.as_ref()).into();
            parts.push(TableSegmentPart::Text(si.into()));
        } else if let Some(tok) = cst.BACKQUOTED_IDENTIFIER() {
            parts.push(TableSegmentPart::Text(backquoted_to_parsed(
                &tok.get_text(),
                interval(tok.as_ref()).into(),
                ctx,
            )));
        }

        for pair in cst.templatePair_all() {
            let parsed = pair
                .simpleIdentifier()
                .map(|si| ParsedSimpleIdentifier::from_cst_with_context(si.as_ref(), ctx))
                .unwrap_or_else(|| {
                    ctx.error("missing template parameter name in table path")
                        .at(pair.as_ref())
                        .emit()
                        .into()
                });
            parts.push(TableSegmentPart::Parameter(parsed));

            if let Some(tok) = pair.IDENTIFIER() {
                let mut si = SimpleIdentifier::with_arc_name(ctx.intern(&tok.get_text()));
                si.span = interval(tok.as_ref()).into();
                parts.push(TableSegmentPart::Text(si.into()));
            } else if let Some(tok) = pair.BACKQUOTED_IDENTIFIER() {
                parts.push(TableSegmentPart::Text(backquoted_to_parsed(
                    &tok.get_text(),
                    interval(tok.as_ref()).into(),
                    ctx,
                )));
            }
        }

        TablePathSegment {
            span: interval(cst.as_ref()).into(),
            parts,
        }
    }
}

/// Parse a raw `BACKQUOTED_IDENTIFIER` token text into a [`ParsedSimpleIdentifier`], applying the
/// same empty-backtick rejection as the shared `simpleIdentifier` CST path and interning through
/// the parse context.
fn backquoted_to_parsed(text: &str, span: Span, ctx: &mut ParseContext) -> ParsedSimpleIdentifier {
    let inner = text.strip_prefix('`').and_then(|s| s.strip_suffix('`'));
    match inner {
        Some(inner) if inner.is_empty() => ctx
            .error("empty backtick identifier")
            .at(&span)
            .emit()
            .into(),
        Some(inner) => {
            let unescaped = inner.replace("``", "`");
            let mut si = SimpleIdentifier::with_arc_name(ctx.intern(&unescaped));
            si.span = span;
            si.into()
        }
        None => ctx
            .error("invalid backtick identifier")
            .at(&span)
            .emit()
            .into(),
    }
}

fn table_path_to_target(
    path_cst: Rc<TablePathContextAll<'static>>,
    ctx: &mut ParseContext,
) -> TableReferenceTarget {
    let segments: Vec<TablePathSegment> = path_cst
        .tablePathSegment_all()
        .into_iter()
        .map(|s| TablePathSegment::from_cst_with_context(s, ctx))
        .collect();

    let is_templated = segments.iter().any(|seg| {
        seg.parts
            .iter()
            .any(|p| matches!(p, TableSegmentPart::Parameter(_)))
    });

    if is_templated {
        return TemplatedTablePath {
            span: interval(path_cst.as_ref()).into(),
            segments,
        }
        .into();
    }

    let mut parts: Vec<SimpleIdentifier> = Vec::with_capacity(segments.len());
    for seg in &segments {
        let mut acc = String::new();
        for p in &seg.parts {
            if let TableSegmentPart::Text(psi) = p {
                match psi.valid_ref() {
                    Ok(si) => acc.push_str(si.name().as_ref()),
                    Err(e) => return ParsedIdentifier::Error(e.clone()).into(),
                }
            }
        }
        if acc.is_empty() {
            let err = ctx
                .error("empty table path segment")
                .at(path_cst.as_ref())
                .emit();
            return ParsedIdentifier::Error(err).into();
        }
        let mut si = SimpleIdentifier::new(acc);
        si.span = seg.span;
        parts.push(si);
    }

    let parsed_id = match parts.as_slice() {
        [] => {
            let err = ctx.error("empty table path").at(path_cst.as_ref()).emit();
            ParsedIdentifier::Error(err)
        }
        [single] => single.clone().into(),
        [first, second, rest @ ..] => {
            CompoundIdentifier::new(first.clone(), second.clone(), rest.to_vec()).into()
        }
    };

    parsed_id.into()
}

impl FromCst<Rc<TableReferenceContextAll<'static>>> for TableReference {
    fn from_cst_with_context(
        cst: Rc<TableReferenceContextAll<'static>>,
        ctx: &mut ParseContext,
    ) -> Self {
        let target = match cst.tablePath() {
            Some(p) => table_path_to_target(p, ctx),
            None => {
                let err = ctx.error("missing table path").at(cst.as_ref()).emit();
                ParsedIdentifier::Error(err).into()
            }
        };

        Self {
            span: interval(cst.as_ref()).into(),
            target,
        }
    }
}

#[derive(Debug, Clone, From, PartialEq, TryUnwrap)]
pub enum FromClause {
    TableAlias(Arc<TableAlias>),
    TableReference(Arc<TableReference>),
}

impl FromClause {
    /// Returns the underlying TableReference, cloning from TableAlias if needed.
    pub fn table_reference(&self) -> TableReference {
        match self {
            FromClause::TableReference(tr) => tr.as_ref().clone(),
            FromClause::TableAlias(ta) => ta.table.clone(),
        }
    }
}

impl Spannable for FromClause {
    fn span(&self) -> Option<RangeInclusive<usize>> {
        match self {
            FromClause::TableAlias(ta) => ta.span(),
            FromClause::TableReference(tr) => tr.span(),
        }
    }
}

impl FromCst<Rc<FromClauseContextAll<'static>>> for FromClause {
    fn from_cst_with_context(
        cst: Rc<FromClauseContextAll<'static>>,
        ctx: &mut ParseContext,
    ) -> Self {
        if let Some(table_alias_ctx) = cst.tableAlias() {
            let alias = table_alias_ctx
                .simpleIdentifier()
                .map(|x| ParsedSimpleIdentifier::from_cst_with_context(x.as_ref(), ctx))
                .unwrap_or_else(|| {
                    ctx.error("missing alias identifier")
                        .at(cst.as_ref())
                        .emit()
                        .into()
                });

            let table_ref_ctx = table_alias_ctx.tableReference();
            let table = match table_ref_ctx.as_ref() {
                Some(tr) => TableReference::from_cst_with_context(tr.clone(), ctx),
                None => TableReference {
                    span: Span::NONE,
                    target: ctx
                        .error("missing table reference")
                        .at(cst.as_ref())
                        .emit()
                        .into(),
                },
            };

            Arc::new(TableAlias {
                span: interval(table_alias_ctx.as_ref()).into(),
                alias,
                table,
            })
            .into()
        } else if let Some(table_ref_ctx) = cst.tableReference() {
            Arc::new(TableReference::from_cst_with_context(table_ref_ctx, ctx)).into()
        } else {
            Arc::new(TableReference {
                span: Span::NONE,
                target: ctx
                    .error("invalid from clause")
                    .at(cst.as_ref())
                    .emit()
                    .into(),
            })
            .into()
        }
    }
}

impl From<TableAlias> for FromClause {
    fn from(value: TableAlias) -> Self {
        Self::from(Arc::new(value))
    }
}

impl From<TableReference> for FromClause {
    fn from(value: TableReference) -> Self {
        Self::from(Arc::new(value))
    }
}

// ============================================================================
// SortExpression
// ============================================================================

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum SortOrder {
    Asc,
    Desc,
}

#[derive(Debug, Clone)]
pub struct SortExpression {
    pub span: Span,
    pub expression: Arc<Expression>,
    pub order: Option<SortOrder>,
}

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

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

impl FromCst<Rc<SortExpressionContextAll<'static>>> for SortExpression {
    fn from_cst_with_context(
        cst: Rc<SortExpressionContextAll<'static>>,
        ctx: &mut ParseContext,
    ) -> Self {
        let expression = if let Some(expr_ctx) = cst.expression() {
            Expression::from_cst_with_context(expr_ctx.clone(), ctx)
        } else {
            let err = ctx
                .error("missing expression in sort")
                .at(cst.as_ref())
                .emit();
            Expression::from_kind(ErrorExpression { error: err.clone() })
        };

        let order = if cst.ASC().is_some() {
            Some(SortOrder::Asc)
        } else if cst.DESC().is_some() {
            Some(SortOrder::Desc)
        } else {
            None
        };

        Self {
            span: interval(cst.as_ref()).into(),
            expression: Arc::new(expression),
            order,
        }
    }
}

// ============================================================================
// Selection
// ============================================================================

#[derive(Debug, Clone)]
pub struct Selection {
    pub span: Span,
    pub identifier: ParsedIdentifier,
}

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

impl Selection {
    pub fn from_identifier(identifier: ParsedIdentifier) -> Self {
        Self {
            span: Span::NONE,
            identifier,
        }
    }
}

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

impl FromCst<Rc<SelectionContextAll<'static>>> for Selection {
    fn from_cst_with_context(
        cst: Rc<SelectionContextAll<'static>>,
        ctx: &mut ParseContext,
    ) -> Self {
        let identifier = cst
            .identifier()
            .map(|id| ParsedIdentifier::from_cst_with_context(id, ctx))
            .unwrap_or_else(|| {
                ctx.error("missing identifier in selection")
                    .at(cst.as_ref())
                    .emit()
                    .into()
            });

        Self {
            span: interval(cst.as_ref()).into(),
            identifier,
        }
    }
}

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

impl Assignment {
    pub fn subelements(&self) -> usize {
        self.expression.subelements() + 1
    }

    pub fn fmt_indented(&self, f: &mut Formatter<'_>, indentation: usize) -> fmt::Result {
        write!(f, "{} = ", self.identifier)?;
        // For simple expressions, hang-indent after "name = " so continuations align.
        // For complex expressions, use the base indentation so block-indented
        // function calls don't get pushed too far right.
        let hang = indentation + self.identifier.to_string().len() + 3;
        let expr_indent = if self.expression.subelements() > 4 {
            indentation
        } else {
            hang
        };
        self.expression.fmt_indented(f, expr_indent)
    }
}

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

impl AssignmentClause {
    pub fn subelements(&self) -> usize {
        match self {
            AssignmentClause::Assignment(a) => a.subelements(),
            AssignmentClause::Error(_) => 0,
        }
    }

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

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

impl Display for TableReference {
    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
        match &self.target {
            TableReferenceTarget::Static(id) => write!(f, "{}", id),
            TableReferenceTarget::Templated(path) => {
                for (si, seg) in path.segments.iter().enumerate() {
                    if si > 0 {
                        write!(f, ".")?;
                    }
                    for part in &seg.parts {
                        match part {
                            TableSegmentPart::Text(t) => write!(f, "{}", t)?,
                            TableSegmentPart::Parameter(p) => {
                                write!(f, "$")?;
                                write!(f, "{{")?;
                                write!(f, "{}", p)?;
                                write!(f, "}}")?;
                            }
                        }
                    }
                }
                Ok(())
            }
        }
    }
}

impl Display for TableAlias {
    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
        write!(f, "{} = {}", self.alias, self.table)
    }
}

impl FromClause {
    pub fn fmt_indented(&self, f: &mut Formatter<'_>, _indentation: usize) -> fmt::Result {
        match self {
            FromClause::TableAlias(ta) => write!(f, "{}", ta),
            FromClause::TableReference(tr) => write!(f, "{}", tr),
        }
    }
}

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

impl Display for SortOrder {
    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
        match self {
            SortOrder::Asc => write!(f, "ASC"),
            SortOrder::Desc => write!(f, "DESC"),
        }
    }
}

impl SortExpression {
    pub fn subelements(&self) -> usize {
        self.expression.subelements() + 1
    }

    pub fn fmt_indented(&self, f: &mut Formatter<'_>, indentation: usize) -> fmt::Result {
        self.expression.fmt_indented(f, indentation)?;
        if let Some(order) = &self.order {
            write!(f, " {}", order)?;
        }
        Ok(())
    }
}

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

impl Display for Selection {
    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
        write!(f, "{}", self.identifier)
    }
}