parser.cpp

/**
 * parser.cpp
 * author: Bao Le
 */

#include <iostream>
#include <vector>
#include <string>
#include "ast.hpp"
#include "parser.hpp"
#include "lexer.hpp"

std::unique_ptr<Program> Parser::parse_program() {
    auto program = std::make_unique<Program>();
    while (tokens[cur_idx].type != TokenType::END_OF_FILE) {
        if (is_type_keyword(tokens[cur_idx].type)) {
            // handle global declaration and initialization outside the main() function
            if (cur_idx + 2 < tokens.size() && tokens[cur_idx + 2].type == TokenType::LPAREN) {
                program->children.push_back(parse_function());    
            }
            else program->children.push_back(parse_declaration());
        }
        else program->children.push_back(parse_statement());
    }
    return program;
}

std::unique_ptr<Function> Parser::parse_function() {
    Token fn_type = expect_type();
    Token func_name = expect(TokenType::IDENTIFIER);
    expect(TokenType::LPAREN);

    // Function args
    std::vector<Param> params;
    while (tokens[cur_idx].type != TokenType::RPAREN) {
        // no void type allowed inside function args
        Token type = expect_primitive_type();
        Token arg = expect(TokenType::IDENTIFIER);
        // ConstantType type, std::string name
        params.push_back({token_type_to_constant_type(type.type), arg.value});
        if (tokens[cur_idx].type == TokenType::COMMA) consume();
    }
    expect(TokenType::RPAREN);                         
    expect(TokenType::LBRACE);                         
    auto body = parse_block();
    expect(TokenType::RBRACE);                         
    return std::make_unique<Function>(func_name.value, token_type_to_constant_type(fn_type.type), std::move(params), std::move(body));
}

// handle function call -> foo(x, y), begin from LPAREN
std::unique_ptr<FunctionCall> Parser::parse_function_call(std::string fn_name) {
    if (peek().type != TokenType::LPAREN)
        throw std::runtime_error("[FUNCTION CALL] expected '(' after function name '" + fn_name + "' at line " + std::to_string(peek().line));
    expect(TokenType::LPAREN);
    
    std::vector<std::unique_ptr<ASTNode>> args;
    while (tokens[cur_idx].type != TokenType::RPAREN) {
        if (tokens[cur_idx].type == TokenType::END_OF_FILE)
            throw std::runtime_error("[FUNCTION CALL] unterminated argument list in call to '" + fn_name + "' at line " + std::to_string(peek().line));
        args.push_back(parse_expression()); // each arg is an expression
        if (tokens[cur_idx].type == TokenType::COMMA) consume();
    }
    expect(TokenType::RPAREN);
    return std::make_unique<FunctionCall>(fn_name, std::move(args));
}

std::unique_ptr<ASTNode> Parser::parse_if() {
    // call consume after validation
    expect(TokenType::IF);
    if (peek().type != TokenType::LPAREN)
        throw std::runtime_error("[IF] expected '(' after 'if' at line " + std::to_string(peek().line));
    expect(TokenType::LPAREN);
    
    auto cond = parse_expression();
    
    if (peek().type != TokenType::RPAREN)
        throw std::runtime_error("[IF] expected ')' after condition at line " + std::to_string(peek().line));
    expect(TokenType::RPAREN);
    
    // body statements
    if (peek().type != TokenType::LBRACE)
        throw std::runtime_error("[IF] expected '{' to open if body at line " + std::to_string(peek().line));
    expect(TokenType::LBRACE);
    auto block = parse_block();
    
    if (peek().type != TokenType::RBRACE)
        throw std::runtime_error("[IF] expected '}' to close if body at line " + std::to_string(peek().line));
    expect(TokenType::RBRACE);
    
    std::optional<std::unique_ptr<StatementBody>> else_body = std::nullopt;
    if (tokens[cur_idx].type == TokenType::ELSE) {
        consume();
        // currently doesn't handle one liner -> else return 0;
        if (peek().type != TokenType::LBRACE)
            throw std::runtime_error("[ELSE] expected '{' to open else body at line " + std::to_string(peek().line));
        expect(TokenType::LBRACE);
        else_body = parse_block();
        if (peek().type != TokenType::RBRACE)
            throw std::runtime_error("[ELSE] expected '}' to close else body at line " + std::to_string(peek().line));
        expect(TokenType::RBRACE);
    }
    return std::make_unique<If>(std::move(cond), std::move(block), std::move(else_body));
}

