#define _POSIX_C_SOURCE 200809L
#include "parser.h"
#include <stdlib.h>
#include <string.h>

extern void _rava_parser_advance(RavaParser_t *parser);
extern bool _rava_parser_check(RavaParser_t *parser, RavaTokenType_e type);
extern bool _rava_parser_match(RavaParser_t *parser, RavaTokenType_e type);
extern bool _rava_parser_expect(RavaParser_t *parser, RavaTokenType_e type, const char *message);
extern RavaASTNode_t* _rava_parser_parse_expression(RavaParser_t *parser);
extern RavaASTNode_t* _rava_parser_parse_type(RavaParser_t *parser);

RavaASTNode_t* _rava_parser_parse_statement(RavaParser_t *parser);
RavaASTNode_t* _rava_parser_parse_block(RavaParser_t *parser);

RavaASTNode_t* _rava_parser_parse_block(RavaParser_t *parser) {
    RavaASTNode_t *block = rava_ast_node_create(RAVA_AST_BLOCK_STMT,
                                                parser->current_token->line,
                                                parser->current_token->column);

    _rava_parser_expect(parser, RAVA_TOKEN_LBRACE, "Expected '{' at start of block");

    while (!_rava_parser_check(parser, RAVA_TOKEN_RBRACE) &&
           !_rava_parser_check(parser, RAVA_TOKEN_EOF)) {
        RavaASTNode_t *stmt = _rava_parser_parse_statement(parser);
        if (stmt) {
            rava_ast_node_add_child(block, stmt);
        }
    }

    _rava_parser_expect(parser, RAVA_TOKEN_RBRACE, "Expected '}' at end of block");

    return block;
}

