#ifndef AST_H
#define AST_H

#include "l_lang.h"
#include <stddef.h>

typedef enum {
    AST_PROGRAM,
    AST_FUNCTION,
    AST_PARAMETER,
    AST_BLOCK,
    AST_LET_STMT,
    AST_ASSIGN_STMT,
    AST_IF_STMT,
    AST_WHILE_STMT,
    AST_RETURN_STMT,
    AST_EXPR_STMT,
    AST_DEFER_STMT,
    AST_LIST_COMP,         // [for x in expr: body]
    AST_BINARY_EXPR,
    AST_UNARY_EXPR,
    AST_CALL_EXPR,
    AST_LITERAL_EXPR,
    AST_IDENT_EXPR,
    AST_STRUCT_DECL,      // struct Point { x: i64, y: i64 }
    AST_STRUCT_INIT,      // Point { x: 10, y: 20 }
    AST_FIELD_ACCESS,     // p.x
    AST_TYPE_ALIAS,       // type Meters = i64
    AST_ENUM_DECL,        // enum Color { Red, Green, Blue }
    AST_ENUM_VARIANT,     // Color::Red
    AST_INDEX_EXPR,       // arr[i]
    AST_ARRAY_ASSIGN_STMT,// arr[i] = expr
    AST_IMPL_BLOCK,       // impl StructName { fn method(...) ... }
    AST_METHOD_CALL,      // receiver.method(args)
    AST_LAMBDA,           // fn(x: T) -> R { body } 匿名函数/闭包
    AST_MOD_DECL,         // mod foo;
    AST_USE_DECL,         // use foo::bar;
    AST_TRAIT_DECL,       // trait Name { fn ... }
} AstKind;

typedef enum {
    OP_ADD, OP_SUB, OP_MUL, OP_DIV, OP_MOD,
    OP_EQ, OP_NE, OP_LT, OP_GT, OP_LE, OP_GE,
    OP_AND, OP_OR,
    OP_NEG, OP_NOT,
} BinaryOp;

// 类型信息（语义分析阶段填充）
typedef struct {
    TypeKind kind;
    const char* struct_name;       // TYPE_STRUCT / TYPE_ENUM
    TypeKind    element_type;       // TYPE_ARRAY: 元素类型的 TypeKind
    const char* element_struct_name; // TYPE_ARRAY: 元素为 struct 时的类型名
    int64_t     array_size;         // TYPE_ARRAY: 固定大小
} TypeInfo;

// AST 节点
struct AstNode {
    AstKind   kind;
    TypeInfo  type;         // 语义分析后填充，默认为 TYPE_UNKNOWN
    SourceLoc loc;          // 源码位置

