use std::ops::RangeInclusive;
use std::rc::Rc;
use std::sync::Arc;

use ordermap::OrderMap;

use crate::tree::ast::identifier::ParsedSimpleIdentifier;
use crate::tree::ast::node::Spannable;
use crate::tree::typed_ast::environment::{TypeEnvironment, UnboundFieldReference};
use crate::{
    antlr::hamelinparser::ExpressionContextAll,
    err::TranslationError,
    func::{
        def::{
            BroadcastResolution, DirectResolution, FunctionDef, FunctionResolution,
            ParameterBinding,
        },
        registry::{FunctionBindingError, InvocationType},
    },
    tree::ast::expression::IntervalUnit,
    tree::ast::{
        expression::{
            ArrayLiteral, BinaryOperator, Cast, Expression, ExpressionKind, FieldLookup,
            FunctionCall, IndexAccess, Lambda, StructLiteral, TsTrunc, TsTruncUnit, TupleLiteral,
            UnaryPostfixOperator, UnaryPrefixOperator,
        },
        identifier::Identifier,
        ParseWithErrors,
    },
    tree::options::ExpressionTypeCheckOptions,
    type_check_expression,
    types::{
        array::Array, matcher::Matcher, range::Range, Type, CALENDAR_INTERVAL, INTERVAL, TIMESTAMP,
        UNKNOWN, VARIANT,
    },
};
use derive_more::derive::{From, TryUnwrap};

use std::fmt::{Display, Formatter};

use crate::tree::ast::identifier::SimpleIdentifier;
use crate::types::decimal_type::Decimal;
use crate::types::function::Function;
use crate::types::struct_type::Struct;
use crate::types::tuple::Tuple;
use crate::types::{BINARY, BOOLEAN, DOUBLE, INT, ROWS, STRING};

use super::context::ExpressionTranslationContext;

/// Algebra for folding over typed expression trees (catamorphism).
///
/// This trait defines how to combine the results of recursively processing
/// an expression's children into a result for the expression itself.
///
/// For tree-to-tree transformations that return untyped AST, use `MapExpressionAlgebra` instead.
pub trait ExpressionAlgebra<T> {
    /// Fold a function application with its folded arguments.
    fn apply(
        &mut self,
        node: &TypedApply,
        expr: &TypedExpression,
        children: ParameterBinding<T>,
    ) -> T;

    /// Fold a broadcast function application with its folded arguments.
    fn broadcast_apply(
        &mut self,
        node: &TypedBroadcastApply,
        expr: &TypedExpression,
        children: ParameterBinding<T>,
    ) -> T;

    /// Fold an array literal with its folded elements.
    fn array_literal(
        &mut self,
        node: &TypedArrayLiteral,
        expr: &TypedExpression,
        children: Vec<T>,
    ) -> T;

    /// Fold a tuple literal with its folded elements.
    fn tuple_literal(
        &mut self,
        node: &TypedTupleLiteral,
        expr: &TypedExpression,
        children: Vec<T>,
    ) -> T;

    /// Fold a struct literal with its folded field values.
    fn struct_literal(
        &mut self,
        node: &TypedStructLiteral,
        expr: &TypedExpression,
        children: Vec<(ParsedSimpleIdentifier, T)>,
    ) -> T;

    /// Fold a variant index access with its folded value.
    fn variant_index_access(
        &mut self,
        node: &TypedVariantIndexAccess,
        expr: &TypedExpression,
        child: T,
    ) -> T;

    /// Fold a field lookup with its folded value.
    fn field_lookup(&mut self, node: &TypedFieldLookup, expr: &TypedExpression, child: T) -> T;

    /// Fold a cast with its folded value.
    fn cast(&mut self, node: &TypedCast, expr: &TypedExpression, child: T) -> T;

    /// Fold a timestamp truncation with its folded expression.
    fn ts_trunc(&mut self, node: &TypedTsTrunc, expr: &TypedExpression, child: T) -> T;

    /// Fold a lambda with its folded body.
    fn lambda(&mut self, node: &TypedLambda, expr: &TypedExpression, body: T) -> T;

    /// Fold a field reference (leaf node).
    fn field_reference(&mut self, node: &TypedFieldReference, expr: &TypedExpression) -> T;

    /// Fold a leaf expression (literals, etc.).
    fn leaf(&mut self, expr: &TypedExpression) -> T;

    /// Fold an error expression.
    fn error(&mut self, node: &TypedErrorExpression, expr: &TypedExpression) -> T;
}

/// Specialized algebra for tree-to-tree transformations that return untyped AST.
///
/// Unlike `ExpressionAlgebra<T>`, this trait has default implementations that
/// rebuild the AST node with transformed children. Override only the variants
/// you need to transform.
///
/// **Key properties:**
/// - Returns `Arc<Expression>` (untyped AST) - caller must re-typecheck
/// - Defaults rebuild the AST node, so you only override variants you care about
/// - Each method receives both the original typed node (for inspection) and
///   the transformed children (for rebuilding)
///
/// **Two sources of children:**
/// - `node` parameter (e.g., `&TypedArrayLiteral`) - contains **original** typed children
/// - `children` parameter (e.g., `Vec<Arc<Expression>>`) - contains **transformed** untyped children
///
/// Use `node` when you need to inspect original types. Use `children` when building the result.
#[allow(unused_variables)]
pub trait MapExpressionAlgebra {
    /// Transform a function application. Default: rebuild AST with transformed children.
    fn apply(
        &mut self,
        node: &TypedApply,
        expr: &TypedExpression,
        children: ParameterBinding<Arc<Expression>>,
    ) -> Arc<Expression> {
        node.replace_children_ast(expr, children)
    }

    /// Transform a broadcast function application. Default: rebuild AST with transformed children.
    fn broadcast_apply(
        &mut self,
        node: &TypedBroadcastApply,
        expr: &TypedExpression,
        children: ParameterBinding<Arc<Expression>>,
    ) -> Arc<Expression> {
        node.replace_children_ast(expr, children)
    }

    /// Transform an array literal. Default: rebuild AST with transformed children.
    fn array_literal(
        &mut self,
        node: &TypedArrayLiteral,
        expr: &TypedExpression,
        children: Vec<Arc<Expression>>,
    ) -> Arc<Expression> {
        node.replace_children_ast(expr, children)
    }

    /// Transform a tuple literal. Default: rebuild AST with transformed children.
    fn tuple_literal(
        &mut self,
        node: &TypedTupleLiteral,
        expr: &TypedExpression,
        children: Vec<Arc<Expression>>,
    ) -> Arc<Expression> {
        node.replace_children_ast(expr, children)
    }

    /// Transform a struct literal. Default: rebuild AST with transformed children.
    fn struct_literal(
        &mut self,
        node: &TypedStructLiteral,
        expr: &TypedExpression,
        children: Vec<(ParsedSimpleIdentifier, Arc<Expression>)>,
    ) -> Arc<Expression> {
        node.replace_children_ast(expr, children)
    }

    /// Transform a variant index access. Default: rebuild AST with transformed child.
    fn variant_index_access(
        &mut self,
        node: &TypedVariantIndexAccess,
        expr: &TypedExpression,
        child: Arc<Expression>,
    ) -> Arc<Expression> {
        node.replace_children_ast(expr, child)
    }

    /// Transform a field lookup. Default: rebuild AST with transformed child.
    fn field_lookup(
        &mut self,
        node: &TypedFieldLookup,
        expr: &TypedExpression,
        child: Arc<Expression>,
    ) -> Arc<Expression> {
        node.replace_children_ast(expr, child)
    }

    /// Transform a cast. Default: rebuild AST with transformed child.
    fn cast(
        &mut self,
        node: &TypedCast,
        expr: &TypedExpression,
        child: Arc<Expression>,
    ) -> Arc<Expression> {
        node.replace_children_ast(expr, child)
    }

    /// Transform a timestamp truncation. Default: rebuild AST with transformed child.
    fn ts_trunc(
        &mut self,
        node: &TypedTsTrunc,
        expr: &TypedExpression,
        child: Arc<Expression>,
    ) -> Arc<Expression> {
        node.replace_children_ast(expr, child)
    }

    /// Transform a lambda. Default: rebuild AST with transformed body.
    fn lambda(
        &mut self,
        node: &TypedLambda,
        expr: &TypedExpression,
        body: Arc<Expression>,
    ) -> Arc<Expression> {
        node.replace_children_ast(expr, body)
    }

    /// Transform a field reference (leaf). Default: keep original AST.
    fn field_reference(
        &mut self,
        node: &TypedFieldReference,
        expr: &TypedExpression,
    ) -> Arc<Expression> {
        expr.ast.clone()
    }

    /// Transform a leaf expression (literals). Default: keep original AST.
    fn leaf(&mut self, expr: &TypedExpression) -> Arc<Expression> {
        expr.ast.clone()
    }

    /// Transform an error expression. Default: keep original AST.
    fn error(&mut self, node: &TypedErrorExpression, expr: &TypedExpression) -> Arc<Expression> {
        expr.ast.clone()
    }
}

/// Blanket implementation: any `MapExpressionAlgebra` is an `ExpressionAlgebra<Arc<Expression>>`.
///
/// This allows using `expr.cata(alg)` with a `MapExpressionAlgebra`, eliminating the need
/// for a separate `map_to_ast` method.
impl<M: MapExpressionAlgebra> ExpressionAlgebra<Arc<Expression>> for M {
    fn apply(
        &mut self,
        node: &TypedApply,
        expr: &TypedExpression,
        children: ParameterBinding<Arc<Expression>>,
    ) -> Arc<Expression> {
        MapExpressionAlgebra::apply(self, node, expr, children)
    }

    fn broadcast_apply(
        &mut self,
        node: &TypedBroadcastApply,
        expr: &TypedExpression,
        children: ParameterBinding<Arc<Expression>>,
    ) -> Arc<Expression> {
        MapExpressionAlgebra::broadcast_apply(self, node, expr, children)
    }

    fn array_literal(
        &mut self,
        node: &TypedArrayLiteral,
        expr: &TypedExpression,
        children: Vec<Arc<Expression>>,
    ) -> Arc<Expression> {
        MapExpressionAlgebra::array_literal(self, node, expr, children)
    }

