#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->type.element_type = 0; n->type.element_struct_name = NULL; \
    n->type.array_size = 0; \
    n->loc = loc

AstNode* ast_make_program(void* alloc, AstNode** fns, size_t fn_count,
                          AstNode** structs, size_t struct_count,
                          AstNode** aliases, size_t alias_count,
                          AstNode** enums, size_t enum_count,
                          AstNode** impls, size_t impl_count, SourceLoc loc) {
    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;
    n->as.program.type_aliases = aliases;
    n->as.program.alias_count = alias_count;
    n->as.program.enums = enums;
    n->as.program.enum_count = enum_count;
    n->as.program.impls = impls;
    n->as.program.impl_count = impl_count;
    return n;
}

AstNode* ast_make_function(void* alloc, const char* name, AstNode** params, size_t pcount,
                           TypeKind ret, const char* ret_struct_name, AstNode* body, SourceLoc loc) {
    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.return_struct_type_name = ret_struct_name;
    n->as.function.body = body;
    return n;
}

AstNode* ast_make_parameter(void* alloc, const char* name, TypeKind type,
                             const char* struct_type_name, SourceLoc loc) {
    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, SourceLoc loc) {
    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,
                      TypeKind annot_elem_type, const char* annot_elem_struct, int64_t annot_array_size, SourceLoc loc) {
    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;
    n->as.let_stmt.annot_element_type = annot_elem_type;
    n->as.let_stmt.annot_element_struct_name = annot_elem_struct;
    n->as.let_stmt.annot_array_size = annot_array_size;
    return n;
}

AstNode* ast_make_assign(void* alloc, const char* name, AstNode* value, SourceLoc loc) {
    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, SourceLoc loc) {
    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, SourceLoc loc) {
    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, SourceLoc loc) {
    NEW(alloc, AST_RETURN_STMT);
    n->as.return_stmt.expr = expr;
    return n;
}

AstNode* ast_make_expr_stmt(void* alloc, AstNode* expr, SourceLoc loc) {
    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, SourceLoc loc) {
    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, SourceLoc loc) {
    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, SourceLoc loc) {
    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, SourceLoc loc) {
    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, SourceLoc loc) {
    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, SourceLoc loc) {
    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_char(void* alloc, int val, SourceLoc loc) {
    NEW(alloc, AST_LITERAL_EXPR);
    n->as.literal.lit_type = TYPE_CHAR; n->as.literal.i64_val = val;
    n->type.kind = TYPE_CHAR;
    return n;
}

AstNode* ast_make_literal_str(void* alloc, const char* val, SourceLoc loc) {
    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, SourceLoc loc) {
    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, SourceLoc loc) {
    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, SourceLoc loc) {
    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,
                                SourceLoc loc) {
    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;
}

AstNode* ast_make_type_alias(void* alloc, const char* name, TypeKind aliased,
                              const char* aliased_struct, SourceLoc loc) {
    NEW(alloc, AST_TYPE_ALIAS);
    n->as.type_alias.name = name;
    n->as.type_alias.aliased_type = aliased;
    n->as.type_alias.aliased_struct_name = aliased_struct;
    return n;
}

// === 枚举相关工厂函数 ===

AstNode* ast_make_enum_decl(void* alloc, const char* name, const char** variants,
                              size_t count, SourceLoc loc) {
    NEW(alloc, AST_ENUM_DECL);
    n->as.enum_decl.name = name;
    n->as.enum_decl.variants = variants;
    n->as.enum_decl.variant_count = count;
    return n;
}

AstNode* ast_make_enum_variant(void* alloc, const char* enum_name,
                                 const char* variant_name, SourceLoc loc) {
    NEW(alloc, AST_ENUM_VARIANT);
    n->as.enum_variant.enum_name = enum_name;
    n->as.enum_variant.variant_name = variant_name;
    n->as.enum_variant.variant_index = -1;
    return n;
}

AstNode* ast_make_index_expr(void* alloc, AstNode* array, AstNode* index, SourceLoc loc) {
    NEW(alloc, AST_INDEX_EXPR);
    n->as.index_expr.array = array;
    n->as.index_expr.index = index;
    return n;
}

AstNode* ast_make_array_assign(void* alloc, const char* name, AstNode* index, AstNode* value, SourceLoc loc) {
    NEW(alloc, AST_ARRAY_ASSIGN_STMT);
    n->as.array_assign.name = name;
    n->as.array_assign.index = index;
    n->as.array_assign.value = value;
    return n;
}

AstNode* ast_make_impl_block(void* alloc, const char* struct_name, AstNode** methods,
                              size_t count, SourceLoc loc) {
    NEW(alloc, AST_IMPL_BLOCK);
    n->as.impl_block.struct_name = struct_name;
    n->as.impl_block.methods = methods;
    n->as.impl_block.method_count = count;
    return n;
}

AstNode* ast_make_method_call(void* alloc, AstNode* receiver, const char* method,
                              AstNode** args, size_t count, SourceLoc loc) {
    NEW(alloc, AST_METHOD_CALL);
    n->as.method_call.receiver = receiver;
    n->as.method_call.method_name = method;
    n->as.method_call.args = args;
    n->as.method_call.arg_count = count;
    return n;
}