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

static void _rava_ir_gen_expression(RavaIRGenerator_t *gen, RavaASTNode_t *expr);
static void _rava_ir_gen_statement(RavaIRGenerator_t *gen, RavaASTNode_t *stmt);

static void _rava_ir_clear_locals(RavaIRGenerator_t *gen) {
    RavaLocalVar_t *current = gen->locals;
    while (current) {
        RavaLocalVar_t *next = current->next;
        free(current->name);
        free(current);
        current = next;
    }
    gen->locals = NULL;
}

static void _rava_ir_add_local(RavaIRGenerator_t *gen, const char *name, size_t index) {
    RavaLocalVar_t *local = calloc(1, sizeof(RavaLocalVar_t));
    local->name = strdup(name);
    local->index = index;
    local->next = gen->locals;
    gen->locals = local;
}

static int _rava_ir_find_local(RavaIRGenerator_t *gen, const char *name) {
    for (RavaLocalVar_t *local = gen->locals; local; local = local->next) {
        if (strcmp(local->name, name) == 0) {
            return (int)local->index;
        }
    }
    return -1;
}

RavaIRGenerator_t* rava_ir_generator_create(RavaSemanticAnalyzer_t *analyzer) {
    RavaIRGenerator_t *gen = calloc(1, sizeof(RavaIRGenerator_t));
    gen->program = rava_program_create();
    gen->current_class = NULL;
    gen->current_method = NULL;
    gen->analyzer = analyzer;
    gen->next_local = 0;
    gen->locals = NULL;
    gen->loop_context = NULL;
    gen->error_message = NULL;
    gen->had_error = false;
    return gen;
}

void rava_ir_generator_destroy(RavaIRGenerator_t *generator) {
    if (!generator) return;
    _rava_ir_clear_locals(generator);
    free(generator->error_message);
    free(generator);
}

static void _rava_ir_emit(RavaIRGenerator_t *gen, RavaInstruction_t instr) {
    if (!gen->current_method) return;
    rava_instruction_list_add(gen->current_method->instructions, instr);
}

static void _rava_ir_gen_binary_expr(RavaIRGenerator_t *gen, RavaASTNode_t *expr) {
    _rava_ir_gen_expression(gen, expr->data.binary.left);
    _rava_ir_gen_expression(gen, expr->data.binary.right);

    RavaInstruction_t instr = {0};
    instr.line = expr->line;

    switch (expr->data.binary.op) {
        case RAVA_BINOP_ADD: instr.opcode = RAVA_OP_ADD; break;
        case RAVA_BINOP_SUB: instr.opcode = RAVA_OP_SUB; break;
        case RAVA_BINOP_MUL: instr.opcode = RAVA_OP_MUL; break;
        case RAVA_BINOP_DIV: instr.opcode = RAVA_OP_DIV; break;
        case RAVA_BINOP_MOD: instr.opcode = RAVA_OP_MOD; break;
        case RAVA_BINOP_EQ: instr.opcode = RAVA_OP_EQ; break;
        case RAVA_BINOP_NE: instr.opcode = RAVA_OP_NE; break;
        case RAVA_BINOP_LT: instr.opcode = RAVA_OP_LT; break;
        case RAVA_BINOP_LE: instr.opcode = RAVA_OP_LE; break;
        case RAVA_BINOP_GT: instr.opcode = RAVA_OP_GT; break;
        case RAVA_BINOP_GE: instr.opcode = RAVA_OP_GE; break;
        case RAVA_BINOP_AND: instr.opcode = RAVA_OP_AND; break;
        case RAVA_BINOP_OR: instr.opcode = RAVA_OP_OR; break;
        default: instr.opcode = RAVA_OP_NOP; break;
    }

    _rava_ir_emit(gen, instr);
}

