rlox/src/frontend/parser.rs

use crate::{
    frontend::{
        ast::{AstNode, Expr, Stmt},
        source_registry::{SourcePosition, SourceSlice},
        tokens::{LiteralValue, LoxFunction, Token, TokenType},
    },
    logging::display_ast::{pretty_print_with_config, PrettyConfig},
    result::{parse_error, runtime_error, LoxResult},
};

pub struct Parser {
    tokens: Vec<Token>,
    current: usize,
}

impl Parser {
    pub fn new(tokens: Vec<Token>) -> Self {
        Self { tokens, current: 0 }
    }

    pub fn parse(&mut self) -> LoxResult<Vec<AstNode<Stmt>>> {
        let mut statements = Vec::new();
        let mut errors = Vec::new();
        println!("\n \n ++++++++++++++++++++++++++++++++++++\nStarting parsing");
        // Con configurazione personalizzata
        let config = PrettyConfig {
            indent: "    ".to_string(),
            max_depth: None,
            show_positions: false,
            ..Default::default()
        };
        while !self.is_at_end() {
            match self.statement() {
                Ok(stmt) => {
                    println!(
                        "Parsed Successfully: {}",
                        pretty_print_with_config(&stmt.node, &config)
                    );

                    statements.push(stmt.clone());
                }
                Err(err) => {
                    println!("Error: {:?}", err);
                    errors.push(err);
                    let old_current = self.current;
                    self.synchronize();
                    println!("Synchronized from {} to {}", old_current, self.current);
                }
            }
        }

        if !errors.is_empty() {
            return Err(errors.into_iter().next().unwrap());
        }

        Ok(statements)
    }

    fn statement(&mut self) -> LoxResult<AstNode<Stmt>> {
        match (&self.peek().token_type, &self.peek_next().token_type) {
            (TokenType::Print, _) => self.print_statement(),
            (TokenType::Return, _) => self.return_statement(),
            (TokenType::StartBlock, _) => self.block_statement(),
            (TokenType::Var, TokenType::Identifier) => {
                self.advance();
                self.var_statement()
            }
            (TokenType::Identifier, TokenType::Colon) => self.var_statement(),
            (TokenType::Identifier, TokenType::Equal) => self.assignment_statement(),
            (TokenType::If, _) => self.if_statement(),
            (TokenType::While, _) => self.while_statement(),
            (TokenType::For, _) => self.for_statement(),
            _ => self.expression_statement(),
        }
    }

    fn for_statement(&mut self) -> LoxResult<AstNode<Stmt>> {
        let start_slice = self.peek().source_slice.clone();
        self.advance(); // consume 'for'

        let variable = self.var_statement()?;
        let condition = self.expression()?;
        self.consume(TokenType::Semicolon, "Expected ';' after for condition")?;
        let increment = self.assignment_statement()?;

        self.consume(TokenType::StartBlock, "Expected 'do' after for range")?;

        let body = self.statement()?;
        self.consume(TokenType::EndBlock, "Expected 'end' after for body")?;

        let end_slice = self.previous().source_slice.clone();
        let combined_slice = SourceSlice::from_positions(
            start_slice.source_id,
            start_slice.start_position,
            end_slice.end_position,
        );

        Ok(AstNode::new(
            Stmt::For {
                variable: Box::new(variable),
                condition: Box::new(condition),
                increment: Box::new(increment),
                body: Box::new(body),
            },
            combined_slice,
        ))
    }

