neutron-engine 0.1.3

/// Iris Parser - Converts tokens into AST

use super::lexer::Token;
use super::token::TokenType;




#[derive(Debug, Clone, PartialEq)]
pub enum Expr {
    Null,
    Bool(bool),
    Number(f64),
    String(String),
    Identifier(String),
    Array(Vec<Expr>),
    Object(Vec<(String, Expr)>),
    Binary {
        left: Box<Expr>,
        op: BinaryOp,
        right: Box<Expr>,
    },
    Unary {
        op: UnaryOp,
        expr: Box<Expr>,
    },
    Call {
        callee: Box<Expr>,
        args: Vec<Expr>,
    },
    Index {
        object: Box<Expr>,
        index: Box<Expr>,
    },
    Member {
        object: Box<Expr>,
        property: String,
    },
    Assign {
        target: Box<Expr>,
        value: Box<Expr>,
    },
    Lambda {
        params: Vec<String>,
        body: Vec<Stmt>,
    },
}

#[derive(Debug, Clone, PartialEq)]
pub enum BinaryOp {
    Add, Sub, Mul, Div, Mod,
    Eq, Neq, Lt, Gt, Lte, Gte,
    And, Or,
}

#[derive(Debug, Clone, PartialEq)]
pub enum UnaryOp {
    Neg, Not,
}

#[derive(Debug, Clone, PartialEq)]
pub enum Stmt {
    Expr(Expr),

    // system iris:"..." directive
    // stores compile-time/runtime traits requested by the user
    SystemIris {
        traits: Vec<String>,
    },