std::unique_ptr<ASTNode> Parser::parse_while() {
    expect(TokenType::WHILE);
        if (peek().type != TokenType::LPAREN) 
            throw std::runtime_error("[WHILE] expected '(' after 'while' at line " + std::to_string(peek().line));
        expect(TokenType::LPAREN);
        auto cond = parse_expression();
        if (peek().type != TokenType::RPAREN)
            throw std::runtime_error("[WHILE] expected ')' after condition at line " + std::to_string(peek().line));
        expect(TokenType::RPAREN);
        if (peek().type != TokenType::LBRACE)
            throw std::runtime_error("[WHILE] expected '{' to open while body at line " + std::to_string(peek().line));
        expect(TokenType::LBRACE);
        auto block = parse_block();
        if (peek().type != TokenType::RBRACE)
            throw std::runtime_error("[WHILE] expected '}' to close while body at line " + std::to_string(peek().line));
        expect(TokenType::RBRACE);
        return std::make_unique<While>(std::move(cond), std::move(block));
}
// handle the block inside {}
std::unique_ptr<StatementBody> Parser::parse_block() {
    // assignment, nested if, while, for statements
    auto block = std::make_unique<StatementBody>();
    while (tokens[cur_idx].type != TokenType::RBRACE && 
        tokens[cur_idx].type != TokenType::END_OF_FILE) {
        block->statements.push_back(parse_statement());
    }
    if (tokens[cur_idx].type == TokenType::END_OF_FILE)
        throw std::runtime_error("[BLOCK] unterminated block — missing '}' at end of file");
    return block;
}

std::unique_ptr<StatementBody> Parser::parse_nested_block() {
    if (peek().type != TokenType::LBRACE)
        throw std::runtime_error("[NESTED BLOCK] expected '{' at line " + std::to_string(peek().line));
    expect(TokenType::LBRACE);
    auto block = parse_block();
    expect(TokenType::RBRACE);
    return block;
}

// redirect assignment, return, nested blocks, and control statements to appropriate parsers
std::unique_ptr<ASTNode> Parser::parse_statement() {
    switch(tokens[cur_idx].type) {
        case TokenType::RETURN:     return parse_return();
        case TokenType::IF:         return parse_if();
        case TokenType::WHILE:     return parse_while();
        case TokenType::LBRACE:     return parse_nested_block();
        case TokenType::IDENTIFIER:    
            // could be assignment or function call
            if (cur_idx + 1 < tokens.size() && tokens[cur_idx + 1].type == TokenType::LPAREN) {
                // advance index to parse from (
                auto call = parse_function_call(consume().value);
                if (peek().type != TokenType::SEMICOLON)
                    throw std::runtime_error("[STATEMENT] expected ';' after function call at line " + std::to_string(peek().line));
                expect(TokenType::SEMICOLON);
                return call;
            }
            return parse_assignment(); 
        default: 
            if (is_primitive_type_keyword(tokens[cur_idx].type)) return parse_declaration();
            throw std::runtime_error("[STATEMENT] unexpected token '" + peek().value + "' at line " + std::to_string(peek().line));
    }
}

// handle declaration + initialization
std::unique_ptr<ASTNode> Parser::parse_declaration() {
    Token type = expect_primitive_type();
    Token identifier = expect(TokenType::IDENTIFIER); // consume the token

    // e.g. int x; -> declaration
    if (tokens[cur_idx].type == TokenType::SEMICOLON) {
        consume();
        return std::make_unique<Declaration>(Variable{identifier.value}, token_type_to_constant_type(type.type), std::nullopt);
    }
    // e.g. int x = 5; -> declaration with initialization
    if (tokens[cur_idx].type == TokenType::ASSIGN) {
        consume();
        auto expr = parse_expression();
        if (peek().type != TokenType::SEMICOLON) throw std::runtime_error("[DECLARATION] expected ';' after initialization of '" + identifier.value + "' at line " + std::to_string(peek().line));
        expect(TokenType::SEMICOLON);
        return std::make_unique<Declaration>(
            Variable{identifier.value}, token_type_to_constant_type(type.type), std::move(expr)
        );
    }
    throw std::runtime_error("[DECLARATION] expected ';' or '=' after identifier '" + identifier.value + "' at line " + std::to_string(peek().line));
}