    // fn for_statement(&mut self) -> LoxResult<AstNode<Stmt>> {
    //     let start_slice = self.peek().source_slice.clone();
    //     self.advance();
    //     let variable = if let Token {
    //         token_type: TokenType::Identifier,
    //         literal: Some(variable_name),
    //         source_slice,
    //         ..
    //     } = self.peek()
    //     {
    //         AstNode {
    //             node: Expr::Literal {
    //                 value: variable_name.clone(),
    //             },
    //             source_slice: source_slice.clone(),
    //         }
    //     } else {
    //         return parse_error(
    //             self.peek().source_slice.clone(),
    //             "Expect variable name after 'for' keyword.",
    //         );
    //     };
    //     self.advance();
    //     self.consume(TokenType::In, "Expect 'in' after for variable.")?;
    //     let iterable = self.expression()?;
    //     let body = self.statement()?;
    //     let end_slice = body.source_slice.clone();
    //     let combined_slice = SourceSlice::from_positions(
    //         start_slice.source_id,
    //         start_slice.start_position,
    //         end_slice.end_position,
    //     );
    //     Ok(AstNode::new(
    //         Stmt::For {
    //             variable: Box::new(variable),
    //             iterable: Box::new(iterable),
    //             body: Box::new(body),
    //         },
    //         combined_slice,
    //     ))
    // }

    fn while_statement(&mut self) -> LoxResult<AstNode<Stmt>> {
        let start_slice = self.peek().source_slice.clone();
        self.advance();
        let condition = self.expression()?;
        let body = self.statement()?;
        let end_slice = body.source_slice.clone();
        let combined_slice = SourceSlice::from_positions(
            start_slice.source_id,
            start_slice.start_position,
            end_slice.end_position,
        );
        Ok(AstNode::new(
            Stmt::While {
                condition: Box::new(condition),
                body: Box::new(body),
            },
            combined_slice,
        ))
    }

    fn if_statement(&mut self) -> LoxResult<AstNode<Stmt>> {
        let start_slice = self.peek().source_slice.clone();
        self.advance(); // consume 'if'
        let condition = self.expression()?;
        self.consume(TokenType::Then, "Expect 'then' after if condition.")?;
        let then_branch = self.statement()?;
        let mut elif_branches = Vec::new();
        let mut last_slice = then_branch.source_slice.clone();

        while self.peek().token_type == TokenType::Elif {
            self.advance(); // consume 'elif'
            let elif_condition = self.expression()?;
            self.consume(TokenType::Then, "Expect 'then' after elif condition.")?;
            let elif_branch = self.statement()?;
            last_slice = elif_branch.source_slice.clone();
            elif_branches.push((Box::new(elif_condition), Box::new(elif_branch)));
        }

        let else_branch = if self.peek().token_type == TokenType::Else {
            self.advance(); // consume 'else'
            let else_stmt = self.statement()?;
            last_slice = else_stmt.source_slice.clone();
            Some(Box::new(else_stmt))
        } else {
            None
        };

        let combined_slice = SourceSlice::from_positions(
            start_slice.source_id,
            start_slice.start_position,
            last_slice.end_position,
        );

        Ok(AstNode::new(
            Stmt::If {
                condition: Box::new(condition),
                then_branch: Box::new(then_branch),
                elif_branch: elif_branches,
                else_branch: else_branch,
            },
            combined_slice,
        ))
    }

    fn var_statement(&mut self) -> LoxResult<AstNode<Stmt>> {
        let start_slice = self.peek().source_slice.clone();
        let name = self.consume(TokenType::Identifier, "Expect variable name.")?;
        let name_lexeme = name.lexeme.clone();

        let mut _type_annotation = LiteralValue::Nil;
        self.consume(TokenType::Colon, "Expect column")?;

        match self.peek().token_type {
            TokenType::Equal | TokenType::Identifier => {
                self.variable_declaration(start_slice, name_lexeme)
            }
            TokenType::Colon => self.function_declaration(start_slice, name_lexeme),
            _ => runtime_error(start_slice, "this is not supposed to be here"),
        }
    }