    fn tuple_literal(
        &mut self,
        node: &TypedTupleLiteral,
        expr: &TypedExpression,
        children: Vec<Arc<Expression>>,
    ) -> Arc<Expression> {
        MapExpressionAlgebra::tuple_literal(self, node, expr, children)
    }

    fn struct_literal(
        &mut self,
        node: &TypedStructLiteral,
        expr: &TypedExpression,
        children: Vec<(ParsedSimpleIdentifier, Arc<Expression>)>,
    ) -> Arc<Expression> {
        MapExpressionAlgebra::struct_literal(self, node, expr, children)
    }

    fn variant_index_access(
        &mut self,
        node: &TypedVariantIndexAccess,
        expr: &TypedExpression,
        child: Arc<Expression>,
    ) -> Arc<Expression> {
        MapExpressionAlgebra::variant_index_access(self, node, expr, child)
    }

    fn field_lookup(
        &mut self,
        node: &TypedFieldLookup,
        expr: &TypedExpression,
        child: Arc<Expression>,
    ) -> Arc<Expression> {
        MapExpressionAlgebra::field_lookup(self, node, expr, child)
    }

    fn cast(
        &mut self,
        node: &TypedCast,
        expr: &TypedExpression,
        child: Arc<Expression>,
    ) -> Arc<Expression> {
        MapExpressionAlgebra::cast(self, node, expr, child)
    }

    fn ts_trunc(
        &mut self,
        node: &TypedTsTrunc,
        expr: &TypedExpression,
        child: Arc<Expression>,
    ) -> Arc<Expression> {
        MapExpressionAlgebra::ts_trunc(self, node, expr, child)
    }

    fn lambda(
        &mut self,
        node: &TypedLambda,
        expr: &TypedExpression,
        body: Arc<Expression>,
    ) -> Arc<Expression> {
        MapExpressionAlgebra::lambda(self, node, expr, body)
    }

    fn field_reference(
        &mut self,
        node: &TypedFieldReference,
        expr: &TypedExpression,
    ) -> Arc<Expression> {
        MapExpressionAlgebra::field_reference(self, node, expr)
    }

    fn leaf(&mut self, expr: &TypedExpression) -> Arc<Expression> {
        MapExpressionAlgebra::leaf(self, expr)
    }

    fn error(&mut self, node: &TypedErrorExpression, expr: &TypedExpression) -> Arc<Expression> {
        MapExpressionAlgebra::error(self, node, expr)
    }
}

#[derive(Debug, Clone)]
pub struct TypedExpression {
    pub ast: Arc<Expression>,
    pub resolved_type: Arc<Type>,
    pub kind: TypedExpressionKind,
}

impl TypedExpression {
    /// Parse a Hamelin expression string into a TypedExpression with context
    pub fn parse_with_context<S: Into<String>>(
        s: S,
        ctx: &mut ExpressionTranslationContext,
    ) -> Self {
        let (ast, _) = Expression::parse_with_errors(s);
        Self::from_ast_with_context(Arc::new(ast), ctx)
    }

    /// Convert a parsed CST (Concrete Syntax Tree) into a TypedExpression with context
    pub fn from_cst_with_context(
        cst: Rc<ExpressionContextAll<'static>>,
        ctx: &mut ExpressionTranslationContext,
    ) -> Self {
        let ast = Expression::from(cst);
        Self::from_ast_with_context(Arc::new(ast), ctx)
    }

    pub fn from_ast_with_context(
        expr: Arc<Expression>,
        ctx: &mut ExpressionTranslationContext,
    ) -> Self {
        let interner = ctx.pipeline_ctx.interner.clone();
        match TypedExpressionKind::from_ast_with_context(&expr, ctx) {
            Ok((kind, resolved_type)) => Self {
                ast: expr,
                resolved_type: interner.intern_type(resolved_type),
                kind,
            },
            Err(error) => Self {
                ast: expr,
                resolved_type: interner.intern_type(UNKNOWN),
                kind: TypedErrorExpression { error }.into(),
            },
        }
    }

    /// Convert an AST expression with type validation against an expected type matcher
    ///
    /// This method type-checks the expression and validates that the resolved type
    /// matches the given matcher. If the type doesn't match, returns an Error variant
    /// with the type mismatch error embedded in the expression itself.
    ///
    /// This is useful for commands that require specific types (e.g., PARSE requires STRING).
    pub fn from_ast_with_expected_type(
        expr: Arc<Expression>,
        ctx: &mut ExpressionTranslationContext,
        expected: &impl Matcher,
    ) -> Self {
        // First, type-check the expression normally
        let typed = Self::from_ast_with_context(expr, ctx);

        // Check if the resolved type matches the expected matcher
        if !expected.matches(&typed.resolved_type) {
            // Type mismatch - create an error expression with the actual resolved type
            let error = ctx
                .error(format!(
                    "Expected type {}, got {}",
                    expected, typed.resolved_type
                ))
                .at(typed.ast.as_ref())
                .emit();

            Self {
                ast: typed.ast,
                resolved_type: typed.resolved_type, // Keep the actual resolved type
                kind: TypedErrorExpression { error }.into(),
            }
        } else {
            // Type matches - return the typed expression as-is
            typed
        }
    }

    /// Find the first expression node matching the predicate.
    ///
    /// Traverses the tree in pre-order, checking the current node first, then
    /// recursively checking children. Returns immediately when a match is found (early exit).
    pub fn find<F>(&self, predicate: &mut F) -> Option<&TypedExpression>
    where
        F: FnMut(&TypedExpression) -> bool,
    {
        // Check current node first (pre-order for early exit)
        if predicate(self) {
            return Some(self);
        }

        // Then check children, returning on first match
        // This match duplicates the fanout logic with fold(), but it's necessary
        // because find needs early-exit semantics
        match &self.kind {
            TypedExpressionKind::Apply(apply) => apply
                .parameter_binding
                .iter()
                .find_map(|arg| arg.find(predicate)),
            TypedExpressionKind::BroadcastApply(ba) => ba
                .parameter_binding
                .iter()
                .find_map(|arg| arg.find(predicate)),
            TypedExpressionKind::VariantIndexAccess(via) => via.value.find(predicate),
            TypedExpressionKind::ArrayLiteral(arr) => {
                arr.elements.iter().find_map(|elem| elem.find(predicate))
            }
            TypedExpressionKind::TupleLiteral(tuple) => {
                tuple.elements.iter().find_map(|elem| elem.find(predicate))
            }
            TypedExpressionKind::StructLiteral(struct_lit) => struct_lit
                .fields
                .iter()
                .map(|(_, v)| v)
                .find_map(|val| val.find(predicate)),
            TypedExpressionKind::FieldLookup(fl) => fl.value.find(predicate),
            TypedExpressionKind::Cast(cast) => cast.value.find(predicate),
            TypedExpressionKind::TsTrunc(ts) => ts.expression.find(predicate),
            TypedExpressionKind::Lambda(lambda) => lambda.body.find(predicate),
            TypedExpressionKind::FieldReference(_) => None,
            TypedExpressionKind::Leaf | TypedExpressionKind::Error(_) => None,
        }
    }

    /// Apply a catamorphism (algebra) to fold this expression tree bottom-up.
    ///
    /// Each node's children are folded first, then the algebra method for that
    /// node type is called with the original node and the folded children.
    pub fn cata<T>(&self, alg: &mut impl ExpressionAlgebra<T>) -> T {
        match &self.kind {
            TypedExpressionKind::Apply(apply) => {
                // Fold each argument, preserving the ParameterBinding structure
                let children = apply
                    .parameter_binding
                    .clone()
                    .try_map(|arg| Ok::<_, std::convert::Infallible>(arg.cata(alg)))
                    .unwrap();
                alg.apply(apply, self, children)
            }
            TypedExpressionKind::BroadcastApply(ba) => {
                let children = ba
                    .parameter_binding
                    .clone()
                    .try_map(|arg| Ok::<_, std::convert::Infallible>(arg.cata(alg)))
                    .unwrap();
                alg.broadcast_apply(ba, self, children)
            }
            TypedExpressionKind::ArrayLiteral(arr) => {
                let children = arr.elements.iter().map(|e| e.cata(alg)).collect();
                alg.array_literal(arr, self, children)
            }
            TypedExpressionKind::TupleLiteral(tuple) => {
                let children = tuple.elements.iter().map(|e| e.cata(alg)).collect();
                alg.tuple_literal(tuple, self, children)
            }
            TypedExpressionKind::StructLiteral(struct_lit) => {
                let children = struct_lit
                    .fields
                    .iter()
                    .map(|(name, expr)| (name.clone(), expr.cata(alg)))
                    .collect();
                alg.struct_literal(struct_lit, self, children)
            }
            TypedExpressionKind::VariantIndexAccess(via) => {
                let child = via.value.cata(alg);
                alg.variant_index_access(via, self, child)
            }
            TypedExpressionKind::FieldLookup(fl) => {
                let child = fl.value.cata(alg);
                alg.field_lookup(fl, self, child)
            }
            TypedExpressionKind::Cast(cast) => {
                let child = cast.value.cata(alg);
                alg.cast(cast, self, child)
            }
            TypedExpressionKind::TsTrunc(ts) => {
                let child = ts.expression.cata(alg);
                alg.ts_trunc(ts, self, child)
            }
            TypedExpressionKind::Lambda(lambda) => {
                let body = lambda.body.cata(alg);
                alg.lambda(lambda, self, body)
            }
            TypedExpressionKind::FieldReference(field_ref) => alg.field_reference(field_ref, self),
            TypedExpressionKind::Leaf => alg.leaf(self),
            TypedExpressionKind::Error(err) => alg.error(err, self),
        }
    }

    /// If this expression is a simple field reference or a chain of struct field lookups (e.g.
    /// `x` or `data.items`), returns the identifier naming the source value to consume after
    /// `UNNEST`. Returns `None` for other expressions or for non-struct field lookups (tuple
    /// indices, ranges, etc.).
    pub fn unnest_source_drop_identifier(&self) -> Option<Identifier> {
        match &self.kind {
            TypedExpressionKind::FieldReference(field_ref) => {
                let field_name = field_ref.field_name.valid_ref().ok()?;
                Some(field_name.clone().into())
            }
            TypedExpressionKind::FieldLookup(fl) => {
                let prefix = fl.value.unnest_source_drop_identifier()?;
                match &fl.access {
                    FieldAccess::StructField(field_id) => {
                        let field_simple = field_id.clone().valid().ok()?;
                        Some(field_simple.prepend(prefix).into())
                    }
                    _ => None,
                }
            }
            _ => None,
        }
    }
}

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

