substrait-explain 0.7.0

use chrono::{DateTime, NaiveDate, NaiveDateTime, NaiveTime};
use substrait::proto::aggregate_rel::Measure;
use substrait::proto::expression::field_reference::{ReferenceType, RootReference, RootType};
use substrait::proto::expression::if_then::IfClause;
use substrait::proto::expression::literal::LiteralType;
use substrait::proto::expression::{
    Cast, FieldReference, IfThen, Literal, ReferenceSegment, RexType, ScalarFunction, cast,
    reference_segment,
};
use substrait::proto::function_argument::ArgType;
use substrait::proto::r#type::{Fp64, I64, Kind, Nullability};
use substrait::proto::{AggregateFunction, Expression, FunctionArgument, Type};

use super::types::get_and_validate_anchor;
use super::{
    MessageParseError, ParsePair, Rule, RuleIter, ScopedParsePair, unescape_string,
    unwrap_single_pair,
};
use crate::extensions::SimpleExtensions;
use crate::extensions::simple::{CompoundName, ExtensionKind};

/// A field index (e.g., parsed from "$0" -> 0).
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct FieldIndex(pub i32);

impl FieldIndex {
    /// Convert this field index to a FieldReference for use in expressions.
    pub fn to_field_reference(self) -> FieldReference {
        // XXX: Why is it so many layers to make a struct field reference? This is
        // surprisingly complex
        FieldReference {
            reference_type: Some(ReferenceType::DirectReference(ReferenceSegment {
                reference_type: Some(reference_segment::ReferenceType::StructField(Box::new(
                    reference_segment::StructField {
                        field: self.0,
                        child: None,
                    },
                ))),
            })),
            root_type: Some(RootType::RootReference(RootReference {})),
        }
    }
}

impl ParsePair for FieldIndex {
    fn rule() -> Rule {
        Rule::reference
    }

    fn message() -> &'static str {
        "FieldIndex"
    }

    fn parse_pair(pair: pest::iterators::Pair<Rule>) -> Self {
        assert_eq!(pair.as_rule(), Self::rule());
        let inner = unwrap_single_pair(pair);
        let index: i32 = inner.as_str().parse().unwrap();
        FieldIndex(index)
    }
}

impl ParsePair for FieldReference {
    fn rule() -> Rule {
        Rule::reference
    }

    fn message() -> &'static str {
        "FieldReference"
    }

    fn parse_pair(pair: pest::iterators::Pair<Rule>) -> Self {
        assert_eq!(pair.as_rule(), Self::rule());

        // TODO: Other types of references.
        FieldIndex::parse_pair(pair).to_field_reference()
    }
}

fn to_int_literal(
    value: pest::iterators::Pair<Rule>,
    typ: Option<Type>,
) -> Result<Literal, MessageParseError> {
    assert_eq!(value.as_rule(), Rule::integer);
    let parsed_value: i64 = value.as_str().parse().unwrap();

    const DEFAULT_KIND: Kind = Kind::I64(I64 {
        type_variation_reference: 0,
        nullability: Nullability::Required as i32,
    });

    // If no type is provided, we assume i64, Nullability::Required.
    let kind = typ.and_then(|t| t.kind).unwrap_or(DEFAULT_KIND);

    let (lit, nullability, tvar) = match &kind {
        // If no type is provided, we assume i64, Nullability::Required.
        Kind::I8(i) => (
            LiteralType::I8(parsed_value as i32),
            i.nullability,
            i.type_variation_reference,
        ),
        Kind::I16(i) => (
            LiteralType::I16(parsed_value as i32),
            i.nullability,
            i.type_variation_reference,
        ),
        Kind::I32(i) => (
            LiteralType::I32(parsed_value as i32),
            i.nullability,
            i.type_variation_reference,
        ),
        Kind::I64(i) => (
            LiteralType::I64(parsed_value),
            i.nullability,
            i.type_variation_reference,
        ),
        k => {
            return Err(MessageParseError::invalid(
                "int_literal_type",
                value.as_span(),
                format!("Invalid type for integer literal: {k:?}"),
            ));
        }
    };

    Ok(Literal {
        literal_type: Some(lit),
        nullable: nullability != Nullability::Required as i32,
        type_variation_reference: tvar,
    })
}

fn to_float_literal(
    value: pest::iterators::Pair<Rule>,
    typ: Option<Type>,
) -> Result<Literal, MessageParseError> {
    assert_eq!(value.as_rule(), Rule::float);
    let parsed_value: f64 = value.as_str().parse().unwrap();

    const DEFAULT_KIND: Kind = Kind::Fp64(Fp64 {
        type_variation_reference: 0,
        nullability: Nullability::Required as i32,
    });

    // If no type is provided, we assume fp64, Nullability::Required.
    let kind = typ.and_then(|t| t.kind).unwrap_or(DEFAULT_KIND);

    let (lit, nullability, tvar) = match &kind {
        Kind::Fp32(f) => (
            LiteralType::Fp32(parsed_value as f32),
            f.nullability,
            f.type_variation_reference,
        ),
        Kind::Fp64(f) => (
            LiteralType::Fp64(parsed_value),
            f.nullability,
            f.type_variation_reference,
        ),
        k => {
            return Err(MessageParseError::invalid(
                "float_literal_type",
                value.as_span(),
                format!("Invalid type for float literal: {k:?}"),
            ));
        }
    };

    Ok(Literal {
        literal_type: Some(lit),
        nullable: nullability != Nullability::Required as i32,
        type_variation_reference: tvar,
    })
}

fn to_boolean_literal(
    value: pest::iterators::Pair<Rule>,
    typ: Option<Type>,
) -> Result<Literal, MessageParseError> {
    assert_eq!(value.as_rule(), Rule::boolean);
    let parsed_value: bool = value.as_str().parse().unwrap();

    let (nullable, tvar) = match typ.and_then(|t| t.kind) {
        Some(Kind::Bool(b)) => (
            b.nullability != Nullability::Required as i32,
            b.type_variation_reference,
        ),
        None => (false, 0),
        Some(k) => {
            return Err(MessageParseError::invalid(
                "bool_literal_type",
                value.as_span(),
                format!("Invalid type for boolean literal: {k:?}"),
            ));
        }
    };

    Ok(Literal {
        literal_type: Some(LiteralType::Boolean(parsed_value)),
        nullable,
        type_variation_reference: tvar,
    })
}