    fn function_declaration(
        &mut self,
        start_slice: SourceSlice,
        name_lexeme: String,
    ) -> LoxResult<AstNode<Stmt>> {
        self.advance();
        self.consume(TokenType::Fn, "Expect 'fn' after '::'")?;
        self.consume(TokenType::LeftParen, "Expect '(' after 'fn' ")?;
        let mut parameters: Vec<(String, String)> = vec![];
        while self.peek().token_type != TokenType::RightParen {
            let expr = self
                .consume(TokenType::Identifier, "Expect identifier after '('")?
                .clone();
            if self.peek().token_type == TokenType::Colon {
                self.advance();
                let type_ =
                    self.consume(TokenType::Identifier, "Expected identifier after ':' ")?;
                parameters.push((expr.lexeme, type_.lexeme.clone()));
            } else {
                parameters.push((expr.lexeme, "Any".to_string()));
            }
            if self.peek().token_type == TokenType::Comma {
                self.advance();
            }
        }
        self.advance();
        let end_position = self.peek().source_slice.clone();
        let combine_position = SourceSlice {
            source_id: start_slice.source_id,
            start_position: start_slice.start_position.clone(),
            end_position: end_position.end_position,
        };
        let mut guard: Option<Box<Expr>> = None;
        if self.peek().token_type == TokenType::RightBrace {
            self.advance();
            guard = Some(Box::new(self.expression()?.node.clone()));
            self.consume(TokenType::RightBrace, "Expected '}' after guard expression")?;
        }

        let body = self.statement()?;
        let node = AstNode {
            node: Expr::Literal {
                value: LiteralValue::Function(LoxFunction {
                    parameters,
                    body,
                    closure: None,
                    guard,
                }),
            },
            source_slice: combine_position.clone(),
        };
        Ok(AstNode::new(
            Stmt::VarDeclaration {
                name: name_lexeme,
                initializer: Some(Box::new(node)),
            },
            combine_position.clone(),
        ))
    }

    fn variable_declaration(
        &mut self,
        start_slice: SourceSlice,
        name_lexeme: String,
    ) -> LoxResult<AstNode<Stmt>> {
        if self.peek().token_type == TokenType::Identifier {
            self.advance();
            if self.peek().token_type == TokenType::Identifier {
                self.advance();
            }
            // todo: make type annotation
        }
        let mut value = AstNode {
            node: Expr::Literal {
                value: LiteralValue::Nil,
            },
            source_slice: start_slice.clone(),
        };
        if self.peek().token_type == TokenType::Equal {
            self.advance();
            value = self.expression()?;
        }
        let semicolon = self.consume(
            TokenType::Semicolon,
            "Expect ';' after variable declaration.",
        )?;
        let end_slice = semicolon.source_slice.clone();
        let combined_slice = SourceSlice::from_positions(
            start_slice.source_id,
            start_slice.start_position,
            end_slice.end_position,
        );
        Ok(AstNode::new(
            Stmt::VarDeclaration {
                name: name_lexeme,
                initializer: Some(Box::new(value)),
            },
            combined_slice,
        ))
    }
    fn assignment_statement(&mut self) -> LoxResult<AstNode<Stmt>> {
        let start_slice = self.peek().source_slice.clone();
        let name = self.consume(TokenType::Identifier, "Expect variable name.")?;
        let name_lexeme = name.lexeme.clone();
        self.consume(TokenType::Equal, "Expect '=' after variable name.")?;
        let value = self.expression()?;
        let semicolon = self.consume(
            TokenType::Semicolon,
            "Expect ';' after variable declaration.",
        )?;
        let end_slice = semicolon.source_slice.clone();
        let combined_slice = SourceSlice::from_positions(
            start_slice.source_id,
            start_slice.start_position,
            end_slice.end_position,
        );
        Ok(AstNode::new(
            Stmt::VarAssigment {
                name: name_lexeme,
                value: Box::new(value),
            },
            combined_slice,
        ))
    }

    fn print_statement(&mut self) -> LoxResult<AstNode<Stmt>> {
        // consume the print keyword
        let start_slice = self.peek().source_slice.clone();
        self.advance();
        let expr = self.expression()?;
        let semicolon = self.consume(TokenType::Semicolon, "Expect ';' after value.")?;
        let end_slice = semicolon.source_slice.clone();
        let combined_slice = SourceSlice::from_positions(
            start_slice.source_id,
            start_slice.start_position,
            end_slice.end_position,
        );
        Ok(AstNode::new(
            Stmt::Print {
                expression: Box::new(expr),
            },
            combined_slice,
        ))
    }