/// Whether a bare identifier resolved to a pipeline column or a scalar `DEF`.
///
/// Resolution order matches type-checking: pipeline/environment bindings win over scalar `DEF`s
/// when both exist for the same name.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub enum FieldReferenceResolution {
    /// Column (or other binding) from the current expression environment.
    Pipeline,
    /// Script-level scalar `DEF` (only when no pipeline binding shadows the name).
    ScalarDef,
}

/// A typed field reference
#[derive(Debug, Clone)]
pub struct TypedFieldReference {
    /// The field name being referenced
    pub field_name: ParsedSimpleIdentifier,
    /// How the name was resolved during type-checking (drives scalar `DEF` inlining in backends).
    pub resolution: FieldReferenceResolution,
}

#[derive(Debug, Clone, From, TryUnwrap)]
pub enum TypedExpressionKind {
    Leaf,
    FieldReference(TypedFieldReference),
    ArrayLiteral(TypedArrayLiteral),
    TupleLiteral(TypedTupleLiteral),
    StructLiteral(TypedStructLiteral),
    Apply(TypedApply),
    BroadcastApply(TypedBroadcastApply),
    VariantIndexAccess(TypedVariantIndexAccess),
    FieldLookup(TypedFieldLookup),
    Cast(TypedCast),
    TsTrunc(TypedTsTrunc),
    Lambda(TypedLambda),
    Error(TypedErrorExpression),
}

impl TypedExpressionKind {
    pub fn from_ast_with_context(
        expr: &Arc<Expression>,
        ctx: &mut ExpressionTranslationContext,
    ) -> Result<(Self, Type), Arc<TranslationError>> {
        match &expr.kind {
            // Simple literals - handle inline since they're trivial
            ExpressionKind::IntLiteral(_) => Ok((TypedExpressionKind::Leaf, INT)),
            ExpressionKind::ScientificLiteral(_) | ExpressionKind::DoubleLiteral(_) => {
                Ok((TypedExpressionKind::Leaf, DOUBLE))
            }
            ExpressionKind::BooleanLiteral(_) => Ok((TypedExpressionKind::Leaf, BOOLEAN)),
            ExpressionKind::StringLiteral(_) => Ok((TypedExpressionKind::Leaf, STRING)),
            ExpressionKind::BinaryLiteral(_) => Ok((TypedExpressionKind::Leaf, BINARY)),
            ExpressionKind::NullLiteral(_) => Ok((TypedExpressionKind::Leaf, UNKNOWN)),
            ExpressionKind::RowsLiteral(_) => Ok((TypedExpressionKind::Leaf, ROWS)),
            ExpressionKind::IntervalLiteral(lit) => {
                let interval_type = match lit.unit {
                    IntervalUnit::Month | IntervalUnit::Quarter | IntervalUnit::Year => {
                        CALENDAR_INTERVAL
                    }
                    _ => INTERVAL,
                };
                Ok((TypedExpressionKind::Leaf, interval_type))
            }
            ExpressionKind::UnboundRangeLiteral(_) => {
                Ok((TypedExpressionKind::Leaf, Range::new(UNKNOWN).into()))
            }

            ExpressionKind::DecimalLiteral(lit) => {
                let decimal_type = Type::Decimal(
                    Decimal::new(lit.precision as i32, lit.scale as i32).map_err(|e| {
                        ctx.error("invalid decimal literal")
                            .at(expr.as_ref())
                            .with_source_boxed(e.into())
                            .emit()
                    })?,
                );
                Ok((TypedExpressionKind::Leaf, decimal_type))
            }

            ExpressionKind::FieldReference(col_ref) => {
                // Extract valid identifier for TypeEnvironment lookup
                let simple_id = col_ref.field_name.clone().valid()?;
                let identifier: Identifier = simple_id.into();

                let (field_type, resolution) = if let Some(ft) = ctx.bindings.lookup(&identifier) {
                    (ft, FieldReferenceResolution::Pipeline)
                } else if let Some(ft) = ctx.pipeline_ctx.lookup_scalar_def_type(&identifier) {
                    (ft, FieldReferenceResolution::ScalarDef)
                } else {
                    let unbound = UnboundFieldReference {
                        field: col_ref.field_name.to_string(),
                        schema: ctx.bindings.as_struct().clone(),
                    };
                    return Err(ctx
                        .error("Field not found")
                        .at(expr.as_ref())
                        .with_source(unbound)
                        .emit());
                };

                Ok((
                    TypedFieldReference {
                        field_name: col_ref.field_name.clone(),
                        resolution,
                    }
                    .into(),
                    (*field_type).clone(),
                ))
            }

            ExpressionKind::TemplateParameter(p) => {
                let simple_id = p.name.clone().valid()?;
                let typ = ctx
                    .pipeline_ctx
                    .lookup_template_param_type(&simple_id)
                    .ok_or_else(|| {
                        ctx.error(format!(
                            "template parameter '{}' is not declared in template_parameters",
                            simple_id
                        ))
                        .at(expr.as_ref())
                        .emit()
                    })?;
                Ok((TypedExpressionKind::Leaf, typ))
            }

            // Delegate complex types to their impl methods
            ExpressionKind::ArrayLiteral(arr_lit) => {
                TypedArrayLiteral::from_ast_with_context(arr_lit, expr, ctx)
            }

            ExpressionKind::TupleLiteral(tuple_lit) => {
                TypedTupleLiteral::from_ast_with_context(tuple_lit, expr, ctx)
            }

            ExpressionKind::PairLiteral(pair_lit) => {
                // Convert pair literal to tuple literal with exactly 2 elements
                let typed_left = TypedExpression::from_ast_with_context(pair_lit.left.clone(), ctx);
                let typed_right =
                    TypedExpression::from_ast_with_context(pair_lit.right.clone(), ctx);

                let left_type = typed_left.resolved_type.as_ref().clone();
                let right_type = typed_right.resolved_type.as_ref().clone();
                let tuple_type = Tuple::new(vec![left_type, right_type]).into();

                Ok((
                    TypedTupleLiteral {
                        elements: vec![Arc::new(typed_left), Arc::new(typed_right)],
                    }
                    .into(),
                    tuple_type,
                ))
            }

            ExpressionKind::StructLiteral(struct_lit) => {
                TypedStructLiteral::from_ast_with_context(struct_lit, ctx)
            }

            ExpressionKind::UnaryPrefixOperator(op) => TypedApply::from_ast_args(
                &op.operator.to_string(),
                vec![op.operand.clone()],
                OrderMap::new(),
                InvocationType::UnaryPrefix,
                expr,
                ctx,
            ),

            ExpressionKind::UnaryPostfixOperator(op) => TypedApply::from_ast_args(
                &op.operator.to_string(),
                vec![op.operand.clone()],
                OrderMap::new(),
                InvocationType::UnaryPostfix,
                expr,
                ctx,
            ),

            ExpressionKind::BinaryOperator(op) => TypedApply::from_ast_args(
                &op.operator.to_string(),
                vec![op.left.clone(), op.right.clone()],
                OrderMap::new(),
                InvocationType::Binary,
                expr,
                ctx,
            ),

            ExpressionKind::FunctionCall(func_call) => {
                let positional_ast: Vec<_> = func_call.positional_args.clone();

                let named_ast: OrderMap<String, Arc<Expression>> = func_call
                    .named_args
                    .iter()
                    .filter_map(|(name, arg)| {
                        let name_str = name.valid_ref().ok()?.as_str().to_string();
                        Some((name_str, arg.clone()))
                    })
                    .collect();

                let function_name = func_call.name.valid_ref()?.as_str();
                TypedApply::from_ast_args(
                    function_name,
                    positional_ast,
                    named_ast,
                    InvocationType::Function,
                    expr.as_ref(),
                    ctx,
                )
            }

            ExpressionKind::IndexAccess(index_access) => {
                // Type check the value to determine how to handle the index
                let typed_value =
                    TypedExpression::from_ast_with_context(index_access.value.clone(), ctx);

                let value_type = typed_value.resolved_type.as_ref();

                match value_type {
                    Type::Variant => {
                        // Variant access requires a literal integer index - becomes TypedVariantIndexAccess
                        match &index_access.index.kind {
                            ExpressionKind::IntLiteral(int_lit) => {
                                if int_lit.int < 0 {
                                    return Err(ctx
                                        .error("Variant index must be non-negative")
                                        .at(index_access.index.as_ref())
                                        .emit());
                                }
                                Ok((
                                    TypedVariantIndexAccess {
                                        value: Arc::new(typed_value),
                                        variant_index: int_lit.int as usize,
                                    }
                                    .into(),
                                    VARIANT,
                                ))
                            }
                            _ => Err(ctx
                                .error("Variant index must be a literal integer")
                                .at(index_access.index.as_ref())
                                .emit()),
                        }
                    }
                    Type::Array(_) | Type::Map(_) => {
                        // Array indexing or map lookup becomes TypedApply with get()
                        // For arrays of maps, this may broadcast: arr["key"] -> get(map, "key") per element
                        let positional =
                            vec![index_access.value.clone(), index_access.index.clone()];
                        let named = OrderMap::new();

                        let registry = ctx.registry.clone();
                        let resolution = registry
                            .bind("get", positional, named, InvocationType::Function, ctx)
                            .map_err(|e| {
                                create_binding_error(index_access.value.as_ref(), e, ctx)
                            })?;

                        match resolution {
                            FunctionResolution::Direct(DirectResolution {
                                function_def,
                                binding,
                                typ,
                            }) => Ok((
                                TypedApply {
                                    function_def,
                                    parameter_binding: binding,
                                }
                                .into(),
                                typ,
                            )),
                            FunctionResolution::Broadcast(BroadcastResolution {
                                function_def,
                                binding,
                                typ,
                                broadcast_position,
                            }) => Ok((
                                TypedBroadcastApply {
                                    function_def,
                                    parameter_binding: binding,
                                    broadcast_position,
                                }
                                .into(),
                                typ,
                            )),
                        }
                    }
                    _ => Err(ctx
                        .error(format!(
                            "Cannot index into type: {}. Expected Array, Map, or Variant",
                            value_type
                        ))
                        .at(index_access.value.as_ref())
                        .emit()),
                }
            }

            ExpressionKind::FieldLookup(field_lookup) => {
                TypedFieldLookup::try_from_ast_with_context(field_lookup, ctx)
            }

            ExpressionKind::Cast(cast_expr) => TypedCast::from_ast_with_context(cast_expr, ctx),

            ExpressionKind::TsTrunc(ts_trunc) => {
                TypedTsTrunc::from_ast_with_context(ts_trunc, expr, ctx)
            }

            ExpressionKind::Error(err_expr) => Ok((
                TypedExpressionKind::Error(TypedErrorExpression {
                    error: err_expr.error.clone(),
                }),
                UNKNOWN,
            )),

            ExpressionKind::Lambda(lambda) => TypedLambda::from_ast_with_context(lambda, ctx),
        }
    }
}

