finx 0.1.0

// Parser for the toy language with recursive descent parsing
// Uses Pratt parsing for expressions to handle operator precedence

use crate::lexer::Token;
use std::collections::HashMap;

#[derive(Debug, PartialEq, Clone)]
pub enum Expr {
    Identifier(String),
    Number(f64),
    String(String),
    Bool(bool),
    Null,
    Binary {
        left: Box<Expr>,
        op: Token,
        right: Box<Expr>,
    },
    Call {
        callee: Box<Expr>,
        args: Vec<Expr>,
    },
    Block(Vec<Stmt>), // Add Block variant for compound statements
}

#[derive(Debug, PartialEq, Clone)]
pub enum Stmt {
    Let {
        name: String,
        value: Expr,
    },
    Assign {
        name: String,
        value: Expr,
    },
    Expr(Expr),
    Fn {
        name: String,
        params: Vec<String>,
        body: Vec<Stmt>,
    },
    If {
        cond: Expr,
        then_branch: Vec<Stmt>,
        else_branch: Option<Box<Stmt>>, // else_branch can be another If or a block
    },
    While {
        cond: Expr,
        body: Vec<Stmt>,
    },
    For {
        var: String,
        start: Expr,
        end: Expr,
        body: Vec<Stmt>,
    },
    Return(Expr),
}

#[derive(Debug)]
pub enum ParseError {
    UnexpectedToken { expected: String, found: Token },
}

impl std::fmt::Display for ParseError {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            ParseError::UnexpectedToken { expected, found } => {
                write!(
                    f,
                    "Unexpected token: expected {}, found {:?}",
                    expected, found
                )
            }
        }
    }
}

pub struct Parser<'a> {
    tokens: &'a [Token],
    pos: usize,
    fn_table: HashMap<String, usize>,
}

impl<'a> Parser<'a> {
    pub fn new(tokens: &'a [Token]) -> Self {
        Parser {
            tokens,
            pos: 0,
            fn_table: HashMap::new(),
        }
    }

    /// Peek at the current token without consuming it
    fn peek(&self) -> Option<&Token> {
        self.tokens.get(self.pos)
    }

    /// Consume and return the current token
    fn next(&mut self) -> Option<&Token> {
        let tok = self.tokens.get(self.pos);
        self.pos += 1;
        tok
    }

    /// Consume a token if it matches the expected type
    fn consume(&mut self, expected: &Token) -> bool {
        if self.peek() == Some(expected) {
            self.next();
            true
        } else {
            false
        }
    }

    /// Parse a block of statements enclosed in braces
    fn parse_block(&mut self) -> Result<Vec<Stmt>, ParseError> {
        if !self.consume(&Token::LBrace) {
            return Ok(Vec::new());
        }

        let mut stmts = Vec::new();
        while let Some(token) = self.peek() {
            if let Token::RBrace = token {
                self.next();
                break;
            }
            match self.parse_stmt()? {
                Some(stmt) => stmts.push(stmt),
                None => {}
            }
        }
        Ok(stmts)
    }

    /// Parse function parameters enclosed in parentheses
    fn parse_function_params(&mut self) -> Result<Vec<String>, ParseError> {
        let mut params = Vec::new();
        if !self.consume(&Token::LParen) {
            return Ok(params);
        }

        while let Some(Token::Identifier(param)) = self.peek() {
            params.push(param.clone());
            self.next();
            if !self.consume(&Token::Comma) {
                break;
            }
        }
        self.consume(&Token::RParen);
        Ok(params)
    }

    /// Parse function call arguments
    fn parse_call_args(&mut self) -> Result<Vec<Expr>, ParseError> {
        let mut args = Vec::new();
        if !self.consume(&Token::LParen) {
            return Ok(args);
        }

        while self.peek().is_some() {
            if self.consume(&Token::RParen) {
                break;
            }
            args.push(self.parse_expression(0)?);
            if !self.consume(&Token::Comma) && self.peek() != Some(&Token::RParen) {
                break;
            }
        }
        self.consume(&Token::RParen);
        Ok(args)
    }

    /// Parse literal values (numbers, strings, booleans, null)
    fn parse_literal(&mut self) -> Result<Option<Expr>, ParseError> {
        match self.peek() {
            Some(Token::NumberLiteral(num)) => {
                let num = *num;
                self.next();
                Ok(Some(Expr::Number(num)))
            }
            Some(Token::StringLiteral(s)) => {
                let s = s.clone();
                self.next();
                Ok(Some(Expr::String(s)))
            }
            Some(Token::True) => {
                self.next();
                Ok(Some(Expr::Bool(true)))
            }
            Some(Token::False) => {
                self.next();
                Ok(Some(Expr::Bool(false)))
            }
            Some(Token::Null) => {
                self.next();
                Ok(Some(Expr::Null))
            }
            _ => Ok(None),
        }
    }

