/*
 * par2serial-cc: ast.h - Abstract Syntax Tree definitions
 * Supports full C subset + parallel extensions
 */
#ifndef P2S_AST_H
#define P2S_AST_H

#include "utils.h"
#include "lexer.h"

/* ── Forward declarations ────────────────────────────────── */
typedef struct ASTNode ASTNode;
typedef struct ASTType ASTType;

VEC_TYPEDEF(ASTNode*, NodeVec);

/* ── Type System ─────────────────────────────────────────── */
typedef enum {
    TYPE_VOID, TYPE_CHAR, TYPE_SHORT, TYPE_INT, TYPE_LONG,
    TYPE_FLOAT, TYPE_DOUBLE,
    TYPE_UNSIGNED_CHAR, TYPE_UNSIGNED_SHORT, TYPE_UNSIGNED_INT, TYPE_UNSIGNED_LONG,
    TYPE_POINTER, TYPE_ARRAY, TYPE_STRUCT, TYPE_UNION, TYPE_ENUM, TYPE_FUNC,
    TYPE_UNKNOWN
} TypeKind;

struct ASTType {
    TypeKind    kind;
    bool        is_const;
    bool        is_volatile;
    bool        is_restrict;
    bool        is_static;
    ASTType    *base;        /* for pointer/array */
    int         array_size;  /* for array; -1 = flexible */
    const char *name;        /* for struct/union/enum */
    /* function type info */
    ASTType   **param_types;
    int         param_count;
    ASTType    *return_type;
};

/* ── AST Node Types ──────────────────────────────────────── */
typedef enum {
    /* Program / Declarations */
    NODE_PROGRAM,            /* top-level: list of declarations */
    NODE_FUNC_DECL,          /* function declaration */
    NODE_VAR_DECL,           /* variable declaration (possibly with init) */
    NODE_PARAM,              /* function parameter */
    NODE_STRUCT_DECL,        /* struct definition */

    /* Statements */
    NODE_BLOCK,              /* compound statement { ... } */
    NODE_EXPR_STMT,          /* expression statement */
    NODE_IF,                 /* if/else */
    NODE_WHILE,              /* while loop */
    NODE_DO_WHILE,           /* do-while loop */
    NODE_FOR,                /* standard for loop */
    NODE_RETURN,             /* return statement */
    NODE_BREAK,
    NODE_CONTINUE,

    /* Expressions */
    NODE_BINARY,             /* a op b */
    NODE_UNARY,              /* op a */
    NODE_POSTFIX,            /* a++ a-- */
    NODE_ASSIGN,             /* a = b (and +=, -=, etc) */
    NODE_TERNARY,            /* a ? b : c */
    NODE_CALL,               /* f(args) */
    NODE_INDEX,              /* a[i] */
    NODE_MEMBER,             /* a.field */
    NODE_ARROW,              /* a->field */
    NODE_CAST,               /* (type)expr */
    NODE_SIZEOF,             /* sizeof(expr/type) */
    NODE_IDENT,              /* identifier reference */
    NODE_INT_LIT,            /* integer literal */
    NODE_FLOAT_LIT,          /* float literal */
    NODE_STRING_LIT,         /* string literal */
    NODE_CHAR_LIT,           /* char literal */
    NODE_COMMA_EXPR,         /* a, b */
    NODE_ADDR_OF,            /* &x */
    NODE_DEREF,              /* *p */

    /* Parallel extensions */
    NODE_PARALLEL_FOR,       /* parallel_for(var, lo, hi) { body } */
    NODE_PARALLEL_REDUCE,    /* parallel_reduce(acc, op, var, lo, hi) { body } */
    NODE_PARALLEL_SCAN,      /* parallel_scan(out, op, var, lo, hi) { body } */
    NODE_PARALLEL_MAP,       /* parallel_map(dst, src, n, func) { body } */
    NODE_BARRIER,            /* barrier() */
    NODE_SHARED_DECL,        /* shared float x[N] */
    NODE_TILE_HINT,          /* tile_hint(name, size) */
    NODE_SIMD_HINT,          /* simd_hint(target) */
    NODE_MEMORY_LAYOUT,      /* memory_layout(var, SOA/AOS) */
    NODE_ATOMIC_ADD,         /* atomic_add(&x, val) */
    NODE_ATOMIC_CAS,         /* atomic_cas(&x, expected, desired) */

    /* Preprocessor (passthrough) */
    NODE_PREPROC,            /* #include, #define, etc. */

    NODE_TYPE_COUNT
} NodeType;