// Helper function to create binding errors
fn create_binding_error(
    expr: &Expression,
    binding_error: FunctionBindingError,
    ctx: &mut ExpressionTranslationContext,
) -> Arc<TranslationError> {
    match binding_error {
        FunctionBindingError::FunctionNotFound { .. } => {
            ctx.error("no function by this name").at(expr).emit()
        }
        FunctionBindingError::NoMatchingSignature {
            function_name: _,
            positional,
            named,
            apply_failure,
        } => {
            // Type the AST expressions for error display to show types in error messages
            // Must do this before creating the builder to avoid borrow conflicts
            let typed_positional: Vec<_> = positional
                .iter()
                .map(|ast| TypedExpression::from_ast_with_context(ast.clone(), ctx))
                .collect();
            let typed_named: Vec<_> = named
                .iter()
                .map(|(_, ast)| TypedExpression::from_ast_with_context(ast.clone(), ctx))
                .collect();

            let mut builder = ctx
                .error("could not find a matching function definition")
                .at(expr)
                .with_source(apply_failure);

            for typed in typed_positional.iter().chain(typed_named.iter()) {
                if let Some(span) = Spannable::span(typed.ast.as_ref()) {
                    builder = builder.add_context(span, &format!("{}", typed.resolved_type));
                }
            }

            builder.emit()
        }
        FunctionBindingError::FunctionNameNotLowercase {
            function_name,
            suggestion,
        } => {
            let msg = match suggestion {
                Some(s) => format!(
                    "function names must be lowercase: '{}' should be '{}'",
                    function_name, s
                ),
                None => format!("function names must be lowercase: '{}'", function_name),
            };
            ctx.error(msg).at(expr).emit()
        }
        FunctionBindingError::AmbiguousBroadcast { function_name } => ctx
            .error(format!(
                "ambiguous broadcast for '{}': multiple array arguments could be broadcast",
                function_name
            ))
            .at(expr)
            .emit(),
        FunctionBindingError::Fatal(error) => ctx
            .error("fatal error during function binding")
            .at(expr)
            .with_source_boxed(error.into())
            .emit(),
    }
}

// Typed expression kind structs
#[derive(Debug, Clone)]
pub struct TypedArrayLiteral {
    pub elements: Vec<Arc<TypedExpression>>,
}

impl TypedArrayLiteral {
    pub fn from_ast_with_context(
        arr_lit: &ArrayLiteral,
        _expr: &Arc<Expression>,
        ctx: &mut ExpressionTranslationContext,
    ) -> Result<(TypedExpressionKind, Type), Arc<TranslationError>> {
        let mut typed_elements = Vec::new();
        let mut element_types = Vec::new();

        for element in &arr_lit.elements {
            let typed_element = TypedExpression::from_ast_with_context(element.clone(), ctx);
            element_types.push(typed_element.resolved_type.as_ref().clone());
            typed_elements.push(Arc::new(typed_element));
        }

        // Determine array element type by merging all element types
        let element_type = if element_types.is_empty() {
            UNKNOWN
        } else {
            let mut merged_type = element_types[0].clone();
            for (i, elt_type) in element_types.iter().enumerate().skip(1) {
                merged_type = merged_type.merge(elt_type.clone()).map_err(|e| {
                    ctx.error("could not determine array element type")
                        .at(arr_lit.elements[i].as_ref())
                        .with_source(e)
                        .emit()
                })?;
            }
            merged_type
        };

        let array_type = Array::new(element_type).into();
        Ok((
            Self {
                elements: typed_elements,
            }
            .into(),
            array_type,
        ))
    }

    /// Rebuild this array literal with new element expressions.
    ///
    /// Preserves the original expression's type and span.
    /// The AST is rebuilt to match the new typed structure.
    pub fn replace_children(
        &self,
        original_expr: &TypedExpression,
        new_elements: Vec<Arc<TypedExpression>>,
    ) -> Arc<TypedExpression> {
        Arc::new(TypedExpression {
            kind: TypedArrayLiteral {
                elements: new_elements.clone(),
            }
            .into(),
            ast: self.replace_children_ast(
                original_expr,
                new_elements.iter().map(|e| e.ast.clone()).collect(),
            ),
            resolved_type: original_expr.resolved_type.clone(),
        })
    }

    /// Rebuild this array literal's AST with new element ASTs.
    ///
    /// Preserves the original expression's span.
    pub fn replace_children_ast(
        &self,
        original_expr: &TypedExpression,
        new_elements: Vec<Arc<Expression>>,
    ) -> Arc<Expression> {
        Arc::new(Expression {
            span: original_expr.ast.span.clone(),
            kind: ArrayLiteral {
                elements: new_elements,
            }
            .into(),
        })
    }
}

#[derive(Debug, Clone)]
pub struct TypedTupleLiteral {
    pub elements: Vec<Arc<TypedExpression>>,
}

impl TypedTupleLiteral {
    pub fn from_ast_with_context(
        tuple_lit: &TupleLiteral,
        expr: &Arc<Expression>,
        ctx: &mut ExpressionTranslationContext,
    ) -> Result<(TypedExpressionKind, Type), Arc<TranslationError>> {
        let mut typed_elements = Vec::new();
        let mut element_types = Vec::new();

        for element in &tuple_lit.elements {
            let typed_element = TypedExpression::from_ast_with_context(element.clone(), ctx);
            element_types.push(typed_element.resolved_type.clone());
            typed_elements.push(Arc::new(typed_element));
        }

        if typed_elements.is_empty() {
            return Err(ctx
                .error("Tuple literal cannot be empty")
                .at(expr.as_ref())
                .emit());
        }

        let tuple_type: Type = Tuple {
            elements: element_types,
        }
        .into();
        Ok((
            Self {
                elements: typed_elements,
            }
            .into(),
            tuple_type,
        ))
    }

    /// Rebuild this tuple literal with new element expressions.
    ///
    /// Preserves the original expression's type and span.
    /// The AST is rebuilt to match the new typed structure.
    pub fn replace_children(
        &self,
        original_expr: &TypedExpression,
        new_elements: Vec<Arc<TypedExpression>>,
    ) -> Arc<TypedExpression> {
        Arc::new(TypedExpression {
            kind: TypedTupleLiteral {
                elements: new_elements.clone(),
            }
            .into(),
            ast: self.replace_children_ast(
                original_expr,
                new_elements.iter().map(|e| e.ast.clone()).collect(),
            ),
            resolved_type: original_expr.resolved_type.clone(),
        })
    }

    /// Rebuild this tuple literal's AST with new element ASTs.
    ///
    /// Preserves the original expression's span.
    pub fn replace_children_ast(
        &self,
        original_expr: &TypedExpression,
        new_elements: Vec<Arc<Expression>>,
    ) -> Arc<Expression> {
        Arc::new(Expression {
            span: original_expr.ast.span.clone(),
            kind: TupleLiteral {
                elements: new_elements,
            }
            .into(),
        })
    }
}

#[derive(Debug, Clone)]
pub struct TypedStructLiteral {
    pub fields: Vec<(ParsedSimpleIdentifier, Arc<TypedExpression>)>,
}

impl TypedStructLiteral {
    pub fn from_ast_with_context(
        struct_lit: &StructLiteral,
        ctx: &mut ExpressionTranslationContext,
    ) -> Result<(TypedExpressionKind, Type), Arc<TranslationError>> {
        let mut type_environment = TypeEnvironment::default();
        let mut fields = Vec::with_capacity(struct_lit.fields.len());
        let mut seen: Vec<SimpleIdentifier> = Vec::new();

        for (maybe_field_name, field_expr) in struct_lit.fields.iter() {
            let field_name = match maybe_field_name.clone().valid() {
                Ok(field) => field,
                Err(_) => continue,
            };

            if seen.contains(&field_name) {
                ctx.error(format!(
                    "duplicate field '{}' in struct literal",
                    field_name
                ))
                .at(maybe_field_name)
                .emit();
                continue;
            }
            seen.push(field_name.clone());

            let typed_expr = Arc::new(TypedExpression::from_ast_with_context(
                field_expr.clone(),
                ctx,
            ));
            type_environment.bind_arc(field_name.clone().into(), typed_expr.resolved_type.clone());
            fields.push((field_name.clone().into(), typed_expr));
        }

        Ok((
            Self { fields: fields }.into(),
            type_environment.as_struct().clone().into(),
        ))
    }

    /// Rebuild this struct literal with new field expressions.
    ///
    /// The new_fields must have the same field names as the original.
    /// Preserves the original expression's type and span.
    /// The AST is rebuilt to match the new typed structure.
    pub fn replace_children(
        &self,
        original_expr: &TypedExpression,
        new_fields: Vec<(ParsedSimpleIdentifier, Arc<TypedExpression>)>,
    ) -> Arc<TypedExpression> {
        Arc::new(TypedExpression {
            kind: TypedStructLiteral {
                fields: new_fields.clone(),
            }
            .into(),
            ast: self.replace_children_ast(
                original_expr,
                new_fields
                    .iter()
                    .map(|(name, expr)| (name.clone(), expr.ast.clone()))
                    .collect(),
            ),
            resolved_type: original_expr.resolved_type.clone(),
        })
    }