fn to_string_literal(
    value: pest::iterators::Pair<Rule>,
    typ: Option<Type>,
) -> Result<Literal, MessageParseError> {
    assert_eq!(value.as_rule(), Rule::string_literal);
    let string_value = unescape_string(value.clone());

    // If no type is provided, default to string
    let Some(typ) = typ else {
        return Ok(Literal {
            literal_type: Some(LiteralType::String(string_value)),
            nullable: false,
            type_variation_reference: 0,
        });
    };

    let Some(kind) = typ.kind else {
        return Ok(Literal {
            literal_type: Some(LiteralType::String(string_value)),
            nullable: false,
            type_variation_reference: 0,
        });
    };

    match &kind {
        Kind::Date(d) => {
            // Parse date in ISO 8601 format: YYYY-MM-DD
            let date_days = parse_date_to_days(&string_value, value.as_span())?;
            Ok(Literal {
                literal_type: Some(LiteralType::Date(date_days)),
                nullable: d.nullability != Nullability::Required as i32,
                type_variation_reference: d.type_variation_reference,
            })
        }
        #[allow(deprecated)]
        Kind::Time(t) => {
            // Parse time in ISO 8601 format: HH:MM:SS[.fff]
            let time_microseconds = parse_time_to_microseconds(&string_value, value.as_span())?;
            Ok(Literal {
                literal_type: Some(LiteralType::Time(time_microseconds)),
                nullable: t.nullability != Nullability::Required as i32,
                type_variation_reference: t.type_variation_reference,
            })
        }
        #[allow(deprecated)]
        Kind::Timestamp(ts) => {
            // Parse timestamp in ISO 8601 format: YYYY-MM-DDTHH:MM:SS[.fff] or YYYY-MM-DD HH:MM:SS[.fff]
            let timestamp_microseconds =
                parse_timestamp_to_microseconds(&string_value, value.as_span())?;
            Ok(Literal {
                literal_type: Some(LiteralType::Timestamp(timestamp_microseconds)),
                nullable: ts.nullability != Nullability::Required as i32,
                type_variation_reference: ts.type_variation_reference,
            })
        }
        _ => {
            // For other types, treat as string
            Ok(Literal {
                literal_type: Some(LiteralType::String(string_value)),
                nullable: false,
                type_variation_reference: 0,
            })
        }
    }
}

fn to_null_literal(
    value: pest::iterators::Pair<Rule>,
    typ: Option<Type>,
) -> Result<Literal, MessageParseError> {
    assert_eq!(value.as_rule(), Rule::null);
    let typ = typ.ok_or_else(|| {
        MessageParseError::invalid(
            "null_literal_type",
            value.as_span(),
            "Null literals require an explicit type annotation, e.g. null:i64?",
        )
    })?;

    Ok(Literal {
        literal_type: Some(LiteralType::Null(typ)),
        nullable: false,
        type_variation_reference: 0,
    })
}

/// Parse a date string using chrono to days since Unix epoch
fn parse_date_to_days(date_str: &str, span: pest::Span) -> Result<i32, MessageParseError> {
    // Try multiple date formats for flexibility
    let formats = ["%Y-%m-%d", "%Y/%m/%d"];

    for format in &formats {
        if let Ok(date) = NaiveDate::parse_from_str(date_str, format) {
            // Calculate days since Unix epoch (1970-01-01)
            let epoch = NaiveDate::from_ymd_opt(1970, 1, 1).unwrap();
            let days = date.signed_duration_since(epoch).num_days();
            return Ok(days as i32);
        }
    }

    Err(MessageParseError::invalid(
        "date_parse_format",
        span,
        format!("Invalid date format: '{date_str}'. Expected YYYY-MM-DD or YYYY/MM/DD"),
    ))
}

/// Parse a time string using chrono to microseconds since midnight
fn parse_time_to_microseconds(time_str: &str, span: pest::Span) -> Result<i64, MessageParseError> {
    // Try multiple time formats for flexibility
    let formats = ["%H:%M:%S%.f", "%H:%M:%S"];

    for format in &formats {
        if let Ok(time) = NaiveTime::parse_from_str(time_str, format) {
            // Convert to microseconds since midnight
            let midnight = NaiveTime::from_hms_opt(0, 0, 0).unwrap();
            let duration = time.signed_duration_since(midnight);
            return Ok(duration.num_microseconds().unwrap_or(0));
        }
    }

    Err(MessageParseError::invalid(
        "time_parse_format",
        span,
        format!("Invalid time format: '{time_str}'. Expected HH:MM:SS or HH:MM:SS.fff"),
    ))
}

/// Parse a timestamp string using chrono to microseconds since Unix epoch
fn parse_timestamp_to_microseconds(
    timestamp_str: &str,
    span: pest::Span,
) -> Result<i64, MessageParseError> {
    // Try multiple timestamp formats for flexibility
    let formats = [
        "%Y-%m-%dT%H:%M:%S%.f", // ISO 8601 with T and fractional seconds
        "%Y-%m-%dT%H:%M:%S",    // ISO 8601 with T
        "%Y-%m-%d %H:%M:%S%.f", // Space separator with fractional seconds
        "%Y-%m-%d %H:%M:%S",    // Space separator
        "%Y/%m/%dT%H:%M:%S%.f", // Alternative date format with T
        "%Y/%m/%dT%H:%M:%S",    // Alternative date format with T
        "%Y/%m/%d %H:%M:%S%.f", // Alternative date format with space
        "%Y/%m/%d %H:%M:%S",    // Alternative date format with space
    ];

    for format in &formats {
        if let Ok(datetime) = NaiveDateTime::parse_from_str(timestamp_str, format) {
            // Calculate microseconds since Unix epoch (1970-01-01 00:00:00)
            let epoch = DateTime::from_timestamp(0, 0).unwrap().naive_utc();
            let duration = datetime.signed_duration_since(epoch);
            return Ok(duration.num_microseconds().unwrap_or(0));
        }
    }

    Err(MessageParseError::invalid(
        "timestamp_parse_format",
        span,
        format!(
            "Invalid timestamp format: '{timestamp_str}'. Expected YYYY-MM-DDTHH:MM:SS or YYYY-MM-DD HH:MM:SS"
        ),
    ))
}