static void _rava_ir_gen_expression(RavaIRGenerator_t *gen, RavaASTNode_t *expr) {
    if (!expr) return;

    RavaInstruction_t instr = {0};
    instr.line = expr->line;

    switch (expr->type) {
        case RAVA_AST_LITERAL_EXPR:
            if (expr->data.literal.literal_type == RAVA_TOKEN_LITERAL_STRING) {
                instr.opcode = RAVA_OP_LOAD_STRING;
                instr.operand.string_value = strdup(expr->data.literal.value.string_value);
                _rava_ir_emit(gen, instr);
            } else if (expr->data.literal.literal_type == RAVA_TOKEN_LITERAL_LONG) {
                instr.opcode = RAVA_OP_LOAD_LONG;
                instr.operand.int_value = expr->data.literal.value.int_value;
                _rava_ir_emit(gen, instr);
            } else if (expr->data.literal.literal_type == RAVA_TOKEN_LITERAL_DOUBLE) {
                instr.opcode = RAVA_OP_LOAD_DOUBLE;
                instr.operand.float_value = expr->data.literal.value.float_value;
                _rava_ir_emit(gen, instr);
            } else {
                instr.opcode = RAVA_OP_LOAD_CONST;
                instr.operand.int_value = expr->data.literal.value.int_value;
                _rava_ir_emit(gen, instr);
            }
            break;

        case RAVA_AST_IDENTIFIER_EXPR: {
            int local_index = _rava_ir_find_local(gen, expr->data.identifier.name);
            if (local_index >= 0) {
                instr.opcode = RAVA_OP_LOAD_LOCAL;
                instr.operand.var.index = (size_t)local_index;
                _rava_ir_emit(gen, instr);
            }
            break;
        }

        case RAVA_AST_BINARY_EXPR:
            _rava_ir_gen_binary_expr(gen, expr);
            break;

        case RAVA_AST_UNARY_EXPR:
            if (expr->data.unary.op == RAVA_UNOP_MINUS) {
                _rava_ir_gen_expression(gen, expr->data.unary.operand);
                instr.opcode = RAVA_OP_NEG;
                _rava_ir_emit(gen, instr);
            } else if (expr->data.unary.op == RAVA_UNOP_NOT) {
                _rava_ir_gen_expression(gen, expr->data.unary.operand);
                instr.opcode = RAVA_OP_NOT;
                _rava_ir_emit(gen, instr);
            } else if (expr->data.unary.op == RAVA_UNOP_PREINC || expr->data.unary.op == RAVA_UNOP_PREDEC) {
                if (expr->data.unary.operand->type == RAVA_AST_IDENTIFIER_EXPR) {
                    int local_index = _rava_ir_find_local(gen, expr->data.unary.operand->data.identifier.name);
                    if (local_index >= 0) {
                        instr.opcode = RAVA_OP_LOAD_LOCAL;
                        instr.operand.var.index = (size_t)local_index;
                        _rava_ir_emit(gen, instr);
                        instr.opcode = RAVA_OP_LOAD_CONST;
                        instr.operand.int_value = 1;
                        _rava_ir_emit(gen, instr);
                        instr.opcode = (expr->data.unary.op == RAVA_UNOP_PREINC) ? RAVA_OP_ADD : RAVA_OP_SUB;
                        _rava_ir_emit(gen, instr);
                        instr.opcode = RAVA_OP_DUP;
                        _rava_ir_emit(gen, instr);
                        instr.opcode = RAVA_OP_STORE_LOCAL;
                        instr.operand.var.index = (size_t)local_index;
                        _rava_ir_emit(gen, instr);
                    }
                }
            } else if (expr->data.unary.op == RAVA_UNOP_POSTINC || expr->data.unary.op == RAVA_UNOP_POSTDEC) {
                if (expr->data.unary.operand->type == RAVA_AST_IDENTIFIER_EXPR) {
                    int local_index = _rava_ir_find_local(gen, expr->data.unary.operand->data.identifier.name);
                    if (local_index >= 0) {
                        instr.opcode = RAVA_OP_LOAD_LOCAL;
                        instr.operand.var.index = (size_t)local_index;
                        _rava_ir_emit(gen, instr);
                        instr.opcode = RAVA_OP_LOAD_LOCAL;
                        instr.operand.var.index = (size_t)local_index;
                        _rava_ir_emit(gen, instr);
                        instr.opcode = RAVA_OP_LOAD_CONST;
                        instr.operand.int_value = 1;
                        _rava_ir_emit(gen, instr);
                        instr.opcode = (expr->data.unary.op == RAVA_UNOP_POSTINC) ? RAVA_OP_ADD : RAVA_OP_SUB;
                        _rava_ir_emit(gen, instr);
                        instr.opcode = RAVA_OP_STORE_LOCAL;
                        instr.operand.var.index = (size_t)local_index;
                        _rava_ir_emit(gen, instr);
                    }
                }
            } else {
                _rava_ir_gen_expression(gen, expr->data.unary.operand);
            }
            break;

        case RAVA_AST_ASSIGN_EXPR:
            if (expr->data.assign.target->type == RAVA_AST_ARRAY_ACCESS_EXPR) {
                _rava_ir_gen_expression(gen, expr->data.assign.target->data.array_access.array);
                _rava_ir_gen_expression(gen, expr->data.assign.target->data.array_access.index);
                _rava_ir_gen_expression(gen, expr->data.assign.value);
                instr.opcode = RAVA_OP_STORE_ARRAY;
                _rava_ir_emit(gen, instr);
            } else if (expr->data.assign.target->type == RAVA_AST_MEMBER_ACCESS_EXPR) {
                _rava_ir_gen_expression(gen, expr->data.assign.target->data.member_access.object);
                _rava_ir_gen_expression(gen, expr->data.assign.value);
                instr.opcode = RAVA_OP_PUT_FIELD;
                instr.operand.field.field_name = strdup(expr->data.assign.target->data.member_access.member);
                _rava_ir_emit(gen, instr);
            } else {
                _rava_ir_gen_expression(gen, expr->data.assign.value);

                if (expr->data.assign.target->type == RAVA_AST_IDENTIFIER_EXPR) {
                    int local_index = _rava_ir_find_local(gen, expr->data.assign.target->data.identifier.name);
                    if (local_index >= 0) {
                        instr.opcode = RAVA_OP_STORE_LOCAL;
                        instr.operand.var.index = (size_t)local_index;
                        _rava_ir_emit(gen, instr);
                    }
                }
            }
            break;

        case RAVA_AST_MEMBER_ACCESS_EXPR:
            _rava_ir_gen_expression(gen, expr->data.member_access.object);
            instr.opcode = RAVA_OP_GET_FIELD;
            instr.operand.field.field_name = strdup(expr->data.member_access.member);
            _rava_ir_emit(gen, instr);
            break;

        case RAVA_AST_ARRAY_ACCESS_EXPR:
            _rava_ir_gen_expression(gen, expr->data.array_access.array);
            _rava_ir_gen_expression(gen, expr->data.array_access.index);
            instr.opcode = RAVA_OP_LOAD_ARRAY;
            _rava_ir_emit(gen, instr);
            break;

        case RAVA_AST_NEW_EXPR:
            if (expr->data.new_expr.type && expr->data.new_expr.type->data.type.is_array) {
                if (expr->data.new_expr.arguments_count > 0) {
                    _rava_ir_gen_expression(gen, expr->data.new_expr.arguments[0]);
                }
                instr.opcode = RAVA_OP_NEW_ARRAY;
                _rava_ir_emit(gen, instr);
            } else if (expr->data.new_expr.type) {
                instr.opcode = RAVA_OP_NEW;
                instr.operand.call.class_name = strdup(expr->data.new_expr.type->data.type.type_name);
                instr.operand.call.arg_count = expr->data.new_expr.arguments_count;
                _rava_ir_emit(gen, instr);
                instr.opcode = RAVA_OP_DUP;
                _rava_ir_emit(gen, instr);
                for (size_t i = 0; i < expr->data.new_expr.arguments_count; i++) {
                    _rava_ir_gen_expression(gen, expr->data.new_expr.arguments[i]);
                }
                instr.opcode = RAVA_OP_CALL_VIRTUAL;
                instr.operand.call.method_name = strdup("<init>");
                _rava_ir_emit(gen, instr);
            }
            break;

        case RAVA_AST_CALL_EXPR: {
            bool is_println = false;
            bool is_print = false;
            bool is_string_length = false;
            bool is_string_charat = false;
            bool is_string_equals = false;
            bool is_file_read = false;
            bool is_file_write = false;
            bool is_file_exists = false;
            bool is_file_delete = false;
            bool is_current_time_millis = false;
            bool is_nano_time = false;
            RavaASTNode_t *string_object = NULL;

            if (expr->data.call.callee->type == RAVA_AST_MEMBER_ACCESS_EXPR) {
                RavaASTNode_t *member = expr->data.call.callee;
                if (member->data.member_access.object->type == RAVA_AST_IDENTIFIER_EXPR &&
                    strcmp(member->data.member_access.object->data.identifier.name, "System") == 0) {
                    if (strcmp(member->data.member_access.member, "currentTimeMillis") == 0) {
                        is_current_time_millis = true;
                    } else if (strcmp(member->data.member_access.member, "nanoTime") == 0) {
                        is_nano_time = true;
                    }
                }
                if (member->data.member_access.object->type == RAVA_AST_MEMBER_ACCESS_EXPR) {
                    RavaASTNode_t *system_out = member->data.member_access.object;
                    if (system_out->data.member_access.object->type == RAVA_AST_IDENTIFIER_EXPR &&
                        strcmp(system_out->data.member_access.object->data.identifier.name, "System") == 0 &&
                        strcmp(system_out->data.member_access.member, "out") == 0) {
                        if (strcmp(member->data.member_access.member, "println") == 0) {
                            is_println = true;
                        } else if (strcmp(member->data.member_access.member, "print") == 0) {
                            is_print = true;
                        }
                    }
                }
                if (member->data.member_access.object->type == RAVA_AST_IDENTIFIER_EXPR &&
                    strcmp(member->data.member_access.object->data.identifier.name, "Files") == 0) {
                    if (strcmp(member->data.member_access.member, "read") == 0 && expr->data.call.arguments_count == 1) {
                        is_file_read = true;
                    } else if (strcmp(member->data.member_access.member, "write") == 0 && expr->data.call.arguments_count == 2) {
                        is_file_write = true;
                    } else if (strcmp(member->data.member_access.member, "exists") == 0 && expr->data.call.arguments_count == 1) {
                        is_file_exists = true;
                    } else if (strcmp(member->data.member_access.member, "delete") == 0 && expr->data.call.arguments_count == 1) {
                        is_file_delete = true;
                    }
                }
                if (strcmp(member->data.member_access.member, "length") == 0 && expr->data.call.arguments_count == 0) {
                    is_string_length = true;
                    string_object = member->data.member_access.object;
                } else if (strcmp(member->data.member_access.member, "charAt") == 0 && expr->data.call.arguments_count == 1) {
                    is_string_charat = true;
                    string_object = member->data.member_access.object;
                } else if (strcmp(member->data.member_access.member, "equals") == 0 && expr->data.call.arguments_count == 1) {
                    is_string_equals = true;
                    string_object = member->data.member_access.object;
                }
            }

            if (is_current_time_millis) {
                instr.opcode = RAVA_OP_CURRENT_TIME_MILLIS;
                _rava_ir_emit(gen, instr);
            } else if (is_nano_time) {
                instr.opcode = RAVA_OP_NANO_TIME;
                _rava_ir_emit(gen, instr);
            } else if (is_file_read) {
                _rava_ir_gen_expression(gen, expr->data.call.arguments[0]);
                instr.opcode = RAVA_OP_FILE_READ;
                _rava_ir_emit(gen, instr);
            } else if (is_file_write) {
                _rava_ir_gen_expression(gen, expr->data.call.arguments[0]);
                _rava_ir_gen_expression(gen, expr->data.call.arguments[1]);
                instr.opcode = RAVA_OP_FILE_WRITE;
                _rava_ir_emit(gen, instr);
            } else if (is_file_exists) {
                _rava_ir_gen_expression(gen, expr->data.call.arguments[0]);
                instr.opcode = RAVA_OP_FILE_EXISTS;
                _rava_ir_emit(gen, instr);
            } else if (is_file_delete) {
                _rava_ir_gen_expression(gen, expr->data.call.arguments[0]);
                instr.opcode = RAVA_OP_FILE_DELETE;
                _rava_ir_emit(gen, instr);
            } else if (is_string_length) {
                _rava_ir_gen_expression(gen, string_object);
                instr.opcode = RAVA_OP_STRING_LENGTH;
                _rava_ir_emit(gen, instr);
            } else if (is_string_charat) {
                _rava_ir_gen_expression(gen, string_object);
                _rava_ir_gen_expression(gen, expr->data.call.arguments[0]);
                instr.opcode = RAVA_OP_STRING_CHARAT;
                _rava_ir_emit(gen, instr);
            } else if (is_string_equals) {
                _rava_ir_gen_expression(gen, string_object);
                _rava_ir_gen_expression(gen, expr->data.call.arguments[0]);
                instr.opcode = RAVA_OP_STRING_EQUALS;
                _rava_ir_emit(gen, instr);
            } else if (is_println || is_print) {
                for (size_t i = 0; i < expr->data.call.arguments_count; i++) {
                    _rava_ir_gen_expression(gen, expr->data.call.arguments[i]);
                }
                if (is_println) {
                    instr.opcode = RAVA_OP_PRINTLN;
                } else {
                    instr.opcode = RAVA_OP_PRINT;
                }
                _rava_ir_emit(gen, instr);
            } else if (expr->data.call.callee->type == RAVA_AST_MEMBER_ACCESS_EXPR) {
                RavaASTNode_t *member = expr->data.call.callee;
                _rava_ir_gen_expression(gen, member->data.member_access.object);
                for (size_t i = 0; i < expr->data.call.arguments_count; i++) {
                    _rava_ir_gen_expression(gen, expr->data.call.arguments[i]);
                }
                instr.opcode = RAVA_OP_CALL_VIRTUAL;
                instr.operand.call.class_name = strdup(gen->current_class->name);
                instr.operand.call.method_name = strdup(member->data.member_access.member);
                instr.operand.call.arg_count = expr->data.call.arguments_count;
                _rava_ir_emit(gen, instr);
            } else {
                for (size_t i = 0; i < expr->data.call.arguments_count; i++) {
                    _rava_ir_gen_expression(gen, expr->data.call.arguments[i]);
                }
                instr.opcode = RAVA_OP_CALL_STATIC;
                if (expr->data.call.callee->type == RAVA_AST_IDENTIFIER_EXPR) {
                    instr.operand.call.class_name = strdup(gen->current_class->name);
                    instr.operand.call.method_name = strdup(expr->data.call.callee->data.identifier.name);
                    instr.operand.call.arg_count = expr->data.call.arguments_count;
                }
                _rava_ir_emit(gen, instr);
            }
            break;
        }

        case RAVA_AST_THIS_EXPR:
            instr.opcode = RAVA_OP_LOAD_THIS;
            _rava_ir_emit(gen, instr);
            break;

        default:
            break;
    }
}