    /// Rebuild this struct literal's AST with new field ASTs.
    ///
    /// Preserves the original expression's span.
    pub fn replace_children_ast(
        &self,
        original_expr: &TypedExpression,
        new_fields: Vec<(ParsedSimpleIdentifier, Arc<Expression>)>,
    ) -> Arc<Expression> {
        Arc::new(Expression {
            span: original_expr.ast.span.clone(),
            kind: StructLiteral { fields: new_fields }.into(),
        })
    }
}

/// A typed function application node - represents all runtime function calls
/// including operators, regular functions, and array/map indexing.
///
/// After type checking, operators like `!x`, `x + y`, function calls like `foo(x, y)`,
/// and runtime indexing like `array[i]` are all semantically equivalent: they're function
/// applications with typed arguments and a known return type.
#[derive(Debug, Clone)]
pub struct TypedApply {
    /// The function being applied (operator, builtin, or user function)
    pub function_def: Arc<dyn FunctionDef>,
    /// The typed arguments bound to the function's parameters
    pub parameter_binding: ParameterBinding<Arc<TypedExpression>>,
}

impl TypedApply {
    /// Try to bind a function call, falling back to broadcast coercion if direct binding fails.
    ///
    /// Receives AST expressions and passes them to the registry for typing.
    /// Returns either a TypedApply (direct call) or TypedBroadcastApply (element-wise broadcast).
    pub fn from_ast_args(
        function_name: &str,
        positional: Vec<Arc<Expression>>,
        named: OrderMap<String, Arc<Expression>>,
        invocation_type: InvocationType,
        expr: &Expression,
        ctx: &mut ExpressionTranslationContext,
    ) -> Result<(TypedExpressionKind, Type), Arc<TranslationError>> {
        let registry = ctx.registry.clone();
        match registry.bind(
            function_name,
            positional.clone(),
            named.clone(),
            invocation_type.clone(),
            ctx,
        ) {
            Ok(resolution) => match resolution {
                FunctionResolution::Direct(DirectResolution {
                    function_def,
                    binding,
                    typ,
                }) => Ok((
                    TypedApply {
                        function_def,
                        parameter_binding: binding,
                    }
                    .into(),
                    typ,
                )),
                FunctionResolution::Broadcast(BroadcastResolution {
                    function_def,
                    binding,
                    typ,
                    broadcast_position,
                }) => Ok((
                    TypedBroadcastApply {
                        function_def,
                        parameter_binding: binding,
                        broadcast_position,
                    }
                    .into(),
                    typ,
                )),
            },
            Err(bind_error) => Err(create_binding_error(expr, bind_error, ctx)),
        }
    }

    /// Rebuild this function application with new argument expressions.
    ///
    /// Preserves the original expression's type, span, and function definition.
    /// The AST is rebuilt to match the new typed structure based on the original AST kind.
    pub fn replace_children(
        &self,
        original_expr: &TypedExpression,
        new_children: ParameterBinding<Arc<TypedExpression>>,
    ) -> Arc<TypedExpression> {
        let children_ast = new_children.clone().map(|e| e.ast.clone());

        Arc::new(TypedExpression {
            kind: TypedApply {
                parameter_binding: new_children,
                function_def: self.function_def.clone(),
            }
            .into(),
            ast: Self::replace_children_ast_inner(original_expr, &children_ast),
            resolved_type: original_expr.resolved_type.clone(),
        })
    }

    /// Rebuild this function application's AST with new argument ASTs.
    ///
    /// Preserves the original expression's span.
    /// The AST is rebuilt to match the structure based on the original AST kind.
    pub fn replace_children_ast(
        &self,
        original_expr: &TypedExpression,
        new_children: ParameterBinding<Arc<Expression>>,
    ) -> Arc<Expression> {
        Self::replace_children_ast_inner(original_expr, &new_children)
    }

    /// Internal helper for AST rebuilding shared by replace_children and replace_children_ast.
    ///
    /// Takes a ParameterBinding so we can look up children by both index and name,
    /// which is necessary to correctly rebuild both positional and named arguments
    /// in the AST.
    fn replace_children_ast_inner(
        original_expr: &TypedExpression,
        new_children: &ParameterBinding<Arc<Expression>>,
    ) -> Arc<Expression> {
        let new_ast = match &original_expr.ast.kind {
            ExpressionKind::BinaryOperator(bin_op) => Expression {
                span: original_expr.ast.span.clone(),
                kind: BinaryOperator {
                    left: new_children
                        .get_by_index(0)
                        .ok()
                        .cloned()
                        .unwrap_or_else(|| bin_op.left.clone()),
                    operator: bin_op.operator.clone(),
                    right: new_children
                        .get_by_index(1)
                        .ok()
                        .cloned()
                        .unwrap_or_else(|| bin_op.right.clone()),
                }
                .into(),
            },
            ExpressionKind::UnaryPrefixOperator(unary) => Expression {
                span: original_expr.ast.span.clone(),
                kind: UnaryPrefixOperator {
                    operator: unary.operator.clone(),
                    operand: new_children
                        .get_by_index(0)
                        .ok()
                        .cloned()
                        .unwrap_or_else(|| unary.operand.clone()),
                }
                .into(),
            },
            ExpressionKind::UnaryPostfixOperator(unary) => Expression {
                span: original_expr.ast.span.clone(),
                kind: UnaryPostfixOperator {
                    operand: new_children
                        .get_by_index(0)
                        .ok()
                        .cloned()
                        .unwrap_or_else(|| unary.operand.clone()),
                    operator: unary.operator.clone(),
                }
                .into(),
            },
            ExpressionKind::FunctionCall(fn_call) => Expression {
                span: original_expr.ast.span.clone(),
                kind: FunctionCall {
                    name: fn_call.name.clone(),
                    positional_args: (0..fn_call.positional_args.len())
                        .map(|i| {
                            new_children
                                .get_by_index(i)
                                .cloned()
                                .unwrap_or_else(|_| fn_call.positional_args[i].clone())
                        })
                        .collect(),
                    named_args: fn_call
                        .named_args
                        .iter()
                        .map(|(key, original_expr)| {
                            let new_expr = key
                                .valid_ref()
                                .ok()
                                .and_then(|k| new_children.get_by_name(k.as_str()).ok().cloned())
                                .unwrap_or_else(|| original_expr.clone());
                            (key.clone(), new_expr)
                        })
                        .collect(),
                }
                .into(),
            },
            ExpressionKind::IndexAccess(idx_access) => Expression {
                span: original_expr.ast.span.clone(),
                kind: IndexAccess {
                    value: new_children
                        .get_by_index(0)
                        .ok()
                        .cloned()
                        .unwrap_or_else(|| idx_access.value.clone()),
                    index: new_children
                        .get_by_index(1)
                        .ok()
                        .cloned()
                        .unwrap_or_else(|| idx_access.index.clone()),
                }
                .into(),
            },
            // For other AST kinds that become Apply, keep original AST
            _ => return original_expr.ast.clone(),
        };

        Arc::new(new_ast)
    }
}

/// A function application that requires broadcasting over an array argument.
///
/// This represents cases like `[1, 2, 3] * 10` where one argument is an array
/// and the operation should be applied element-wise. During lowering, this
/// becomes `transform(array, x -> function(x, other_args...))`.
///
/// Unlike `TypedApply`, this preserves the original AST structure while recording
/// the broadcast decision for later lowering. This keeps the TypedAST honest about
/// what the user wrote while still capturing the semantic transformation needed.
///
/// The resolved_type of the containing TypedExpression is `Array(element_result_type)`,
/// so the element result type can be recovered by unwrapping the array.
#[derive(Debug, Clone)]
pub struct TypedBroadcastApply {
    /// The function being broadcast (e.g., Add, Multiply, Upper)
    pub function_def: Arc<dyn FunctionDef>,
    /// The parameter binding with the original arguments as the user supplied them.
    /// The argument at `broadcast_position` is the array that will be iterated over.
    pub parameter_binding: ParameterBinding<Arc<TypedExpression>>,
    /// Which positional argument position contains the array being broadcast over.
    pub broadcast_position: usize,
}

impl TypedBroadcastApply {
    /// Rebuild this broadcast apply with new argument expressions.
    ///
    /// Preserves the original expression's type, span, function definition, and broadcast position.
    /// The AST is rebuilt to match the new typed structure based on the original AST kind.
    pub fn replace_children(
        &self,
        original_expr: &TypedExpression,
        new_children: ParameterBinding<Arc<TypedExpression>>,
    ) -> Arc<TypedExpression> {
        let children_ast = new_children.clone().map(|e| e.ast.clone());

        Arc::new(TypedExpression {
            kind: TypedBroadcastApply {
                parameter_binding: new_children,
                function_def: self.function_def.clone(),
                broadcast_position: self.broadcast_position,
            }
            .into(),
            ast: TypedApply::replace_children_ast_inner(original_expr, &children_ast),
            resolved_type: original_expr.resolved_type.clone(),
        })
    }

    /// Rebuild this broadcast apply's AST with new argument ASTs.
    ///
    /// Preserves the original expression's span.
    /// Uses the same AST rebuilding logic as TypedApply.
    pub fn replace_children_ast(
        &self,
        original_expr: &TypedExpression,
        new_children: ParameterBinding<Arc<Expression>>,
    ) -> Arc<Expression> {
        TypedApply::replace_children_ast_inner(original_expr, &new_children)
    }
}

/// A variant index access with compile-time constant index.
///
/// This is distinct from `Apply` because the index is lifted to a compile-time
/// constant during type checking, not evaluated at runtime.
///
/// Example: `union_value[0]` where `0` is extracted as a constant `usize`.
#[derive(Debug, Clone)]
pub struct TypedVariantIndexAccess {
    /// The union/variant value being indexed
    pub value: Arc<TypedExpression>,
    /// The compile-time constant variant index (lifted from a literal at type-check time)
    pub variant_index: usize,
}

impl TypedVariantIndexAccess {
    /// Rebuild this variant index access with a new value expression.
    ///
    /// Preserves the original expression's type, span, and variant index.
    /// The AST is rebuilt to match the new typed structure.
    pub fn replace_children(
        &self,
        original_expr: &TypedExpression,
        new_value: Arc<TypedExpression>,
    ) -> Arc<TypedExpression> {
        Arc::new(TypedExpression {
            kind: TypedVariantIndexAccess {
                value: new_value.clone(),
                variant_index: self.variant_index,
            }
            .into(),
            ast: self.replace_children_ast(original_expr, new_value.ast.clone()),
            resolved_type: original_expr.resolved_type.clone(),
        })
    }