impl ScopedParsePair for Literal {
    fn rule() -> Rule {
        Rule::literal
    }

    fn message() -> &'static str {
        "Literal"
    }

    fn parse_pair(
        extensions: &SimpleExtensions,
        pair: pest::iterators::Pair<Rule>,
    ) -> Result<Self, MessageParseError> {
        assert_eq!(pair.as_rule(), Self::rule());
        let mut pairs = pair.into_inner();
        let value = pairs.next().unwrap(); // First item is always the value
        let typ = pairs.next(); // Second item is optional type
        assert!(pairs.next().is_none());
        let typ = match typ {
            Some(t) => Some(Type::parse_pair(extensions, t)?),
            None => None,
        };
        match value.as_rule() {
            Rule::integer => to_int_literal(value, typ),
            Rule::float => to_float_literal(value, typ),
            Rule::boolean => to_boolean_literal(value, typ),
            Rule::string_literal => to_string_literal(value, typ),
            Rule::null => to_null_literal(value, typ),
            _ => unreachable!("Literal unexpected rule: {:?}", value.as_rule()),
        }
    }
}

impl ScopedParsePair for ScalarFunction {
    fn rule() -> Rule {
        Rule::function_call
    }

    fn message() -> &'static str {
        "ScalarFunction"
    }

    fn parse_pair(
        extensions: &SimpleExtensions,
        pair: pest::iterators::Pair<Rule>,
    ) -> Result<Self, MessageParseError> {
        assert_eq!(pair.as_rule(), Self::rule());
        let span = pair.as_span();
        let mut iter = RuleIter::from(pair.into_inner());

        // Parse compound function name (required) — e.g. "equal" or "equal:any_any"
        let name = iter.parse_next::<CompoundName>();

        // Parse optional anchor (e.g., #1)
        let anchor = iter
            .try_pop(Rule::anchor)
            .map(|n| unwrap_single_pair(n).as_str().parse::<u32>().unwrap());

        // Parse optional URN anchor (e.g., @1)
        let _urn_anchor = iter
            .try_pop(Rule::urn_anchor)
            .map(|n| unwrap_single_pair(n).as_str().parse::<u32>().unwrap());

        // Parse argument list (required)
        let argument_list = iter.pop(Rule::argument_list);
        let mut arguments = Vec::new();
        for e in argument_list.into_inner() {
            arguments.push(FunctionArgument {
                arg_type: Some(ArgType::Value(Expression::parse_pair(extensions, e)?)),
            });
        }

        // Parse required output type (e.g., :i64). pop is safe here because
        // the grammar guarantees the type token is always present.
        let output_type = Some(Type::parse_pair(extensions, iter.pop(Rule::r#type))?);

        iter.done();
        let anchor = get_and_validate_anchor(
            extensions,
            ExtensionKind::Function,
            anchor,
            name.full(),
            span,
        )?;
        Ok(ScalarFunction {
            function_reference: anchor,
            arguments,
            options: vec![], // TODO: Function Options
            output_type,
            #[allow(deprecated)]
            args: vec![],
        })
    }
}

impl ScopedParsePair for Cast {
    fn rule() -> Rule {
        Rule::cast_expression
    }

    fn message() -> &'static str {
        "Cast"
    }

    fn parse_pair(
        extensions: &SimpleExtensions,
        pair: pest::iterators::Pair<Rule>,
    ) -> Result<Self, MessageParseError> {
        assert_eq!(pair.as_rule(), Self::rule());
        let mut pairs = pair.into_inner();

        let expr_pair = pairs.next().unwrap();

        // Optional failure behavior prefix: ? = RETURN_NULL, ! = THROW_EXCEPTION
        let next = pairs.next().unwrap();
        let (failure_behavior, type_pair) = if next.as_rule() == Rule::cast_failure_behavior {
            let fb = match next.as_str() {
                "?" => cast::FailureBehavior::ReturnNull as i32,
                "!" => cast::FailureBehavior::ThrowException as i32,
                _ => unreachable!("Grammar guarantees cast_failure_behavior is ? or !"),
            };
            (fb, pairs.next().unwrap())
        } else {
            (cast::FailureBehavior::Unspecified as i32, next)
        };

        assert!(pairs.next().is_none());

        let input = Expression::parse_pair(extensions, expr_pair)?;
        let target_type = Type::parse_pair(extensions, type_pair)?;

        Ok(Cast {
            r#type: Some(target_type),
            input: Some(Box::new(input)),
            failure_behavior,
        })
    }
}

impl ScopedParsePair for Expression {
    fn rule() -> Rule {
        Rule::expression
    }

    fn message() -> &'static str {
        "Expression"
    }

    fn parse_pair(
        extensions: &SimpleExtensions,
        pair: pest::iterators::Pair<Rule>,
    ) -> Result<Self, MessageParseError> {
        assert_eq!(pair.as_rule(), Self::rule());
        let inner = unwrap_single_pair(pair);
        match inner.as_rule() {
            Rule::literal => Ok(Expression {
                rex_type: Some(RexType::Literal(Literal::parse_pair(extensions, inner)?)),
            }),
            Rule::function_call => Ok(Expression {
                rex_type: Some(RexType::ScalarFunction(ScalarFunction::parse_pair(
                    extensions, inner,
                )?)),
            }),
            Rule::reference => Ok(Expression {
                rex_type: Some(RexType::Selection(Box::new(FieldReference::parse_pair(
                    inner,
                )))),
            }),
            Rule::if_then => Ok(Expression {
                rex_type: Some(RexType::IfThen(Box::new(IfThen::parse_pair(
                    extensions, inner,
                )?))),
            }),
            Rule::cast_expression => Ok(Expression {
                rex_type: Some(RexType::Cast(Box::new(Cast::parse_pair(
                    extensions, inner,
                )?))),
            }),
            _ => unreachable!(
                "Grammar guarantees expression can only be literal, function_call, reference, if_then, or cast_expression, got: {:?}",
                inner.as_rule()
            ),
        }
    }
}