static void _rava_ir_gen_statement(RavaIRGenerator_t *gen, RavaASTNode_t *stmt) {
    if (!stmt) return;

    RavaInstruction_t instr = {0};
    instr.line = stmt->line;

    switch (stmt->type) {
        case RAVA_AST_BLOCK_STMT:
            for (size_t i = 0; i < stmt->children_count; i++) {
                _rava_ir_gen_statement(gen, stmt->children[i]);
            }
            break;

        case RAVA_AST_IF_STMT: {
            _rava_ir_gen_expression(gen, stmt->data.if_stmt.condition);

            int else_label = rava_instruction_list_new_label(gen->current_method->instructions);
            int end_label = rava_instruction_list_new_label(gen->current_method->instructions);

            instr.opcode = RAVA_OP_JUMP_IF_FALSE;
            instr.operand.label_id = else_label;
            _rava_ir_emit(gen, instr);

            _rava_ir_gen_statement(gen, stmt->data.if_stmt.then_stmt);

            instr.opcode = RAVA_OP_JUMP;
            instr.operand.label_id = end_label;
            _rava_ir_emit(gen, instr);

            instr.opcode = RAVA_OP_LABEL;
            instr.operand.label_id = else_label;
            _rava_ir_emit(gen, instr);

            if (stmt->data.if_stmt.else_stmt) {
                _rava_ir_gen_statement(gen, stmt->data.if_stmt.else_stmt);
            }

            instr.opcode = RAVA_OP_LABEL;
            instr.operand.label_id = end_label;
            _rava_ir_emit(gen, instr);
            break;
        }

        case RAVA_AST_WHILE_STMT: {
            int start_label = rava_instruction_list_new_label(gen->current_method->instructions);
            int end_label = rava_instruction_list_new_label(gen->current_method->instructions);

            RavaLoopContext_t loop_ctx = { .break_label = end_label, .continue_label = start_label, .parent = gen->loop_context };
            gen->loop_context = &loop_ctx;

            instr.opcode = RAVA_OP_LABEL;
            instr.operand.label_id = start_label;
            _rava_ir_emit(gen, instr);

            _rava_ir_gen_expression(gen, stmt->data.while_stmt.condition);

            instr.opcode = RAVA_OP_JUMP_IF_FALSE;
            instr.operand.label_id = end_label;
            _rava_ir_emit(gen, instr);

            _rava_ir_gen_statement(gen, stmt->data.while_stmt.body);

            instr.opcode = RAVA_OP_JUMP;
            instr.operand.label_id = start_label;
            _rava_ir_emit(gen, instr);

            instr.opcode = RAVA_OP_LABEL;
            instr.operand.label_id = end_label;
            _rava_ir_emit(gen, instr);

            gen->loop_context = loop_ctx.parent;
            break;
        }

        case RAVA_AST_FOR_STMT: {
            if (stmt->data.for_stmt.init) {
                if (stmt->data.for_stmt.init->type == RAVA_AST_VAR_DECL) {
                    _rava_ir_gen_statement(gen, stmt->data.for_stmt.init);
                } else {
                    _rava_ir_gen_expression(gen, stmt->data.for_stmt.init);
                    instr.opcode = RAVA_OP_POP;
                    _rava_ir_emit(gen, instr);
                }
            }

            int start_label = rava_instruction_list_new_label(gen->current_method->instructions);
            int continue_label = rava_instruction_list_new_label(gen->current_method->instructions);
            int end_label = rava_instruction_list_new_label(gen->current_method->instructions);

            RavaLoopContext_t loop_ctx = { .break_label = end_label, .continue_label = continue_label, .parent = gen->loop_context };
            gen->loop_context = &loop_ctx;

            instr.opcode = RAVA_OP_LABEL;
            instr.operand.label_id = start_label;
            _rava_ir_emit(gen, instr);

            if (stmt->data.for_stmt.condition) {
                _rava_ir_gen_expression(gen, stmt->data.for_stmt.condition);
                instr.opcode = RAVA_OP_JUMP_IF_FALSE;
                instr.operand.label_id = end_label;
                _rava_ir_emit(gen, instr);
            }

            _rava_ir_gen_statement(gen, stmt->data.for_stmt.body);

            instr.opcode = RAVA_OP_LABEL;
            instr.operand.label_id = continue_label;
            _rava_ir_emit(gen, instr);

            if (stmt->data.for_stmt.update) {
                _rava_ir_gen_expression(gen, stmt->data.for_stmt.update);
                instr.opcode = RAVA_OP_POP;
                _rava_ir_emit(gen, instr);
            }

            instr.opcode = RAVA_OP_JUMP;
            instr.operand.label_id = start_label;
            _rava_ir_emit(gen, instr);

            instr.opcode = RAVA_OP_LABEL;
            instr.operand.label_id = end_label;
            _rava_ir_emit(gen, instr);

            gen->loop_context = loop_ctx.parent;
            break;
        }

        case RAVA_AST_RETURN_STMT:
            if (stmt->data.return_stmt.value) {
                _rava_ir_gen_expression(gen, stmt->data.return_stmt.value);
                instr.opcode = RAVA_OP_RETURN;
            } else {
                instr.opcode = RAVA_OP_RETURN_VOID;
            }
            _rava_ir_emit(gen, instr);
            break;

        case RAVA_AST_EXPR_STMT:
            if (stmt->children_count > 0) {
                _rava_ir_gen_expression(gen, stmt->children[0]);
                instr.opcode = RAVA_OP_POP;
                _rava_ir_emit(gen, instr);
            }
            break;

        case RAVA_AST_VAR_DECL: {
            _rava_ir_add_local(gen, stmt->data.var_decl.name, gen->next_local);
            gen->next_local++;
            if (stmt->data.var_decl.initializer) {
                _rava_ir_gen_expression(gen, stmt->data.var_decl.initializer);
                instr.opcode = RAVA_OP_STORE_LOCAL;
                instr.operand.var.index = gen->next_local - 1;
                _rava_ir_emit(gen, instr);
            }
            break;
        }

        case RAVA_AST_BREAK_STMT:
            if (gen->loop_context) {
                instr.opcode = RAVA_OP_JUMP;
                instr.operand.label_id = gen->loop_context->break_label;
                _rava_ir_emit(gen, instr);
            }
            break;

        case RAVA_AST_CONTINUE_STMT:
            if (gen->loop_context) {
                instr.opcode = RAVA_OP_JUMP;
                instr.operand.label_id = gen->loop_context->continue_label;
                _rava_ir_emit(gen, instr);
            }
            break;

        default:
            break;
    }
}

