linkrs-query 0.1.0

//! Tool Statement Parsing Module
//!
//! Responsible for parsing statements related to tool functions, including USE, SHOW, EXPLAIN, FETCH, LOOKUP, UNWIND, RETURN, WITH, YIELD, SET, REMOVE, etc.

use crate::parser::ast::stmt::*;
use crate::parser::ast::types::OrderDirection;
use crate::parser::core::error::{ParseError, ParseErrorKind};
use crate::parser::parsing::clause_parser::ClauseParser;
use crate::parser::parsing::expr_parser::parse_expression_with_context;
use crate::parser::parsing::parse_context::ParseContext;
use crate::parser::TokenKind;
use linkrs_core::types::expr::contextual::ContextualExpression;

/// Tool Syntax Parser
pub struct UtilStmtParser;

impl UtilStmtParser {
    pub fn new() -> Self {
        Self
    }

    /// Analysis of the USE statement
    pub fn parse_use_statement(&mut self, ctx: &mut ParseContext) -> Result<Stmt, ParseError> {
        let start_span = ctx.current_span();
        ctx.expect_token(TokenKind::Use)?;

        // `USE GRAPH <name>` is a namespace alias for `USE <name>`.
        // `SPACE` is intentionally not consumed here: it stays a legal
        // space name in USE position, unlike CREATE/DROP where SPACE is
        // the object keyword (see ddl_parser::consume_space_or_graph).
        // A lone `USE graph` names the space `graph`: GRAPH only acts as
        // the alias keyword when another identifier follows it.
        if ctx.check_keyword("GRAPH") {
            let graph_word = ctx.current_token().lexeme.clone();
            ctx.consume_keyword("GRAPH")?;
            match ctx.expect_identifier() {
                Ok(space) => {
                    return Ok(Stmt::Use(UseStmt {
                        span: start_span,
                        space,
                    }));
                }
                Err(_) => {
                    return Ok(Stmt::Use(UseStmt {
                        span: start_span,
                        space: graph_word,
                    }));
                }
            }
        }

        let space = ctx.expect_identifier()?;

        Ok(Stmt::Use(UseStmt {
            span: start_span,
            space,
        }))
    }

    /// Analyzing the internal methods of the SHOW CREATE statement
    pub fn parse_show_create_internal(
        &mut self,
        ctx: &mut ParseContext,
        start_span: crate::parser::ast::types::Span,
    ) -> Result<Stmt, ParseError> {
        ctx.expect_token(TokenKind::Create)?;

        let target = if ctx.match_token(TokenKind::Space) {
            ShowCreateTarget::Space(ctx.expect_identifier()?)
        } else if ctx.match_token(TokenKind::Tag) {
            ShowCreateTarget::Tag(ctx.expect_identifier()?)
        } else if ctx.match_token(TokenKind::Edge) {
            ShowCreateTarget::Edge(ctx.expect_identifier()?)
        } else if ctx.match_token(TokenKind::Index) {
            ShowCreateTarget::Index(ctx.expect_identifier()?)
        } else {
            return Err(ParseError::new(
                ParseErrorKind::UnexpectedToken,
                "Expected SPACE, TAG, EDGE, or INDEX after SHOW CREATE".to_string(),
                ctx.current_position(),
            ));
        };

        let span = ctx.span_since(start_span.start);

        Ok(Stmt::ShowCreate(ShowCreateStmt { span, target }))
    }

    /// Analysis of the internal methods of the SHOW ROLES command
    pub fn parse_show_roles_internal(
        &mut self,
        ctx: &mut ParseContext,
        start_span: crate::parser::ast::types::Span,
    ) -> Result<Stmt, ParseError> {
        ctx.expect_token(TokenKind::Roles)?;

        // Optional IN <space_name> clause; space names accept the same
        // hyphenated/quoted forms as other DCL object names.
        let space_name = if ctx.match_token(TokenKind::In) {
            Some(ctx.expect_dcl_name()?)
        } else {
            None
        };

        let span = ctx.span_since(start_span.start);

        Ok(Stmt::ShowRoles(ShowRolesStmt { span, space_name }))
    }

    /// Analyzing the EXPLAIN statement
    pub fn parse_explain_statement(&mut self, ctx: &mut ParseContext) -> Result<Stmt, ParseError> {
        let _start_span = ctx.current_span();
        ctx.expect_token(TokenKind::Explain)?;

        // After “EXPLAIN”, a substatement needs to be parsed.
        // Here, we need to invoke the main parser, but due to circular dependency issues, we are returning a placeholder.
        // The actual parsing will be handled within the StmtParser.
        Err(ParseError::new(
            ParseErrorKind::SyntaxError,
            "EXPLAIN should be handled by main parser".to_string(),
            ctx.current_position(),
        ))
    }