    // 节点特有数据（按 kind 解释）
    union {
        // AST_PROGRAM
        struct { struct AstNode** functions; size_t fn_count;
                 struct AstNode** structs;   size_t struct_count;
                 struct AstNode** type_aliases; size_t alias_count;
                 struct AstNode** enums;     size_t enum_count;
                 struct AstNode** impls;     size_t impl_count; } program;
        // AST_FUNCTION
        struct { const char* name; struct AstNode** params; size_t param_count;
                 TypeKind return_type; const char* return_struct_type_name;
                 struct AstNode* body; bool is_pub;
                 const char** type_params; size_t type_param_count;
                 TypeKind* multi_ret_types; const char** multi_ret_snames; size_t multi_ret_count;
                 const char** captured; TypeKind* cap_types; size_t cap_count; } function;
        // AST_PARAMETER (也用作结构体字段: name + type)
        struct { const char* name; TypeKind type; const char* struct_type_name; bool is_out;
                 TypeKind arr_elem_type; const char* arr_elem_struct; int64_t arr_size; } parameter;
        // AST_BLOCK
        struct { struct AstNode** stmts; size_t stmt_count; } block;
        // AST_LET_STMT
        struct { const char* name; TypeKind annot_type; bool has_type_annot; bool is_mut; bool is_const; struct AstNode* init;
                 const char* struct_type_name;
                 TypeKind annot_element_type; const char* annot_element_struct_name; int64_t annot_array_size; } let_stmt;
        // AST_ASSIGN_STMT
        struct { const char* name; struct AstNode* value; } assign_stmt;
        // AST_IF_STMT
        struct { struct AstNode* cond; struct AstNode* then_block; struct AstNode* else_block; } if_stmt;
        // AST_WHILE_STMT
        struct { struct AstNode* cond; struct AstNode* body; } while_stmt;
        // AST_RETURN_STMT
        struct { struct AstNode* expr; } return_stmt;
        // AST_EXPR_STMT
        struct { struct AstNode* expr; } expr_stmt;
        // AST_DEFER_STMT
        struct { struct AstNode* body; } defer_stmt;
        // AST_LIST_COMP
        struct { const char* var_name; struct AstNode* array; struct AstNode* map_expr; } list_comp;
        // AST_BINARY_EXPR
        struct { BinaryOp op; struct AstNode* left; struct AstNode* right; } binary;
        // AST_UNARY_EXPR
        struct { BinaryOp op; struct AstNode* operand; } unary;
        // AST_CALL_EXPR
        struct { const char* name; struct AstNode** args; const char** arg_names; size_t arg_count; } call;
        // AST_LITERAL_EXPR
        struct { TypeKind lit_type; union { int64_t i64_val; double f64_val; bool bool_val; const char* str_val; }; } literal;
        // AST_IDENT_EXPR
        struct { const char* name; } ident;
        // AST_STRUCT_DECL
        struct { const char* name; struct AstNode** fields; size_t field_count; } struct_decl;
        // AST_STRUCT_INIT
        struct { const char* type_name; const char** field_names;
                 struct AstNode** field_values; size_t field_count; } struct_init;
        // AST_FIELD_ACCESS
        struct { struct AstNode* object; const char* field; int field_index; } field_access;
        // AST_TYPE_ALIAS
        struct { const char* name; TypeKind aliased_type; const char* aliased_struct_name; } type_alias;
        // AST_ENUM_DECL
        struct { const char* name; const char** variants;
                 TypeKind* variant_payload_types;           // 每个变体的 payload 类型(TYPE_VOID=无)
                 const char** variant_payload_struct_names;  // payload 为 struct 时的类型名
                 size_t variant_count; } enum_decl;
        // AST_ENUM_VARIANT
        struct { const char* enum_name; const char* variant_name;
                 AstNode* payload;   // payload 表达式 (NULL=无 payload)
                 int variant_index; } enum_variant;
        // AST_INDEX_EXPR
        struct { struct AstNode* array; struct AstNode* index; } index_expr;
        // AST_ARRAY_ASSIGN_STMT
        struct { const char* name; struct AstNode* index; struct AstNode* value; } array_assign;
        // AST_IMPL_BLOCK
        struct { const char* struct_name; struct AstNode** methods; size_t method_count; } impl_block;
        // AST_METHOD_CALL
        struct { struct AstNode* receiver; const char* method_name; struct AstNode** args; const char** arg_names; size_t arg_count; } method_call;
        // AST_LAMBDA
        struct { struct AstNode** params; size_t param_count;
                 TypeKind return_type; const char* return_struct_type_name;
                 struct AstNode* body;
                 const char* generated_name;   // 自动生成的顶层函数名
                 const char** captured; size_t captured_count; } lambda;
        // AST_MOD_DECL
        struct { const char* name; struct AstNode* ast; } mod_decl;
        // AST_USE_DECL
        struct { const char* path; const char* item; } use_decl;
        // AST_TRAIT_DECL
        struct { const char* name; struct AstNode** methods; size_t method_count; } trait_decl;
    } as;
};

// 创建节点的辅助函数（内存来自 arena，通过 void* 传递避免循环依赖）
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);
AstNode* ast_make_function(void* alloc, const char* name, AstNode** params, size_t pcount,
                           TypeKind ret, const char* ret_struct_name, AstNode* body,
                           bool is_pub, const char** type_params, size_t tp_count,
                           SourceLoc loc);