    fn block_statement(&mut self) -> LoxResult<AstNode<Stmt>> {
        let start_slice = self.peek().source_slice.clone();
        self.advance();
        let mut statements = Vec::new();
        while self.peek().token_type != TokenType::EndBlock && !self.is_at_end() {
            let stmt = self.statement()?;
            let should_break = matches!(stmt.node, Stmt::Expression { .. });
            statements.push(stmt);
            if should_break {
                break;
            }
        }
        let end_token = self.consume(TokenType::EndBlock, "Expect 'end' after block.")?;
        let end_slice = end_token.source_slice.clone();
        let combined_slice = SourceSlice::from_positions(
            start_slice.source_id,
            start_slice.start_position,
            end_slice.end_position,
        );
        Ok(AstNode::new(
            Stmt::Block {
                statements: Box::new(statements),
            },
            combined_slice,
        ))
    }

    fn expression_statement(&mut self) -> LoxResult<AstNode<Stmt>> {
        let start_slice = self.peek().source_slice.clone();
        let expr = self.expression()?;
        if self.peek().token_type == TokenType::Semicolon {
            let semicolon = self.advance();
            let end_slice = semicolon.source_slice.clone();
            let combined_slice = SourceSlice::from_positions(
                start_slice.source_id,
                start_slice.start_position,
                end_slice.end_position,
            );
            return Ok(AstNode::new(
                Stmt::Stmt {
                    expression: Box::new(expr),
                },
                combined_slice,
            ));
        }
        // Use the expression's source slice for expression statements without semicolon
        let expr_slice = expr.source_slice.clone();
        Ok(AstNode::new(
            Stmt::Expression {
                expression: Box::new(expr),
            },
            expr_slice,
        ))
    }

    fn return_statement(&mut self) -> LoxResult<AstNode<Stmt>> {
        let start_slice = self.peek().source_slice.clone();
        self.advance();
        let expr = self.expression()?;
        let semicolon = self.consume(TokenType::Semicolon, "Expect ';' after return value.")?;
        let end_slice = semicolon.source_slice.clone();
        let combined_slice = SourceSlice::from_positions(
            start_slice.source_id,
            start_slice.start_position,
            end_slice.end_position,
        );
        Ok(AstNode::new(
            Stmt::Return {
                expression: Box::new(expr),
            },
            combined_slice,
        ))
    }

    fn expression(&mut self) -> LoxResult<AstNode<Expr>> {
        self.or_and()
    }

    fn or_and(&mut self) -> LoxResult<AstNode<Expr>> {
        let mut expr = self.equality()?;

        while [TokenType::Or, TokenType::And].contains(&self.peek().token_type) {
            let start_slice = expr.source_slice.clone();
            let operator = self.peek().token_type.clone();
            self.advance();
            let right = self.equality()?;
            let end_slice = right.source_slice.clone();
            let combined_slice = SourceSlice::from_positions(
                start_slice.source_id,
                start_slice.start_position,
                end_slice.end_position,
            );
            expr = AstNode::new(
                Expr::Binary {
                    left: Box::new(expr),
                    operator,
                    right: Box::new(right),
                },
                combined_slice,
            );
        }

        Ok(expr)
    }
    fn equality(&mut self) -> LoxResult<AstNode<Expr>> {
        let mut expr = self.comparison()?;

        while [TokenType::BangEqual, TokenType::EqualEqual].contains(&self.peek().token_type) {
            let start_slice = expr.source_slice.clone();
            let operator = self.peek().token_type.clone();
            self.advance();
            let right = self.comparison()?;
            let end_slice = right.source_slice.clone();
            let combined_slice = SourceSlice::from_positions(
                start_slice.source_id,
                start_slice.start_position,
                end_slice.end_position,
            );
            expr = AstNode::new(
                Expr::Binary {
                    left: Box::new(expr),
                    operator,
                    right: Box::new(right),
                },
                combined_slice,
            );
        }

        Ok(expr)
    }

