#ifndef AST_H

#define AST_H

#include "common.h"

typedef enum {

TYPE_BOOL,

TYPE_INT,

TYPE_FLOAT,

TYPE_STR,

TYPE_TENSOR,

TYPE_MAP,

TYPE_FUNC,

TYPE_THREAD,

TYPE_UNKNOWN

} DeclType;

typedef struct Expr Expr;

typedef struct Stmt Stmt;

typedef struct Param {

DeclType type;

int num_base; // 0 = parent int/float, 2..64 = named base

char *name;

bool coerced;

Expr *default_value; // optional

Expr *schema_expr; // NEW: map schema expression, NULL if none

} Param;

typedef struct ParamList {

Param *items;

size_t count;

size_t capacity;

} ParamList;

typedef enum {

EXPR_BOOL,

EXPR_INT,

EXPR_FLOAT,

EXPR_STR,

EXPR_PTR,

EXPR_IDENT,

EXPR_TYPED_IDENT,

EXPR_CALL,

EXPR_ASYNC,

EXPR_TENSOR,

EXPR_MAP,

EXPR_INDEX,

EXPR_RANGE,

EXPR_WILDCARD,

EXPR_LAMBDA,

EXPR_FMT_STR

} ExprType;

typedef struct {

Expr **items;

size_t count;

size_t capacity;

} ExprList;

struct Expr {

ExprType type;

int line;

int column;

union {

bool bool_value;

struct {

int64_t value;

int base;

} int_value;

struct {

double value;

int base;

int base_is_nan;

} flt_value;

char *str_value;

char *ident;

struct {

DeclType decl_type;

int decl_base; // 0 = parent int/float, 2..64 = named base

char *name;

Expr *schema_expr; // NEW: map schema expression, NULL if none

} typed_ident;

struct {

Expr *callee;

ExprList args;

char **kw_names;

ExprList kw_args;

size_t kw_count;

size_t kw_capacity;

} call;

struct {

Stmt *block;

} async;

struct {

ExprList parts; // alternating str parts and embedded expressions

} fmt_str;

struct {

ParamList params;

DeclType return_type;

int return_base; // 0 = parent int/float, 2..64 = named base

Expr *return_schema_expr; // NEW: map return schema, NULL if none

Stmt *body;

} lambda;

ExprList tns_items;

struct {

ExprList keys;

ExprList values;

} map_items;

struct {

Expr *target;

ExprList indices;

bool is_map;

} index;

struct {

Expr *start;

Expr *end;

} range;

char *ptr_name;

} as;

};

typedef enum {

STMT_BLOCK,

STMT_ASYNC,

STMT_EXPR,

STMT_ASSIGN,

STMT_DECL,

STMT_IF,

STMT_WHILE,

STMT_FOR,

STMT_PARFOR,

STMT_FUNC,

STMT_RETURN,

STMT_BREAK,

STMT_CONTINUE,

STMT_THREAD,

STMT_POP,

STMT_TRY,

STMT_GOTO,

STMT_GOTOPOINT

} StmtType;

typedef struct {

Stmt **items;

size_t count;

size_t capacity;

} StmtList;