static void _rava_ir_gen_method(RavaIRGenerator_t *gen, RavaASTNode_t *method_node) {
    if (method_node->type != RAVA_AST_METHOD_DECL) return;

    RavaType_t *return_type = rava_type_from_name(
        method_node->data.method_decl.return_type->data.type.type_name);

    RavaMethod_t *method = rava_method_create(method_node->data.method_decl.name, return_type);
    gen->current_method = method;
    gen->next_local = 0;
    _rava_ir_clear_locals(gen);

    rava_symbol_table_enter_scope(gen->analyzer->symbol_table, method_node->data.method_decl.name);

    for (size_t i = 0; i < method_node->children_count; i++) {
        RavaASTNode_t *child = method_node->children[i];

        if (child->type == RAVA_AST_PARAM_DECL) {
            _rava_ir_add_local(gen, child->data.var_decl.name, gen->next_local);
            gen->next_local++;
        } else if (child->type == RAVA_AST_BLOCK_STMT) {
            _rava_ir_gen_statement(gen, child);
        }
    }

    rava_symbol_table_exit_scope(gen->analyzer->symbol_table);

    method->local_count = gen->next_local;
    rava_class_add_method(gen->current_class, method);
    gen->current_method = NULL;
}

static void _rava_ir_gen_constructor(RavaIRGenerator_t *gen, RavaASTNode_t *ctor_node) {
    if (ctor_node->type != RAVA_AST_CONSTRUCTOR_DECL) return;

    RavaType_t *void_type = rava_type_from_name("void");
    RavaMethod_t *method = rava_method_create("<init>", void_type);
    gen->current_method = method;
    gen->next_local = 0;
    _rava_ir_clear_locals(gen);

    for (size_t i = 0; i < ctor_node->children_count; i++) {
        RavaASTNode_t *child = ctor_node->children[i];
        if (child->type == RAVA_AST_PARAM_DECL) {
            _rava_ir_add_local(gen, child->data.var_decl.name, gen->next_local);
            gen->next_local++;
        } else if (child->type == RAVA_AST_BLOCK_STMT) {
            _rava_ir_gen_statement(gen, child);
        }
    }

    RavaInstruction_t ret_instr = {0};
    ret_instr.opcode = RAVA_OP_RETURN_VOID;
    _rava_ir_emit(gen, ret_instr);

    method->local_count = gen->next_local;
    rava_class_add_method(gen->current_class, method);
    gen->current_method = NULL;
}