std::unique_ptr<Return> Parser::parse_return() {
    expect(TokenType::RETURN);
    // void return 
    if (tokens[cur_idx].type == TokenType::SEMICOLON) {
        consume();
        return std::make_unique<Return>(std::nullopt);
    }
    auto expr = parse_expression();
    if (peek().type != TokenType::SEMICOLON)
        throw std::runtime_error("[RETURN] expected ';' after return value at line " + std::to_string(peek().line));
    expect(TokenType::SEMICOLON);
    return std::make_unique<Return>(std::move(expr));
}

// variable assignment and reassignment
std::unique_ptr<Assignment> Parser::parse_assignment() {
    Token identifier = expect(TokenType::IDENTIFIER);
    if (peek().type != TokenType::ASSIGN)
        throw std::runtime_error("[ASSIGNMENT] expected '=' after identifier '" + identifier.value + "' at line " + std::to_string(peek().line));
    // =
    expect(TokenType::ASSIGN);               
    auto expr = parse_expression();
    if (peek().type != TokenType::SEMICOLON)
        throw std::runtime_error("[ASSIGNMENT] expected ';' after assignment to '" + identifier.value + "' at line " + std::to_string(peek().line));
    expect(TokenType::SEMICOLON);               
    return std::make_unique<Assignment>(
        std::make_unique<Variable>(identifier.value), std::move(expr)
    );   
}

// handle +, -, <, >, ==, !=
std::unique_ptr<ASTNode> Parser::parse_expression() {
    std::unique_ptr<ASTNode> node = parse_term();
    while (is_binary_op(tokens[cur_idx].type)){
        // the current binary operator being passed awaiting rhs eval recursively
        BinaryOperator op = token_to_op(consume().type);
        // parse_factor will handle exceptions
        std::unique_ptr<ASTNode> right = parse_term(); 
        node = std::make_unique<BinaryOperation>(op, std::move(node), std::move(right)); 
    }
    return node;
}

// handle * and / 
std::unique_ptr<ASTNode> Parser::parse_term() {
    // build a single binary operation from the left side x + (y * z)
    std::unique_ptr<ASTNode> node = parse_factor();
    while (tokens[cur_idx].type == TokenType::STAR || tokens[cur_idx].type == TokenType::SLASH) {
        BinaryOperator op = (tokens[cur_idx].type == TokenType::STAR) 
            ? BinaryOperator::Star
            : BinaryOperator::Slash; 
        consume();
        std::unique_ptr<ASTNode> right = parse_factor();
        node = std::make_unique<BinaryOperation>(op, std::move(node), std::move(right));
    }
    return node;
}

// handle literals and identifiers and advance to the next index
std::unique_ptr<ASTNode> Parser::parse_factor() {
    // handle ()
    if (tokens[cur_idx].type == TokenType::LPAREN) {
        consume();
        auto expr = parse_expression();
        if (peek().type != TokenType::RPAREN)
            throw std::runtime_error("[EXPRESSION] expected ')' to close expression at line " + std::to_string(peek().line));
        expect(TokenType::RPAREN);
        return expr;
    }
    // e.g. abc
    if (tokens[cur_idx].type == TokenType::IDENTIFIER) {
        Token t = consume();    
        // function call
        if (tokens[cur_idx].type == TokenType::LPAREN) {
            return parse_function_call(t.value);
        }
        // normal variable
        return std::make_unique<Variable>(t.value);
    }
    // e.g. 123
    if (tokens[cur_idx].type == TokenType::INT_LITERAL) {
        Token t = consume();
        return std::make_unique<Constant>(ConstantType::Int, t.value); 
    }
    // e.g. 'a'
    if (tokens[cur_idx].type == TokenType::CHAR_LITERAL) {
        Token t = consume();
        return std::make_unique<Constant>(ConstantType::Char, t.value); 
    }
    // e.g. true
    if (tokens[cur_idx].type == TokenType::BOOL_LITERAL) {
        Token t = consume();
        return std::make_unique<Constant>(ConstantType::Bool, t.value); 
    }
    if (tokens[cur_idx].type == TokenType::STRING_LITERAL) {
        Token t = consume();
        return std::make_unique<Constant>(ConstantType::String, t.value);
    }
    throw std::runtime_error("[FACTOR] unexpected token '" + tokens[cur_idx].value + "' at line " + std::to_string(peek().line) + " — expected identifier, literal, or '('");
}