    /// Rebuild this variant index access's AST with a new value AST.
    ///
    /// Preserves the original expression's span and index.
    pub fn replace_children_ast(
        &self,
        original_expr: &TypedExpression,
        new_value: Arc<Expression>,
    ) -> Arc<Expression> {
        if let ExpressionKind::IndexAccess(idx) = &original_expr.ast.kind {
            Arc::new(Expression {
                span: original_expr.ast.span.clone(),
                kind: IndexAccess {
                    value: new_value,
                    index: idx.index.clone(),
                }
                .into(),
            })
        } else {
            original_expr.ast.clone()
        }
    }
}

/// Describes what kind of field access this is, allowing for optimized evaluation
#[derive(Debug, Clone)]
pub enum FieldAccess {
    /// Regular struct field access
    StructField(ParsedSimpleIdentifier),
    /// Tuple element access by position (e.g., .f0, .f1)
    TupleElement(usize),
    /// Variant field access
    VariantField(ParsedSimpleIdentifier),
    /// Range .begin access (lower bound)
    RangeBegin,
    /// Range .end access (upper bound)
    RangeEnd,
    /// Broadcast struct-field access over an array: produced when the value
    /// is `Array<Struct>`. Every backend translates this to an element-wise
    /// struct field extraction without a transform/explode detour.
    BroadcastStructField(ParsedSimpleIdentifier),
    /// Broadcast tuple-element access over an array: produced when the value
    /// is `Array<Tuple>`.
    BroadcastTupleElement(usize),
    /// Broadcast variant-field access over an array: produced when the value
    /// is `Array<Variant>`.
    BroadcastVariantField(ParsedSimpleIdentifier),
}

#[derive(Debug, Clone)]
pub struct TypedFieldLookup {
    pub value: Arc<TypedExpression>,
    pub access: FieldAccess,
}

impl TypedFieldLookup {
    pub fn try_from_ast_with_context(
        field_lookup: &FieldLookup,
        ctx: &mut ExpressionTranslationContext,
    ) -> Result<(TypedExpressionKind, Type), Arc<TranslationError>> {
        let typed_value = TypedExpression::from_ast_with_context(field_lookup.value.clone(), ctx);

        let value_type = typed_value.resolved_type.as_ref();
        let field_id = &field_lookup.field_identifier;

        let (access, result_type) = match value_type {
            Type::Struct(struct_type) => {
                let field_type = struct_type
                    .lookup(&field_id.clone().valid()?.into())
                    .cloned()
                    .ok_or_else(|| {
                        ctx.error(format!("Field {} not found in struct", field_id))
                            .at(field_lookup.value.as_ref())
                            .emit()
                    })?;
                (FieldAccess::StructField(field_id.clone()), field_type)
            }
            Type::Tuple(tuple) => {
                let index = Self::parse_tuple_field_index(field_id, field_lookup, tuple, ctx)?;
                (
                    FieldAccess::TupleElement(index),
                    tuple.elements[index].as_ref().clone(),
                )
            }
            Type::Variant => (FieldAccess::VariantField(field_id.clone()), VARIANT),
            Type::Array(array) => Self::resolve_array_field(array, field_id, field_lookup, ctx)?,
            Type::Range(range) | Type::RangeInclusive(range) => {
                let field_name = field_id.valid_ref()?;
                let access = match field_name.as_str() {
                    "begin" => FieldAccess::RangeBegin,
                    "end" => FieldAccess::RangeEnd,
                    _ => {
                        return Err(ctx
                            .error(format!(
                                "Range type only has 'begin' and 'end' fields, got: {}",
                                field_name
                            ))
                            .at(&field_lookup.field_identifier)
                            .emit());
                    }
                };
                (access, range.of.as_ref().clone())
            }
            Type::Unknown => {
                // For error resilience, treat unknown as struct field
                (FieldAccess::StructField(field_id.clone()), UNKNOWN)
            }
            _ => {
                return Err(ctx
                    .error(format!(
                        "Cannot access field on type: {}. Expected Struct, Tuple, Array, Range, or Variant",
                        value_type
                    ))
                    .at(field_lookup.value.as_ref())
                    .emit());
            }
        };

        Ok((
            Self {
                value: Arc::new(typed_value),
                access,
            }
            .into(),
            result_type,
        ))
    }

    /// Resolve both the `FieldAccess` variant and the result type for a
    /// broadcast field lookup over an `Array<T>`. One match, one source of
    /// truth for "which element kinds support field access."
    fn resolve_array_field(
        array: &Array,
        field_id: &ParsedSimpleIdentifier,
        field_lookup: &FieldLookup,
        ctx: &mut ExpressionTranslationContext,
    ) -> Result<(FieldAccess, Type), Arc<TranslationError>> {
        match array.element_type.as_ref() {
            Type::Struct(inner_struct) => {
                let field_simple = field_id.clone().valid()?;
                let lookup_id = field_simple.into();
                let element_type = inner_struct.lookup(&lookup_id).cloned().ok_or_else(|| {
                    ctx.error(format!(
                        "Field {} not found in array element struct",
                        field_lookup.field_identifier
                    ))
                    .at(&field_lookup.field_identifier)
                    .emit()
                })?;
                Ok((
                    FieldAccess::BroadcastStructField(field_id.clone()),
                    Array::new(element_type).into(),
                ))
            }
            Type::Tuple(tuple) => {
                let index = Self::parse_tuple_field_index(field_id, field_lookup, tuple, ctx)?;
                let element_type = tuple.elements[index].as_ref().clone();
                Ok((
                    FieldAccess::BroadcastTupleElement(index),
                    Array::new(element_type).into(),
                ))
            }
            Type::Variant => Ok((
                FieldAccess::BroadcastVariantField(field_id.clone()),
                Array::new(VARIANT).into(),
            )),
            _ => Err(ctx
                .error(format!(
                    "Cannot access field on array element type: {}. Array elements must be Struct, Tuple, or Variant",
                    array.element_type
                ))
                .at(field_lookup.value.as_ref())
                .emit()),
        }
    }

    /// Parse a tuple field identifier (`f0`, `f1`, ...) into its index.
    ///
    /// Used for both direct tuple access and broadcast tuple access over
    /// `Array<Tuple>`. Validates that the name has an `fN` shape and that
    /// `N` is in bounds for the tuple's arity.
    fn parse_tuple_field_index(
        field_id: &ParsedSimpleIdentifier,
        field_lookup: &FieldLookup,
        tuple: &Tuple,
        ctx: &mut ExpressionTranslationContext,
    ) -> Result<usize, Arc<TranslationError>> {
        let field_name = field_id.valid_ref()?;
        let index = if let Some(rest) = field_name.as_str().strip_prefix('f') {
            rest.parse::<usize>().map_err(|_| {
                ctx.error(format!(
                    "Invalid tuple field: {}. Tuple fields must be 'f0', 'f1', 'f2', etc.",
                    field_id
                ))
                .at(&field_lookup.field_identifier)
                .emit()
            })?
        } else {
            return Err(ctx
                .error(format!(
                    "Invalid tuple field: {}. Tuple fields must be 'f0', 'f1', 'f2', etc.",
                    field_id
                ))
                .at(&field_lookup.field_identifier)
                .emit());
        };

        if index >= tuple.elements.len() {
            return Err(ctx
                .error(format!(
                    "Tuple index {} out of bounds. Tuple has {} elements",
                    index,
                    tuple.elements.len()
                ))
                .at(&field_lookup.field_identifier)
                .emit());
        }

        Ok(index)
    }

    /// Rebuild this field lookup with a new value expression.
    ///
    /// Preserves the original expression's type, span, and field access kind.
    /// The AST is rebuilt to match the new typed structure.
    pub fn replace_children(
        &self,
        original_expr: &TypedExpression,
        new_value: Arc<TypedExpression>,
    ) -> Arc<TypedExpression> {
        Arc::new(TypedExpression {
            kind: TypedFieldLookup {
                value: new_value.clone(),
                access: self.access.clone(),
            }
            .into(),
            ast: self.replace_children_ast(original_expr, new_value.ast.clone()),
            resolved_type: original_expr.resolved_type.clone(),
        })
    }

    /// Rebuild this field lookup's AST with a new value AST.
    ///
    /// Preserves the original expression's span and field identifier.
    pub fn replace_children_ast(
        &self,
        original_expr: &TypedExpression,
        new_value: Arc<Expression>,
    ) -> Arc<Expression> {
        if let ExpressionKind::FieldLookup(fl) = &original_expr.ast.kind {
            Arc::new(Expression {
                span: original_expr.ast.span.clone(),
                kind: FieldLookup {
                    value: new_value,
                    field_identifier: fl.field_identifier.clone(),
                }
                .into(),
            })
        } else {
            original_expr.ast.clone()
        }
    }
}

/// Describes how to cast to/from Variant for variant-compatible types.
///
/// Scalar variants here match Hamelin types that a variant can hold as-is
/// (Int/Double/Decimal/String/Boolean), plus Timestamp for JSON serialization,
/// and the additional scalar targets reachable from a string variant via the
/// double-cast rule (Timestamp only for `variant → T`).
/// Array/Struct/Map recurse on their inner types.
#[derive(Debug, Clone, PartialEq)]
pub enum VariantCastKind {
    Int,
    Double,
    /// Decimal target carries precision/scale so backend translators can build
    /// the correct typed SQL (`DECIMAL(p, s)`) rather than degrading to DOUBLE.
    Decimal {
        precision: i32,
        scale: i32,
    },
    String,
    Boolean,
    Timestamp,
    Unknown,
    Variant,
    Array(Box<VariantCastKind>),
    Struct(Vec<(std::string::String, VariantCastKind)>),
    Map(Box<VariantCastKind>, Box<VariantCastKind>),
}

impl VariantCastKind {
    /// Variant-compatible types for `T → variant` serialization.
    ///
    /// Timestamps serialize the same way as in JSON/`to_json_string`: as an ISO-8601
    /// string, including when nested inside structs, arrays, or maps.
    pub fn try_from_type(ty: &Type) -> Option<Self> {
        Self::try_from_type_with(ty, leaf_for_to_variant)
    }

