substrait-explain 0.7.0

use pest::iterators::Pair;
use substrait::proto::r#type::{Kind, Nullability, Parameter};
use substrait::proto::{self, Type};

use super::{ParsePair, Rule, ScopedParsePair, iter_pairs, unwrap_single_pair};
use crate::extensions::SimpleExtensions;
use crate::extensions::simple::{ExtensionKind, MissingReference};
use crate::parser::MessageParseError;

/// Given a name (plain or compound) and an optional anchor, resolve and
/// validate the extension anchor.
///
/// For `Function` kind, delegates to [`SimpleExtensions::resolve_function`]
/// which encapsulates the full base-name fallback logic.
///
/// For other kinds, performs a direct anchor/name validation.
pub(crate) fn get_and_validate_anchor(
    extensions: &SimpleExtensions,
    kind: ExtensionKind,
    anchor: Option<u32>,
    name: &str,
    span: pest::Span,
) -> Result<u32, MessageParseError> {
    if kind == ExtensionKind::Function {
        return extensions
            .resolve_function(name, anchor)
            .map(|r| r.anchor)
            .map_err(|e| {
                MessageParseError::lookup(kind.name(), e, span, "Error resolving function")
            });
    }
    // For non-function kinds, validate the anchor/name pair directly.
    match anchor {
        Some(a) => match extensions.find_by_anchor(kind, a) {
            Err(e) => Err(MessageParseError::lookup(
                kind.name(),
                e,
                span,
                "Error matching name to anchor",
            )),
            Ok((_, stored)) => {
                if stored.full() == name || stored.base() == name {
                    Ok(a)
                } else {
                    Err(MessageParseError::lookup(
                        kind.name(),
                        MissingReference::Mismatched(kind, name.to_string(), a),
                        span,
                        "Error matching name to anchor",
                    ))
                }
            }
        },
        None => extensions.find_by_name(kind, name).map_err(|e| {
            MessageParseError::lookup(kind.name(), e, span, "Error finding extension for name")
        }),
    }
}

impl ParsePair for Nullability {
    fn rule() -> Rule {
        Rule::nullability
    }

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

    fn parse_pair(pair: Pair<Rule>) -> Self {
        assert_eq!(pair.as_rule(), Rule::nullability);
        match pair.as_str() {
            "?" => Nullability::Nullable,
            "" => Nullability::Required,
            "⁉" => Nullability::Unspecified,
            _ => panic!("Invalid nullability: {}", pair.as_str()),
        }
    }
}

impl ScopedParsePair for Parameter {
    fn rule() -> Rule {
        Rule::parameter
    }

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

    fn parse_pair(
        extensions: &SimpleExtensions,
        pair: Pair<Rule>,
    ) -> Result<Self, MessageParseError> {
        assert_eq!(pair.as_rule(), Rule::parameter);
        let inner = unwrap_single_pair(pair);
        match inner.as_rule() {
            Rule::r#type => Ok(Parameter {
                parameter: Some(proto::r#type::parameter::Parameter::DataType(
                    Type::parse_pair(extensions, inner)?,
                )),
            }),
            _ => unimplemented!("{:?}", inner.as_rule()),
        }
    }
}