RavaASTNode_t* _rava_parser_parse_statement(RavaParser_t *parser) {
    if (_rava_parser_check(parser, RAVA_TOKEN_LBRACE)) {
        return _rava_parser_parse_block(parser);
    }

    if (_rava_parser_match(parser, RAVA_TOKEN_KEYWORD_IF)) {
        RavaASTNode_t *node = rava_ast_node_create(RAVA_AST_IF_STMT,
                                                   parser->current_token->line,
                                                   parser->current_token->column);

        _rava_parser_expect(parser, RAVA_TOKEN_LPAREN, "Expected '(' after 'if'");
        node->data.if_stmt.condition = _rava_parser_parse_expression(parser);
        _rava_parser_expect(parser, RAVA_TOKEN_RPAREN, "Expected ')' after if condition");

        node->data.if_stmt.then_stmt = _rava_parser_parse_statement(parser);

        if (_rava_parser_match(parser, RAVA_TOKEN_KEYWORD_ELSE)) {
            node->data.if_stmt.else_stmt = _rava_parser_parse_statement(parser);
        } else {
            node->data.if_stmt.else_stmt = NULL;
        }

        return node;
    }

    if (_rava_parser_match(parser, RAVA_TOKEN_KEYWORD_WHILE)) {
        RavaASTNode_t *node = rava_ast_node_create(RAVA_AST_WHILE_STMT,
                                                   parser->current_token->line,
                                                   parser->current_token->column);

        _rava_parser_expect(parser, RAVA_TOKEN_LPAREN, "Expected '(' after 'while'");
        node->data.while_stmt.condition = _rava_parser_parse_expression(parser);
        _rava_parser_expect(parser, RAVA_TOKEN_RPAREN, "Expected ')' after while condition");

        node->data.while_stmt.body = _rava_parser_parse_statement(parser);

        return node;
    }

    if (_rava_parser_match(parser, RAVA_TOKEN_KEYWORD_FOR)) {
        RavaASTNode_t *node = rava_ast_node_create(RAVA_AST_FOR_STMT,
                                                   parser->current_token->line,
                                                   parser->current_token->column);

        _rava_parser_expect(parser, RAVA_TOKEN_LPAREN, "Expected '(' after 'for'");

        if (!_rava_parser_check(parser, RAVA_TOKEN_SEMICOLON)) {
            bool is_type_decl = false;
            if (_rava_parser_check(parser, RAVA_TOKEN_KEYWORD_INT) ||
                _rava_parser_check(parser, RAVA_TOKEN_KEYWORD_LONG) ||
                _rava_parser_check(parser, RAVA_TOKEN_KEYWORD_DOUBLE) ||
                _rava_parser_check(parser, RAVA_TOKEN_KEYWORD_BOOLEAN) ||
                _rava_parser_check(parser, RAVA_TOKEN_KEYWORD_CHAR)) {
                is_type_decl = true;
            } else if (_rava_parser_check(parser, RAVA_TOKEN_IDENTIFIER)) {
                if (parser->peek_token && parser->peek_token->type == RAVA_TOKEN_IDENTIFIER) {
                    is_type_decl = true;
                }
            }
            if (is_type_decl) {
                RavaASTNode_t *type = _rava_parser_parse_type(parser);
                RavaASTNode_t *var_decl = rava_ast_node_create(RAVA_AST_VAR_DECL,
                                                               parser->current_token->line,
                                                               parser->current_token->column);
                var_decl->data.var_decl.type = type;
                if (_rava_parser_check(parser, RAVA_TOKEN_IDENTIFIER)) {
                    var_decl->data.var_decl.name = strdup(parser->current_token->lexeme);
                    _rava_parser_advance(parser);
                }
                if (_rava_parser_match(parser, RAVA_TOKEN_ASSIGN)) {
                    var_decl->data.var_decl.initializer = _rava_parser_parse_expression(parser);
                }
                node->data.for_stmt.init = var_decl;
            } else {
                node->data.for_stmt.init = _rava_parser_parse_expression(parser);
            }
        } else {
            node->data.for_stmt.init = NULL;
        }
        _rava_parser_expect(parser, RAVA_TOKEN_SEMICOLON, "Expected ';' after for init");

        if (!_rava_parser_check(parser, RAVA_TOKEN_SEMICOLON)) {
            node->data.for_stmt.condition = _rava_parser_parse_expression(parser);
        } else {
            node->data.for_stmt.condition = NULL;
        }
        _rava_parser_expect(parser, RAVA_TOKEN_SEMICOLON, "Expected ';' after for condition");

        if (!_rava_parser_check(parser, RAVA_TOKEN_RPAREN)) {
            node->data.for_stmt.update = _rava_parser_parse_expression(parser);
        } else {
            node->data.for_stmt.update = NULL;
        }
        _rava_parser_expect(parser, RAVA_TOKEN_RPAREN, "Expected ')' after for update");

        node->data.for_stmt.body = _rava_parser_parse_statement(parser);

        return node;
    }

    if (_rava_parser_match(parser, RAVA_TOKEN_KEYWORD_RETURN)) {
        RavaASTNode_t *node = rava_ast_node_create(RAVA_AST_RETURN_STMT,
                                                   parser->current_token->line,
                                                   parser->current_token->column);

        if (!_rava_parser_check(parser, RAVA_TOKEN_SEMICOLON)) {
            node->data.return_stmt.value = _rava_parser_parse_expression(parser);
        } else {
            node->data.return_stmt.value = NULL;
        }

        _rava_parser_expect(parser, RAVA_TOKEN_SEMICOLON, "Expected ';' after return statement");

        return node;
    }

    if (_rava_parser_match(parser, RAVA_TOKEN_KEYWORD_BREAK)) {
        RavaASTNode_t *node = rava_ast_node_create(RAVA_AST_BREAK_STMT,
                                                   parser->current_token->line,
                                                   parser->current_token->column);
        _rava_parser_expect(parser, RAVA_TOKEN_SEMICOLON, "Expected ';' after break");
        return node;
    }

    if (_rava_parser_match(parser, RAVA_TOKEN_KEYWORD_CONTINUE)) {
        RavaASTNode_t *node = rava_ast_node_create(RAVA_AST_CONTINUE_STMT,
                                                   parser->current_token->line,
                                                   parser->current_token->column);
        _rava_parser_expect(parser, RAVA_TOKEN_SEMICOLON, "Expected ';' after continue");
        return node;
    }

    bool is_type_decl = false;
    if (_rava_parser_check(parser, RAVA_TOKEN_KEYWORD_INT) ||
        _rava_parser_check(parser, RAVA_TOKEN_KEYWORD_BOOLEAN) ||
        _rava_parser_check(parser, RAVA_TOKEN_KEYWORD_CHAR) ||
        _rava_parser_check(parser, RAVA_TOKEN_KEYWORD_BYTE) ||
        _rava_parser_check(parser, RAVA_TOKEN_KEYWORD_SHORT) ||
        _rava_parser_check(parser, RAVA_TOKEN_KEYWORD_LONG) ||
        _rava_parser_check(parser, RAVA_TOKEN_KEYWORD_FLOAT) ||
        _rava_parser_check(parser, RAVA_TOKEN_KEYWORD_DOUBLE)) {
        is_type_decl = true;
    } else if (_rava_parser_check(parser, RAVA_TOKEN_IDENTIFIER)) {
        if (parser->peek_token && parser->peek_token->type == RAVA_TOKEN_IDENTIFIER) {
            is_type_decl = true;
        }
    }

    if (is_type_decl) {
        RavaASTNode_t *type = _rava_parser_parse_type(parser);
        RavaASTNode_t *var_decl = rava_ast_node_create(RAVA_AST_VAR_DECL,
                                                       parser->current_token->line,
                                                       parser->current_token->column);
        var_decl->data.var_decl.type = type;

        if (_rava_parser_check(parser, RAVA_TOKEN_IDENTIFIER)) {
            var_decl->data.var_decl.name = strdup(parser->current_token->lexeme);
            _rava_parser_advance(parser);
        }

        if (_rava_parser_match(parser, RAVA_TOKEN_ASSIGN)) {
            var_decl->data.var_decl.initializer = _rava_parser_parse_expression(parser);
        } else {
            var_decl->data.var_decl.initializer = NULL;
        }

        _rava_parser_expect(parser, RAVA_TOKEN_SEMICOLON, "Expected ';' after variable declaration");

        return var_decl;
    }

    RavaASTNode_t *expr = _rava_parser_parse_expression(parser);
    _rava_parser_expect(parser, RAVA_TOKEN_SEMICOLON, "Expected ';' after expression");

    RavaASTNode_t *stmt = rava_ast_node_create(RAVA_AST_EXPR_STMT,
                                               parser->current_token->line,
                                               parser->current_token->column);
    rava_ast_node_add_child(stmt, expr);

    return stmt;
}