    /// Parse identifier or function call
    fn parse_identifier_or_call(&mut self) -> Result<Expr, ParseError> {
        if let Some(Token::Identifier(name)) = self.peek() {
            let name = name.clone();
            self.next();

            // Check for function call
            if self.peek() == Some(&Token::LParen) {
                let args = self.parse_call_args()?;
                Ok(Expr::Call {
                    callee: Box::new(Expr::Identifier(name)),
                    args,
                })
            } else {
                Ok(Expr::Identifier(name))
            }
        } else {
            Err(ParseError::UnexpectedToken {
                expected: "identifier".to_string(),
                found: self.peek().cloned().unwrap_or(Token::Error),
            })
        }
    }

    /// Parse primary expressions (literals, identifiers, parenthesized expressions, blocks)
    fn parse_primary(&mut self) -> Result<Expr, ParseError> {
        // Handle unary minus
        if self.consume(&Token::Minus) {
            let expr = self.parse_primary()?;
            if let Expr::Number(num) = expr {
                return Ok(Expr::Number(-num));
            }

            return Ok(Expr::Binary {
                left: Box::new(Expr::Number(0.0)),
                op: Token::Minus,
                right: Box::new(expr),
            });
        }

        // Try to parse literals first
        if let Some(literal) = self.parse_literal()? {
            return Ok(literal);
        }

        // Parse identifier or function call
        if let Some(Token::Identifier(_)) = self.peek() {
            return self.parse_identifier_or_call();
        }

        // Parse parenthesized expression
        if self.consume(&Token::LParen) {
            let expr = self.parse_expression(0)?;
            self.consume(&Token::RParen);
            return Ok(expr);
        }

        // Parse block expression
        if self.peek() == Some(&Token::LBrace) {
            let stmts = self.parse_block()?;
            return Ok(Expr::Block(stmts));
        }

        Err(ParseError::UnexpectedToken {
            expected: "primary expression".to_string(),
            found: self.peek().cloned().unwrap_or(Token::Error),
        })
    }

    /// Parse a let statement
    fn parse_let_stmt(&mut self) -> Result<Stmt, ParseError> {
        self.next(); // consume 'let'

        let name = if let Some(Token::Identifier(name)) = self.peek() {
            let name = name.clone();
            self.next();
            name
        } else {
            return Err(ParseError::UnexpectedToken {
                expected: "identifier".to_string(),
                found: self.peek().cloned().unwrap_or(Token::Error),
            });
        };

        if !self.consume(&Token::Assign) {
            return Err(ParseError::UnexpectedToken {
                expected: "=".to_string(),
                found: self.peek().cloned().unwrap_or(Token::Error),
            });
        }

        let value = self.parse_expression(0)?;
        self.consume(&Token::Semicolon);
        Ok(Stmt::Let { name, value })
    }

    /// Parse a return statement
    fn parse_return_stmt(&mut self) -> Result<Stmt, ParseError> {
        self.next(); // consume 'return'
        let expr = self.parse_expression(0)?;
        self.consume(&Token::Semicolon);
        Ok(Stmt::Return(expr))
    }

    /// Parse an if statement
    fn parse_if_stmt(&mut self) -> Result<Stmt, ParseError> {
        self.next(); // consume 'if'
        let cond = self.parse_expression(0)?;
        let then_branch = self.parse_block()?;
        let else_branch = self.parse_else_chain()?;
        Ok(Stmt::If {
            cond,
            then_branch,
            else_branch,
        })
    }

    /// Parse a function definition
    fn parse_fn_stmt(&mut self) -> Result<Stmt, ParseError> {
        self.next(); // consume 'fn'

        let name = if let Some(Token::Identifier(name)) = self.peek() {
            let name = name.clone();
            self.next();
            name
        } else {
            return Err(ParseError::UnexpectedToken {
                expected: "function name".to_string(),
                found: self.peek().cloned().unwrap_or(Token::Error),
            });
        };

        let params = self.parse_function_params()?;
        self.fn_table.insert(name.clone(), params.len());
        let body = self.parse_block()?;

        Ok(Stmt::Fn { name, params, body })
    }

    /// Parse assignment or expression statement starting with an identifier
    fn parse_identifier_stmt(&mut self) -> Result<Stmt, ParseError> {
        let save_pos = self.pos;

        // Check if this is an assignment
        if let Some(Token::Assign) = self.tokens.get(self.pos + 1) {
            let name = if let Some(Token::Identifier(name)) = self.peek() {
                let name = name.clone();
                self.next();
                name
            } else {
                return Err(ParseError::UnexpectedToken {
                    expected: "identifier".to_string(),
                    found: self.peek().cloned().unwrap_or(Token::Error),
                });
            };

            self.next(); // consume '='
            let value = self.parse_expression(0)?;
            self.consume(&Token::Semicolon);
            return Ok(Stmt::Assign { name, value });
        }

        // Otherwise, parse as expression statement
        let expr = match self.parse_expression(0) {
            Ok(expr) => expr,
            Err(e) => {
                self.pos = save_pos + 1;
                return Err(e);
            }
        };

        self.consume(&Token::Semicolon);
        Ok(Stmt::Expr(expr))
    }