fn parse_simple_type(pair: Pair<Rule>) -> Type {
    assert_eq!(pair.as_rule(), Rule::simple_type);
    let mut iter = iter_pairs(pair.into_inner());
    let name = iter.pop(Rule::simple_type_name).as_str();
    let nullability = iter.parse_next::<Nullability>();
    iter.done();

    let kind = match name {
        "boolean" => Kind::Bool(proto::r#type::Boolean {
            nullability: nullability.into(),
            type_variation_reference: 0,
        }),
        "i64" => Kind::I64(proto::r#type::I64 {
            nullability: nullability.into(),
            type_variation_reference: 0,
        }),
        "i32" => Kind::I32(proto::r#type::I32 {
            nullability: nullability.into(),
            type_variation_reference: 0,
        }),
        "i16" => Kind::I16(proto::r#type::I16 {
            nullability: nullability.into(),
            type_variation_reference: 0,
        }),
        "i8" => Kind::I8(proto::r#type::I8 {
            nullability: nullability.into(),
            type_variation_reference: 0,
        }),
        "fp32" => Kind::Fp32(proto::r#type::Fp32 {
            nullability: nullability.into(),
            type_variation_reference: 0,
        }),
        "fp64" => Kind::Fp64(proto::r#type::Fp64 {
            nullability: nullability.into(),
            type_variation_reference: 0,
        }),
        "string" => Kind::String(proto::r#type::String {
            nullability: nullability.into(),
            type_variation_reference: 0,
        }),
        "binary" => Kind::Binary(proto::r#type::Binary {
            nullability: nullability.into(),
            type_variation_reference: 0,
        }),
        #[allow(deprecated)]
        "timestamp" => Kind::Timestamp(proto::r#type::Timestamp {
            nullability: nullability.into(),
            type_variation_reference: 0,
        }),
        #[allow(deprecated)]
        "timestamp_tz" => Kind::TimestampTz(proto::r#type::TimestampTz {
            nullability: nullability.into(),
            type_variation_reference: 0,
        }),
        "date" => Kind::Date(proto::r#type::Date {
            nullability: nullability.into(),
            type_variation_reference: 0,
        }),
        #[allow(deprecated)]
        "time" => Kind::Time(proto::r#type::Time {
            nullability: nullability.into(),
            type_variation_reference: 0,
        }),
        "interval_year" => Kind::IntervalYear(proto::r#type::IntervalYear {
            nullability: nullability.into(),
            type_variation_reference: 0,
        }),
        "uuid" => Kind::Uuid(proto::r#type::Uuid {
            nullability: nullability.into(),
            type_variation_reference: 0,
        }),
        _ => unreachable!("Type {} exists in parser but not implemented in code", name),
    };
    Type { kind: Some(kind) }
}

fn parse_compound_type(
    extensions: &SimpleExtensions,
    pair: Pair<Rule>,
) -> Result<Type, MessageParseError> {
    assert_eq!(pair.as_rule(), Rule::compound_type);
    let inner = unwrap_single_pair(pair);
    match inner.as_rule() {
        Rule::list_type => parse_list_type(extensions, inner),
        // Rule::map_type => parse_map_type(inner),
        // Rule::struct_type => parse_struct_type(inner),
        Rule::precision_timestamp_tz_type
        | Rule::precision_timestamp_type
        | Rule::precision_time_type => parse_precision_type(inner),
        _ => unimplemented!("{:?}", inner.as_rule()),
    }
}

fn parse_precision_type(pair: Pair<Rule>) -> Result<Type, MessageParseError> {
    let rule = pair.as_rule();
    let mut iter = iter_pairs(pair.into_inner());
    let nullability = iter.parse_next::<Nullability>();
    let precision_pair = iter.pop(Rule::integer);
    let precision_span = precision_pair.as_span();
    let precision = precision_pair.as_str().parse::<i32>().unwrap();
    if !(0..=12).contains(&precision) {
        return Err(MessageParseError::invalid(
            "precision time type",
            precision_span,
            format!("precision must be between 0 and 12, got {precision}"),
        ));
    }
    iter.done();
    let kind = match rule {
        Rule::precision_timestamp_type => {
            Kind::PrecisionTimestamp(proto::r#type::PrecisionTimestamp {
                precision,
                nullability: nullability.into(),
                type_variation_reference: 0,
            })
        }
        Rule::precision_timestamp_tz_type => {
            Kind::PrecisionTimestampTz(proto::r#type::PrecisionTimestampTz {
                precision,
                nullability: nullability.into(),
                type_variation_reference: 0,
            })
        }
        Rule::precision_time_type => Kind::PrecisionTime(proto::r#type::PrecisionTime {
            precision,
            nullability: nullability.into(),
            type_variation_reference: 0,
        }),
        _ => unreachable!("parse_precision_type called with rule {:?}", rule),
    };
    Ok(Type { kind: Some(kind) })
}