    /// Variant-compatible target types for `variant → T` unboxing. Uses the same
    /// scalar leaf set as [`Self::try_from_type`]. Runtime unboxing to targets like
    /// `timestamp` still follows the boxed-type rule (e.g. variant-string parsed like
    /// a native string cast).
    pub fn try_from_target_type(ty: &Type) -> Option<Self> {
        Self::try_from_type_with(ty, leaf_for_to_variant)
    }

    /// Single recursive walk shared by `try_from_type` and `try_from_target_type`.
    /// `leaf` decides which scalar leaf types are accepted; Array/Struct/Map
    /// recursion is identical between the two callers.
    fn try_from_type_with(ty: &Type, leaf: fn(&Type) -> Option<VariantCastKind>) -> Option<Self> {
        if let Some(kind) = leaf(ty) {
            return Some(kind);
        }
        match ty {
            Type::Array(arr) => {
                let inner = Self::try_from_type_with(&arr.element_type, leaf)?;
                Some(VariantCastKind::Array(Box::new(inner)))
            }
            Type::Struct(s) => {
                let fields: Option<Vec<_>> = s
                    .iter()
                    .map(|(name, field_type)| {
                        let kind = Self::try_from_type_with(field_type, leaf)?;
                        Some((name.name().to_string(), kind))
                    })
                    .collect();
                Some(VariantCastKind::Struct(fields?))
            }
            Type::Map(m) => {
                let key = Self::try_from_type_with(&m.key_type, leaf)?;
                let value = Self::try_from_type_with(&m.value_type, leaf)?;
                Some(VariantCastKind::Map(Box::new(key), Box::new(value)))
            }
            _ => None,
        }
    }
}

/// Scalar leaves accepted in both `T → variant` and `variant → T` cast kinds
/// (the structural walk is the same; runtime unboxing rules differ per `CastKind`).
fn leaf_for_to_variant(ty: &Type) -> Option<VariantCastKind> {
    match ty {
        Type::Int => Some(VariantCastKind::Int),
        Type::Double => Some(VariantCastKind::Double),
        Type::Decimal(d) => Some(VariantCastKind::Decimal {
            precision: d.precision,
            scale: d.scale,
        }),
        Type::String => Some(VariantCastKind::String),
        Type::Boolean => Some(VariantCastKind::Boolean),
        Type::Timestamp => Some(VariantCastKind::Timestamp),
        Type::Unknown => Some(VariantCastKind::Unknown),
        Type::Variant => Some(VariantCastKind::Variant),
        _ => None,
    }
}

/// The kind of cast operation to perform
#[derive(Debug, Clone, PartialEq)]
pub enum CastKind {
    /// Identity cast - no conversion needed
    Identity,

    /// Cast a null literal to a specific type
    NullToType,

    // Numeric casts
    IntToDouble,
    IntToDecimal,
    DoubleToInt,
    DoubleToDecimal,
    DecimalToInt,
    DecimalToDouble,
    DecimalToDecimal,

    // Boolean casts
    IntToBoolean,
    BooleanToInt,
    BooleanToDouble,
    BooleanToDecimal,
    DoubleToBoolean,
    DecimalToBoolean,

    // String casts - to string
    ToStringFromInt,
    ToStringFromDouble,
    ToStringFromBoolean,
    ToStringFromTimestamp,
    ToStringFromBinary,
    ToStringFromDecimal,
    ToStringFromInterval,
    ToStringFromCalendarInterval,

    // String casts - from string
    StringToInt,
    StringToDouble,
    StringToBoolean,
    StringToTimestamp,
    StringToDecimal,

    // Variant casts
    ToVariant(VariantCastKind),
    FromVariant(VariantCastKind),

    // Collection casts
    ArrayElementCast(Box<CastKind>),
    TupleToStruct(Vec<(String, CastKind)>),

    /// Project/align a struct to a target struct schema.
    /// - Fields in target that exist in source (by name): extracted and cast
    /// - Fields in target that don't exist in source: NULL
    /// - Fields in source that don't exist in target: dropped (not an error)
    /// Each entry is `(field_name, cast_kind)` where:
    /// - `Identity` or other cast - field exists in source, apply the cast
    /// - `NullToType` - field doesn't exist in source, output null
    /// Field order matches target struct. Supports nested expansion via recursive CastKind.
    StructExpansion(Vec<(SimpleIdentifier, CastKind)>),

    // Range casts
    RangeElementCast(Box<CastKind>),
    IntervalToTimestampRange,
    TimestampToTimestampRange,
    IntervalRangeToTimestampRange,
}

#[derive(Debug, Clone)]
pub struct TypedCast {
    pub value: Arc<TypedExpression>,
    pub target_type: Type,
    pub cast_kind: CastKind,
}

impl TypedCast {
    pub fn from_ast_with_context(
        cast_expr: &Cast,
        ctx: &mut ExpressionTranslationContext,
    ) -> Result<(TypedExpressionKind, Type), Arc<TranslationError>> {
        let typed_value = TypedExpression::from_ast_with_context(cast_expr.expression.clone(), ctx);
        let target_type = (*cast_expr.target_type).clone();

        let cast_kind =
            CastKind::from_ast_with_context(&typed_value, &target_type).ok_or_else(|| {
                ctx.error(format!(
                    "Cannot cast from {} to {}",
                    typed_value.resolved_type, target_type
                ))
                .at(cast_expr.expression.as_ref())
                .emit()
            })?;

        Ok((
            Self {
                value: Arc::new(typed_value),
                target_type: target_type.clone(),
                cast_kind,
            }
            .into(),
            target_type,
        ))
    }

    /// Rebuild this cast with a new value expression.
    ///
    /// Preserves the original expression's type, span, target type, and cast kind.
    /// The AST is rebuilt to match the new typed structure.
    pub fn replace_children(
        &self,
        original_expr: &TypedExpression,
        new_value: Arc<TypedExpression>,
    ) -> Arc<TypedExpression> {
        Arc::new(TypedExpression {
            kind: TypedCast {
                value: new_value.clone(),
                target_type: self.target_type.clone(),
                cast_kind: self.cast_kind.clone(),
            }
            .into(),
            ast: self.replace_children_ast(original_expr, new_value.ast.clone()),
            resolved_type: original_expr.resolved_type.clone(),
        })
    }

    /// Rebuild this cast's AST with a new value AST.
    ///
    /// Preserves the original expression's span and target type.
    pub fn replace_children_ast(
        &self,
        original_expr: &TypedExpression,
        new_value: Arc<Expression>,
    ) -> Arc<Expression> {
        Arc::new(Expression {
            span: original_expr.ast.span.clone(),
            kind: Cast {
                expression: new_value,
                target_type: Arc::new(self.target_type.clone()),
            }
            .into(),
        })
    }
}

impl CastKind {
    /// Determine the cast kind from a typed expression to a target type.
    /// Returns None if the cast is not valid.
    pub fn from_ast_with_context(from: &TypedExpression, to: &Type) -> Option<CastKind> {
        Self::from_types(from.resolved_type.as_ref(), to)
    }

    /// Determine cast kind from types only (for recursive cases like array elements).
    /// Does not handle null literals since we don't have AST context.
    fn from_types(from: &Type, to: &Type) -> Option<CastKind> {
        if from == to {
            return Some(CastKind::Identity);
        }

        // Unknown can be cast to any type - covers null literals and error recovery
        if matches!(from, Type::Unknown) {
            return Some(CastKind::NullToType);
        }

        match (from, to) {
            (Type::Int, Type::Double) => Some(CastKind::IntToDouble),
            (Type::Int, Type::Decimal(_)) => Some(CastKind::IntToDecimal),
            (Type::Double, Type::Int) => Some(CastKind::DoubleToInt),
            (Type::Double, Type::Decimal(_)) => Some(CastKind::DoubleToDecimal),
            (Type::Decimal(_), Type::Int) => Some(CastKind::DecimalToInt),
            (Type::Decimal(_), Type::Double) => Some(CastKind::DecimalToDouble),
            (Type::Decimal(_), Type::Decimal(_)) => Some(CastKind::DecimalToDecimal),
            (Type::Int, Type::Boolean) => Some(CastKind::IntToBoolean),
            (Type::Boolean, Type::Int) => Some(CastKind::BooleanToInt),
            (Type::Boolean, Type::Double) => Some(CastKind::BooleanToDouble),
            (Type::Boolean, Type::Decimal(_)) => Some(CastKind::BooleanToDecimal),
            (Type::Double, Type::Boolean) => Some(CastKind::DoubleToBoolean),
            (Type::Decimal(_), Type::Boolean) => Some(CastKind::DecimalToBoolean),
            (Type::Int, Type::String) => Some(CastKind::ToStringFromInt),
            (Type::Double, Type::String) => Some(CastKind::ToStringFromDouble),
            (Type::Boolean, Type::String) => Some(CastKind::ToStringFromBoolean),
            (Type::Timestamp, Type::String) => Some(CastKind::ToStringFromTimestamp),
            (Type::Binary, Type::String) => Some(CastKind::ToStringFromBinary),
            (Type::Decimal(_), Type::String) => Some(CastKind::ToStringFromDecimal),
            (Type::Interval, Type::String) => Some(CastKind::ToStringFromInterval),
            (Type::CalendarInterval, Type::String) => Some(CastKind::ToStringFromCalendarInterval),
            (Type::String, Type::Int) => Some(CastKind::StringToInt),
            (Type::String, Type::Double) => Some(CastKind::StringToDouble),
            (Type::String, Type::Boolean) => Some(CastKind::StringToBoolean),
            (Type::String, Type::Timestamp) => Some(CastKind::StringToTimestamp),
            (Type::String, Type::Decimal(_)) => Some(CastKind::StringToDecimal),
            // Variant casts - only allowed for Variant-compatible types
            (_, Type::Variant) => {
                let kind = VariantCastKind::try_from_type(from)?;
                Some(CastKind::ToVariant(kind))
            }
            (Type::Variant, _) => {
                let kind = VariantCastKind::try_from_target_type(to)?;
                Some(CastKind::FromVariant(kind))
            }
            (Type::Array(from_arr), Type::Array(to_arr)) => {
                let element_cast_kind =
                    CastKind::from_types(&from_arr.element_type, &to_arr.element_type)?;
                Some(CastKind::ArrayElementCast(Box::new(element_cast_kind)))
            }
            (Type::Tuple(tuple), Type::Struct(struct_type)) => {
                if tuple.elements.len() != struct_type.len() {
                    return None;
                }
                let field_casts: Option<Vec<(String, CastKind)>> = struct_type
                    .iter()
                    .enumerate()
                    .map(|(i, (name, field_type))| {
                        let element_type = tuple.elements.get(i)?;
                        let cast_kind = CastKind::from_types(element_type, field_type)?;
                        Some((name.name().to_string(), cast_kind))
                    })
                    .collect();
                Some(CastKind::TupleToStruct(field_casts?))
            }
            (
                Type::Interval | Type::CalendarInterval,
                Type::Range(range) | Type::RangeInclusive(range),
            ) if matches!(range.of.as_ref(), Type::Timestamp) => {
                Some(CastKind::IntervalToTimestampRange)
            }
            (Type::Timestamp, Type::Range(range) | Type::RangeInclusive(range))
                if matches!(range.of.as_ref(), Type::Timestamp) =>
            {
                Some(CastKind::TimestampToTimestampRange)
            }
            (
                Type::Range(from_range) | Type::RangeInclusive(from_range),
                Type::Range(to_range) | Type::RangeInclusive(to_range),
            ) => {
                if matches!(
                    from_range.of.as_ref(),
                    Type::Interval | Type::CalendarInterval
                ) && matches!(to_range.of.as_ref(), Type::Timestamp)
                {
                    return Some(CastKind::IntervalRangeToTimestampRange);
                }
                let element_cast = CastKind::from_types(&from_range.of, &to_range.of)?;
                Some(CastKind::RangeElementCast(Box::new(element_cast)))
            }
            // Struct expansion/projection: cast source struct to target struct schema.
            // - Fields in target that exist in source (by name): use source value with cast
            // - Fields in target that don't exist in source: NULL
            // - Fields in source that don't exist in target: dropped
            (Type::Struct(from_struct), Type::Struct(to_struct)) => {
                // Build field casts in target field order
                let field_casts: Option<Vec<(SimpleIdentifier, CastKind)>> = to_struct
                    .iter()
                    .map(|(name, to_type)| {
                        match from_struct.lookup(name) {
                            Some(from_type) => {
                                // Field exists in source - compute cast kind
                                let cast_kind = CastKind::from_types(from_type, to_type)?;
                                Some((name.clone(), cast_kind))
                            }
                            None => {
                                // Field doesn't exist in source - use NullToType
                                Some((name.clone(), CastKind::NullToType))
                            }
                        }
                    })
                    .collect();

                Some(CastKind::StructExpansion(field_casts?))
            }
            _ => None,
        }
    }
}