/* ── Operator types ──────────────────────────────────────── */
typedef enum {
    OP_ADD, OP_SUB, OP_MUL, OP_DIV, OP_MOD,
    OP_AND, OP_OR, OP_XOR, OP_NOT, OP_LNOT,
    OP_LSHIFT, OP_RSHIFT,
    OP_EQ, OP_NEQ, OP_LT, OP_GT, OP_LE, OP_GE,
    OP_LAND, OP_LOR,
    OP_NEG, OP_POS,        /* unary - + */
    OP_INC, OP_DEC,        /* ++ -- (prefix) */
    OP_POST_INC, OP_POST_DEC,
    OP_ASSIGN,
    OP_ADD_ASSIGN, OP_SUB_ASSIGN, OP_MUL_ASSIGN, OP_DIV_ASSIGN,
    OP_MOD_ASSIGN, OP_AND_ASSIGN, OP_OR_ASSIGN, OP_XOR_ASSIGN,
    OP_LSHIFT_ASSIGN, OP_RSHIFT_ASSIGN,
    OP_COMMA,
} OpType;

/* ── Reduction operator for parallel constructs ──────────── */
typedef enum {
    REDUCE_ADD, REDUCE_MUL, REDUCE_MIN, REDUCE_MAX,
    REDUCE_AND, REDUCE_OR, REDUCE_XOR
} ReduceOp;

/* ── AST Node ────────────────────────────────────────────── */
struct ASTNode {
    NodeType     type;
    SourceLoc    loc;
    ASTType     *expr_type;   /* resolved type (filled by semantic analysis) */

    union {
        /* NODE_PROGRAM */
        struct { NodeVec decls; } program;

        /* NODE_FUNC_DECL */
        struct {
            const char *name;
            ASTType    *return_type;
            NodeVec     params;
            ASTNode    *body;       /* NULL = forward decl */
            bool        is_inline;
            bool        is_static;
        } func;

        /* NODE_VAR_DECL */
        struct {
            const char *name;
            ASTType    *var_type;
            ASTNode    *init;       /* NULL if no initializer */
        } var;

        /* NODE_PARAM */
        struct {
            const char *name;
            ASTType    *param_type;
        } param;

        /* NODE_STRUCT_DECL */
        struct {
            const char *name;
            NodeVec     fields;
            bool        is_union;
        } struc;

        /* NODE_BLOCK */
        struct { NodeVec stmts; } block;

        /* NODE_IF */
        struct {
            ASTNode *cond;
            ASTNode *then_body;
            ASTNode *else_body;     /* NULL if no else */
        } if_stmt;

        /* NODE_WHILE / NODE_DO_WHILE */
        struct {
            ASTNode *cond;
            ASTNode *body;
        } while_stmt;

        /* NODE_FOR */
        struct {
            ASTNode *init;
            ASTNode *cond;
            ASTNode *step;
            ASTNode *body;
        } for_stmt;

        /* NODE_RETURN */
        struct { ASTNode *value; } ret;

        /* NODE_EXPR_STMT */
        struct { ASTNode *expr; } expr_stmt;

        /* NODE_BINARY */
        struct {
            OpType   op;
            ASTNode *left;
            ASTNode *right;
        } binary;

        /* NODE_UNARY */
        struct {
            OpType   op;
            ASTNode *operand;
        } unary;

        /* NODE_POSTFIX */
        struct {
            OpType   op;
            ASTNode *operand;
        } postfix;

        /* NODE_ASSIGN */
        struct {
            OpType   op;
            ASTNode *target;
            ASTNode *value;
        } assign;