struct Stmt {

StmtType type;

int line;

int column;

char *src_text;

union {

StmtList block;

struct {

Expr *expr;

} expr_stmt;

struct {

bool has_type;

DeclType decl_type;

int decl_base; // 0 = parent int/float, 2..64 = named base

char *name;

Expr *target;

Expr *value;

Expr *schema_expr; // NEW: map schema expression, NULL if none

} assign;

struct {

DeclType decl_type;

int decl_base; // 0 = parent int/float, 2..64 = named base

char *name;

Expr *schema_expr; // NEW: map schema expression, NULL if none

} decl;

struct {

Expr *condition;

Stmt *then_branch;

ExprList elif_conditions;

StmtList elif_blocks;

Stmt *else_branch; // optional

} if_stmt;

struct {

Expr *condition;

Stmt *body;

} while_stmt;

struct {

char *counter;

Expr *target;

Stmt *body;

} for_stmt;

struct {

char *counter;

Expr *target;

Stmt *body;

} parfor_stmt;

struct {

char *name;

ParamList params;

DeclType return_type;

int return_base; // 0 = parent int/float, 2..64 = named base

Expr *return_schema_expr; // NEW: map return schema, NULL if none

Stmt *body;

} func_stmt;

struct {

Expr *value;

} return_stmt;

struct {

Expr *value;

} break_stmt;

struct {

Stmt *body;

} async_stmt;

struct {

char *name;

Stmt *body;

} thread_stmt;

struct {

Stmt *try_block;

char *catch_name;

Stmt *catch_block;

} try_stmt;

struct {

Expr *target;

} goto_stmt;

struct {

Expr *expr;

} pop_stmt;

struct {

Expr *target;

} gotopoint_stmt;

} as;

};

Expr *expr_bool(bool value, int line, int column);

Expr *expr_int(int64_t value, int base, int line, int column);

Expr *expr_float(double value, int base, int base_is_nan, int line, int column);

Expr *expr_str(char *value, int line, int column);

Expr *expr_ptr(char *name, int line, int column);

Expr *expr_ident(char *name, int line, int column);

Expr *expr_typed_ident(DeclType decl_type, int decl_base, char *name, Expr *schema_expr, int line, int column);

Expr *expr_call(Expr *callee, int line, int column);

void call_kw_add(Expr *call, char *name, Expr *value);

Expr *expr_tensor(int line, int column);

Expr *expr_async(Stmt *block, int line, int column);

Expr *expr_map(int line, int column);

Expr *expr_index(Expr *target, int line, int column, bool is_map);

Expr *expr_range(Expr *start, Expr *end, int line, int column);

Expr *expr_wildcard(int line, int column);

Expr *expr_lambda(ParamList params, DeclType return_type, int return_base, Expr *return_schema_expr, Stmt *body,

int line, int column);

Expr *expr_fmt_str(ExprList parts, int line, int column);

void expr_list_add(ExprList *list, Expr *expr);

Stmt *stmt_block(int line, int column);

Stmt *stmt_async(Stmt *body, int line, int column);

Stmt *stmt_expr(Expr *expr, int line, int column);

Stmt *stmt_assign(bool has_type, DeclType decl_type, int decl_base, char *name, Expr *target, Expr *value,

Expr *schema_expr, int line, int column);

Stmt *stmt_decl(DeclType decl_type, int decl_base, char *name, Expr *schema_expr, int line, int column);

Stmt *stmt_if(Expr *cond, Stmt *then_branch, int line, int column);

Stmt *stmt_while(Expr *cond, Stmt *body, int line, int column);

Stmt *stmt_for(char *counter, Expr *target, Stmt *body, int line, int column);

Stmt *stmt_parfor(char *counter, Expr *target, Stmt *body, int line, int column);

Stmt *stmt_func(char *name, DeclType ret, int return_base, Expr *return_schema_expr, Stmt *body, int line, int column);

Stmt *stmt_return(Expr *value, int line, int column);

Stmt *stmt_pop(Expr *expr, int line, int column);

Stmt *stmt_break(Expr *value, int line, int column);

Stmt *stmt_continue(int line, int column);

Stmt *stmt_thread(char *name, Stmt *body, int line, int column);

Stmt *stmt_try(Stmt *try_block, char *catch_name, Stmt *catch_block, int line, int column);

Stmt *stmt_goto(Expr *target, int line, int column);

Stmt *stmt_gotopoint(Expr *target, int line, int column);

void stmt_list_add(StmtList *list, Stmt *stmt);

void param_list_add(ParamList *list, Param param);

// Attach original source text (single line) to a statement node.

void stmt_set_src(Stmt *stmt, const char *src);

void free_expr(Expr *expr);

void free_stmt(Stmt *stmt);

#endif // AST_H