    Let {
        name: String,
        value: Expr,
    },
    Const {
        name: String,
        value: Expr,
    },
    Fn {
        name: String,
        params: Vec<String>,
        body: Vec<Stmt>,
    },
    Block(Vec<Stmt>),
    If {
        condition: Expr,
        then_branch: Vec<Stmt>,
        else_branch: Option<Vec<Stmt>>,
    },
    While {
        condition: Expr,
        body: Vec<Stmt>,
    },
    For {
        var: String,
        iterable: Expr,
        body: Vec<Stmt>,
    },
    Return(Option<Expr>),
    Break,
    Continue,
    Import {
        path: String,
    },
}

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

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

    pub fn parse(&mut self) -> Result<Vec<Stmt>, String> {
        let mut stmts = Vec::new();
        while !self.is_at_end() {
            if self.match_token(&[TokenType::Newline]) || self.match_token(&[TokenType::Comment(String::new())]) {
                continue;
            }
            stmts.push(self.statement()?);
        }
        Ok(stmts)
    }

    fn statement(&mut self) -> Result<Stmt, String> {
        if self.match_token(&[TokenType::Import]) {
            self.import_statement()
        } else if self.match_token(&[TokenType::System]) {
            self.system_iris_statement()
        } else if self.match_token(&[TokenType::Let]) {
            self.let_statement()
        } else if self.match_token(&[TokenType::Const]) {
            self.const_statement()
        } else if self.match_token(&[TokenType::Fn]) {
            self.fn_statement()
        } else if self.match_token(&[TokenType::If]) {
            self.if_statement()
        } else if self.match_token(&[TokenType::Elif]) {
            Err("Iris does not use 'elif'. Use nested if logic: else { if condition { ... } }".to_string())
        } else if self.match_token(&[TokenType::While]) {
            self.while_statement()
        } else if self.match_token(&[TokenType::For]) {
            self.for_statement()
        } else if self.match_token(&[TokenType::Return]) {
            self.return_statement()
        } else if self.match_token(&[TokenType::Break]) {
            Ok(Stmt::Break)
        } else if self.match_token(&[TokenType::Continue]) {
            Ok(Stmt::Continue)
        } else if self.check(&TokenType::LBrace) {
            self.block()
        } else {
            let expr = self.expression()?;
            self.consume_semicolon();
            Ok(Stmt::Expr(expr))
        }
    }

    fn system_iris_statement(&mut self) -> Result<Stmt, String> {
        // Syntax supported:
        // system iris:"borrow checker"
        // system iris:"borrow checker","ownership"
        // system iris:"borrow checker", "ownership"

        // Expect identifier `iris`
        let namespace = self.consume_identifier("Expected identifier after 'system' (expected 'iris')")?;
        if namespace != "iris" {
            return Err(format!("Expected 'iris' after 'system' but got '{}'", namespace));
        }

        self.consume(&TokenType::Colon, "Expected ':' after 'iris'")?;

        let mut traits = Vec::new();
        loop {
            match self.peek_token_type() {
                Some(TokenType::String(s)) => {
                    traits.extend(
                        s.split(',')
                            .map(str::trim)
                            .filter(|trait_name| !trait_name.is_empty())
                            .map(str::to_string),
                    );
                    self.advance();
                }
                _ => {
                    return Err(format!("Expected string trait after 'system iris:' at line {}", self.peek().line));
                }
            }

            if self.match_token(&[TokenType::Comma]) {
                continue;
            } else {
                break;
            }
        }

        self.consume_semicolon();
        Ok(Stmt::SystemIris { traits })
    }

    fn import_statement(&mut self) -> Result<Stmt, String> {
        let path = match self.peek_token_type() {
            Some(TokenType::String(s)) => {
                let s = s.clone();
                self.advance();
                s
            }
            _ => return Err(format!("Expected string path after 'import' at line {}", self.peek().line)),
        };
        self.consume_semicolon();
        Ok(Stmt::Import { path })
    }

    fn let_statement(&mut self) -> Result<Stmt, String> {
        let name = self.consume_identifier("Expected variable name after 'let'")?;
        self.consume(&TokenType::Assign, "Expected '=' after variable name")?;
        let value = self.expression()?;
        self.consume_semicolon();
        Ok(Stmt::Let { name, value })
    }

    fn const_statement(&mut self) -> Result<Stmt, String> {
        let name = self.consume_identifier("Expected constant name after 'const'")?;
        self.consume(&TokenType::Assign, "Expected '=' after constant name")?;
        let value = self.expression()?;
        self.consume_semicolon();
        Ok(Stmt::Const { name, value })
    }

    fn fn_statement(&mut self) -> Result<Stmt, String> {
        let name = self.consume_identifier("Expected function name after 'fn'")?;
        self.consume(&TokenType::LParen, "Expected '(' after function name")?;
        let mut params = Vec::new();
        if !self.check(&TokenType::RParen) {
            loop {
                params.push(self.consume_identifier("Expected parameter name")?);
                if !self.match_token(&[TokenType::Comma]) {
                    break;
                }
            }
        }
        self.consume(&TokenType::RParen, "Expected ')' after parameters")?;
        let body = self.block()?;
        let body = match body {
            Stmt::Block(stmts) => stmts,
            _ => vec![body],
        };
        Ok(Stmt::Fn { name, params, body })
    }

    fn if_statement(&mut self) -> Result<Stmt, String> {
        let condition = self.expression()?;
        let then_branch = match self.block()? {
            Stmt::Block(stmts) => stmts,
            s => vec![s],
        };
        let else_branch = if self.match_token(&[TokenType::Else]) {
            if self.check(&TokenType::If) {
                return Err("Iris does not use 'else if'. Use nested if logic: else { if condition { ... } }".to_string());
            }
            Some(match self.block()? {
                Stmt::Block(stmts) => stmts,
                s => vec![s],
            })
        } else {
            None
        };
        Ok(Stmt::If { condition, then_branch, else_branch })
    }

    fn while_statement(&mut self) -> Result<Stmt, String> {
        let condition = self.expression()?;
        let body = match self.block()? {
            Stmt::Block(stmts) => stmts,
            s => vec![s],
        };
        Ok(Stmt::While { condition, body })
    }

    fn for_statement(&mut self) -> Result<Stmt, String> {
        let var = self.consume_identifier("Expected variable name after 'for'")?;
        self.consume(&TokenType::In, "Expected 'in' after for variable")?;
        let iterable = self.expression()?;
        let body = match self.block()? {
            Stmt::Block(stmts) => stmts,
            s => vec![s],
        };
        Ok(Stmt::For { var, iterable, body })
    }

    fn return_statement(&mut self) -> Result<Stmt, String> {
        let value = if self.check(&TokenType::Semicolon) || self.check(&TokenType::Newline) || self.check(&TokenType::RBrace) || self.is_at_end() {
            None
        } else {
            Some(self.expression()?)
        };
        self.consume_semicolon();
        Ok(Stmt::Return(value))
    }

    fn block(&mut self) -> Result<Stmt, String> {
        self.consume(&TokenType::LBrace, "Expected '{'")?;
        let mut stmts = Vec::new();
        while !self.check(&TokenType::RBrace) && !self.is_at_end() {
            if self.match_token(&[TokenType::Newline]) || self.match_token(&[TokenType::Comment(String::new())]) {
                continue;
            }
            stmts.push(self.statement()?);
        }
        self.consume(&TokenType::RBrace, "Expected '}' after block")?;
        Ok(Stmt::Block(stmts))
    }

    fn expression(&mut self) -> Result<Expr, String> {
        self.assignment()
    }

    fn assignment(&mut self) -> Result<Expr, String> {
        let expr = self.or()?;
        if self.match_token(&[TokenType::Assign, TokenType::PlusAssign, TokenType::MinusAssign, TokenType::StarAssign, TokenType::SlashAssign]) {
            let op = self.previous().token_type.clone();
            let value = self.assignment()?;
            match expr {
                Expr::Identifier(name) => {
                    let value = match op {
                        TokenType::PlusAssign => Expr::Binary { left: Box::new(Expr::Identifier(name.clone())), op: BinaryOp::Add, right: Box::new(value) },
                        TokenType::MinusAssign => Expr::Binary { left: Box::new(Expr::Identifier(name.clone())), op: BinaryOp::Sub, right: Box::new(value) },
                        TokenType::StarAssign => Expr::Binary { left: Box::new(Expr::Identifier(name.clone())), op: BinaryOp::Mul, right: Box::new(value) },
                        TokenType::SlashAssign => Expr::Binary { left: Box::new(Expr::Identifier(name.clone())), op: BinaryOp::Div, right: Box::new(value) },
                        _ => value,
                    };
                    Ok(Expr::Assign { target: Box::new(Expr::Identifier(name)), value: Box::new(value) })
                }
                Expr::Member { object, property } => {
                    Ok(Expr::Assign { target: Box::new(Expr::Member { object, property }), value: Box::new(value) })
                }
                Expr::Index { object, index } => {
                    Ok(Expr::Assign { target: Box::new(Expr::Index { object, index }), value: Box::new(value) })
                }
                _ => Err("Invalid assignment target".to_string()),
            }
        } else {
            Ok(expr)
        }
    }

    fn or(&mut self) -> Result<Expr, String> {
        let mut expr = self.and()?;
        while self.match_token(&[TokenType::Or]) {
            let right = self.and()?;
            expr = Expr::Binary { left: Box::new(expr), op: BinaryOp::Or, right: Box::new(right) };
        }
        Ok(expr)
    }

    fn and(&mut self) -> Result<Expr, String> {
        let mut expr = self.equality()?;
        while self.match_token(&[TokenType::And]) {
            let right = self.equality()?;
            expr = Expr::Binary { left: Box::new(expr), op: BinaryOp::And, right: Box::new(right) };
        }
        Ok(expr)
    }

    fn equality(&mut self) -> Result<Expr, String> {
        let mut expr = self.comparison()?;
        while self.match_token(&[TokenType::Eq, TokenType::Neq]) {
            let op = match self.previous().token_type {
                TokenType::Eq => BinaryOp::Eq,
                TokenType::Neq => BinaryOp::Neq,
                _ => unreachable!(),
            };
            let right = self.comparison()?;
            expr = Expr::Binary { left: Box::new(expr), op, right: Box::new(right) };
        }
        Ok(expr)
    }

    fn comparison(&mut self) -> Result<Expr, String> {
        let mut expr = self.term()?;
        while self.match_token(&[TokenType::Lt, TokenType::Gt, TokenType::Lte, TokenType::Gte]) {
            let op = match self.previous().token_type {
                TokenType::Lt => BinaryOp::Lt,
                TokenType::Gt => BinaryOp::Gt,
                TokenType::Lte => BinaryOp::Lte,
                TokenType::Gte => BinaryOp::Gte,
                _ => unreachable!(),
            };
            let right = self.term()?;
            expr = Expr::Binary { left: Box::new(expr), op, right: Box::new(right) };
        }
        Ok(expr)
    }

    fn term(&mut self) -> Result<Expr, String> {
        let mut expr = self.factor()?;
        while self.match_token(&[TokenType::Plus, TokenType::Minus]) {
            let op = match self.previous().token_type {
                TokenType::Plus => BinaryOp::Add,
                TokenType::Minus => BinaryOp::Sub,
                _ => unreachable!(),
            };
            let right = self.factor()?;
            expr = Expr::Binary { left: Box::new(expr), op, right: Box::new(right) };
        }
        Ok(expr)
    }

    fn factor(&mut self) -> Result<Expr, String> {
        let mut expr = self.unary()?;
        while self.match_token(&[TokenType::Star, TokenType::Slash, TokenType::Percent]) {
            let op = match self.previous().token_type {
                TokenType::Star => BinaryOp::Mul,
                TokenType::Slash => BinaryOp::Div,
                TokenType::Percent => BinaryOp::Mod,
                _ => unreachable!(),
            };
            let right = self.unary()?;
            expr = Expr::Binary { left: Box::new(expr), op, right: Box::new(right) };
        }
        Ok(expr)
    }