impl ScopedParsePair for IfClause {
    fn rule() -> Rule {
        Rule::if_clause
    }

    fn message() -> &'static str {
        "IfClause"
    }

    fn parse_pair(
        extensions: &SimpleExtensions,
        pair: pest::iterators::Pair<Rule>,
    ) -> Result<Self, MessageParseError> {
        assert_eq!(pair.as_rule(), Self::rule());
        let mut pairs = pair.into_inner(); // should have 2 children, 2 expressions

        let condition = pairs.next().unwrap();
        let result = pairs.next().unwrap();
        assert!(pairs.next().is_none());

        let ex1 = Some(Expression::parse_pair(extensions, condition)?);
        let ex2 = Some(Expression::parse_pair(extensions, result)?);

        Ok(IfClause {
            r#if: ex1,
            then: ex2,
        })
    }
}

impl ScopedParsePair for IfThen {
    fn rule() -> Rule {
        Rule::if_then
    }
    fn message() -> &'static str {
        "IfThen"
    }

    fn parse_pair(
        extensions: &SimpleExtensions,
        pair: pest::iterators::Pair<Rule>,
    ) -> Result<Self, MessageParseError> {
        assert_eq!(pair.as_rule(), Self::rule());

        let mut iter = RuleIter::from(pair.into_inner()); // should have 2 or more children

        let mut ifs: Vec<IfClause> = Vec::new();

        // gets all of the if clauses
        while let Some(p) = iter.try_pop(Rule::if_clause) {
            let if_clause = IfClause::parse_pair(extensions, p)?;
            ifs.push(if_clause);
        }

        let pair = iter.try_pop(Rule::expression).unwrap(); // should be else expression
        iter.done();
        let else_clause = Some(Box::new(Expression::parse_pair(extensions, pair)?));

        Ok(IfThen {
            ifs,
            r#else: else_clause,
        })
    }
}
pub struct Name(pub String);

impl ParsePair for Name {
    fn rule() -> Rule {
        Rule::name
    }

    fn message() -> &'static str {
        "Name"
    }

    fn parse_pair(pair: pest::iterators::Pair<Rule>) -> Self {
        assert_eq!(pair.as_rule(), Self::rule());
        let inner = unwrap_single_pair(pair);
        match inner.as_rule() {
            Rule::identifier => Name(inner.as_str().to_string()),
            Rule::quoted_name => Name(unescape_string(inner)),
            _ => unreachable!("Name unexpected rule: {:?}", inner.as_rule()),
        }
    }
}

impl ParsePair for CompoundName {
    fn rule() -> Rule {
        Rule::function_signature
    }

    fn message() -> &'static str {
        "CompoundName"
    }

    fn parse_pair(pair: pest::iterators::Pair<Rule>) -> Self {
        assert_eq!(pair.as_rule(), Self::rule());
        CompoundName::new(pair.as_str())
    }
}

impl ScopedParsePair for Measure {
    fn rule() -> Rule {
        Rule::aggregate_measure
    }

    fn message() -> &'static str {
        "Measure"
    }

    fn parse_pair(
        extensions: &SimpleExtensions,
        pair: pest::iterators::Pair<Rule>,
    ) -> Result<Self, MessageParseError> {
        assert_eq!(pair.as_rule(), Self::rule());

        // Extract the inner function_call from aggregate_measure
        let function_call_pair = unwrap_single_pair(pair);
        assert_eq!(function_call_pair.as_rule(), Rule::function_call);

        // Parse as ScalarFunction, then convert to AggregateFunction
        let scalar = ScalarFunction::parse_pair(extensions, function_call_pair)?;
        Ok(Measure {
            measure: Some(AggregateFunction {
                function_reference: scalar.function_reference,
                arguments: scalar.arguments,
                options: scalar.options,
                output_type: scalar.output_type,
                invocation: 0, // TODO: support invocation (ALL, DISTINCT, etc.)
                phase: 0, // TODO: support phase (INITIAL_TO_RESULT, PARTIAL_TO_INTERMEDIATE, etc.)
                sorts: vec![], // TODO: support sorts for ordered aggregates
                #[allow(deprecated)]
                args: scalar.args,
            }),
            filter: None, // TODO: support filter conditions on aggregate measures
        })
    }
}

#[cfg(test)]
mod tests {
    use pest::Parser as PestParser;

    use super::*;
    use crate::parser::ExpressionParser;

