#ifndef TEST_UTILS_H
#define TEST_UTILS_H

#include "unittest.h"
#include "../lexer/lexer.h"
#include "../parser/parser.h"
#include "../semantic/semantic.h"
#include "../ir/ir.h"
#include "../ir/ir_gen.h"
#include "../runtime/runtime.h"
#include <stdio.h>
#include <stdlib.h>
#include <string.h>

typedef struct {
    RavaLexer_t *lexer;
    RavaParser_t *parser;
    RavaASTNode_t *ast;
    RavaSemanticAnalyzer_t *analyzer;
    RavaIRGenerator_t *ir_gen;
    RavaProgram_t *program;
    RavaVM_t *vm;
    bool parse_ok;
    bool semantic_ok;
    bool ir_ok;
    bool execute_ok;
    int32_t return_value;
    char *error_stage;
    char *error_message;
} RavaTestContext_t;

static inline void rava_test_context_init(RavaTestContext_t *ctx) {
    memset(ctx, 0, sizeof(RavaTestContext_t));
}

static inline void rava_test_context_cleanup(RavaTestContext_t *ctx) {
    if (ctx->vm) rava_vm_destroy(ctx->vm);
    if (ctx->program) rava_program_destroy(ctx->program);
    if (ctx->ir_gen) rava_ir_generator_destroy(ctx->ir_gen);
    if (ctx->analyzer) rava_semantic_analyzer_destroy(ctx->analyzer);
    if (ctx->ast) rava_ast_node_destroy(ctx->ast);
    if (ctx->parser) rava_parser_destroy(ctx->parser);
    if (ctx->lexer) rava_lexer_destroy(ctx->lexer);
}

static inline bool rava_test_compile_and_run(RavaTestContext_t *ctx, const char *source,
                                              const char *class_name, const char *method_name) {
    ctx->lexer = rava_lexer_create(source);
    ctx->parser = rava_parser_create(ctx->lexer);
    ctx->ast = rava_parser_parse(ctx->parser);

    if (!ctx->ast || ctx->parser->had_error) {
        ctx->parse_ok = false;
        ctx->error_stage = "parse";
        ctx->error_message = ctx->parser->error_message ? ctx->parser->error_message : "unknown parse error";
        return false;
    }
    ctx->parse_ok = true;

    ctx->analyzer = rava_semantic_analyzer_create();
    if (!rava_semantic_analyze(ctx->analyzer, ctx->ast)) {
        ctx->semantic_ok = false;
        ctx->error_stage = "semantic";
        ctx->error_message = ctx->analyzer->error_message ? ctx->analyzer->error_message : "unknown semantic error";
        return false;
    }
    ctx->semantic_ok = true;

    ctx->ir_gen = rava_ir_generator_create(ctx->analyzer);
    ctx->program = rava_ir_generate(ctx->ir_gen, ctx->ast);

    if (!ctx->program) {
        ctx->ir_ok = false;
        ctx->error_stage = "ir";
        ctx->error_message = "IR generation failed";
        return false;
    }
    ctx->ir_ok = true;

    ctx->vm = rava_vm_create(ctx->program);

    if (!rava_vm_execute(ctx->vm, class_name, method_name)) {
        ctx->execute_ok = false;
        ctx->error_stage = "runtime";
        ctx->error_message = ctx->vm->error_message ? ctx->vm->error_message : "unknown runtime error";
        return false;
    }
    ctx->execute_ok = true;

    RavaValue_t result_val = rava_vm_get_result(ctx->vm);
    ctx->return_value = rava_value_as_int(result_val);

    return true;
}

static inline char* rava_test_read_file(const char *filename) {
    FILE *file = fopen(filename, "r");
    if (!file) return NULL;
    fseek(file, 0, SEEK_END);
    long size = ftell(file);
    fseek(file, 0, SEEK_SET);
    char *content = malloc(size + 1);
    size_t read_bytes = fread(content, 1, size, file);
    content[read_bytes] = '\0';
    fclose(file);
    return content;
}

#define RAVA_TEST_RUN(result, source, class_name, method_name, expected, msg) \
    do { \
        RavaTestContext_t _ctx; \
        rava_test_context_init(&_ctx); \
        bool _success = rava_test_compile_and_run(&_ctx, source, class_name, method_name); \
        if (!_success) { \
            char _err_buf[256]; \
            snprintf(_err_buf, sizeof(_err_buf), "%s: %s error: %s", msg, _ctx.error_stage, _ctx.error_message); \
            UNITTEST_ASSERT_TRUE(result, false, _err_buf); \
        } else { \
            UNITTEST_ASSERT_EQUAL_INT(result, expected, _ctx.return_value, msg); \
        } \
        rava_test_context_cleanup(&_ctx); \
    } while(0)

#define RAVA_TEST_EXPECT_PARSE_ERROR(result, source, msg) \
    do { \
        RavaTestContext_t _ctx; \
        rava_test_context_init(&_ctx); \
        rava_test_compile_and_run(&_ctx, source, "Test", "main"); \
        bool _is_parse_error = !_ctx.parse_ok; \
        UNITTEST_ASSERT_TRUE(result, _is_parse_error, msg); \
        rava_test_context_cleanup(&_ctx); \
    } while(0)