    /// Analysis of the FETCH statement
    pub fn parse_fetch_statement(&mut self, ctx: &mut ParseContext) -> Result<Stmt, ParseError> {
        let start_span = ctx.current_span();
        ctx.expect_token(TokenKind::Fetch)?;

        // The FETCH PROP ON <tag> <ids> syntax is supported.
        let _with_props = ctx.match_token(TokenKind::Prop);

        // `->` is an edge arrow inside FETCH, not the JSON access operator
        // (`expr -> 'key'` postfix JSON get), so expression parsing must not
        // consume `src -> dst` as a JSON get.
        let target = ctx.with_edge_syntax_mode(|ctx| {
            if ctx.match_token(TokenKind::On) {
                // Check for * (all tags)
                let tag_name = if ctx.match_token(TokenKind::Star) {
                    None // None means all tags
                } else {
                    Some(ctx.expect_identifier()?)
                };

                let first_expr = self.parse_expression(ctx)?;
                if ctx.check_token(TokenKind::Arrow) {
                    ctx.expect_token(TokenKind::Arrow)?;
                    let dst = self.parse_expression(ctx)?;
                    let rank = if ctx.match_token(TokenKind::At) {
                        Some(self.parse_expression(ctx)?)
                    } else {
                        None
                    };
                    let edge_type = tag_name.ok_or_else(|| {
                        ParseError::new(
                            ParseErrorKind::SyntaxError,
                            "FETCH edge form requires an edge type; '*' is only valid for vertices"
                                .to_string(),
                            ctx.current_position(),
                        )
                    })?;
                    Ok(FetchTarget::Edges {
                        src: first_expr,
                        dst,
                        edge_type,
                        rank,
                        properties: None,
                    })
                } else {
                    let mut ids = vec![first_expr];
                    while ctx.match_token(TokenKind::Comma) {
                        ids.push(self.parse_expression(ctx)?);
                    }
                    Ok(FetchTarget::Vertices {
                        tag_name,
                        ids,
                        properties: None,
                    })
                }
            } else if ctx.match_token(TokenKind::Tag) {
                let tag_name = Some(ctx.expect_identifier()?);
                let ids = self.parse_expression_list(ctx)?;
                Ok(FetchTarget::Vertices {
                    tag_name,
                    ids,
                    properties: None,
                })
            } else if ctx.match_token(TokenKind::Edge) {
                let edge_type = ctx.expect_identifier()?;
                let src = self.parse_expression(ctx)?;
                ctx.expect_token(TokenKind::Arrow)?;
                let dst = self.parse_expression(ctx)?;
                let rank = if ctx.match_token(TokenKind::At) {
                    Some(self.parse_expression(ctx)?)
                } else {
                    None
                };
                Ok(FetchTarget::Edges {
                    src,
                    dst,
                    edge_type,
                    rank,
                    properties: None,
                })
            } else {
                let ids = self.parse_expression_list(ctx)?;
                Ok(FetchTarget::Vertices {
                    tag_name: None,
                    ids,
                    properties: None,
                })
            }
        })?;

        let yield_clause = if ctx.match_token(TokenKind::Yield) {
            Some(ClauseParser::new().parse_yield_clause(ctx)?)
        } else {
            None
        };

        Ok(Stmt::Fetch(FetchStmt {
            span: start_span,
            target,
            yield_clause,
        }))
    }

    /// Analysis of the LOOKUP statement
    pub fn parse_lookup_statement(&mut self, ctx: &mut ParseContext) -> Result<Stmt, ParseError> {
        let start_span = ctx.current_span();
        ctx.expect_token(TokenKind::Lookup)?;

        let target = if ctx.match_token(TokenKind::On) {
            // Check if it's an EDGE or TAG keyword followed by identifier
            if ctx.current_token().kind == TokenKind::Edge {
                // LOOKUP ON EDGE <name>
                ctx.next_token(); // consume EDGE
                let name = ctx.expect_identifier()?;
                LookupTarget::Edge(name)
            } else if ctx.current_token().kind == TokenKind::Tag {
                // LOOKUP ON TAG <name>
                ctx.next_token(); // consume TAG
                let name = ctx.expect_identifier()?;
                LookupTarget::Tag(name)
            } else {
                // LOOKUP ON <name> - type will be resolved during validation
                let name = ctx.expect_identifier()?;
                LookupTarget::Unspecified(name)
            }
        } else {
            LookupTarget::Unspecified(String::new())
        };

        let where_clause = if ctx.match_token(TokenKind::Where) {
            ctx.recover_clause(|_| Ok(None), |c| self.parse_expression(c).map(Some))?
        } else {
            None
        };

        let yield_clause = if ctx.match_token(TokenKind::Yield) {
            Some(ClauseParser::new().parse_yield_clause(ctx)?)
        } else {
            None
        };

        Ok(Stmt::Lookup(LookupStmt {
            span: start_span,
            target,
            where_clause,
            yield_clause,
        }))
    }