    fn parse_exact(rule: Rule, input: &'_ str) -> pest::iterators::Pair<'_, Rule> {
        let mut pairs = ExpressionParser::parse(rule, input).unwrap();
        assert_eq!(pairs.as_str(), input);
        let pair = pairs.next().unwrap();
        assert_eq!(pairs.next(), None);
        pair
    }

    fn assert_parses_to<T: ParsePair + PartialEq + std::fmt::Debug>(input: &str, expected: T) {
        let pair = parse_exact(T::rule(), input);
        let actual = T::parse_pair(pair);
        assert_eq!(actual, expected);
    }

    fn assert_parses_with<T: ScopedParsePair + PartialEq + std::fmt::Debug>(
        ext: &SimpleExtensions,
        input: &str,
        expected: T,
    ) {
        let pair = parse_exact(T::rule(), input);
        let actual = T::parse_pair(ext, pair).unwrap();
        assert_eq!(actual, expected);
    }

    #[test]
    fn test_parse_field_reference() {
        assert_parses_to("$1", FieldIndex(1).to_field_reference());
    }

    #[test]
    fn test_parse_integer_literal() {
        let extensions = SimpleExtensions::default();
        let expected = Literal {
            literal_type: Some(LiteralType::I64(1)),
            nullable: false,
            type_variation_reference: 0,
        };
        assert_parses_with(&extensions, "1", expected);
    }

    #[test]
    fn test_parse_float_literal() {
        // First test that the grammar can parse floats
        let pairs = ExpressionParser::parse(Rule::float, "3.82").unwrap();
        let parsed_text = pairs.as_str();
        assert_eq!(parsed_text, "3.82");

        let extensions = SimpleExtensions::default();
        let expected = Literal {
            literal_type: Some(LiteralType::Fp64(3.82)),
            nullable: false,
            type_variation_reference: 0,
        };
        assert_parses_with(&extensions, "3.82", expected);
    }

    #[test]
    fn test_parse_negative_float_literal() {
        let extensions = SimpleExtensions::default();
        let expected = Literal {
            literal_type: Some(LiteralType::Fp64(-2.5)),
            nullable: false,
            type_variation_reference: 0,
        };
        assert_parses_with(&extensions, "-2.5", expected);
    }

    #[test]
    fn test_parse_boolean_true_literal() {
        let extensions = SimpleExtensions::default();
        let expected = Literal {
            literal_type: Some(LiteralType::Boolean(true)),
            nullable: false,
            type_variation_reference: 0,
        };
        assert_parses_with(&extensions, "true", expected);
    }

    #[test]
    fn test_parse_boolean_false_literal() {
        let extensions = SimpleExtensions::default();
        let expected = Literal {
            literal_type: Some(LiteralType::Boolean(false)),
            nullable: false,
            type_variation_reference: 0,
        };
        assert_parses_with(&extensions, "false", expected);
    }

    #[test]
    fn test_parse_nullable_boolean_literal() {
        let extensions = SimpleExtensions::default();
        let expected_true = Literal {
            literal_type: Some(LiteralType::Boolean(true)),
            nullable: true,
            type_variation_reference: 0,
        };
        let expected_false = Literal {
            literal_type: Some(LiteralType::Boolean(false)),
            nullable: true,
            type_variation_reference: 0,
        };
        assert_parses_with(&extensions, "true:boolean?", expected_true);
        assert_parses_with(&extensions, "false:boolean?", expected_false);
    }

    #[test]
    fn test_parse_nullable_integer_literal() {
        let extensions = SimpleExtensions::default();
        let expected_i32 = Literal {
            literal_type: Some(LiteralType::I32(78)),
            nullable: true,
            type_variation_reference: 0,
        };
        let expected_i64 = Literal {
            literal_type: Some(LiteralType::I64(42)),
            nullable: true,
            type_variation_reference: 0,
        };
        assert_parses_with(&extensions, "78:i32?", expected_i32);
        assert_parses_with(&extensions, "42:i64?", expected_i64);
    }

    #[test]
    fn test_parse_nullable_float_literal() {
        let extensions = SimpleExtensions::default();
        let expected_fp64 = Literal {
            literal_type: Some(LiteralType::Fp64(3.19)),
            nullable: true,
            type_variation_reference: 0,
        };
        assert_parses_with(&extensions, "3.19:fp64?", expected_fp64);
    }

    #[test]
    fn test_parse_float_literal_with_fp32_type() {
        let extensions = SimpleExtensions::default();
        let pair = parse_exact(Rule::literal, "3.82:fp32");
        let result = Literal::parse_pair(&extensions, pair).unwrap();

        match result.literal_type {
            Some(LiteralType::Fp32(val)) => assert!((val - 3.82).abs() < f32::EPSILON),
            _ => panic!("Expected Fp32 literal type"),
        }
    }

    #[test]
    fn test_parse_date_literal() {
        let extensions = SimpleExtensions::default();
        let pair = parse_exact(Rule::literal, "'2023-12-25':date");
        let result = Literal::parse_pair(&extensions, pair).unwrap();

        match result.literal_type {
            Some(LiteralType::Date(days)) => {
                // 2023-12-25 should be a positive number of days since 1970-01-01
                assert!(
                    days > 0,
                    "Expected positive days since epoch, got: {}",
                    days
                );
            }
            _ => panic!("Expected Date literal type, got: {:?}", result.literal_type),
        }
    }

    #[test]
    fn test_parse_time_literal() {
        let extensions = SimpleExtensions::default();
        let pair = parse_exact(Rule::literal, "'14:30:45':time");
        let result = Literal::parse_pair(&extensions, pair).unwrap();

        match result.literal_type {
            #[allow(deprecated)]
            Some(LiteralType::Time(microseconds)) => {
                // 14:30:45 = (14*3600 + 30*60 + 45) * 1_000_000 microseconds
                let expected = (14 * 3600 + 30 * 60 + 45) * 1_000_000;
                assert_eq!(microseconds, expected);
            }
            _ => panic!("Expected Time literal type, got: {:?}", result.literal_type),
        }
    }

    #[test]
    fn test_parse_timestamp_literal_with_t() {
        let extensions = SimpleExtensions::default();
        let pair = parse_exact(Rule::literal, "'2023-01-01T12:00:00':timestamp");
        let result = Literal::parse_pair(&extensions, pair).unwrap();

        match result.literal_type {
            #[allow(deprecated)]
            Some(LiteralType::Timestamp(microseconds)) => {
                assert!(
                    microseconds > 0,
                    "Expected positive microseconds since epoch"
                );
            }
            _ => panic!(
                "Expected Timestamp literal type, got: {:?}",
                result.literal_type
            ),
        }
    }

    #[test]
    fn test_parse_timestamp_literal_with_space() {
        let extensions = SimpleExtensions::default();
        let pair = parse_exact(Rule::literal, "'2023-01-01 12:00:00':timestamp");
        let result = Literal::parse_pair(&extensions, pair).unwrap();

        match result.literal_type {
            #[allow(deprecated)]
            Some(LiteralType::Timestamp(microseconds)) => {
                assert!(
                    microseconds > 0,
                    "Expected positive microseconds since epoch"
                );
            }
            _ => panic!(
                "Expected Timestamp literal type, got: {:?}",
                result.literal_type
            ),
        }
    }

    /// Helper function to create a literal boolean expression
    fn make_literal_bool(value: bool) -> Expression {
        Expression {
            rex_type: Some(RexType::Literal(Literal {
                literal_type: Some(LiteralType::Boolean(value)),
                nullable: false,
                type_variation_reference: 0,
            })),
        }
    }

    #[test]
    fn test_parse_if_then_single_clause() {
        let extensions = SimpleExtensions::default();
        let input = "if_then(true -> 42, _ -> 0)";
        let pair = parse_exact(Rule::if_then, input);
        let result = IfThen::parse_pair(&extensions, pair).unwrap();

        assert_eq!(result.ifs.len(), 1);
        assert!(result.r#else.is_some());
    }

    #[test]
    fn test_parse_if_then_with_typed_literals() {
        let extensions = SimpleExtensions::default();
        let input = "if_then(true -> 100:i32, _ -> -100:i32)";
        let pair = parse_exact(Rule::if_then, input);
        let result = IfThen::parse_pair(&extensions, pair).unwrap();

        assert_eq!(result.ifs.len(), 1);
        assert!(result.r#else.is_some());
    }

    #[test]
    fn test_parse_if_then_with_date_literals() {
        let extensions = SimpleExtensions::default();
        let input = "if_then(true -> '2023-12-25':date, _ -> '1970-01-01':date)";
        let pair = parse_exact(Rule::if_then, input);
        let result = IfThen::parse_pair(&extensions, pair).unwrap();

        assert_eq!(result.ifs.len(), 1);
        assert!(result.r#else.is_some());
    }

    #[test]
    fn test_parse_if_then_with_time_literals() {
        let extensions = SimpleExtensions::default();
        let input = "if_then(true -> '14:30:45':time, _ -> '00:00:00':time)";
        let pair = parse_exact(Rule::if_then, input);
        let result = IfThen::parse_pair(&extensions, pair).unwrap();

        assert_eq!(result.ifs.len(), 1);
        assert!(result.r#else.is_some());
    }

    #[test]
    fn test_parse_if_then_with_timestamp_literals() {
        let extensions = SimpleExtensions::default();
        let input = "if_then(true -> '2023-01-01T12:00:00':timestamp, _ -> '1970-01-01T00:00:00':timestamp)";
        let pair = parse_exact(Rule::if_then, input);
        let result = IfThen::parse_pair(&extensions, pair).unwrap();

        assert_eq!(result.ifs.len(), 1);
        assert!(result.r#else.is_some());
    }

    #[test]
    fn test_parse_if_clause_with_whitespace_variations() {
        let extensions = SimpleExtensions::default();

        // Test with various whitespace patterns
        let inputs = vec!["true->false", "true -> false", "true  ->  false"];

        for input in inputs {
            let pair = parse_exact(Rule::if_clause, input);
            let result = IfClause::parse_pair(&extensions, pair).unwrap();
            assert!(result.r#if.is_some());
            assert!(result.then.is_some());
        }
    }

    #[test]
    fn test_if_clause_structure() {
        let extensions = SimpleExtensions::default();
        let pair = parse_exact(Rule::if_clause, "42 -> 100");
        let result = IfClause::parse_pair(&extensions, pair).unwrap();

        // Verify the if clause has both condition and result
        let if_expr = result.r#if.as_ref().unwrap();
        let then_expr = result.then.as_ref().unwrap();

        // Check that they are literal expressions
        match (&if_expr.rex_type, &then_expr.rex_type) {
            (Some(RexType::Literal(_)), Some(RexType::Literal(_))) => {
                // Success - both are literals as expected
            }
            _ => panic!("Expected both if and then to be literals"),
        }
    }

    #[test]
    fn test_if_then_structure() {
        let extensions = SimpleExtensions::default();
        let input = "if_then(true -> 1, false -> 2, _ -> 0)";
        let pair = parse_exact(Rule::if_then, input);
        let result = IfThen::parse_pair(&extensions, pair).unwrap();

        // Verify structure
        assert_eq!(result.ifs.len(), 2);

        // Check each if clause
        for clause in &result.ifs {
            assert!(clause.r#if.is_some(), "If clause condition should exist");
            assert!(clause.then.is_some(), "If clause result should exist");
        }

        // Check else clause
        assert!(result.r#else.is_some(), "Else clause should exist");
    }

    #[test]
    fn test_parse_if_then_mixed_types_in_conditions() {
        let extensions = SimpleExtensions::default();
        // Different types in conditions (not results)
        let input = "if_then(true -> 1, true -> 'yes', 'yes' -> true, 42 -> 2, $0 -> 3, _ -> 0)";
        let pair = parse_exact(Rule::if_then, input);
        let result = IfThen::parse_pair(&extensions, pair).unwrap();

        assert_eq!(result.ifs.len(), 5);
        assert!(result.r#else.is_some());
    }

    #[test]
    fn test_if_then_preserves_clause_order() {
        let extensions = SimpleExtensions::default();
        let input = "if_then(1 -> 10, 2 -> 20, 3 -> 30, _ -> 0)";
        let pair = parse_exact(Rule::if_then, input);
        let result = IfThen::parse_pair(&extensions, pair).unwrap();

        assert_eq!(result.ifs.len(), 3);

        // Verify the clauses are in order by checking the literal values
        for (i, clause) in result.ifs.iter().enumerate() {
            if let Some(Expression {
                rex_type: Some(RexType::Literal(lit)),
            }) = &clause.r#if
                && let Some(LiteralType::I64(val)) = &lit.literal_type
            {
                assert_eq!(*val, (i as i64) + 1);
            }
        }
    }

    #[test]
    fn test_parse_if_then() {
        let extensions = SimpleExtensions::default();

        let c1 = IfClause {
            r#if: Some(make_literal_bool(true)),
            then: Some(make_literal_bool(true)),
        };

        let c2 = IfClause {
            r#if: Some(make_literal_bool(false)),
            then: Some(make_literal_bool(false)),
        };

        let if_clause = IfThen {
            ifs: vec![c1, c2],
            r#else: Some(Box::new(make_literal_bool(false))),
        };
        assert_parses_with(
            &extensions,
            "if_then(true -> true , false -> false, _ -> false)",
            if_clause,
        );
    }

    // ---- Tests for function_signature grammar rule ----

    fn parse_function_signature(input: &str) -> CompoundName {
        let pair = parse_exact(Rule::function_signature, input);
        CompoundName::parse_pair(pair)
    }

    #[test]
    fn test_compound_name_plain() {
        assert_eq!(parse_function_signature("add").full(), "add");
    }

    #[test]
    fn test_compound_name_full_zero_arg_type_signature() {
        // A Full name whose type signature encodes zero argument types (nothing after the colon).
        let n = parse_function_signature("add:");
        assert_eq!(n.full(), "add:");
        assert_eq!(n.base(), "add");
        assert!(n.matches("add:"));
        assert!(!n.matches("add:i64_i64"));
        assert!(n.matches("add"));
    }

    #[test]
    fn test_compound_name_with_signature() {
        assert_eq!(
            parse_function_signature("equal:any_any").full(),
            "equal:any_any"
        );
        assert_eq!(
            parse_function_signature("regexp_match_substring:str_str_i64").full(),
            "regexp_match_substring:str_str_i64"
        );
        assert_eq!(
            parse_function_signature("add:i64_i64").full(),
            "add:i64_i64"
        );
    }

    #[test]
    fn test_compound_name_trailing_colon_grammar() {
        // "count:" (trailing colon, zero-arg type signature) parses as a compound name with
        // an empty signature suffix: base "count", has_signature true, full "count:".
        let name = parse_function_signature("count:");
        assert_eq!(name.base(), "count");
        assert_eq!(name.full(), "count:");
        assert!(
            name.has_signature(),
            "trailing colon must set has_signature"
        );
    }

    #[test]
    fn test_compound_name_stops_at_opening_paren() {
        // In a function call, the function_signature must stop before the '('.
        let pairs = ExpressionParser::parse(Rule::function_signature, "equal:any_any").unwrap();
        assert_eq!(pairs.as_str(), "equal:any_any");
    }

    // ---- Tests for ScalarFunction parsing with compound names ----

    fn make_extensions_for_fn_tests() -> SimpleExtensions {
        let mut exts = SimpleExtensions::default();
        exts.add_extension_urn("urn".to_string(), 1).unwrap();
        exts.add_extension(
            crate::extensions::simple::ExtensionKind::Function,
            1,
            1,
            "equal:any_any".to_string(),
        )
        .unwrap();
        exts.add_extension(
            crate::extensions::simple::ExtensionKind::Function,
            1,
            2,
            "equal:str_str".to_string(),
        )
        .unwrap();
        exts.add_extension(
            crate::extensions::simple::ExtensionKind::Function,
            1,
            3,
            "add:i64_i64".to_string(),
        )
        .unwrap();
        exts
    }

    #[test]
    fn test_scalar_function_full_compound_name() {
        // Full compound name without anchor
        let exts = make_extensions_for_fn_tests();
        let pair = parse_exact(Rule::function_call, "equal:any_any($0, $1):boolean");
        let f = ScalarFunction::parse_pair(&exts, pair).unwrap();
        assert_eq!(f.function_reference, 1);
        assert_eq!(f.arguments.len(), 2);
        assert!(
            f.output_type.is_some(),
            "output_type must be set after parsing"
        );
    }

    #[test]
    fn test_scalar_function_second_overload() {
        let exts = make_extensions_for_fn_tests();
        let pair = parse_exact(Rule::function_call, "equal:str_str($0, $1):boolean");
        let f = ScalarFunction::parse_pair(&exts, pair).unwrap();

        assert_eq!(f.arguments.len(), 2);
        assert_eq!(f.function_reference, 2);
    }

    #[test]
    fn test_scalar_function_base_name_unique_overload() {
        // "add" has only one overload; base-name lookup should succeed
        let exts = make_extensions_for_fn_tests();
        let pair = parse_exact(Rule::function_call, "add($0, $1):i64");
        let f = ScalarFunction::parse_pair(&exts, pair).unwrap();

        assert_eq!(f.arguments.len(), 2);
        assert_eq!(f.function_reference, 3);
        assert!(
            f.output_type.is_some(),
            "output_type must be set after parsing"
        );
    }

    #[test]
    fn test_scalar_function_base_name_ambiguous_fails() {
        // "equal" has two overloads; base-name lookup should fail
        let exts = make_extensions_for_fn_tests();
        let pair = parse_exact(Rule::function_call, "equal($0, $1):boolean");
        let result = ScalarFunction::parse_pair(&exts, pair);
        assert!(result.is_err(), "ambiguous base name should fail");
    }

    #[test]
    fn test_scalar_function_compound_name_with_anchor() {
        let exts = make_extensions_for_fn_tests();
        let pair = parse_exact(Rule::function_call, "equal:any_any#1($0, $1):boolean");
        let f = ScalarFunction::parse_pair(&exts, pair).unwrap();
        assert_eq!(f.function_reference, 1);
        assert_eq!(f.arguments.len(), 2);
    }

    #[test]
    fn test_scalar_function_base_name_with_anchor() {
        // Base name + explicit anchor should resolve (anchor 1 stores equal:any_any)
        let exts = make_extensions_for_fn_tests();
        let pair = parse_exact(Rule::function_call, "equal#1($0, $1):boolean");
        let f = ScalarFunction::parse_pair(&exts, pair).unwrap();
        assert_eq!(f.function_reference, 1);
        assert_eq!(f.arguments.len(), 2);
    }

    #[test]
    fn test_scalar_function_wrong_name_for_anchor_fails() {
        let exts = make_extensions_for_fn_tests();
        let pair = parse_exact(Rule::function_call, "like#1($0):boolean");
        let result = ScalarFunction::parse_pair(&exts, pair);
        assert!(result.is_err(), "mismatched name/anchor should fail");
    }

    #[test]
    fn test_scalar_function_user_defined_type_in_signature() {
        // u!-prefixed type segments in function signatures parse and resolve.
        let mut exts = SimpleExtensions::default();
        exts.add_extension_urn("urn".to_string(), 1).unwrap();
        exts.add_extension(
            crate::extensions::simple::ExtensionKind::Function,
            1,
            10,
            "json_extract_path:u!json_str".to_string(),
        )
        .unwrap();

        let pair = parse_exact(
            Rule::function_call,
            "json_extract_path:u!json_str($0, $1):string",
        );
        let f = ScalarFunction::parse_pair(&exts, pair).unwrap();
        assert_eq!(f.function_reference, 10);
        assert_eq!(f.arguments.len(), 2);
    }

    #[test]
    fn test_scalar_function_missing_type_fails_to_parse() {
        // The grammar requires a type annotation; "add($0, $1)" without ":i64" must fail.
        let result = ExpressionParser::parse(Rule::function_call, "add($0, $1)");
        assert!(
            result.is_err(),
            "function call without type annotation should fail to parse"
        );
    }

    #[test]
    fn test_parse_cast_expression_basic() {
        let extensions = SimpleExtensions::default();
        let pair = parse_exact(Rule::cast_expression, "(78:i32)::i16");
        let result = Cast::parse_pair(&extensions, pair).unwrap();

        // Input should be 78:i32
        let input = result.input.as_ref().unwrap();
        match &input.rex_type {
            Some(RexType::Literal(lit)) => match &lit.literal_type {
                Some(LiteralType::I32(v)) => assert_eq!(*v, 78),
                other => panic!("Expected I32 literal, got: {:?}", other),
            },
            other => panic!("Expected literal, got: {:?}", other),
        }

        // Target type should be i16
        let target = result.r#type.as_ref().unwrap();
        match &target.kind {
            Some(substrait::proto::r#type::Kind::I16(_)) => {}
            other => panic!("Expected i16 type, got: {:?}", other),
        }

        assert_eq!(result.failure_behavior, 0);
    }

    #[test]
    fn test_parse_cast_expression_via_expression_rule() {
        let extensions = SimpleExtensions::default();
        let pair = parse_exact(Rule::expression, "(78:i32)::i16");
        let result = Expression::parse_pair(&extensions, pair).unwrap();

        match result.rex_type {
            Some(RexType::Cast(_)) => {}
            other => panic!("Expected Cast rex type, got: {:?}", other),
        }
    }

    #[test]
    fn test_parse_cast_expression_nested() {
        let extensions = SimpleExtensions::default();
        let pair = parse_exact(Rule::cast_expression, "((78:i32)::i16)::i32");
        let result = Cast::parse_pair(&extensions, pair).unwrap();

        // Input should itself be a Cast
        let input = result.input.as_ref().unwrap();
        match &input.rex_type {
            Some(RexType::Cast(inner)) => {
                let inner_input = inner.input.as_ref().unwrap();
                match &inner_input.rex_type {
                    Some(RexType::Literal(lit)) => match &lit.literal_type {
                        Some(LiteralType::I32(v)) => assert_eq!(*v, 78),
                        other => panic!("Expected I32 literal, got: {:?}", other),
                    },
                    other => panic!("Expected literal, got: {:?}", other),
                }
            }
            other => panic!("Expected inner Cast, got: {:?}", other),
        }

        match &result.r#type.as_ref().unwrap().kind {
            Some(substrait::proto::r#type::Kind::I32(_)) => {}
            other => panic!("Expected i32 outer type, got: {:?}", other),
        }
    }

    #[test]
    fn test_parse_cast_expression_with_boolean() {
        let extensions = SimpleExtensions::default();
        let pair = parse_exact(Rule::cast_expression, "(true)::i32");
        let result = Cast::parse_pair(&extensions, pair).unwrap();

        let input = result.input.as_ref().unwrap();
        match &input.rex_type {
            Some(RexType::Literal(lit)) => match &lit.literal_type {
                Some(LiteralType::Boolean(v)) => assert!(*v),
                other => panic!("Expected Boolean literal, got: {:?}", other),
            },
            other => panic!("Expected literal, got: {:?}", other),
        }
    }

    #[test]
    fn test_parse_cast_expression_with_whitespace() {
        let extensions = SimpleExtensions::default();
        // Grammar allows optional whitespace around the expression and ::
        let pair = parse_exact(Rule::cast_expression, "( 78:i32 ) :: i16");
        let result = Cast::parse_pair(&extensions, pair).unwrap();
        assert!(result.input.is_some());
        assert!(result.r#type.is_some());
    }

    #[test]
    fn test_parse_cast_unspecified_failure_behavior() {
        let extensions = SimpleExtensions::default();
        let pair = parse_exact(Rule::cast_expression, "(78:i32)::i16");
        let result = Cast::parse_pair(&extensions, pair).unwrap();
        assert_eq!(
            result.failure_behavior,
            cast::FailureBehavior::Unspecified as i32
        );
    }

    #[test]
    fn test_parse_cast_return_null_failure_behavior() {
        let extensions = SimpleExtensions::default();
        let pair = parse_exact(Rule::cast_expression, "(78:i32)::?i16");
        let result = Cast::parse_pair(&extensions, pair).unwrap();
        assert_eq!(
            result.failure_behavior,
            cast::FailureBehavior::ReturnNull as i32
        );
    }

    #[test]
    fn test_parse_cast_throw_exception_failure_behavior() {
        let extensions = SimpleExtensions::default();
        let pair = parse_exact(Rule::cast_expression, "(78:i32)::!i16");
        let result = Cast::parse_pair(&extensions, pair).unwrap();
        assert_eq!(
            result.failure_behavior,
            cast::FailureBehavior::ThrowException as i32
        );
    }

    #[test]
    fn test_parse_cast_to_user_defined_type_with_u_prefix() {
        // Cast target type is a u!-prefixed UDT; exercises the user_defined_type rule in the cast path.
        let mut extensions = SimpleExtensions::default();
        extensions.add_extension_urn("urn".to_string(), 1).unwrap();
        extensions
            .add_extension(
                crate::extensions::simple::ExtensionKind::Type,
                1,
                5,
                "u!json".to_string(),
            )
            .unwrap();

        let pair = parse_exact(Rule::cast_expression, "($0)::u!json");
        let result = Cast::parse_pair(&extensions, pair).unwrap();
        match result.r#type.as_ref().unwrap().kind.as_ref().unwrap() {
            substrait::proto::r#type::Kind::UserDefined(u) => {
                assert_eq!(u.type_reference, 5);
            }
            other => panic!("expected UserDefined, got {other:?}"),
        }
    }

    #[test]
    fn test_function_call_u_prefix_base_name_rejected() {
        // u! is not valid in a function call base name. The grammar's function_signature
        // rule uses `identifier` as the base, which cannot match "u!" + identifier.
        assert!(
            ExpressionParser::parse(Rule::function_call, "u!json_get($0)").is_err(),
            "u! prefix in function call base name must be rejected by the grammar"
        );
    }
}