// retrieve the current token and then advance to the next index
Token Parser::consume() {
    Token c = peek();
    advance();
    return c;
}

// check for token's expected type, return the current token, and advance to next index
Token Parser::expect(TokenType type) {
    if (peek().type != type) 
        throw std::runtime_error("[EXPECT] unexpected token '" + peek().value + "' at line " + std::to_string(peek().line));
    return consume();
}

// for function return types — allows void
Token Parser::expect_type() {
    if (peek().type != TokenType::INT_KEYWORD  &&
        peek().type != TokenType::CHAR_KEYWORD &&
        peek().type != TokenType::BOOL_KEYWORD &&
        peek().type != TokenType::VOID_KEYWORD) {
        throw std::runtime_error("[TYPE] expected type keyword (int, char, bool, void) but got '" + peek().value + "' at line " + std::to_string(peek().line));
    }
    return consume();
}

// for variable declarations — void not allowed
Token Parser::expect_primitive_type() {
    if (peek().type != TokenType::INT_KEYWORD && 
        peek().type != TokenType::CHAR_KEYWORD && 
        peek().type != TokenType::BOOL_KEYWORD) {
        throw std::runtime_error("[PRIMITIVE TYPE] expected primitive type (int, char, bool) but got '" + peek().value + "' at line " + std::to_string(peek().line));
    }
    return consume();
}

BinaryOperator Parser::token_to_op(TokenType type) {
    switch(type) {
        case TokenType::PLUS:         return BinaryOperator::Plus;
        case TokenType::MINUS:        return BinaryOperator::Minus;
        case TokenType::LESS_THAN:    return BinaryOperator::LessThan;
        case TokenType::GREATER_THAN: return BinaryOperator::GreaterThan;
        case TokenType::EQUAL:        return BinaryOperator::Equal;
        case TokenType::NOT_EQUAL:    return BinaryOperator::NotEqual;
        default: throw std::runtime_error("[BINARY OP] unknown operator token '" + std::to_string(static_cast<int>(type)) + "'");
    }
}

bool Parser::is_binary_op(TokenType type) {
    return type == TokenType::PLUS || type == TokenType::MINUS ||
           type == TokenType::LESS_THAN || type == TokenType::GREATER_THAN ||
           type == TokenType::EQUAL || type == TokenType::NOT_EQUAL;
}

// only if statement for now
bool Parser::is_statement(TokenType type) {
    return type == TokenType::IF || type == TokenType::WHILE;
}

bool Parser::is_type_keyword(TokenType type) {
    return type == TokenType::INT_KEYWORD  ||
           type == TokenType::VOID_KEYWORD ||
           type == TokenType::CHAR_KEYWORD ||
           type == TokenType::BOOL_KEYWORD;
}

bool Parser::is_primitive_type_keyword(TokenType type) {
    return type == TokenType::INT_KEYWORD  ||
           type == TokenType::CHAR_KEYWORD ||
           type == TokenType::BOOL_KEYWORD;
}

ConstantType Parser::token_type_to_constant_type(TokenType type) {
    switch(type) {
        case TokenType::INT_KEYWORD:  return ConstantType::Int;
        case TokenType::CHAR_KEYWORD: return ConstantType::Char;
        case TokenType::BOOL_KEYWORD: return ConstantType::Bool;
        case TokenType::VOID_KEYWORD: return ConstantType::Void;
        default: throw std::runtime_error("[TYPE] cannot convert token to ConstantType");
    }
}