    fn comparison(&mut self) -> LoxResult<AstNode<Expr>> {
        let mut expr = self.term()?;

        while [
            TokenType::Greater,
            TokenType::GreaterEqual,
            TokenType::Less,
            TokenType::LessEqual,
        ]
        .contains(&self.peek().token_type)
        {
            let start_slice = expr.source_slice.clone();
            let operator = self.peek().token_type.clone();
            self.advance();
            let right = self.term()?;
            let end_slice = right.source_slice.clone();
            let combined_slice = SourceSlice::from_positions(
                start_slice.source_id,
                start_slice.start_position,
                end_slice.end_position,
            );
            expr = AstNode::new(
                Expr::Binary {
                    left: Box::new(expr),
                    operator,
                    right: Box::new(right),
                },
                combined_slice,
            );
        }

        Ok(expr)
    }

    fn term(&mut self) -> LoxResult<AstNode<Expr>> {
        let mut expr = self.factor()?;

        while [TokenType::Plus, TokenType::Minus].contains(&self.peek().token_type) {
            let start_slice = expr.source_slice.clone();
            let operator = self.peek().token_type.clone();
            self.advance();
            let right = self.factor()?;
            let end_slice = right.source_slice.clone();
            let combined_slice = SourceSlice::from_positions(
                start_slice.source_id,
                start_slice.start_position,
                end_slice.end_position,
            );
            expr = AstNode::new(
                Expr::Binary {
                    left: Box::new(expr),
                    operator,
                    right: Box::new(right),
                },
                combined_slice,
            );
        }

        Ok(expr)
    }

    fn factor(&mut self) -> LoxResult<AstNode<Expr>> {
        let mut expr = self.unary()?;

        while [TokenType::Slash, TokenType::Star].contains(&self.peek().token_type) {
            let start_slice = expr.source_slice.clone();
            let operator = self.peek().token_type.clone();
            self.advance();
            let right = self.unary()?;
            let end_slice = right.source_slice.clone();
            let combined_slice = SourceSlice::from_positions(
                start_slice.source_id,
                start_slice.start_position,
                end_slice.end_position,
            );
            expr = AstNode::new(
                Expr::Binary {
                    left: Box::new(expr),
                    operator,
                    right: Box::new(right),
                },
                combined_slice,
            );
        }

        Ok(expr)
    }

    fn unary(&mut self) -> LoxResult<AstNode<Expr>> {
        if [TokenType::Bang, TokenType::Minus].contains(&self.peek().token_type) {
            let operator = self.peek().token_type.clone();
            let source_slice = self.peek().source_slice.clone();
            self.advance();
            let right = self.unary()?;
            return Ok(AstNode::new(
                Expr::Unary {
                    operator,
                    operand: Box::new(right),
                },
                source_slice,
            ));
        }

        self.call()
    }

    fn call(&mut self) -> LoxResult<AstNode<Expr>> {
        let mut expr = self.primary()?;

        loop {
            if self.peek().token_type == TokenType::LeftParen {
                expr = self.finish_call(expr)?;
            } else {
                break;
            }
        }

        Ok(expr)
    }
    fn finish_call(&mut self, callee: AstNode<Expr>) -> LoxResult<AstNode<Expr>> {
        let start_slice = callee.source_slice.start_position.clone();
        self.advance(); // Consume '('

        let mut arguments = Vec::new();

        if self.peek().token_type != TokenType::RightParen {
            loop {
                arguments.push(self.expression()?);
                if self.peek().token_type == TokenType::Comma {
                    self.advance();
                } else {
                    break;
                }
            }
        }

        self.consume(TokenType::RightParen, "Expect ')' after arguments.")?;

        // Estendi il source_slice per includere le parentesi di chiusura
        let end_slice = self.previous().source_slice.end_position.clone();
        let full_slice = SourceSlice {
            source_id: callee.source_slice.source_id.clone(),
            start_position: start_slice.clone(),
            end_position: end_slice.clone(),
        };

        Ok(AstNode::new(
            Expr::Call {
                callee: Box::new(callee),
                arguments,
            },
            full_slice,
        ))
    }