        /* NODE_TERNARY */
        struct {
            ASTNode *cond;
            ASTNode *then_expr;
            ASTNode *else_expr;
        } ternary;

        /* NODE_CALL */
        struct {
            ASTNode *func;
            NodeVec  args;
        } call;

        /* NODE_INDEX */
        struct {
            ASTNode *array;
            ASTNode *index;
        } index;

        /* NODE_MEMBER / NODE_ARROW */
        struct {
            ASTNode    *object;
            const char *field;
        } member;

        /* NODE_CAST */
        struct {
            ASTType *target_type;
            ASTNode *expr;
        } cast;

        /* NODE_SIZEOF */
        struct {
            ASTType *of_type;     /* sizeof(type) */
            ASTNode *of_expr;     /* sizeof expr (one or the other) */
        } size_of;

        /* NODE_IDENT */
        struct { const char *name; } ident;

        /* NODE_INT_LIT */
        struct { int64_t value; } int_lit;

        /* NODE_FLOAT_LIT */
        struct { double value; } float_lit;

        /* NODE_STRING_LIT */
        struct { const char *value; } string_lit;

        /* NODE_CHAR_LIT */
        struct { int64_t value; } char_lit;

        /* NODE_ADDR_OF / NODE_DEREF */
        struct { ASTNode *operand; } unary_expr;

        /* ── Parallel Nodes ─────────────────────────────── */

        /* NODE_PARALLEL_FOR */
        struct {
            const char *iter_var;   /* loop variable name */
            ASTNode    *lo;         /* lower bound */
            ASTNode    *hi;         /* upper bound (exclusive) */
            ASTNode    *step;       /* stride (NULL = 1) */
            ASTNode    *body;       /* loop body */
        } par_for;

        /* NODE_PARALLEL_REDUCE */
        struct {
            const char *accum_var;  /* accumulator variable name */
            ReduceOp    op;         /* reduction operator */
            const char *iter_var;
            ASTNode    *lo;
            ASTNode    *hi;
            ASTNode    *body;
            ASTNode    *init_val;   /* initial value for accumulator */
        } par_reduce;

        /* NODE_PARALLEL_SCAN */
        struct {
            const char *output_var;
            ReduceOp    op;
            const char *iter_var;
            ASTNode    *lo;
            ASTNode    *hi;
            ASTNode    *body;
        } par_scan;

        /* NODE_PARALLEL_MAP */
        struct {
            const char *dst;
            const char *src;
            ASTNode    *count;
            const char *func_name;
            ASTNode    *body;
        } par_map;

        /* NODE_SHARED_DECL */
        struct {
            const char *name;
            ASTType    *var_type;
            ASTNode    *size;
        } shared;

        /* NODE_TILE_HINT */
        struct {
            const char *target;
            int         tile_size;
        } tile_hint;

        /* NODE_SIMD_HINT */
        struct {
            const char *target;     /* "sse", "avx", "avx2", "avx512", "neon" */
        } simd_hint;

        /* NODE_MEMORY_LAYOUT */
        struct {
            const char *target;
            bool        soa;        /* true=SOA, false=AOS */
        } mem_layout;

        /* NODE_ATOMIC_ADD */
        struct {
            ASTNode *target;
            ASTNode *value;
        } atomic_add;

        /* NODE_ATOMIC_CAS */
        struct {
            ASTNode *target;
            ASTNode *expected;
            ASTNode *desired;
        } atomic_cas;

        /* NODE_PREPROC */
        struct { const char *text; } preproc;
    };
};

/* ── AST construction helpers ────────────────────────────── */
ASTNode *ast_new(Arena *a, NodeType type, SourceLoc loc);
ASTType *ast_new_type(Arena *a, TypeKind kind);
ASTType *ast_ptr_type(Arena *a, ASTType *base);
ASTType *ast_array_type(Arena *a, ASTType *base, int size);

/* ── AST printing (debug) ────────────────────────────────── */
void ast_print(ASTNode *node, int indent);
const char *node_type_str(NodeType t);
const char *op_type_str(OpType op);

#endif /* P2S_AST_H */