#include "test_utils.h"
#include "parser.h"
#include "lexer.h"
#include "arena.h"
#include <string.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; }
    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);
    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);
}

// === 新增测试 ===

void test_struct_decl_and_init() {
    AstNode* ast = parse_string(
        "struct Point { x: i64, y: i64 }"
        "fn main() -> i64 { let p = Point { x: 1, y: 2 }; return p.x; }");
    ASSERT(ast != NULL);
    ASSERT(ast->as.program.struct_count == 1);
    AstNode* sd = ast->as.program.structs[0];
    ASSERT(sd->kind == AST_STRUCT_DECL);
    ASSERT(strcmp(sd->as.struct_decl.name, "Point") == 0);
    ASSERT(sd->as.struct_decl.field_count == 2);
}

void test_literals() {
    // i64
    AstNode* a1 = parse_string("fn main() -> i64 { return 42; }");
    ASSERT(a1 != NULL);
    // f64
    AstNode* a2 = parse_string("fn main() -> f64 { return 3.14; }");
    ASSERT(a2 != NULL);
    // bool true/false
    AstNode* a3 = parse_string("fn main() -> bool { return true; }");
    ASSERT(a3 != NULL);
    AstNode* a4 = parse_string("fn main() -> bool { return false; }");
    ASSERT(a4 != NULL);
    // str
    AstNode* a5 = parse_string("fn main() -> str { return \"hello\"; }");
    ASSERT(a5 != NULL);
}

void test_comparison_chaining() {
    AstNode* ast = parse_string("fn main() -> bool { return 1 + 2 == 3; }");
    ASSERT(ast != NULL);
    // 1 + 2 == 3 应解析为 (1+2) == 3, 因 == 优先级(40) < +(50)
    AstNode* body = ast->as.program.functions[0]->as.function.body;
    AstNode* ret_expr = body->as.block.stmts[0]->as.return_stmt.expr;
    ASSERT(ret_expr->kind == AST_BINARY_EXPR);
    ASSERT(ret_expr->as.binary.op == OP_EQ);
    ASSERT(ret_expr->as.binary.left->kind == AST_BINARY_EXPR);
    ASSERT(ret_expr->as.binary.left->as.binary.op == OP_ADD);
}

void test_guard_desugar() {
    AstNode* ast = parse_string("fn main() { guard x >= 0 else { return -1; } return 0; }");
    ASSERT(ast != NULL);
    // guard 去糖为 if !(x >= 0) { return -1; }
    AstNode* body = ast->as.program.functions[0]->as.function.body;
    AstNode* first = body->as.block.stmts[0];
    ASSERT(first->kind == AST_IF_STMT);
    ASSERT(first->as.if_stmt.cond->kind == AST_UNARY_EXPR);
    ASSERT(first->as.if_stmt.cond->as.unary.op == OP_NOT);
}

void test_named_args() {
    AstNode* ast = parse_string(
        "fn draw(x: i64, y: i64) { }"
        "fn main() { draw(x: 10, y: 20); return 0; }");
    ASSERT(ast != NULL);
    AstNode* body = ast->as.program.functions[1]->as.function.body;
    AstNode* call = body->as.block.stmts[0]->as.expr_stmt.expr;
    ASSERT(call->kind == AST_CALL_EXPR);
    ASSERT(call->as.call.arg_count == 2);
    // 命名参数应有 name_arr
}

void test_field_access() {
    AstNode* ast = parse_string(
        "struct Pt { x: i64 }"
        "fn main() -> i64 { let p = Pt { x: 42 }; return p.x; }");
    ASSERT(ast != NULL);
    AstNode* body = ast->as.program.functions[0]->as.function.body;
    AstNode* ret_expr = body->as.block.stmts[1]->as.return_stmt.expr;
    ASSERT(ret_expr->kind == AST_FIELD_ACCESS);
}

void test_method_call_parse() {
    AstNode* ast = parse_string(
        "struct Pt { x: i64 }"
        "extend Pt { fn get(self: Pt) -> i64 { return self.x; } }"
        "fn main() -> i64 { let p = Pt { x: 10 }; return p.get(); }");
    ASSERT(ast != NULL);
}

void test_match_parse() {
    AstNode* ast = parse_string(
        "fn main() -> i64 { match 1 { 1 => { return 10; }, _ => { return 0; } } }");
    ASSERT(ast != NULL);
}

void test_enum_decl() {
    AstNode* ast = parse_string("enum Color { Red, Green, Blue } fn main() -> i64 { return 0; }");
    ASSERT(ast != NULL);
    ASSERT(ast->as.program.enum_count == 1);
    AstNode* ed = ast->as.program.enums[0];
    ASSERT(ed->kind == AST_ENUM_DECL);
    ASSERT(ed->as.enum_decl.variant_count == 3);
}

void test_for_desugar() {
    AstNode* ast = parse_string(
        "fn main() { for i in 0 to 5 { print_i64(i); } return 0; }");
    ASSERT(ast != NULL);
    // for 去糖为 { var i = 0; while i < 5 { ... i = i + 1; } }
    AstNode* body = ast->as.program.functions[0]->as.function.body;
    ASSERT(body->as.block.stmt_count == 2);  // for-block + return 0
    // 第一个语句是 for 脱糖块
    AstNode* for_block = body->as.block.stmts[0];
    ASSERT(for_block->kind == AST_BLOCK);
    ASSERT(for_block->as.block.stmt_count == 2);  // let + while
}

void test_pipe_operator() {
    AstNode* ast = parse_string(
        "fn main() -> i64 { return 10 |> double(5); }");
    ASSERT(ast != NULL);
}

void test_type_alias_parse() {
    AstNode* ast = parse_string("type Meters = i64; fn main() -> i64 { return 0; }");
    ASSERT(ast != NULL);
}

void test_trait_decl() {
    AstNode* ast = parse_string(
        "trait Show { fn show(self: Self) -> void; }"
        "fn main() -> i64 { return 0; }");
    ASSERT(ast != NULL);
}

void test_array_type_parse() {
    AstNode* ast = parse_string(
        "fn main() -> i64 { var a: i64[10] = a; a[0] = 42; return a[0]; }");
    ASSERT(ast != NULL);
}

void test_if_expr_parse() {
    AstNode* ast = parse_string(
        "fn main() -> i64 { let x = if true { 10; } else { 20; }; return x; }");
    ASSERT(ast != NULL);
}

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);
    TEST_RUN(test_struct_decl_and_init);
    TEST_RUN(test_literals);
    TEST_RUN(test_comparison_chaining);
    TEST_RUN(test_guard_desugar);
    TEST_RUN(test_named_args);
    TEST_RUN(test_field_access);
    TEST_RUN(test_method_call_parse);
    TEST_RUN(test_match_parse);
    TEST_RUN(test_enum_decl);
    TEST_RUN(test_for_desugar);
    TEST_RUN(test_pipe_operator);
    TEST_RUN(test_type_alias_parse);
    TEST_RUN(test_trait_decl);
    TEST_RUN(test_array_type_parse);
    TEST_RUN(test_if_expr_parse);
    return test_summary();
}