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

void test_type_error() {
    Arena a = arena_create(1);
    size_t tc; ErrorInfo lex_err = {0};
    Token* toks = lex(&a, "fn main() { let x: i64 = 1; let y: i64 = x + true; return; }",
                      "test", &tc, &lex_err);
    ASSERT(toks != NULL);
    ErrorInfo parse_err = {0};
    AstNode* ast = parse(&a, toks, tc, "test", &parse_err);
    ASSERT(ast != NULL);

    ErrorList errors; error_init(&errors);
    sema_analyze(ast, &errors, &a);
    ASSERT(errors.count > 0);
    arena_destroy(&a);
}

void test_undefined_var() {
    Arena a = arena_create(1);
    size_t tc; ErrorInfo lex_err = {0};
    Token* toks = lex(&a, "fn main() { let x: i64 = y; return; }", "test", &tc, &lex_err);
    ASSERT(toks != NULL);
    ErrorInfo parse_err = {0};
    AstNode* ast = parse(&a, toks, tc, "test", &parse_err);
    ASSERT(ast != NULL);

    ErrorList errors; error_init(&errors);
    sema_analyze(ast, &errors, &a);
    ASSERT(errors.count > 0);
    arena_destroy(&a);
}

void test_simple_ok() {
    Arena a = arena_create(1);
    size_t tc; ErrorInfo lex_err = {0};
    Token* toks = lex(&a, "fn main() { let x: i64 = 42; print_i64(x); return; }",
                      "test", &tc, &lex_err);
    ASSERT(toks != NULL);
    ErrorInfo parse_err = {0};
    AstNode* ast = parse(&a, toks, tc, "test", &parse_err);
    ASSERT(ast != NULL);

    ErrorList errors; error_init(&errors);
    sema_analyze(ast, &errors, &a);
    ASSERT(errors.count == 0);
    arena_destroy(&a);
}

void test_let_mut_assign_ok() {
    Arena a = arena_create(1);
    size_t tc; ErrorInfo lex_err = {0};
    Token* toks = lex(&a,
        "fn main() { let mut x: i64 = 0; x = 42; print_i64(x); return; }",
        "test", &tc, &lex_err);
    ASSERT(toks != NULL);
    ErrorInfo parse_err = {0};
    AstNode* ast = parse(&a, toks, tc, "test", &parse_err);
    ASSERT(ast != NULL);
    ErrorList errors; error_init(&errors);
    sema_analyze(ast, &errors, &a);
    ASSERT(errors.count == 0);
    arena_destroy(&a);
}

void test_assign_immutable_error() {
    Arena a = arena_create(1);
    size_t tc; ErrorInfo lex_err = {0};
    Token* toks = lex(&a,
        "fn main() { let x: i64 = 0; x = 42; return; }",
        "test", &tc, &lex_err);
    ASSERT(toks != NULL);
    ErrorInfo parse_err = {0};
    AstNode* ast = parse(&a, toks, tc, "test", &parse_err);
    ASSERT(ast != NULL);
    ErrorList errors; error_init(&errors);
    sema_analyze(ast, &errors, &a);
    ASSERT(errors.count > 0);  // 不可变变量赋值应报错
    arena_destroy(&a);
}

void test_str_type_ok() {
    Arena a = arena_create(1);
    size_t tc; ErrorInfo lex_err = {0};
    Token* toks = lex(&a,
        "fn main() { let msg: str = \"hello\"; print_str(msg); return; }",
        "test", &tc, &lex_err);
    ASSERT(toks != NULL);
    ErrorInfo parse_err = {0};
    AstNode* ast = parse(&a, toks, tc, "test", &parse_err);
    ASSERT(ast != NULL);
    ErrorList errors; error_init(&errors);
    sema_analyze(ast, &errors, &a);
    ASSERT(errors.count == 0);
    arena_destroy(&a);
}

void test_str_concat_type_ok() {
    Arena a = arena_create(1);
    size_t tc; ErrorInfo lex_err = {0};
    Token* toks = lex(&a,
        "fn main() { let a: str = \"a\"; let b: str = \"b\"; let c: str = a + b; print_str(c); return; }",
        "test", &tc, &lex_err);
    ASSERT(toks != NULL);
    ErrorInfo parse_err = {0};
    AstNode* ast = parse(&a, toks, tc, "test", &parse_err);
    ASSERT(ast != NULL);
    ErrorList errors; error_init(&errors);
    sema_analyze(ast, &errors, &a);
    ASSERT(errors.count == 0);
    arena_destroy(&a);
}

/* === struct 类型检查测试 === */

