#include "test_utils.h"
#include "parser.h"
#include "lexer.h"
#include "arena.h"

static AstNode* parse_string(const char* src) {
    Arena* a = malloc(sizeof(Arena));
    *a = arena_create(1);
    size_t tcount;
    ErrorInfo lex_err = {0};
    Token* tokens = lex(a, src, "test", &tcount, &lex_err);
    if (!tokens) { arena_destroy(a); free(a); return NULL; }
    ErrorInfo parse_err = {0};
    AstNode* ast = parse(a, tokens, tcount, "test", &parse_err);
    if (!ast) { arena_destroy(a); free(a); return NULL; }
    // NOTE: arena and tokens must stay alive for AST - leak intentionally in test
    return ast;
}

void test_simple_function() {
    AstNode* ast = parse_string("fn main() { return 42; }");
    ASSERT(ast != NULL);
    ASSERT(ast->kind == AST_PROGRAM);
    ASSERT(ast->as.program.fn_count == 1);
    AstNode* fn = ast->as.program.functions[0];
    ASSERT(fn->kind == AST_FUNCTION);
}

void test_arithmetic_expr() {
    AstNode* ast = parse_string("fn main() { return 1 + 2 * 3; }");
    ASSERT(ast != NULL);
    AstNode* body = ast->as.program.functions[0]->as.function.body;
    AstNode* ret = body->as.block.stmts[0];
    ASSERT(ret->kind == AST_RETURN_STMT);
    AstNode* expr = ret->as.return_stmt.expr;
    ASSERT(expr->kind == AST_BINARY_EXPR);
    ASSERT(expr->as.binary.op == OP_ADD);
    // 1 + (2 * 3): right should be *, left should be 1
    ASSERT(expr->as.binary.right->kind == AST_BINARY_EXPR);
    ASSERT(expr->as.binary.right->as.binary.op == OP_MUL);
}

void test_if_statement() {
    AstNode* ast = parse_string("fn main() { if true { return 1; } else { return 0; } }");
    ASSERT(ast != NULL);
}

void test_while_loop() {
    AstNode* ast = parse_string("fn main() { while true { return; } }");
    ASSERT(ast != NULL);
}

void test_function_with_params() {
    AstNode* ast = parse_string("fn add(a: i64, b: i64) -> i64 { return a + b; }");
    ASSERT(ast != NULL);
    AstNode* fn = ast->as.program.functions[0];
    ASSERT(fn->as.function.param_count == 2);
    ASSERT(fn->as.function.return_type == TYPE_I64);
}

int main(void) {
    TEST_RUN(test_simple_function);
    TEST_RUN(test_arithmetic_expr);
    TEST_RUN(test_if_statement);
    TEST_RUN(test_while_loop);
    TEST_RUN(test_function_with_params);
    return test_summary();
}