    /// Pratt parser for expressions with operator precedence
    fn parse_expression(&mut self, min_prec: u8) -> Result<Expr, ParseError> {
        let mut left = self.parse_primary()?;

        loop {
            let op_prec = match self.peek() {
                Some(Token::Plus) | Some(Token::Minus) => 10,
                Some(Token::Star) | Some(Token::Slash) | Some(Token::Percent) => 20,
                Some(Token::EqEq) | Some(Token::BangEq) | Some(Token::Gt) | Some(Token::Lt)
                | Some(Token::GtEq) | Some(Token::LtEq) => 5,
                _ => break,
            };

            if op_prec < min_prec {
                break;
            }

            let op = self.next().unwrap().clone();
            let right = self.parse_expression(op_prec + 1)?;
            left = Expr::Binary {
                left: Box::new(left),
                op,
                right: Box::new(right),
            };
        }
        Ok(left)
    }

    /// Parse a single statement
    fn parse_stmt(&mut self) -> Result<Option<Stmt>, ParseError> {
        match self.peek() {
            Some(Token::If) => Ok(Some(self.parse_if_stmt()?)),
            Some(Token::Let) => Ok(Some(self.parse_let_stmt()?)),
            Some(Token::Return) => Ok(Some(self.parse_return_stmt()?)),
            Some(Token::Fn) => Ok(Some(self.parse_fn_stmt()?)),
            Some(Token::While) => Ok(Some(self.parse_while_stmt()?)),
            Some(Token::For) => Ok(Some(self.parse_for_stmt()?)),
            Some(Token::Identifier(_)) => Ok(Some(self.parse_identifier_stmt()?)),
            Some(_) => {
                // Try to parse as expression statement (e.g., "2 + 3", function calls, etc.)
                if let Ok(expr) = self.parse_expression(0) {
                    Ok(Some(Stmt::Expr(expr)))
                } else {
                    self.next(); // Skip unknown tokens
                    Ok(None)
                }
            }
            None => Ok(None),
        }
    }

    /// Parse the entire program into a list of statements
    pub fn parse(&mut self) -> Result<Vec<Stmt>, ParseError> {
        let mut stmts = Vec::new();
        while self.pos < self.tokens.len() {
            match self.parse_stmt()? {
                Some(stmt) => stmts.push(stmt),
                None => {}
            }
        }
        Ok(stmts)
    }
    /// Parse an optional else or else-if chain
    fn parse_else_chain(&mut self) -> Result<Option<Box<Stmt>>, ParseError> {
        if !self.consume(&Token::Else) {
            return Ok(None);
        }

        // Check for else if
        if self.consume(&Token::If) {
            let cond = self.parse_expression(0)?;
            let then_branch = self.parse_block()?;
            let else_branch = self.parse_else_chain()?;
            return Ok(Some(Box::new(Stmt::If {
                cond,
                then_branch,
                else_branch,
            })));
        }

        // Parse else block
        if self.peek() == Some(&Token::LBrace) {
            let else_block = self.parse_block()?;

            // If single statement, return it directly; otherwise wrap in Block expression
            if else_block.len() == 1 {
                Ok(Some(Box::new(else_block.into_iter().next().unwrap())))
            } else {
                Ok(Some(Box::new(Stmt::Expr(Expr::Block(else_block)))))
            }
        } else {
            Ok(None)
        }
    }

    /// Parse a while statement
    fn parse_while_stmt(&mut self) -> Result<Stmt, ParseError> {
        self.next(); // consume 'while'
        let cond = self.parse_expression(0)?;
        let body = self.parse_block()?;
        Ok(Stmt::While { cond, body })
    }

    /// Parse a for statement (for var start..end syntax)
    fn parse_for_stmt(&mut self) -> Result<Stmt, ParseError> {
        self.next(); // consume 'for'

        let var = if let Some(Token::Identifier(name)) = self.peek() {
            let name = name.clone();
            self.next();
            name
        } else {
            return Err(ParseError::UnexpectedToken {
                expected: "loop variable identifier".to_string(),
                found: self.peek().cloned().unwrap_or(Token::Error),
            });
        };

        let start = self.parse_expression(0)?;

        if !self.consume(&Token::Range) {
            return Err(ParseError::UnexpectedToken {
                expected: "..".to_string(),
                found: self.peek().cloned().unwrap_or(Token::Error),
            });
        }

        let end = self.parse_expression(0)?;
        let body = self.parse_block()?;

        Ok(Stmt::For {
            var,
            start,
            end,
            body,
        })
    }
}