void test_struct_field_type_mismatch() {
    Arena a = arena_create(1);
    size_t tc; ErrorInfo lex_err = {0};
    Token* toks = lex(&a,
        "struct Point { x: i64, y: i64 } fn main() { let p: Point = Point { x: 10, y: true }; return; }",
        "test", &tc, &lex_err);
    ASSERT(toks != NULL);
    ErrorInfo parse_err = {0};
    AstNode* ast = parse(&a, toks, tc, "test", &parse_err);
    ASSERT(ast != NULL);

    ErrorList errors; error_init(&errors);
    sema_analyze(ast, &errors, &a);
    ASSERT(errors.count > 0);  // y 字段类型不匹配: true 是 bool, 不是 i64
    arena_destroy(&a);
}

void test_struct_undefined() {
    Arena a = arena_create(1);
    size_t tc; ErrorInfo lex_err = {0};
    Token* toks = lex(&a,
        "fn main() { let p: Unknown = Unknown { x: 1 }; return; }",
        "test", &tc, &lex_err);
    ASSERT(toks != NULL);
    ErrorInfo parse_err = {0};
    AstNode* ast = parse(&a, toks, tc, "test", &parse_err);
    ASSERT(ast != NULL);

    ErrorList errors; error_init(&errors);
    sema_analyze(ast, &errors, &a);
    ASSERT(errors.count > 0);  // Unknown 未定义
    arena_destroy(&a);
}

void test_struct_field_count_mismatch() {
    Arena a = arena_create(1);
    size_t tc; ErrorInfo lex_err = {0};
    Token* toks = lex(&a,
        "struct Point { x: i64, y: i64 } fn main() { let p: Point = Point { x: 10 }; return; }",
        "test", &tc, &lex_err);
    ASSERT(toks != NULL);
    ErrorInfo parse_err = {0};
    AstNode* ast = parse(&a, toks, tc, "test", &parse_err);
    ASSERT(ast != NULL);

    ErrorList errors; error_init(&errors);
    sema_analyze(ast, &errors, &a);
    ASSERT(errors.count > 0);  // 缺少字段 'y'
    arena_destroy(&a);
}

/* === 类型别名测试 === */

void test_type_alias_ok() {
    Arena a = arena_create(1);
    size_t tc; ErrorInfo lex_err = {0};
    Token* toks = lex(&a,
        "type Meters = i64; fn main() { let x: Meters = 100; print_i64(x); return; }",
        "test", &tc, &lex_err);
    ASSERT(toks != NULL);
    ErrorInfo parse_err = {0};
    AstNode* ast = parse(&a, toks, tc, "test", &parse_err);
    ASSERT(ast != NULL);
    ASSERT(ast->as.program.alias_count == 1);

    ErrorList errors; error_init(&errors);
    sema_analyze(ast, &errors, &a);
    ASSERT(errors.count == 0);
    arena_destroy(&a);
}

void test_type_alias_struct() {
    Arena a = arena_create(1);
    size_t tc; ErrorInfo lex_err = {0};
    Token* toks = lex(&a,
        "struct Point { x: i64, y: i64 } type P = Point; fn main() { let p: P = Point { x: 1, y: 2 }; return; }",
        "test", &tc, &lex_err);
    ASSERT(toks != NULL);
    ErrorInfo parse_err = {0};
    AstNode* ast = parse(&a, toks, tc, "test", &parse_err);
    ASSERT(ast != NULL);
    ASSERT(ast->as.program.alias_count == 1);

    ErrorList errors; error_init(&errors);
    sema_analyze(ast, &errors, &a);
    ASSERT(errors.count == 0);
    arena_destroy(&a);
}

void test_struct_nested_type_ok() {
    Arena a = arena_create(1);
    size_t tc; ErrorInfo lex_err = {0};
    Token* toks = lex(&a,
        "struct Point { x: i64, y: i64 } struct Rect { tl: Point, br: Point } fn main() { let r: Rect = Rect { tl: Point { x: 0, y: 0 }, br: Point { x: 1, y: 1 } }; print_i64(r.tl.x); return; }",
        "test", &tc, &lex_err);
    ASSERT(toks != NULL);
    ErrorInfo parse_err = {0};
    AstNode* ast = parse(&a, toks, tc, "test", &parse_err);
    ASSERT(ast != NULL);

    ErrorList errors; error_init(&errors);
    sema_analyze(ast, &errors, &a);
    ASSERT(errors.count == 0);  // 嵌套结构体类型检查应通过
    arena_destroy(&a);
}

/* === enum 语义分析测试 === */

