substrait-explain 0.9.0

use std::collections::HashMap;
use std::{fmt, thread};

use pest::error::{Error as PestError, ErrorVariant};
use pest::iterators::{Pair, Pairs};
use pest::{Parser as PestParser, Span};
use pest_derive::Parser as PestDeriveParser;
use substrait::proto::sort_field::SortDirection;
use thiserror::Error;

use crate::extensions::SimpleExtensions;
use crate::extensions::simple::MissingReference;

#[derive(PestDeriveParser)]
#[grammar = "parser/expression_grammar.pest"] // Path relative to src
pub(crate) struct ExpressionParser;

/// An error that occurs when parsing a message within a specific line. Contains
/// context pointing at that specific error.
#[derive(Error, Debug, Clone)]
#[error("{kind} Error parsing {message}:\n{error}")]
pub struct MessageParseError {
    message: &'static str,
    kind: ErrorKind,
    #[source]
    error: Box<PestError<Rule>>,
}

#[derive(Debug, Clone)]
pub(crate) enum ErrorKind {
    Syntax,
    InvalidValue,
    Lookup(MissingReference),
}

impl MessageParseError {
    pub(crate) fn invalid(message: &'static str, span: Span, description: impl ToString) -> Self {
        let error = PestError::new_from_span(
            ErrorVariant::CustomError {
                message: description.to_string(),
            },
            span,
        );
        Self::new(message, ErrorKind::InvalidValue, Box::new(error))
    }

    pub(crate) fn lookup(
        message: &'static str,
        missing: MissingReference,
        span: Span,
        description: impl ToString,
    ) -> Self {
        let error = PestError::new_from_span(
            ErrorVariant::CustomError {
                message: description.to_string(),
            },
            span,
        );
        Self::new(message, ErrorKind::Lookup(missing), Box::new(error))
    }
}

impl MessageParseError {
    pub(crate) fn new(message: &'static str, kind: ErrorKind, error: Box<PestError<Rule>>) -> Self {
        Self {
            message,
            kind,
            error,
        }
    }
}

impl fmt::Display for ErrorKind {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            ErrorKind::Syntax => write!(f, "Syntax"),
            ErrorKind::InvalidValue => write!(f, "Invalid value"),
            ErrorKind::Lookup(e) => write!(f, "Invalid reference ({e})"),
        }
    }
}

pub(crate) fn unwrap_single_pair(pair: Pair<Rule>) -> Pair<Rule> {
    let mut pairs = pair.into_inner();
    let pair = pairs.next().unwrap();
    assert_eq!(pairs.next(), None);
    pair
}

/// Unescapes a quoted string literal, handling escape sequences.
///
/// # Arguments
/// * `pair` - The pest pair containing the string to unescape (must be Rule::string_literal or Rule::quoted_name).
///
/// # Returns
/// * `String` with the unescaped contents.
///
/// # Panics
/// Panics if the rule is not `string_literal` or `quoted_name` (this should never happen
/// if the pest grammar is working correctly).
///
pub(crate) fn unescape_string(pair: Pair<Rule>) -> String {
    let s = pair.as_str();

    // Determine opener/closer based on rule type
    let (opener, closer) = match pair.as_rule() {
        Rule::string_literal => ('\'', '\''),
        Rule::quoted_name => ('"', '"'),
        _ => panic!(
            "unescape_string called with unexpected rule: {:?}",
            pair.as_rule()
        ),
    };

    let mut result = String::new();
    let mut chars = s.chars();
    let first = chars.next().expect("Empty string literal");

    assert_eq!(
        first, opener,
        "Expected opening quote '{opener}', got '{first}'"
    );

    // Skip the opening quote
    while let Some(c) = chars.next() {
        match c {
            c if c == closer => {
                // Skip the closing quote, and assert that there are no more characters.
                assert_eq!(
                    chars.next(),
                    None,
                    "Unexpected characters after closing quote"
                );
                break;
            }
            '\\' => {
                let next = chars
                    .next()
                    .expect("Incomplete escape sequence at end of string");
                match next {
                    'n' => result.push('\n'),
                    't' => result.push('\t'),
                    'r' => result.push('\r'),
                    // For all other characters (especially `"`, `'`, and `\`), we just
                    // push the character.
                    _ => result.push(next),
                }
            }
            _ => result.push(c),
        }
    }
    result
}

// A trait for converting a pest::iterators::Pair<Rule> into a Rust type. This
// is used to convert from the uniformly structured nesting
// pest::iterators::Pair<Rule> into more structured types.
pub(crate) trait ParsePair: Sized {
    // The rule that this type is parsed from.
    fn rule() -> Rule;

    // The name of the protobuf message type that this type corresponds to.
    fn message() -> &'static str;