AstNode* ast_make_parameter(void* alloc, const char* name, TypeKind type, const char* struct_type_name, bool is_out,
                           TypeKind arr_elem, const char* arr_elem_sname, int64_t arr_size, SourceLoc loc);
AstNode* ast_make_block(void* alloc, AstNode** stmts, size_t count, SourceLoc loc);
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);
AstNode* ast_make_assign(void* alloc, const char* name, AstNode* value, SourceLoc loc);
AstNode* ast_make_if(void* alloc, AstNode* cond, AstNode* then_b, AstNode* else_b, SourceLoc loc);
AstNode* ast_make_while(void* alloc, AstNode* cond, AstNode* body, SourceLoc loc);
AstNode* ast_make_return(void* alloc, AstNode* expr, SourceLoc loc);
AstNode* ast_make_expr_stmt(void* alloc, AstNode* expr, SourceLoc loc);
AstNode* ast_make_defer_stmt(void* alloc, AstNode* expr, SourceLoc loc);
AstNode* ast_make_list_comp(void* alloc, const char* var, AstNode* arr, AstNode* map, SourceLoc loc);
AstNode* ast_make_binary(void* alloc, BinaryOp op, AstNode* left, AstNode* right, SourceLoc loc);
AstNode* ast_make_unary(void* alloc, BinaryOp op, AstNode* operand, SourceLoc loc);
AstNode* ast_make_call(void* alloc, const char* name, AstNode** args, const char** arg_names, size_t count, SourceLoc loc);
AstNode* ast_make_literal_i64(void* alloc, int64_t val, SourceLoc loc);
AstNode* ast_make_literal_f64(void* alloc, double val, SourceLoc loc);
AstNode* ast_make_literal_bool(void* alloc, bool val, SourceLoc loc);
AstNode* ast_make_literal_char(void* alloc, int val, SourceLoc loc);
AstNode* ast_make_literal_str(void* alloc, const char* val, SourceLoc loc);
AstNode* ast_make_ident(void* alloc, const char* name, SourceLoc loc);
AstNode* ast_make_struct_decl(void* alloc, const char* name, AstNode** fields, size_t count, SourceLoc loc);
AstNode* ast_make_struct_init(void* alloc, const char* type_name, const char** fnames, AstNode** fvals, size_t count, SourceLoc loc);
AstNode* ast_make_field_access(void* alloc, AstNode* object, const char* field, SourceLoc loc);
AstNode* ast_make_type_alias(void* alloc, const char* name, TypeKind aliased,
                              const char* aliased_struct, SourceLoc loc);
AstNode* ast_make_enum_decl(void* alloc, const char* name, const char** variants,
                             TypeKind* payload_types, const char** payload_struct_names,
                             size_t count, SourceLoc loc);
AstNode* ast_make_enum_variant(void* alloc, const char* enum_name, const char* variant_name,
                                AstNode* payload, SourceLoc loc);
AstNode* ast_make_index_expr(void* alloc, AstNode* array, AstNode* index, SourceLoc loc);
AstNode* ast_make_array_assign(void* alloc, const char* name, AstNode* index, AstNode* value, SourceLoc loc);
AstNode* ast_make_impl_block(void* alloc, const char* struct_name, AstNode** methods, size_t count, SourceLoc loc);
AstNode* ast_make_method_call(void* alloc, AstNode* receiver, const char* method, AstNode** args, const char** arg_names, size_t count, SourceLoc loc);
AstNode* ast_make_lambda(void* alloc, AstNode** params, size_t pcount,
                         TypeKind ret, const char* ret_struct_name,
                         AstNode* body, SourceLoc loc);
AstNode* ast_make_mod_decl(void* alloc, const char* name, AstNode* sub_ast, SourceLoc loc);
AstNode* ast_make_use_decl(void* alloc, const char* path, const char* item, SourceLoc loc);
AstNode* ast_make_trait_decl(void* alloc, const char* name, AstNode** methods, size_t count, SourceLoc loc);

#endif