fn parse_list_type(
    extensions: &SimpleExtensions,
    pair: Pair<Rule>,
) -> Result<Type, MessageParseError> {
    assert_eq!(pair.as_rule(), Rule::list_type);
    let mut iter = iter_pairs(pair.into_inner());
    let nullability = iter.parse_next::<Nullability>();
    let inner = iter.parse_next_scoped::<Type>(extensions)?;
    iter.done();

    Ok(Type {
        kind: Some(Kind::List(Box::new(proto::r#type::List {
            nullability: nullability.into(),
            r#type: Some(Box::new(inner)),
            type_variation_reference: 0,
        }))),
    })
}

fn parse_parameters(
    extensions: &SimpleExtensions,
    pair: Pair<Rule>,
) -> Result<Vec<Parameter>, MessageParseError> {
    assert_eq!(pair.as_rule(), Rule::parameters);
    let mut iter = iter_pairs(pair.into_inner());
    let mut params = Vec::new();
    while let Some(param) = iter.parse_if_next_scoped::<Parameter>(extensions) {
        params.push(param?);
    }
    iter.done();
    Ok(params)
}

fn parse_user_defined_type(
    extensions: &SimpleExtensions,
    pair: Pair<Rule>,
) -> Result<Type, MessageParseError> {
    let span = pair.as_span();
    assert_eq!(pair.as_rule(), Rule::user_defined_type);
    let mut iter = iter_pairs(pair.into_inner());
    // TODO: quoted names not yet supported — see grammar TODO at user_defined_type
    let name = iter.pop(Rule::identifier).as_str().to_string();
    let anchor = iter
        .try_pop(Rule::anchor)
        .map(|n| unwrap_single_pair(n).as_str().parse::<u32>().unwrap());

    // TODO: Handle urn_anchor; validate that it matches the anchor
    let _urn_anchor = iter
        .try_pop(Rule::urn_anchor)
        .map(|n| unwrap_single_pair(n).as_str().parse::<u32>().unwrap());

    let nullability = iter.parse_next::<Nullability>();
    let parameters = match iter.try_pop(Rule::parameters) {
        Some(p) => parse_parameters(extensions, p)?,
        None => Vec::new(),
    };
    iter.done();

    let anchor = get_and_validate_anchor(extensions, ExtensionKind::Type, anchor, &name, span)?;

    Ok(Type {
        kind: Some(Kind::UserDefined(proto::r#type::UserDefined {
            type_reference: anchor,
            nullability: nullability.into(),
            type_parameters: parameters,
            type_variation_reference: 0,
        })),
    })
}

impl ScopedParsePair for Type {
    fn rule() -> Rule {
        Rule::r#type
    }

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

    fn parse_pair(
        extensions: &SimpleExtensions,
        pair: Pair<Rule>,
    ) -> Result<Self, MessageParseError> {
        assert_eq!(pair.as_rule(), Rule::r#type);
        let inner = unwrap_single_pair(pair);
        match inner.as_rule() {
            Rule::simple_type => Ok(parse_simple_type(inner)),
            Rule::compound_type => parse_compound_type(extensions, inner),
            Rule::user_defined_type => parse_user_defined_type(extensions, inner),
            _ => unreachable!(
                "Grammar guarantees type can only be simple_type, compound_type, or user_defined_type, got: {:?}",
                inner.as_rule()
            ),
        }
    }
}

#[cfg(test)]
mod tests {
    use pest::Parser;
    use substrait::proto::r#type::{I64, Kind, Nullability};

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