    // Parse a single instance of this type from a pest::iterators::Pair<Rule>.
    // The input must match the rule returned by `rule`; otherwise, a panic is
    // expected.
    fn parse_pair(pair: Pair<Rule>) -> Self;

    fn parse_str(s: &str) -> Result<Self, MessageParseError> {
        let mut pairs = <ExpressionParser as PestParser<Rule>>::parse(Self::rule(), s)
            .map_err(|e| MessageParseError::new(Self::message(), ErrorKind::Syntax, Box::new(e)))?;
        assert_eq!(pairs.as_str(), s);
        let pair = pairs.next().unwrap();
        assert_eq!(pairs.next(), None);
        Ok(Self::parse_pair(pair))
    }
}

/// A trait for types that are parsed from a `pest::iterators::Pair<Rule>` that
/// depends on the context - e.g. extension lookups or other contextual
/// information. This is used for types that are not directly parsed from the
/// grammar, but rather require additional context to parse correctly.
pub(crate) trait ScopedParsePair: Sized {
    // The rule that this type is parsed from.
    fn rule() -> Rule;

    // The name of the protobuf message type that this type corresponds to.
    fn message() -> &'static str;

    // Parse a single instance of this type from a `pest::iterators::Pair<Rule>`.
    // The input must match the rule returned by `rule`; otherwise, a panic is
    // expected.
    fn parse_pair(
        extensions: &SimpleExtensions,
        pair: Pair<Rule>,
    ) -> Result<Self, MessageParseError>;
}

pub(crate) fn iter_pairs(pair: Pairs<'_, Rule>) -> RuleIter<'_> {
    RuleIter {
        iter: pair,
        done: false,
    }
}

pub(crate) struct RuleIter<'a> {
    iter: Pairs<'a, Rule>,
    // Set to true when done is called, so destructor doesn't panic
    done: bool,
}

impl<'a> From<Pairs<'a, Rule>> for RuleIter<'a> {
    fn from(iter: Pairs<'a, Rule>) -> Self {
        RuleIter { iter, done: false }
    }
}

impl<'a> RuleIter<'a> {
    pub(crate) fn peek(&self) -> Option<Pair<'a, Rule>> {
        self.iter.peek()
    }

    // Pop the next pair if it matches the rule. Returns None if not.
    pub(crate) fn try_pop(&mut self, rule: Rule) -> Option<Pair<'a, Rule>> {
        match self.peek() {
            Some(pair) if pair.as_rule() == rule => {
                self.iter.next();
                Some(pair)
            }
            _ => None,
        }
    }

    // Pop the next pair, asserting it matches the given rule. Panics if not.
    pub(crate) fn pop(&mut self, rule: Rule) -> Pair<'a, Rule> {
        let pair = self.iter.next().expect("expected another pair");
        assert_eq!(
            pair.as_rule(),
            rule,
            "expected rule {:?}, got {:?}",
            rule,
            pair.as_rule()
        );
        pair
    }

    // Parse the next pair if it matches the rule. Returns None if not.
    pub(crate) fn parse_if_next<T: ParsePair>(&mut self) -> Option<T> {
        match self.peek() {
            Some(pair) if pair.as_rule() == T::rule() => {
                self.iter.next();
                Some(T::parse_pair(pair))
            }
            _ => None,
        }
    }

    // Parse the next pair if it matches the rule. Returns None if not.
    pub(crate) fn parse_if_next_scoped<T: ScopedParsePair>(
        &mut self,
        extensions: &SimpleExtensions,
    ) -> Option<Result<T, MessageParseError>> {
        match self.peek() {
            Some(pair) if pair.as_rule() == T::rule() => {
                self.iter.next();
                Some(T::parse_pair(extensions, pair))
            }
            _ => None,
        }
    }

    // Parse the next pair, assuming it matches the rule. Panics if not.
    pub(crate) fn parse_next<T: ParsePair>(&mut self) -> T {
        let pair = self.iter.next().unwrap();
        T::parse_pair(pair)
    }

    // Parse the next pair, assuming it matches the rule. Panics if not.
    pub(crate) fn parse_next_scoped<T: ScopedParsePair>(
        &mut self,
        extensions: &SimpleExtensions,
    ) -> Result<T, MessageParseError> {
        let pair = self.iter.next().unwrap();
        T::parse_pair(extensions, pair)
    }

    pub(crate) fn done(mut self) {
        self.done = true;
        // A rule may end with the `EOI` marker to force full input consumption
        // (e.g. `virtual_read_relation`). That marker carries no data, so it is
        // not leftover content — skip it before asserting the iterator is empty.
        let next = match self.iter.next() {
            Some(pair) if pair.as_rule() == Rule::EOI => self.iter.next(),
            other => other,
        };
        assert_eq!(next, None);
    }
}