fn unary(&mut self) -> Result<Expr, String> {
    if self.match_token(&[TokenType::Minus, TokenType::Not, TokenType::Bang]) {
        let op = match self.previous().token_type {
            TokenType::Minus => UnaryOp::Neg,
            TokenType::Not => UnaryOp::Not,
            TokenType::Bang => UnaryOp::Not, // ! as logical not or macro prefix
            _ => unreachable!(),
        };
        let expr = self.unary()?;
        Ok(Expr::Unary { op, expr: Box::new(expr) })
    } else {
        self.call()
    }
}

    fn call(&mut self) -> Result<Expr, String> {
        let mut expr = self.primary()?;
        loop {
            if self.match_token(&[TokenType::LParen]) {
                let mut args = Vec::new();
                if !self.check(&TokenType::RParen) {
                    loop {
                        args.push(self.expression()?);
                        if !self.match_token(&[TokenType::Comma]) {
                            break;
                        }
                    }
                }
                self.consume(&TokenType::RParen, "Expected ')' after arguments")?;
                expr = Expr::Call { callee: Box::new(expr), args };
            } else if self.match_token(&[TokenType::Dot, TokenType::Scope]) {
                let property = self.consume_identifier("Expected property name after namespace access")?;
                expr = Expr::Member { object: Box::new(expr), property };
            } else if self.match_token(&[TokenType::LBracket]) {
                let index = self.expression()?;
                self.consume(&TokenType::RBracket, "Expected ']' after index")?;
                expr = Expr::Index { object: Box::new(expr), index: Box::new(index) };
            } else {
                break;
            }
        }
        Ok(expr)
    }

    fn primary(&mut self) -> Result<Expr, String> {
        if self.match_token(&[TokenType::Null]) {
            Ok(Expr::Null)
        } else if self.match_token(&[TokenType::Bool(true)]) {
            Ok(Expr::Bool(true))
        } else if self.match_token(&[TokenType::Bool(false)]) {
            Ok(Expr::Bool(false))
        } else if let Some(TokenType::Number(n)) = self.peek_token_type() {
            let n = *n;
            self.advance();
            Ok(Expr::Number(n))
        } else if let Some(TokenType::String(s)) = self.peek_token_type() {
            let s = s.clone();
            self.advance();
            Ok(Expr::String(s))
        } else if self.match_token(&[TokenType::LBracket]) {
            let mut elements = Vec::new();
            if !self.check(&TokenType::RBracket) {
                loop {
                    elements.push(self.expression()?);
                    if !self.match_token(&[TokenType::Comma]) {
                        break;
                    }
                }
            }
            self.consume(&TokenType::RBracket, "Expected ']' after array elements")?;
            Ok(Expr::Array(elements))
        } else if self.match_token(&[TokenType::LBrace]) {
            let mut pairs = Vec::new();
            if !self.check(&TokenType::RBrace) {
                loop {
                    let key = self.consume_identifier("Expected object key")?;
                    self.consume(&TokenType::Colon, "Expected ':' after object key")?;
                    let value = self.expression()?;
                    pairs.push((key, value));
                    if !self.match_token(&[TokenType::Comma]) {
                        break;
                    }
                }
            }
            self.consume(&TokenType::RBrace, "Expected '}' after object pairs")?;
            Ok(Expr::Object(pairs))
        } else if self.match_token(&[TokenType::Fn]) {
            self.consume(&TokenType::LParen, "Expected '(' after 'fn'")?;
            let mut params = Vec::new();
            if !self.check(&TokenType::RParen) {
                loop {
                    params.push(self.consume_identifier("Expected parameter name")?);
                    if !self.match_token(&[TokenType::Comma]) {
                        break;
                    }
                }
            }
            self.consume(&TokenType::RParen, "Expected ')' after parameters")?;
            let body = match self.block()? {
                Stmt::Block(stmts) => stmts,
                s => vec![s],
            };
            Ok(Expr::Lambda { params, body })
        } else if let Some(TokenType::Identifier(name)) = self.peek_token_type() {
            let name = name.clone();
            self.advance();
            Ok(Expr::Identifier(name))
        } else {
            Err(format!("Unexpected token: {:?}", self.peek()))
        }
    }

    // Helper methods
    fn match_token(&mut self, types: &[TokenType]) -> bool {
        for t in types {
            if self.check(t) {
                self.advance();
                return true;
            }
        }
        false
    }

    fn check(&self, token_type: &TokenType) -> bool {
        if self.is_at_end() {
            return false;
        }
        token_matches(&self.peek().token_type, token_type)
    }

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

    fn is_at_end(&self) -> bool {
        matches!(self.peek().token_type, TokenType::Eof)
    }

    fn peek(&self) -> &Token {
        &self.tokens[self.current]
    }

    fn peek_token_type(&self) -> Option<&TokenType> {
        if self.is_at_end() {
            None
        } else {
            Some(&self.peek().token_type)
        }
    }

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

    fn consume(&mut self, token_type: &TokenType, message: &str) -> Result<(), String> {
        if self.check(token_type) {
            self.advance();
            Ok(())
        } else {
            Err(format!("{} at line {}, column {} (got {:?})", 
                message, self.peek().line, self.peek().column, self.peek().token_type))
        }
    }

    fn consume_identifier(&mut self, message: &str) -> Result<String, String> {
        if let Some(TokenType::Identifier(name)) = self.peek_token_type() {
            let name = name.clone();
            self.advance();
            Ok(name)
        } else {
            Err(format!("{} at line {}, column {} (got {:?})", 
                message, self.peek().line, self.peek().column, self.peek().token_type))
        }
    }

    fn consume_semicolon(&mut self) {
        self.match_token(&[TokenType::Semicolon]);
    }
}

fn token_matches(actual: &TokenType, expected: &TokenType) -> bool {
    match (actual, expected) {
        (TokenType::Comment(_), TokenType::Comment(_)) => true,
        (TokenType::LParen, TokenType::LeftParen) | (TokenType::LeftParen, TokenType::LParen) => true,
        (TokenType::RParen, TokenType::RightParen) | (TokenType::RightParen, TokenType::RParen) => true,
        (TokenType::LBrace, TokenType::LeftBrace) | (TokenType::LeftBrace, TokenType::LBrace) => true,
        (TokenType::RBrace, TokenType::RightBrace) | (TokenType::RightBrace, TokenType::RBrace) => true,
        (a, b) => std::mem::discriminant(a) == std::mem::discriminant(b),
    }
}

pub fn parse(tokens: &[Token]) -> Result<Vec<Stmt>, String> {
    let mut parser = Parser::new(tokens.to_vec());
    parser.parse()
}