#include "test_utils.h"
#include "codegen.h"
#include "ast.h"
#include "arena.h"
#include <llvm-c/Analysis.h>

void test_codegen_simple_function() {
    Arena a = arena_create(1);

    // 构造 AST: fn main() -> i64 { return 42; }
    AstNode* ret = ast_make_return(&a, ast_make_literal_i64(&a, 42, 1, 1), 1, 1);
    AstNode* stmts[] = { ret };
    AstNode* body = ast_make_block(&a, stmts, 1, 1, 1);
    AstNode* fn = ast_make_function(&a, "main", NULL, 0, TYPE_I64, body, 1, 1);
    AstNode* fns[] = { fn };
    AstNode* prog = ast_make_program(&a, fns, 1, NULL, 0, 1, 1);

    const char* err = NULL;
    LLVMModuleRef mod = codegen_module(prog, &a, "test_mod", &err);
    ASSERT(mod != NULL);
    ASSERT(err == NULL);

    // 验证模块
    char* verify_err = NULL;
    int failed = LLVMVerifyModule(mod, LLVMReturnStatusAction, &verify_err);
    ASSERT(!failed);

    LLVMDisposeModule(mod);
    arena_destroy(&a);
}

void test_codegen_if_else() {
    Arena a = arena_create(1);

    // fn main() -> i64 { if true { return 1; } else { return 0; } }
    AstNode* then_ret = ast_make_return(&a, ast_make_literal_i64(&a, 1, 1, 1), 1, 1);
    AstNode* then_stmts[] = { then_ret };
    AstNode* then_block = ast_make_block(&a, then_stmts, 1, 1, 1);
    AstNode* else_ret = ast_make_return(&a, ast_make_literal_i64(&a, 0, 1, 1), 1, 1);
    AstNode* else_stmts[] = { else_ret };
    AstNode* else_block = ast_make_block(&a, else_stmts, 1, 1, 1);
    AstNode* if_stmt = ast_make_if(&a,
        ast_make_literal_bool(&a, true, 1, 1), then_block, else_block, 1, 1);
    AstNode* stmts[] = { if_stmt };
    AstNode* body = ast_make_block(&a, stmts, 1, 1, 1);
    AstNode* fn = ast_make_function(&a, "main", NULL, 0, TYPE_I64, body, 1, 1);
    AstNode* fns[] = { fn };
    AstNode* prog = ast_make_program(&a, fns, 1, NULL, 0, 1, 1);

    const char* err = NULL;
    LLVMModuleRef mod = codegen_module(prog, &a, "test_mod2", &err);
    ASSERT(mod != NULL);

    char* verify_err = NULL;
    int failed = LLVMVerifyModule(mod, LLVMReturnStatusAction, &verify_err);
    ASSERT(!failed);

    LLVMDisposeModule(mod);
    arena_destroy(&a);
}

void test_codegen_binary_ops() {
    Arena a = arena_create(1);

    // fn main() -> i64 { return 1 + 2 * 3; }
    AstNode* expr = ast_make_binary(&a, OP_ADD,
        ast_make_literal_i64(&a, 1, 1, 1),
        ast_make_binary(&a, OP_MUL,
            ast_make_literal_i64(&a, 2, 1, 1),
            ast_make_literal_i64(&a, 3, 1, 1), 1, 1),
        1, 1);
    AstNode* ret = ast_make_return(&a, expr, 1, 1);
    AstNode* stmts[] = { ret };
    AstNode* body = ast_make_block(&a, stmts, 1, 1, 1);
    AstNode* fn = ast_make_function(&a, "main", NULL, 0, TYPE_I64, body, 1, 1);
    AstNode* fns[] = { fn };
    AstNode* prog = ast_make_program(&a, fns, 1, NULL, 0, 1, 1);

    const char* err = NULL;
    LLVMModuleRef mod = codegen_module(prog, &a, "test_mod3", &err);
    ASSERT(mod != NULL);

    char* verify_err = NULL;
    int failed = LLVMVerifyModule(mod, LLVMReturnStatusAction, &verify_err);
    ASSERT(!failed);

    LLVMDisposeModule(mod);
    arena_destroy(&a);
}

void test_codegen_while_loop() {
    Arena a = arena_create(1);
    // fn main() -> i64 { while true { return 0; } return 1; }
    AstNode* while_body_stmts[] = {
        ast_make_return(&a, ast_make_literal_i64(&a, 0, 1, 1), 1, 1)
    };
    AstNode* while_body = ast_make_block(&a, while_body_stmts, 1, 1, 1);
    AstNode* while_stmt = ast_make_while(&a,
        ast_make_literal_bool(&a, true, 1, 1), while_body, 1, 1);
    AstNode* ret = ast_make_return(&a, ast_make_literal_i64(&a, 1, 1, 1), 1, 1);
    AstNode* stmts[] = { while_stmt, ret };
    AstNode* fn_body = ast_make_block(&a, stmts, 2, 1, 1);
    AstNode* fn = ast_make_function(&a, "main", NULL, 0, TYPE_I64, fn_body, 1, 1);
    AstNode* fns[] = { fn };
    AstNode* prog = ast_make_program(&a, fns, 1, NULL, 0, 1, 1);

    const char* err = NULL;
    LLVMModuleRef mod = codegen_module(prog, &a, "test_while", &err);
    ASSERT(mod != NULL);
    char* verify_err = NULL;
    ASSERT(!LLVMVerifyModule(mod, LLVMReturnStatusAction, &verify_err));
    LLVMDisposeModule(mod);
    arena_destroy(&a);
}

int main(void) {
    TEST_RUN(test_codegen_simple_function);
    TEST_RUN(test_codegen_if_else);
    TEST_RUN(test_codegen_binary_ops);
    TEST_RUN(test_codegen_while_loop);
    return test_summary();
}