#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_array_initializer(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)) {
        if (parser->had_error) {
            break;
        }
        RavaASTNode_t *stmt = _rava_parser_parse_statement(parser);
        if (stmt) {
            rava_ast_node_add_child(block, stmt);
        } else if (parser->had_error) {
            break;
        }
    }

    _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_DO)) {
        RavaASTNode_t *node = rava_ast_node_create(RAVA_AST_DO_WHILE_STMT,
                                                   parser->current_token->line,
                                                   parser->current_token->column);

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

        _rava_parser_expect(parser, RAVA_TOKEN_KEYWORD_WHILE, "Expected 'while' after do body");
        _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 do-while condition");
        _rava_parser_expect(parser, RAVA_TOKEN_SEMICOLON, "Expected ';' after do-while");

        return node;
    }

    if (_rava_parser_match(parser, RAVA_TOKEN_KEYWORD_FOR)) {
        int line = parser->current_token->line;
        int column = 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);
                char *var_name = NULL;
                if (_rava_parser_check(parser, RAVA_TOKEN_IDENTIFIER)) {
                    var_name = strdup(parser->current_token->lexeme);
                    _rava_parser_advance(parser);
                }

                if (_rava_parser_match(parser, RAVA_TOKEN_COLON)) {
                    RavaASTNode_t *node = rava_ast_node_create(RAVA_AST_ENHANCED_FOR_STMT, line, column);
                    node->data.enhanced_for.var_type = type;
                    node->data.enhanced_for.var_name = var_name;
                    node->data.enhanced_for.iterable = _rava_parser_parse_expression(parser);
                    _rava_parser_expect(parser, RAVA_TOKEN_RPAREN, "Expected ')' after enhanced for");
                    node->data.enhanced_for.body = _rava_parser_parse_statement(parser);
                    return node;
                }

                RavaASTNode_t *node = rava_ast_node_create(RAVA_AST_FOR_STMT, line, column);
                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;
                var_decl->data.var_decl.name = var_name;
                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;
                _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;
            } else {
                RavaASTNode_t *node = rava_ast_node_create(RAVA_AST_FOR_STMT, line, column);
                node->data.for_stmt.init = _rava_parser_parse_expression(parser);
                _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;
            }
        } else {
            RavaASTNode_t *node = rava_ast_node_create(RAVA_AST_FOR_STMT, line, column);
            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_SWITCH)) {
        RavaASTNode_t *node = rava_ast_node_create(RAVA_AST_SWITCH_STMT,
                                                   parser->current_token->line,
                                                   parser->current_token->column);

        _rava_parser_expect(parser, RAVA_TOKEN_LPAREN, "Expected '(' after 'switch'");
        node->data.switch_stmt.expression = _rava_parser_parse_expression(parser);
        _rava_parser_expect(parser, RAVA_TOKEN_RPAREN, "Expected ')' after switch expression");
        _rava_parser_expect(parser, RAVA_TOKEN_LBRACE, "Expected '{' after switch");

        while (!_rava_parser_check(parser, RAVA_TOKEN_RBRACE) &&
               !_rava_parser_check(parser, RAVA_TOKEN_EOF) &&
               !parser->had_error) {

            if (_rava_parser_match(parser, RAVA_TOKEN_KEYWORD_CASE)) {
                RavaASTNode_t *case_node = rava_ast_node_create(RAVA_AST_CASE_STMT,
                                                                parser->current_token->line,
                                                                parser->current_token->column);
                case_node->data.case_stmt.is_default = false;
                case_node->data.case_stmt.value = _rava_parser_parse_expression(parser);
                _rava_parser_expect(parser, RAVA_TOKEN_COLON, "Expected ':' after case value");

                while (!_rava_parser_check(parser, RAVA_TOKEN_KEYWORD_CASE) &&
                       !_rava_parser_check(parser, RAVA_TOKEN_KEYWORD_DEFAULT) &&
                       !_rava_parser_check(parser, RAVA_TOKEN_RBRACE) &&
                       !_rava_parser_check(parser, RAVA_TOKEN_EOF) &&
                       !parser->had_error) {
                    RavaASTNode_t *stmt = _rava_parser_parse_statement(parser);
                    if (stmt) {
                        rava_ast_node_add_child(case_node, stmt);
                    } else if (parser->had_error) {
                        break;
                    }
                }

                rava_ast_node_add_child(node, case_node);

            } else if (_rava_parser_match(parser, RAVA_TOKEN_KEYWORD_DEFAULT)) {
                RavaASTNode_t *default_node = rava_ast_node_create(RAVA_AST_CASE_STMT,
                                                                   parser->current_token->line,
                                                                   parser->current_token->column);
                default_node->data.case_stmt.is_default = true;
                default_node->data.case_stmt.value = NULL;
                _rava_parser_expect(parser, RAVA_TOKEN_COLON, "Expected ':' after default");

                while (!_rava_parser_check(parser, RAVA_TOKEN_KEYWORD_CASE) &&
                       !_rava_parser_check(parser, RAVA_TOKEN_KEYWORD_DEFAULT) &&
                       !_rava_parser_check(parser, RAVA_TOKEN_RBRACE) &&
                       !_rava_parser_check(parser, RAVA_TOKEN_EOF) &&
                       !parser->had_error) {
                    RavaASTNode_t *stmt = _rava_parser_parse_statement(parser);
                    if (stmt) {
                        rava_ast_node_add_child(default_node, stmt);
                    } else if (parser->had_error) {
                        break;
                    }
                }

                rava_ast_node_add_child(node, default_node);
            } else {
                break;
            }
        }

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

        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;
    }

    if (_rava_parser_match(parser, RAVA_TOKEN_KEYWORD_THROW)) {
        RavaASTNode_t *node = rava_ast_node_create(RAVA_AST_THROW_STMT,
                                                   parser->current_token->line,
                                                   parser->current_token->column);
        node->data.throw_stmt.expression = _rava_parser_parse_expression(parser);
        _rava_parser_expect(parser, RAVA_TOKEN_SEMICOLON, "Expected ';' after throw");
        return node;
    }

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

        node->data.try_stmt.try_block = _rava_parser_parse_block(parser);
        node->data.try_stmt.catch_clauses = NULL;
        node->data.try_stmt.catch_count = 0;
        node->data.try_stmt.finally_block = NULL;

        size_t capacity = 4;
        node->data.try_stmt.catch_clauses = malloc(capacity * sizeof(RavaASTNode_t*));

        while (_rava_parser_match(parser, RAVA_TOKEN_KEYWORD_CATCH)) {
            RavaASTNode_t *catch_node = rava_ast_node_create(RAVA_AST_CATCH_CLAUSE,
                                                              parser->current_token->line,
                                                              parser->current_token->column);

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

            if (_rava_parser_check(parser, RAVA_TOKEN_IDENTIFIER)) {
                catch_node->data.catch_clause.exception_type = strdup(parser->current_token->lexeme);
                _rava_parser_advance(parser);
            }

            if (_rava_parser_check(parser, RAVA_TOKEN_IDENTIFIER)) {
                catch_node->data.catch_clause.exception_name = strdup(parser->current_token->lexeme);
                _rava_parser_advance(parser);
            }

            _rava_parser_expect(parser, RAVA_TOKEN_RPAREN, "Expected ')' after catch parameter");

            catch_node->data.catch_clause.body = _rava_parser_parse_block(parser);

            if (node->data.try_stmt.catch_count >= capacity) {
                capacity *= 2;
                node->data.try_stmt.catch_clauses = realloc(
                    node->data.try_stmt.catch_clauses,
                    capacity * sizeof(RavaASTNode_t*));
            }
            node->data.try_stmt.catch_clauses[node->data.try_stmt.catch_count++] = catch_node;
        }

        if (_rava_parser_match(parser, RAVA_TOKEN_KEYWORD_FINALLY)) {
            node->data.try_stmt.finally_block = _rava_parser_parse_block(parser);
        }

        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)) {
            if (_rava_parser_check(parser, RAVA_TOKEN_LBRACE)) {
                var_decl->data.var_decl.initializer = _rava_parser_parse_array_initializer(parser);
            } else {
                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;
}