    fn primary(&mut self) -> LoxResult<AstNode<Expr>> {
        match self.peek().token_type {
            TokenType::False => {
                let source_slice = self.peek().source_slice.clone();
                self.advance();
                Ok(AstNode::new(
                    Expr::Literal {
                        value: LiteralValue::Boolean(false),
                    },
                    source_slice,
                ))
            }
            TokenType::True => {
                let source_slice = self.peek().source_slice.clone();
                self.advance();
                Ok(AstNode::new(
                    Expr::Literal {
                        value: LiteralValue::Boolean(true),
                    },
                    source_slice,
                ))
            }
            TokenType::Nil => {
                let source_slice = self.peek().source_slice.clone();
                self.advance();
                Ok(AstNode::new(
                    Expr::Literal {
                        value: LiteralValue::Nil,
                    },
                    source_slice,
                ))
            }
            TokenType::Number => {
                let source_slice = self.peek().source_slice.clone();
                let literal = self.peek().literal.clone();
                self.advance();
                if let Some(literal) = literal {
                    Ok(AstNode::new(Expr::Literal { value: literal }, source_slice))
                } else {
                    parse_error(self.peek().source_slice.clone(), "Expected number literal")
                }
            }
            TokenType::String => {
                let source_slice = self.peek().source_slice.clone();
                let literal = self.peek().literal.clone();
                self.advance();
                if let Some(literal) = literal {
                    Ok(AstNode::new(Expr::Literal { value: literal }, source_slice))
                } else {
                    parse_error(self.peek().source_slice.clone(), "Expected string literal")
                }
            }
            TokenType::LeftParen => {
                let source_slice = self.peek().source_slice.clone();
                self.advance();
                let expr = self.expression()?;
                self.consume(TokenType::RightParen, "Expect ')' after expression.")?;
                Ok(AstNode::new(
                    Expr::Grouping {
                        expression: Box::new(expr),
                    },
                    source_slice,
                ))
            }
            TokenType::Identifier => {
                let name = self.peek().lexeme.clone();
                let source_slice = self.peek().source_slice.clone();
                self.advance();
                Ok(AstNode::new(Expr::Identifier { name }, source_slice))
            }
            _ => parse_error(self.peek().source_slice.clone(), "Expect expression."),
        }
    }

    fn is_at_end(&self) -> bool {
        if self.current >= self.tokens.len() {
            true
        } else {
            self.tokens[self.current].token_type == TokenType::Eof
        }
    }

    fn advance(&mut self) -> &Token {
        if !self.is_at_end() {
            self.current += 1;
        }
        self.previous()
    }

    fn peek(&self) -> &Token {
        if self.is_at_end() {
            &self.tokens.last().unwrap()
        } else {
            &self.tokens[self.current]
        }
    }

    fn peek_next(&self) -> &Token {
        if self.is_at_end() {
            &self.peek()
        } else {
            &self.tokens[self.current + 1]
        }
    }

    fn previous(&self) -> &Token {
        if self.current == 0 {
            &self.tokens[0]
        } else {
            &self.tokens[self.current - 1]
        }
    }

    fn consume(&mut self, token_type: TokenType, message: &str) -> LoxResult<&Token> {
        if self.peek().token_type == token_type {
            self.advance();
            Ok(self.previous())
        } else {
            return parse_error(self.peek().source_slice.clone(), message);
        }
    }

    fn synchronize(&mut self) {
        self.advance();

        while !self.is_at_end() {
            if self.previous().token_type == TokenType::Semicolon {
                return;
            }

            match self.peek().token_type {
                TokenType::Class
                | TokenType::Fun
                | TokenType::For
                | TokenType::If
                | TokenType::While
                | TokenType::Print
                | TokenType::Return => return,
                _ => {}
            }

            self.advance();
        }
    }
}