void test_enum_ok() {
    Arena a = arena_create(1);
    size_t tc; ErrorInfo lex_err = {0};
    Token* toks = lex(&a,
        "enum Color { Red, Green, Blue } fn main() { let c = Color::Green; print_i64(c); return; }",
        "test", &tc, &lex_err);
    ASSERT(toks != NULL);
    ErrorInfo parse_err = {0};
    AstNode* ast = parse(&a, toks, tc, "test", &parse_err);
    ASSERT(ast != NULL);

    ErrorList errors; error_init(&errors);
    sema_analyze(ast, &errors, &a);
    ASSERT(errors.count == 0);
    arena_destroy(&a);
}

void test_enum_bad_variant() {
    Arena a = arena_create(1);
    size_t tc; ErrorInfo lex_err = {0};
    Token* toks = lex(&a,
        "enum Color { Red, Green, Blue } fn main() { let c = Color::Yellow; return; }",
        "test", &tc, &lex_err);
    ASSERT(toks != NULL);
    ErrorInfo parse_err = {0};
    AstNode* ast = parse(&a, toks, tc, "test", &parse_err);
    ASSERT(ast != NULL);

    ErrorList errors; error_init(&errors);
    sema_analyze(ast, &errors, &a);
    ASSERT(errors.count > 0);  // Yellow 不存在
    arena_destroy(&a);
}

/* === 数组语义分析测试 === */

void test_array_ok() {
    Arena a = arena_create(1);
    size_t tc; ErrorInfo lex_err = {0};
    Token* toks = lex(&a,
        "fn main() { let arr: [i64; 3] = arr; arr[0]; return; }",
        "test", &tc, &lex_err);
    ASSERT(toks != NULL);
    ErrorInfo parse_err = {0};
    AstNode* ast = parse(&a, toks, tc, "test", &parse_err);
    ASSERT(ast != NULL);

    ErrorList errors; error_init(&errors);
    sema_analyze(ast, &errors, &a);
    ASSERT(errors.count == 0);
    arena_destroy(&a);
}

void test_array_index_type_error() {
    Arena a = arena_create(1);
    size_t tc; ErrorInfo lex_err = {0};
    Token* toks = lex(&a,
        "fn main() { let arr: [i64; 3] = arr; arr[true]; return; }",
        "test", &tc, &lex_err);
    ASSERT(toks != NULL);
    ErrorInfo parse_err = {0};
    AstNode* ast = parse(&a, toks, tc, "test", &parse_err);
    ASSERT(ast != NULL);

    ErrorList errors; error_init(&errors);
    sema_analyze(ast, &errors, &a);
    ASSERT(errors.count > 0);  // true 不是 i64
    arena_destroy(&a);
}

void test_array_not_indexable() {
    Arena a = arena_create(1);
    size_t tc; ErrorInfo lex_err = {0};
    Token* toks = lex(&a,
        "fn main() { let x: i64 = 0; x[0]; return; }",
        "test", &tc, &lex_err);
    ASSERT(toks != NULL);
    ErrorInfo parse_err = {0};
    AstNode* ast = parse(&a, toks, tc, "test", &parse_err);
    ASSERT(ast != NULL);

    ErrorList errors; error_init(&errors);
    sema_analyze(ast, &errors, &a);
    ASSERT(errors.count > 0);  // i64 不是数组
    arena_destroy(&a);
}

void test_array_assign_ok() {
    Arena a = arena_create(1);
    size_t tc; ErrorInfo lex_err = {0};
    Token* toks = lex(&a,
        "fn main() { let arr: [i64; 3] = arr; arr[0] = 42; return; }",
        "test", &tc, &lex_err);
    ASSERT(toks != NULL);
    ErrorInfo parse_err = {0};
    AstNode* ast = parse(&a, toks, tc, "test", &parse_err);
    ASSERT(ast != NULL);

    ErrorList errors; error_init(&errors);
    sema_analyze(ast, &errors, &a);
    ASSERT(errors.count == 0);
    arena_destroy(&a);
}

/* === 方法调用语义分析测试 === */

void test_method_call_ok() {
    Arena a = arena_create(1);
    size_t tc; ErrorInfo lex_err = {0};
    Token* toks = lex(&a,
        "struct Point { x: i64, y: i64 } impl Point { fn get_x(self: Point) -> i64 { return self.x; } } fn main() -> i64 { let p: Point = Point { x: 42, y: 0 }; return p.get_x(); }",
        "test", &tc, &lex_err);
    ASSERT(toks != NULL);
    ErrorInfo parse_err = {0};
    AstNode* ast = parse(&a, toks, tc, "test", &parse_err);
    ASSERT(ast != NULL);

    ErrorList errors; error_init(&errors);
    sema_analyze(ast, &errors, &a);
    ASSERT(errors.count == 0);
    arena_destroy(&a);
}