/// Make sure that the iterator was completely consumed when the iterator is
/// dropped - that we didn't leave any partially-parsed tokens.
///
/// This is not strictly necessary, but it's a good way to catch bugs.
impl Drop for RuleIter<'_> {
    fn drop(&mut self) {
        if self.done || thread::panicking() {
            return;
        }
        // If the iterator is not done, something probably went wrong.
        assert_eq!(self.iter.next(), None);
    }
}

/// A collection of named arguments (`name=value` pairs) extracted from a
/// named-argument-list rule, keyed by name with duplicate-name rejection.
pub(crate) struct ParsedNamedArgs<'a> {
    map: HashMap<&'a str, Pair<'a, Rule>>,
}

impl<'a> ParsedNamedArgs<'a> {
    pub(crate) fn new(pairs: Pairs<'a, Rule>, rule: Rule) -> Result<Self, MessageParseError> {
        let mut map = HashMap::new();
        for pair in pairs {
            assert_eq!(pair.as_rule(), rule);
            let mut inner = pair.clone().into_inner();
            let name_pair = inner.next().unwrap();
            let value_pair = inner.next().unwrap();
            assert_eq!(inner.next(), None);
            let name = name_pair.as_str();
            if map.contains_key(name) {
                return Err(MessageParseError::invalid(
                    "NamedArg",
                    name_pair.as_span(),
                    format!("Duplicate argument: {name}"),
                ));
            }
            map.insert(name, value_pair);
        }
        Ok(Self { map })
    }

    // Returns the pair if it exists and matches the rule, otherwise None.
    pub(crate) fn pop(mut self, name: &str, rule: Rule) -> (Self, Option<Pair<'a, Rule>>) {
        let pair = self.map.remove(name).inspect(|pair| {
            assert_eq!(pair.as_rule(), rule, "Rule mismatch for argument {name}");
        });
        (self, pair)
    }

    // Returns an error if there are any unused arguments.
    pub(crate) fn done(self) -> Result<(), MessageParseError> {
        if let Some((name, pair)) = self.map.iter().next() {
            return Err(MessageParseError::invalid(
                "NamedArgExtractor",
                // No span available for all unused args; use default.
                pair.as_span(),
                format!("Unknown argument: {name}"),
            ));
        }
        Ok(())
    }
}

/// Map a sort-direction enum identifier (without the leading `&`) to a
/// [`SortDirection`]. Shared by the `Sort` relation's `sort_field` parser and
/// the window function's `order=` parser, which reach it from different grammar
/// rules (`sort_direction` vs a generic `enum_value`) but accept the same set
/// of variant names. Lives in `common` so neither `relations` nor
/// `expressions` depends on the other for it.
pub(crate) fn sort_direction_from_str(
    name: &str,
    span: pest::Span,
) -> Result<SortDirection, MessageParseError> {
    match name {
        "AscNullsFirst" => Ok(SortDirection::AscNullsFirst),
        "AscNullsLast" => Ok(SortDirection::AscNullsLast),
        "DescNullsFirst" => Ok(SortDirection::DescNullsFirst),
        "DescNullsLast" => Ok(SortDirection::DescNullsLast),
        other => Err(MessageParseError::invalid(
            "SortDirection",
            span,
            format!("Unknown sort direction: {other}"),
        )),
    }
}

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

    use super::{ErrorKind, ExpressionParser, MessageParseError, ParsePair, ScopedParsePair};
    use crate::extensions::SimpleExtensions;

    /// Test-only adapter for parsing individual grammar fragments from strings.
    ///
    /// Production parsing goes through [`ParsePair`] and the structural [`Parser`](crate::Parser).
    pub(crate) trait Parse {
        fn parse(input: &str) -> Result<Self, MessageParseError>
        where
            Self: Sized;
    }

    impl<T: ParsePair> Parse for T {
        fn parse(input: &str) -> Result<Self, MessageParseError> {
            T::parse_str(input)
        }
    }

    /// Test-only adapter for parsing context-dependent grammar fragments from strings.
    ///
    /// Production parsing goes through [`ScopedParsePair`] and the structural
    /// [`Parser`](crate::Parser).
    pub(crate) trait ScopedParse: Sized {
        fn parse(extensions: &SimpleExtensions, input: &str) -> Result<Self, MessageParseError>
        where
            Self: Sized;
    }

    impl<T: ScopedParsePair> ScopedParse for T {
        fn parse(extensions: &SimpleExtensions, input: &str) -> Result<Self, MessageParseError> {
            let mut pairs = ExpressionParser::parse(Self::rule(), input).map_err(|e| {
                MessageParseError::new(Self::message(), ErrorKind::Syntax, Box::new(e))
            })?;
            assert_eq!(pairs.as_str(), input);
            let pair = pairs.next().unwrap();
            assert_eq!(pairs.next(), None);
            Self::parse_pair(extensions, pair)
        }
    }
}