    /// Analysis of the UNWIND statement
    pub fn parse_unwind_statement(&mut self, ctx: &mut ParseContext) -> Result<Stmt, ParseError> {
        let start_span = ctx.current_span();
        ctx.expect_token(TokenKind::Unwind)?;

        let expression = self.parse_expression(ctx)?;

        ctx.match_token(TokenKind::As);

        let variable = ctx.expect_identifier()?;

        let return_clause = if ctx.match_token(TokenKind::Return) {
            Some(ClauseParser::new().parse_return_clause(ctx)?)
        } else {
            None
        };

        let order_by = None;
        let limit = None;
        let skip = None;

        Ok(Stmt::Unwind(UnwindStmt {
            span: start_span,
            expression,
            variable,
            return_clause,
            order_by,
            limit,
            skip,
        }))
    }

    /// Analysis of the RETURN statement
    pub fn parse_return_statement(&mut self, ctx: &mut ParseContext) -> Result<Stmt, ParseError> {
        let start_span = ctx.current_span();
        ctx.expect_token(TokenKind::Return)?;

        let return_clause = ClauseParser::new().parse_return_clause(ctx)?;

        Ok(Stmt::Return(ReturnStmt {
            span: start_span,
            items: return_clause.items,
            distinct: return_clause.distinct,
            order_by: return_clause.order_by,
            skip: return_clause.skip,
            limit: return_clause.limit,
        }))
    }

    /// Analysis of the WITH statement
    pub fn parse_with_statement(&mut self, ctx: &mut ParseContext) -> Result<Stmt, ParseError> {
        let start_span = ctx.current_span();
        ctx.expect_token(TokenKind::With)?;

        let recursive = ctx.match_token(TokenKind::Recursive);

        // Common table expressions: `name AS (<query>) [, ...]`.
        // A `IDENT AS (` sequence can never start a plain WITH item (an
        // alias after AS must be an identifier), so detection is unambiguous.
        let mut ctes = Vec::new();
        while let Some(cte) = Self::try_parse_cte_def(ctx)? {
            ctes.push(cte);
            if !ctx.match_token(TokenKind::Comma) {
                break;
            }
        }

        let mut items = Vec::new();
        let distinct = ctx.match_token(TokenKind::Distinct);

        loop {
            let expr = self.parse_expression(ctx)?;
            let alias = if ctx.match_token(TokenKind::As) {
                Some(ctx.expect_identifier()?)
            } else {
                None
            };
            items.push(ReturnItem::Expression {
                expression: expr,
                alias,
            });
            if !ctx.match_token(TokenKind::Comma) {
                break;
            }
        }

        let where_clause = if ctx.match_token(TokenKind::Where) {
            ctx.recover_clause(|_| Ok(None), |c| self.parse_expression(c).map(Some))?
        } else {
            None
        };

        // Explanation of `ORDER BY`
        let order_by = if ctx.match_token(TokenKind::Order) {
            ctx.recover_clause(
                |_| Ok(None),
                |c| {
                    c.expect_token(TokenKind::By)?;
                    self.parse_order_by_clause(c).map(Some)
                },
            )?
        } else {
            None
        };

        // Analysis of SKIP
        let skip = if ctx.match_token(TokenKind::Skip) {
            ctx.recover_clause(
                |_| Ok(None),
                |c| {
                    let count = c.expect_integer_literal()? as usize;
                    Ok(Some(crate::parser::ast::types::SkipClause {
                        span: c.current_span(),
                        count,
                    }))
                },
            )?
        } else {
            None
        };

        // Analysis of the LIMIT clause
        let limit = if ctx.match_token(TokenKind::Limit) {
            ctx.recover_clause(
                |_| Ok(None),
                |c| {
                    let count = c.expect_integer_literal()? as usize;
                    Ok(Some(crate::parser::ast::types::LimitClause {
                        span: c.current_span(),
                        count,
                    }))
                },
            )?
        } else {
            None
        };

        Ok(Stmt::With(WithStmt {
            span: start_span,
            items,
            where_clause,
            distinct,
            order_by,
            skip,
            limit,
            recursive,
            ctes,
        }))
    }