void test_method_undefined() {
    Arena a = arena_create(1);
    size_t tc; ErrorInfo lex_err = {0};
    Token* toks = lex(&a,
        "struct Point { x: i64, y: i64 } fn main() -> i64 { let p: Point = Point { x: 42, y: 0 }; return p.nope(); }",
        "test", &tc, &lex_err);
    ASSERT(toks != NULL);
    ErrorInfo parse_err = {0};
    AstNode* ast = parse(&a, toks, tc, "test", &parse_err);
    ASSERT(ast != NULL);

    ErrorList errors; error_init(&errors);
    sema_analyze(ast, &errors, &a);
    ASSERT(errors.count > 0);  // nope 不是 Point 的方法
    arena_destroy(&a);
}

/* === match 语义分析测试 === */

void test_match_enum_sema_ok() {
    Arena a = arena_create(1);
    size_t tc; ErrorInfo lex_err = {0};
    Token* toks = lex(&a,
        "enum Color { Red, Green, Blue } fn main() -> i64 { let c = Color::Green; match c { Color::Red => { return 1; } Color::Green => { return 2; } _ => { return 0; } } return 0; }",
        "test", &tc, &lex_err);
    ASSERT(toks != NULL);
    ErrorInfo parse_err = {0};
    AstNode* ast = parse(&a, toks, tc, "test", &parse_err);
    ASSERT(ast != NULL);

    ErrorList errors; error_init(&errors);
    sema_analyze(ast, &errors, &a);
    ASSERT(errors.count == 0);  // match 去糖后应通过类型检查
    arena_destroy(&a);
}

void test_match_int_sema_ok() {
    Arena a = arena_create(1);
    size_t tc; ErrorInfo lex_err = {0};
    Token* toks = lex(&a,
        "fn main() -> i64 { let x: i64 = 42; match x { 1 => { return 10; } 42 => { return 20; } _ => { return 0; } } return 0; }",
        "test", &tc, &lex_err);
    ASSERT(toks != NULL);
    ErrorInfo parse_err = {0};
    AstNode* ast = parse(&a, toks, tc, "test", &parse_err);
    ASSERT(ast != NULL);

    ErrorList errors; error_init(&errors);
    sema_analyze(ast, &errors, &a);
    ASSERT(errors.count == 0);
    arena_destroy(&a);
}

void test_match_wildcard_only_sema_ok() {
    Arena a = arena_create(1);
    size_t tc; ErrorInfo lex_err = {0};
    Token* toks = lex(&a,
        "fn main() -> i64 { let x: i64 = 42; match x { _ => { return 99; } } return 0; }",
        "test", &tc, &lex_err);
    ASSERT(toks != NULL);
    ErrorInfo parse_err = {0};
    AstNode* ast = parse(&a, toks, tc, "test", &parse_err);
    ASSERT(ast != NULL);

    ErrorList errors; error_init(&errors);
    sema_analyze(ast, &errors, &a);
    ASSERT(errors.count == 0);  // 纯通配符 match 应通过
    arena_destroy(&a);
}

int main(void) {
    TEST_RUN(test_type_error);
    TEST_RUN(test_undefined_var);
    TEST_RUN(test_simple_ok);
    TEST_RUN(test_let_mut_assign_ok);
    TEST_RUN(test_assign_immutable_error);
    TEST_RUN(test_str_type_ok);
    TEST_RUN(test_str_concat_type_ok);
    TEST_RUN(test_struct_field_type_mismatch);
    TEST_RUN(test_struct_undefined);
    TEST_RUN(test_struct_field_count_mismatch);
    TEST_RUN(test_struct_nested_type_ok);
    TEST_RUN(test_type_alias_ok);
    TEST_RUN(test_type_alias_struct);
    TEST_RUN(test_enum_ok);
    TEST_RUN(test_enum_bad_variant);
    TEST_RUN(test_array_ok);
    TEST_RUN(test_array_index_type_error);
    TEST_RUN(test_array_not_indexable);
    TEST_RUN(test_array_assign_ok);
    TEST_RUN(test_method_call_ok);
    TEST_RUN(test_method_undefined);
    TEST_RUN(test_match_enum_sema_ok);
    TEST_RUN(test_match_int_sema_ok);
    TEST_RUN(test_match_wildcard_only_sema_ok);
    return test_summary();
}