#[derive(Debug, Clone)]
pub struct TypedTsTrunc {
    pub expression: Arc<TypedExpression>,
    pub unit: TsTruncUnit,
}

impl TypedTsTrunc {
    pub fn from_ast_with_context(
        ts_trunc: &TsTrunc,
        expr: &Arc<Expression>,
        ctx: &mut ExpressionTranslationContext,
    ) -> Result<(TypedExpressionKind, Type), Arc<TranslationError>> {
        if let TsTruncUnit::Template(param) = &ts_trunc.unit {
            let simple_id = param.name.clone().valid()?;
            let param_type = ctx
                .pipeline_ctx
                .lookup_template_param_type(&simple_id)
                .ok_or_else(|| {
                    ctx.error(format!(
                        "template parameter '{}' is not declared in template_parameters",
                        simple_id
                    ))
                    .at(expr.as_ref())
                    .emit()
                })?;
            if param_type != INTERVAL {
                return Err(ctx
                    .error(format!(
                        "template parameter '{}' in @ position must be declared with primitive type interval",
                        simple_id
                    ))
                    .at(expr.as_ref())
                    .emit());
            }
        }

        let typed_expr = TypedExpression::from_ast_with_context(ts_trunc.expression.clone(), ctx);
        let input_type = &typed_expr.resolved_type;

        // Validate input type
        if input_type.as_ref() != &TIMESTAMP
            && input_type.as_ref() != &INTERVAL
            && input_type.as_ref() != &CALENDAR_INTERVAL
        {
            return Err(ctx
                .error(format!(
                    "TsTrunc input must be TIMESTAMP, INTERVAL, or CALENDAR_INTERVAL, got: {}",
                    input_type
                ))
                .at(ts_trunc.expression.as_ref())
                .emit());
        }

        Ok((
            TypedExpressionKind::TsTrunc(TypedTsTrunc {
                expression: Arc::new(typed_expr),
                unit: ts_trunc.unit.clone(),
            }),
            TIMESTAMP,
        ))
    }

    /// Rebuild this timestamp truncation with a new expression.
    ///
    /// Preserves the original expression's type, span, and truncation unit.
    /// The AST is rebuilt to match the new typed structure.
    pub fn replace_children(
        &self,
        original_expr: &TypedExpression,
        new_expression: Arc<TypedExpression>,
    ) -> Arc<TypedExpression> {
        Arc::new(TypedExpression {
            kind: TypedTsTrunc {
                expression: new_expression.clone(),
                unit: self.unit.clone(),
            }
            .into(),
            ast: self.replace_children_ast(original_expr, new_expression.ast.clone()),
            resolved_type: original_expr.resolved_type.clone(),
        })
    }

    /// Rebuild this timestamp truncation's AST with a new expression AST.
    ///
    /// Preserves the original expression's span and truncation unit.
    pub fn replace_children_ast(
        &self,
        original_expr: &TypedExpression,
        new_expression: Arc<Expression>,
    ) -> Arc<Expression> {
        Arc::new(Expression {
            span: original_expr.ast.span.clone(),
            kind: TsTrunc {
                expression: new_expression,
                unit: self.unit.clone(),
            }
            .into(),
        })
    }
}

/// Fully typed lambda - all parameters have concrete types and body is type-checked.
#[derive(Debug, Clone)]
pub struct TypedLambda {
    pub parameters: Vec<TypedLambdaParameter>,
    pub body: Arc<TypedExpression>,
}

impl TypedLambda {
    /// Type-check a lambda expression from AST when all parameters have type hints.
    ///
    /// Returns `None` if any parameter is missing a type hint.
    /// Returns `Err` if the body fails to type-check.
    pub fn from_ast_with_context(
        lambda: &Lambda,
        ctx: &mut ExpressionTranslationContext,
    ) -> Result<(TypedExpressionKind, Type), Arc<TranslationError>> {
        // Check if all params have type hints
        if let Some(bad_param) = lambda.parameters.iter().find(|p| p.type_hint.is_none()) {
            return Err(ctx
                .error("Missing type hint for lambda parameter")
                .at(&bad_param.name)
                .emit());
        }

        // All hints present: type-check the body
        let typed_params: Vec<TypedLambdaParameter> = lambda
            .parameters
            .iter()
            .map(|p| TypedLambdaParameter {
                name: p.name.clone().into(),
                typ: p.type_hint.clone().unwrap_or_else(|| Arc::new(UNKNOWN)),
            })
            .collect();

        // Build lambda parameter bindings, then shadow outer env
        let mut lambda_params = Struct::default();
        for (param, ast_param) in typed_params.iter().zip(lambda.parameters.iter()) {
            let id: Identifier = param.name.clone().into();
            if lambda_params.lookup(&param.name).is_some() {
                return Err(ctx
                    .error(format!("Duplicate lambda parameter: {}", param.name))
                    .at(&ast_param.name)
                    .emit());
            }
            lambda_params.insert_arc(id, param.typ.clone());
        }
        let extended_env =
            TypeEnvironment::new(ctx.bindings.as_struct().prepend_overwrite(&lambda_params));

        // Type-check body with extended bindings and no specials (no aggregates in lambdas - D9)
        let mut body_ctx = ctx.without_specials(Arc::new(extended_env));
        let typed_body = TypedExpression::from_ast_with_context(lambda.body.clone(), &mut body_ctx);

        let function_type = Function {
            params: typed_params.iter().map(|p| p.typ.clone()).collect(),
            return_type: typed_body.resolved_type.clone(),
        };

        Ok((
            TypedLambda {
                parameters: typed_params,
                body: Arc::new(typed_body),
            }
            .into(),
            function_type.into(),
        ))
    }

    /// Rebuild this typed lambda with a new body expression.
    ///
    /// Preserves the original expression's type, span, and parameters.
    /// The AST is rebuilt to match the new typed structure.
    pub fn replace_children(
        &self,
        original_expr: &TypedExpression,
        new_body: Arc<TypedExpression>,
    ) -> Arc<TypedExpression> {
        Arc::new(TypedExpression {
            kind: TypedLambda {
                parameters: self.parameters.clone(),
                body: new_body.clone(),
            }
            .into(),
            ast: self.replace_children_ast(original_expr, new_body.ast.clone()),
            resolved_type: original_expr.resolved_type.clone(),
        })
    }

    /// Rebuild this typed lambda's AST with a new body AST.
    ///
    /// Preserves the original expression's span and parameters.
    pub fn replace_children_ast(
        &self,
        original_expr: &TypedExpression,
        new_body: Arc<Expression>,
    ) -> Arc<Expression> {
        if let ExpressionKind::Lambda(lambda) = &original_expr.ast.kind {
            Arc::new(Expression {
                span: original_expr.ast.span.clone(),
                kind: Lambda {
                    parameters: lambda.parameters.clone(),
                    body: new_body,
                }
                .into(),
            })
        } else {
            original_expr.ast.clone()
        }
    }
}

#[derive(Debug, Clone)]
pub struct TypedLambdaParameter {
    pub name: SimpleIdentifier,
    pub typ: Arc<Type>,
}

#[derive(Debug, Clone)]
pub struct TypedErrorExpression {
    pub error: Arc<TranslationError>,
}

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

/// Convert an AST Expression to a TypedExpression wrapped in Rc.
/// Uses type_check_expression with default options.
/// This is used for default parameter values in function definitions.
impl From<Expression> for Arc<TypedExpression> {
    fn from(expr: Expression) -> Self {
        Arc::new(type_check_expression(expr, ExpressionTypeCheckOptions::default()).output)
    }
}