    #[test]
    fn test_parse_simple_type() {
        let mut pairs = ExpressionParser::parse(Rule::simple_type, "i64").unwrap();
        let pair = pairs.next().unwrap();
        assert_eq!(pairs.next(), None);
        let t = parse_simple_type(pair);
        assert_eq!(
            t,
            Type {
                kind: Some(Kind::I64(I64 {
                    nullability: Nullability::Required as i32,
                    type_variation_reference: 0,
                })),
            }
        );

        let mut pairs = ExpressionParser::parse(Rule::simple_type, "string?").unwrap();
        let pair = pairs.next().unwrap();
        assert_eq!(pairs.next(), None);
        let t = parse_simple_type(pair);
        assert_eq!(
            t,
            Type {
                kind: Some(Kind::String(proto::r#type::String {
                    nullability: Nullability::Nullable as i32,
                    type_variation_reference: 0,
                })),
            }
        );
    }

    #[test]
    fn test_parse_type() {
        let extensions = SimpleExtensions::default();
        let mut pairs = ExpressionParser::parse(Rule::r#type, "i64").unwrap();
        let pair = pairs.next().unwrap();
        assert_eq!(pairs.next(), None);
        let t = Type::parse_pair(&extensions, pair).unwrap();
        assert_eq!(
            t,
            Type {
                kind: Some(Kind::I64(I64 {
                    nullability: Nullability::Required as i32,
                    type_variation_reference: 0,
                }))
            }
        );
    }

    #[test]
    fn test_parse_list_type() {
        let extensions = SimpleExtensions::default();
        let mut pairs = ExpressionParser::parse(Rule::list_type, "list<i64>").unwrap();
        let pair = pairs.next().unwrap();
        assert_eq!(pairs.next(), None);
        let t = parse_list_type(&extensions, pair).unwrap();
        assert_eq!(
            t,
            Type {
                kind: Some(Kind::List(Box::new(proto::r#type::List {
                    nullability: Nullability::Required as i32,
                    r#type: Some(Box::new(Type {
                        kind: Some(Kind::I64(I64 {
                            nullability: Nullability::Required as i32,
                            type_variation_reference: 0,
                        }))
                    })),
                    type_variation_reference: 0,
                })))
            }
        );
    }

    #[test]
    fn test_parse_parameters() {
        let extensions = SimpleExtensions::default();
        let mut pairs = ExpressionParser::parse(Rule::parameters, "<i64?,string>").unwrap();
        let pair = pairs.next().unwrap();
        assert_eq!(pairs.next(), None);
        let t = parse_parameters(&extensions, pair).unwrap();
        assert_eq!(
            t,
            vec![
                Parameter {
                    parameter: Some(proto::r#type::parameter::Parameter::DataType(Type {
                        kind: Some(Kind::I64(proto::r#type::I64 {
                            nullability: Nullability::Nullable as i32,
                            type_variation_reference: 0,
                        })),
                    })),
                },
                Parameter {
                    parameter: Some(proto::r#type::parameter::Parameter::DataType(Type {
                        kind: Some(Kind::String(proto::r#type::String {
                            nullability: Nullability::Required as i32,
                            type_variation_reference: 0,
                        })),
                    })),
                },
            ]
        );
    }

    #[test]
    fn test_udts() {
        let mut extensions = SimpleExtensions::default();
        extensions
            .add_extension_urn("some_source".to_string(), 4)
            .unwrap();
        extensions
            .add_extension(ExtensionKind::Type, 4, 42, "udt".to_string())
            .unwrap();
        let mut pairs = ExpressionParser::parse(Rule::user_defined_type, "udt#42<i64?>").unwrap();
        let pair = pairs.next().unwrap();
        assert_eq!(pairs.next(), None);

        let t = parse_user_defined_type(&extensions, pair).unwrap();
        assert_eq!(
            t,
            Type {
                kind: Some(Kind::UserDefined(proto::r#type::UserDefined {
                    type_reference: 42,
                    type_variation_reference: 0,
                    nullability: Nullability::Required as i32,
                    type_parameters: vec![Parameter {
                        parameter: Some(proto::r#type::parameter::Parameter::DataType(Type {
                            kind: Some(Kind::I64(proto::r#type::I64 {
                                nullability: Nullability::Nullable as i32,
                                type_variation_reference: 0,
                            })),
                        })),
                    }],
                }))
            }
        );
    }

    #[test]
    fn test_udts_with_u_prefix() {
        // Extensions declared with "u!" prefix (e.g. "# 7 @ 4: u!json") normalize to
        // bare name at storage time. Both "u!json" and "json" in the plan resolve to the
        // same anchor.
        let mut extensions = SimpleExtensions::default();
        extensions
            .add_extension_urn("some_source".to_string(), 4)
            .unwrap();
        extensions
            .add_extension(ExtensionKind::Type, 4, 7, "u!json".to_string())
            .unwrap();

        // Plan uses u! prefix with explicit anchor
        let pair = ExpressionParser::parse(Rule::user_defined_type, "u!json#7")
            .unwrap()
            .next()
            .unwrap();
        let t = parse_user_defined_type(&extensions, pair).unwrap();
        assert_eq!(
            t,
            Type {
                kind: Some(Kind::UserDefined(proto::r#type::UserDefined {
                    type_reference: 7,
                    type_variation_reference: 0,
                    nullability: Nullability::Required as i32,
                    type_parameters: vec![],
                }))
            }
        );

        // Plan uses bare name — also resolves to the same anchor
        let pair = ExpressionParser::parse(Rule::user_defined_type, "json?")
            .unwrap()
            .next()
            .unwrap();
        let t = parse_user_defined_type(&extensions, pair).unwrap();
        assert_eq!(
            t,
            Type {
                kind: Some(Kind::UserDefined(proto::r#type::UserDefined {
                    type_reference: 7,
                    type_variation_reference: 0,
                    nullability: Nullability::Nullable as i32,
                    type_parameters: vec![],
                }))
            }
        );

        // Plan uses u! prefix with no anchor — anchor-free lookup must strip u! and find the type
        let pair = ExpressionParser::parse(Rule::user_defined_type, "u!json")
            .unwrap()
            .next()
            .unwrap();
        let t = parse_user_defined_type(&extensions, pair).unwrap();
        assert_eq!(
            t,
            Type {
                kind: Some(Kind::UserDefined(proto::r#type::UserDefined {
                    type_reference: 7,
                    type_variation_reference: 0,
                    nullability: Nullability::Required as i32,
                    type_parameters: vec![],
                }))
            }
        );
    }

    #[test]
    fn test_udt_u_prefix_matches_plain_registration() {
        // Extension registered without "u!" and plan uses "u!json" — must succeed.
        let mut extensions = SimpleExtensions::default();
        extensions.add_extension_urn("src".to_string(), 1).unwrap();
        extensions
            .add_extension(ExtensionKind::Type, 1, 3, "json".to_string())
            .unwrap();

        let pair = ExpressionParser::parse(Rule::user_defined_type, "u!json")
            .unwrap()
            .next()
            .unwrap();
        let t = parse_user_defined_type(&extensions, pair).unwrap();
        assert_eq!(
            t.kind,
            Some(Kind::UserDefined(proto::r#type::UserDefined {
                type_reference: 3,
                type_variation_reference: 0,
                nullability: Nullability::Required as i32,
                type_parameters: vec![],
            }))
        );
    }

    #[test]
    fn test_udt_plain_matches_u_prefix_registration() {
        // Extension registered with "u!" prefix and plan uses bare "json" — must succeed.
        let mut extensions = SimpleExtensions::default();
        extensions.add_extension_urn("src".to_string(), 1).unwrap();
        extensions
            .add_extension(ExtensionKind::Type, 1, 3, "u!json".to_string())
            .unwrap();

        let pair = ExpressionParser::parse(Rule::user_defined_type, "json")
            .unwrap()
            .next()
            .unwrap();
        let t = parse_user_defined_type(&extensions, pair).unwrap();
        assert_eq!(
            t.kind,
            Some(Kind::UserDefined(proto::r#type::UserDefined {
                type_reference: 3,
                type_variation_reference: 0,
                nullability: Nullability::Required as i32,
                type_parameters: vec![],
            }))
        );
    }
}