#include "ast.h"
#include <stdlib.h>

// 使用宏简化节点创建
#define NEW(alloc, k) \
    AstNode* n = (AstNode*)arena_alloc_impl(alloc, sizeof(AstNode)); \
    if (!n) return NULL; \
    n->kind = (k); n->type.kind = TYPE_UNKNOWN; n->type.struct_name = NULL; \
    n->line = line; n->col = col

AstNode* ast_make_program(void* alloc, AstNode** fns, size_t fn_count,
                          AstNode** structs, size_t struct_count, int line, int col) {
    NEW(alloc, AST_PROGRAM);
    n->as.program.functions = fns;
    n->as.program.fn_count = fn_count;
    n->as.program.structs = structs;
    n->as.program.struct_count = struct_count;
    return n;
}

AstNode* ast_make_function(void* alloc, const char* name, AstNode** params, size_t pcount,
                           TypeKind ret, AstNode* body, int line, int col) {
    NEW(alloc, AST_FUNCTION);
    n->as.function.name = name; n->as.function.params = params;
    n->as.function.param_count = pcount; n->as.function.return_type = ret;
    n->as.function.body = body;
    return n;
}

AstNode* ast_make_parameter(void* alloc, const char* name, TypeKind type,
                             const char* struct_type_name, int line, int col) {
    NEW(alloc, AST_PARAMETER);
    n->as.parameter.name = name; n->as.parameter.type = type;
    n->as.parameter.struct_type_name = struct_type_name;
    return n;
}

AstNode* ast_make_block(void* alloc, AstNode** stmts, size_t count, int line, int col) {
    NEW(alloc, AST_BLOCK);
    n->as.block.stmts = stmts; n->as.block.stmt_count = count;
    return n;
}

AstNode* ast_make_let(void* alloc, const char* name, TypeKind annot_type, bool has_type_annot,
                      bool is_mut, AstNode* init, const char* struct_type_name, int line, int col) {
    NEW(alloc, AST_LET_STMT);
    n->as.let_stmt.name = name; n->as.let_stmt.annot_type = annot_type;
    n->as.let_stmt.has_type_annot = has_type_annot; n->as.let_stmt.is_mut = is_mut;
    n->as.let_stmt.init = init;
    n->as.let_stmt.struct_type_name = struct_type_name;
    return n;
}

AstNode* ast_make_assign(void* alloc, const char* name, AstNode* value, int line, int col) {
    NEW(alloc, AST_ASSIGN_STMT);
    n->as.assign_stmt.name = name; n->as.assign_stmt.value = value;
    return n;
}

AstNode* ast_make_if(void* alloc, AstNode* cond, AstNode* then_b, AstNode* else_b, int line, int col) {
    NEW(alloc, AST_IF_STMT);
    n->as.if_stmt.cond = cond; n->as.if_stmt.then_block = then_b;
    n->as.if_stmt.else_block = else_b;
    return n;
}

AstNode* ast_make_while(void* alloc, AstNode* cond, AstNode* body, int line, int col) {
    NEW(alloc, AST_WHILE_STMT);
    n->as.while_stmt.cond = cond; n->as.while_stmt.body = body;
    return n;
}

AstNode* ast_make_return(void* alloc, AstNode* expr, int line, int col) {
    NEW(alloc, AST_RETURN_STMT);
    n->as.return_stmt.expr = expr;
    return n;
}

AstNode* ast_make_expr_stmt(void* alloc, AstNode* expr, int line, int col) {
    NEW(alloc, AST_EXPR_STMT);
    n->as.expr_stmt.expr = expr;
    return n;
}

AstNode* ast_make_binary(void* alloc, BinaryOp op, AstNode* left, AstNode* right, int line, int col) {
    NEW(alloc, AST_BINARY_EXPR);
    n->as.binary.op = op; n->as.binary.left = left; n->as.binary.right = right;
    return n;
}

AstNode* ast_make_unary(void* alloc, BinaryOp op, AstNode* operand, int line, int col) {
    NEW(alloc, AST_UNARY_EXPR);
    n->as.unary.op = op; n->as.unary.operand = operand;
    return n;
}

AstNode* ast_make_call(void* alloc, const char* name, AstNode** args, size_t count, int line, int col) {
    NEW(alloc, AST_CALL_EXPR);
    n->as.call.name = name; n->as.call.args = args; n->as.call.arg_count = count;
    return n;
}

AstNode* ast_make_literal_i64(void* alloc, int64_t val, int line, int col) {
    NEW(alloc, AST_LITERAL_EXPR);
    n->as.literal.lit_type = TYPE_I64; n->as.literal.i64_val = val;
    n->type.kind = TYPE_I64;
    return n;
}

AstNode* ast_make_literal_f64(void* alloc, double val, int line, int col) {
    NEW(alloc, AST_LITERAL_EXPR);
    n->as.literal.lit_type = TYPE_F64; n->as.literal.f64_val = val;
    n->type.kind = TYPE_F64;
    return n;
}

AstNode* ast_make_literal_bool(void* alloc, bool val, int line, int col) {
    NEW(alloc, AST_LITERAL_EXPR);
    n->as.literal.lit_type = TYPE_BOOL; n->as.literal.bool_val = val;
    n->type.kind = TYPE_BOOL;
    return n;
}

AstNode* ast_make_literal_str(void* alloc, const char* val, int line, int col) {
    NEW(alloc, AST_LITERAL_EXPR);
    n->as.literal.lit_type = TYPE_STR; n->as.literal.str_val = val;
    n->type.kind = TYPE_STR;
    return n;
}

AstNode* ast_make_ident(void* alloc, const char* name, int line, int col) {
    NEW(alloc, AST_IDENT_EXPR);
    n->as.ident.name = name;
    return n;
}

// === 结构体相关工厂函数 ===

AstNode* ast_make_struct_decl(void* alloc, const char* name, AstNode** fields,
                               size_t count, int line, int col) {
    NEW(alloc, AST_STRUCT_DECL);
    n->as.struct_decl.name = name;
    n->as.struct_decl.fields = fields;
    n->as.struct_decl.field_count = count;
    return n;
}

AstNode* ast_make_struct_init(void* alloc, const char* type_name,
                               const char** fnames, AstNode** fvals,
                               size_t count, int line, int col) {
    NEW(alloc, AST_STRUCT_INIT);
    n->as.struct_init.type_name = type_name;
    n->as.struct_init.field_names = fnames;
    n->as.struct_init.field_values = fvals;
    n->as.struct_init.field_count = count;
    return n;
}

AstNode* ast_make_field_access(void* alloc, AstNode* object, const char* field,
                                int line, int col) {
    NEW(alloc, AST_FIELD_ACCESS);
    n->as.field_access.object = object;
    n->as.field_access.field = field;
    n->as.field_access.field_index = -1;
    return n;
}