#define RAVA_TEST_EXPECT_SEMANTIC_ERROR(result, source, msg) \
    do { \
        RavaTestContext_t _ctx; \
        rava_test_context_init(&_ctx); \
        rava_test_compile_and_run(&_ctx, source, "Test", "main"); \
        bool _is_semantic_error = _ctx.parse_ok && !_ctx.semantic_ok; \
        UNITTEST_ASSERT_TRUE(result, _is_semantic_error, msg); \
        rava_test_context_cleanup(&_ctx); \
    } while(0)

#define RAVA_TEST_EXPECT_RUNTIME_ERROR(result, source, class_name, method_name, msg) \
    do { \
        RavaTestContext_t _ctx; \
        rava_test_context_init(&_ctx); \
        rava_test_compile_and_run(&_ctx, source, class_name, method_name); \
        bool _is_runtime_error = _ctx.parse_ok && _ctx.semantic_ok && _ctx.ir_ok && !_ctx.execute_ok; \
        UNITTEST_ASSERT_TRUE(result, _is_runtime_error, msg); \
        rava_test_context_cleanup(&_ctx); \
    } while(0)

#define RAVA_TEST_FILE_EXECUTES(result, filename, class_name, method_name, msg) \
    do { \
        char *_source = rava_test_read_file(filename); \
        UNITTEST_ASSERT_NOT_NULL(result, _source, "Failed to read file: " filename); \
        if (_source) { \
            RavaTestContext_t _ctx; \
            rava_test_context_init(&_ctx); \
            bool _success = rava_test_compile_and_run(&_ctx, _source, class_name, method_name); \
            if (!_success) { \
                char _err_buf[256]; \
                snprintf(_err_buf, sizeof(_err_buf), "%s: %s error: %s", msg, _ctx.error_stage, _ctx.error_message); \
                UNITTEST_ASSERT_TRUE(result, false, _err_buf); \
            } else { \
                UNITTEST_ASSERT_TRUE(result, true, msg); \
            } \
            rava_test_context_cleanup(&_ctx); \
            free(_source); \
        } \
    } while(0)

#define RAVA_TEST_LEXER_OK(result, source, msg) \
    do { \
        RavaLexer_t *_lexer = rava_lexer_create(source); \
        RavaToken_t *_token; \
        bool _has_error = false; \
        do { \
            _token = rava_lexer_next_token(_lexer); \
            if (_token->type == RAVA_TOKEN_ERROR) { \
                _has_error = true; \
                rava_token_destroy(_token); \
                break; \
            } \
            RavaTokenType_e _type = _token->type; \
            rava_token_destroy(_token); \
            if (_type == RAVA_TOKEN_EOF) break; \
        } while (1); \
        UNITTEST_ASSERT_TRUE(result, !_has_error, msg); \
        rava_lexer_destroy(_lexer); \
    } while(0)

#define RAVA_TEST_PARSER_OK(result, source, msg) \
    do { \
        RavaLexer_t *_lexer = rava_lexer_create(source); \
        RavaParser_t *_parser = rava_parser_create(_lexer); \
        RavaASTNode_t *_ast = rava_parser_parse(_parser); \
        bool _parse_ok = _ast && !_parser->had_error; \
        UNITTEST_ASSERT_TRUE(result, _parse_ok, msg); \
        if (_ast) rava_ast_node_destroy(_ast); \
        rava_parser_destroy(_parser); \
        rava_lexer_destroy(_lexer); \
    } while(0)

#define RAVA_TEST_SEMANTIC_OK(result, source, msg) \
    do { \
        RavaLexer_t *_lexer = rava_lexer_create(source); \
        RavaParser_t *_parser = rava_parser_create(_lexer); \
        RavaASTNode_t *_ast = rava_parser_parse(_parser); \
        bool _semantic_ok = false; \
        if (_ast && !_parser->had_error) { \
            RavaSemanticAnalyzer_t *_analyzer = rava_semantic_analyzer_create(); \
            _semantic_ok = rava_semantic_analyze(_analyzer, _ast); \
            rava_semantic_analyzer_destroy(_analyzer); \
        } \
        UNITTEST_ASSERT_TRUE(result, _semantic_ok, msg); \
        if (_ast) rava_ast_node_destroy(_ast); \
        rava_parser_destroy(_parser); \
        rava_lexer_destroy(_lexer); \
    } while(0)

#define RAVA_TEST_IR_OK(result, source, msg) \
    do { \
        RavaLexer_t *_lexer = rava_lexer_create(source); \
        RavaParser_t *_parser = rava_parser_create(_lexer); \
        RavaASTNode_t *_ast = rava_parser_parse(_parser); \
        bool _ir_ok = false; \
        RavaSemanticAnalyzer_t *_analyzer = NULL; \
        RavaIRGenerator_t *_ir_gen = NULL; \
        RavaProgram_t *_program = NULL; \
        if (_ast && !_parser->had_error) { \
            _analyzer = rava_semantic_analyzer_create(); \
            if (rava_semantic_analyze(_analyzer, _ast)) { \
                _ir_gen = rava_ir_generator_create(_analyzer); \
                _program = rava_ir_generate(_ir_gen, _ast); \
                _ir_ok = (_program != NULL); \
            } \
        } \
        UNITTEST_ASSERT_TRUE(result, _ir_ok, msg); \
        if (_program) rava_program_destroy(_program); \
        if (_ir_gen) rava_ir_generator_destroy(_ir_gen); \
        if (_analyzer) rava_semantic_analyzer_destroy(_analyzer); \
        if (_ast) rava_ast_node_destroy(_ast); \
        rava_parser_destroy(_parser); \
        rava_lexer_destroy(_lexer); \
    } while(0)

#endif