    /// Try to parse one `name AS (<query>)` CTE definition.
    ///
    /// Returns `Ok(None)` (without consuming input) when the upcoming tokens
    /// do not form a CTE header.
    fn try_parse_cte_def(ctx: &mut ParseContext) -> Result<Option<CteDef>, ParseError> {
        let ckpt = ctx.checkpoint();
        let name = match ctx.expect_identifier() {
            Ok(name) => name,
            Err(_) => {
                ctx.restore(ckpt);
                return Ok(None);
            }
        };
        // A CTE definition is always `name AS (`; `name (` starts a function
        // call item instead (e.g. `WITH count(*) AS total`), so back off and
        // let the expression parser handle it.
        if ctx.check_token(TokenKind::LParen) {
            ctx.restore(ckpt);
            return Ok(None);
        }
        if !ctx.check_keyword("AS") {
            ctx.restore(ckpt);
            return Ok(None);
        }
        ctx.consume_keyword("AS")?;
        if !ctx.check_token(TokenKind::LParen) {
            ctx.restore(ckpt);
            return Ok(None);
        }
        let cte_span = ctx.current_span();
        ctx.expect_token(TokenKind::LParen)?;
        let body = crate::parser::parsing::StmtParser::parse_statement(ctx).map_err(|e| {
            ParseError::new(
                ParseErrorKind::SyntaxError,
                format!("Invalid CTE body for '{}': {}", name, e),
                ctx.current_position(),
            )
        })?;
        ctx.expect_token(TokenKind::RParen)?;
        Ok(Some(CteDef {
            span: cte_span,
            name,
            body: Box::new(body),
        }))
    }

    /// Analysis of the YIELD statement
    pub fn parse_yield_statement(&mut self, ctx: &mut ParseContext) -> Result<Stmt, ParseError> {
        let start_span = ctx.current_span();
        ctx.expect_token(TokenKind::Yield)?;

        let yield_clause = ClauseParser::new().parse_yield_clause(ctx)?;

        Ok(Stmt::Yield(YieldStmt {
            span: start_span,
            items: yield_clause.items,
            where_clause: yield_clause.where_clause,
            distinct: yield_clause.distinct,
            order_by: yield_clause.order_by,
            skip: yield_clause.skip,
            limit: yield_clause.limit,
            sample: yield_clause.sample,
        }))
    }

    /// Analyzing the SET statement
    pub fn parse_set_statement(&mut self, ctx: &mut ParseContext) -> Result<Stmt, ParseError> {
        let start_span = ctx.current_span();
        ctx.expect_token(TokenKind::Set)?;

        let set_clause = ClauseParser::new().parse_set_clause(ctx)?;

        Ok(Stmt::Set(SetStmt {
            span: start_span,
            assignments: set_clause.assignments,
        }))
    }

    /// Analysis of the REMOVE statement
    pub fn parse_remove_statement(&mut self, ctx: &mut ParseContext) -> Result<Stmt, ParseError> {
        let start_span = ctx.current_span();
        ctx.expect_token(TokenKind::Remove)?;

        let mut items = Vec::new();
        loop {
            items.push(self.parse_expression(ctx)?);
            if !ctx.match_token(TokenKind::Comma) {
                break;
            }
        }

        Ok(Stmt::Remove(RemoveStmt {
            span: start_span,
            items,
        }))
    }

    /// Parse the list of expressions
    fn parse_expression_list(
        &mut self,
        ctx: &mut ParseContext,
    ) -> Result<Vec<ContextualExpression>, ParseError> {
        let mut expressions = Vec::new();

        loop {
            expressions.push(self.parse_expression(ctx)?);
            if !ctx.match_token(TokenKind::Comma) {
                break;
            }
        }

        Ok(expressions)
    }

    /// Analyzing the expression
    fn parse_expression(
        &mut self,
        ctx: &mut ParseContext,
    ) -> Result<ContextualExpression, ParseError> {
        parse_expression_with_context(ctx, ctx.expression_context_clone())
    }

    /// Analysis of the ORDER BY clause
    fn parse_order_by_clause(
        &mut self,
        ctx: &mut ParseContext,
    ) -> Result<OrderByClause, ParseError> {
        let span = ctx.current_span();
        let mut items = Vec::new();

        loop {
            let expr = self.parse_expression(ctx)?;
            let direction = if ctx.match_token(TokenKind::Asc) {
                OrderDirection::Asc
            } else if ctx.match_token(TokenKind::Desc) {
                OrderDirection::Desc
            } else {
                OrderDirection::Asc
            };
            items.push(OrderByItem {
                expression: expr,
                direction,
            });
            if !ctx.match_token(TokenKind::Comma) {
                break;
            }
        }

        Ok(OrderByClause { span, items })
    }
}

impl Default for UtilStmtParser {
    fn default() -> Self {
        Self::new()
    }
}