static void _rava_ir_gen_class(RavaIRGenerator_t *gen, RavaASTNode_t *class_node) {
    if (class_node->type != RAVA_AST_CLASS_DECL) return;

    RavaClass_t *class = rava_class_create(class_node->data.class_decl.name);
    if (class_node->data.class_decl.superclass) {
        class->superclass = strdup(class_node->data.class_decl.superclass);
    }
    gen->current_class = class;

    rava_symbol_table_enter_scope(gen->analyzer->symbol_table, class_node->data.class_decl.name);

    for (size_t i = 0; i < class_node->children_count; i++) {
        RavaASTNode_t *child = class_node->children[i];

        if (child->type == RAVA_AST_METHOD_DECL) {
            _rava_ir_gen_method(gen, child);
        } else if (child->type == RAVA_AST_CONSTRUCTOR_DECL) {
            _rava_ir_gen_constructor(gen, child);
        }
    }

    rava_symbol_table_exit_scope(gen->analyzer->symbol_table);

    rava_program_add_class(gen->program, class);
    gen->current_class = NULL;
}

RavaProgram_t* rava_ir_generate(RavaIRGenerator_t *generator, RavaASTNode_t *root) {
    if (!generator || !root) return NULL;

    if (root->type == RAVA_AST_COMPILATION_UNIT) {
        for (size_t i = 0; i < root->children_count; i++) {
            RavaASTNode_t *child = root->children[i];
            if (child->type == RAVA_AST_CLASS_DECL) {
                _rava_ir_gen_class(generator, child);
            }
        }
    }

    return generator->program;
}