#ifndef AST_H
#define AST_H

#include <stdio.h>

#include "limits.h"
#include "symtab.h"
#include "types.h"

/* String representations of value types. */
extern const char *type_names[];
/* String representations of node classes. */
extern const char *node_names[];

/* Span into parser's pointer pool. Used instead of embedded arrays
 * to keep node_t small. Access via parser_get_ptrs(). */
typedef struct {
    u32 idx; /* Starting index into parser_t.ptrs */
    u16 len; /* Number of elements */
    u16 cap; /* Capacity (for growable spans) */
} nodespan_t;

/* Variable declaration */
typedef struct {
    struct node_t *ident;
    struct node_t *type;
    struct node_t *value;
    struct node_t *align;
    bool mutable;
} var_decl_t;

/* Constant declaration */
typedef struct {
    struct node_t *ident;
    struct node_t *type;
    struct node_t *value;
} const_decl_t;

typedef struct {
    struct node_t *ident;
    struct node_t *type;
    struct node_t *value;
} static_decl_t;

/* Record type declaration. */
typedef struct {
    struct node_t *name;    /* Name of the record */
    nodespan_t     fields;  /* Fields of the record */
    struct node_t *attribs; /* Attributes (e.g. pub) */
    bool           tuple;   /* Unlabeled fields */
} record_decl_t;

typedef struct {
    nodespan_t fields; /* Fields of the record */
} record_type_t;

typedef struct {
    struct node_t *name;  /* Field name */
    struct node_t *value; /* Field value expression */
} record_lit_field_t;

/* Union variant definition. */
typedef struct {
    struct node_t *name;       /* Name of the variant */
    struct node_t *type;       /* Optional payload type */
    i32            value;      /* Value or tag for the variant */
    struct node_t *value_expr; /* Literal expression when explicitly assigned */
} union_variant_t;

/* Union type declaration. */
typedef struct {
    struct node_t *name;     /* Name of the union */
    nodespan_t     variants; /* Variants of the union */
    struct node_t *attribs;  /* Attributes (e.g. pub) */
} union_decl_t;

/* Function definition. */
typedef struct {
    struct node_t *ident;
    nodespan_t     params;
    struct node_t *return_type;
    nodespan_t     throws;
    struct node_t *body; /* Will be NULL for `extern` functions */
    struct node_t *attribs;
} fn_decl_t;

/* Node type. */
typedef enum {
    /* Literals. */
    NODE_NUMBER,
    NODE_CHAR,
    NODE_STRING,
    NODE_BOOL,
    NODE_NIL,              /* Nil literal */
    NODE_UNDEF,            /* Undefined literal */
    NODE_RECORD_LIT,       /* Record literal (e.g. Foo { x: 1, y: 2 }) */
    NODE_RECORD_LIT_FIELD, /* Record literal field */
    NODE_ARRAY_LIT,        /* Array literal (e.g. [1, 2, 3]) */
    NODE_ARRAY_REPEAT_LIT, /* Array repeat literal (e.g. [0; 24]) */
    NODE_ARRAY_INDEX,      /* Array indexing (e.g. arr[0]) */
    NODE_RANGE,            /* Range expression (e.g. 0..5 or ..) */

    /* Expressions. */
    NODE_IDENT,
    NODE_SUPER, /* `super` path segment */
    NODE_BINOP,
    NODE_UNOP, /* Unary operation (e.g., not x) */
    NODE_BLOCK,
    NODE_CALL,
    NODE_BUILTIN,
    NODE_CALL_ARG, /* Function call argument, optionally labeled (e.g. x: 1) */
    NODE_ACCESS,   /* (eg. foo.bar) */
    NODE_SCOPE,    /* (eg. foo::bar) */
    NODE_REF,      /* (eg., &foo) */
    NODE_AS,       /* (eg. x as i32) */

    /* Statements. */
    NODE_MOD,      /* Module declaration: mod foo; */
    NODE_MOD_BODY, /* Module body content */
    NODE_RETURN,
    NODE_THROW,
    NODE_WHILE,
    NODE_WHILE_LET, /* While let statement */
    NODE_LOOP,
    NODE_TRY,
    NODE_IF,
    NODE_IF_LET,     /* If let statement (eg. if let x in (opt) { ... }) */
    NODE_IF_CASE,    /* If case statement (eg. if case Foo::Bar(x) = expr) */
    NODE_GUARD_CASE, /* let case else statement */
    NODE_GUARD_LET,  /* let else statement */
    NODE_MATCH,      /* Switch statement (eg. switch (x) { ... }) */
    NODE_MATCH_CASE, /* Switch case (eg. case 1 => { ... }) */
    NODE_CATCH,
    NODE_FN,
    NODE_VAR, /* Variable declaration */
    NODE_CONST,
    NODE_STATIC,
    NODE_PARAM, /* Function parameter */
    NODE_BREAK,
    NODE_FOR,
    NODE_ASSIGN,
    NODE_EXPR_STMT,
    NODE_USE,   /* Module use declaration (e.g. use std.net) */
    NODE_PANIC, /* Panic statement */

    /* Type declarations. */
    NODE_TYPE,
    NODE_PTR, /* Pointer type */
    NODE_ALIGN,
    NODE_ATTRIBUTE,
    NODE_RECORD,       /* Record type declaration. */
    NODE_RECORD_FIELD, /* Field in record declaration */
    NODE_RECORD_TYPE,  /* Anonymous record type */
    NODE_UNION,
    NODE_UNION_VARIANT,
    NODE_PLACEHOLDER, /* Placeholder `_` for unused bindings */
} nodeclass_t;

/* Compiler built-ins */
typedef enum {
    BUILTIN_SIZE_OF,
    BUILTIN_ALIGN_OF,
    BUILTIN_SLICE_OF
} builtin_kind_t;

/* Binary operators. */
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_BAND, /* Bitwise AND (&) */
    OP_BOR,  /* Bitwise OR (|) */
    OP_XOR,  /* Bitwise XOR (^) */
    OP_SHL,  /* Shift left (<<) */
    OP_SHR,  /* Shift right (>>) */
} binop_t;

/* Unary operators. */
typedef enum {
    OP_NOT,
    OP_NEG, /* Numeric negation */
    OP_DEREF,
    OP_BNOT, /* Bitwise NOT (~) */
} unop_t;

/* A node in the abstract syntax tree. */
typedef struct node_t {
    /* Fields set by parser. */
    nodeclass_t cls;    /* Node class. */
    u32         offset; /* Byte offset into source code for this node. */
    u32         length; /* Length of source code for this node. */
    const char *file;   /* Source file this node was parsed from. */

    /* Fields set by resolver. */
    symbol_t      *sym;  /* Symbol table entry, if any. */
    struct type_t *type; /* Type context, if any. */

    /* Node value, set during parsing. */
    union {
        /* Boolean literal. */
        bool bool_lit;
        /* Character literal. */
        char char_lit;

        struct {
            const char *data;
            u16         length;
        } string_lit;

        /* Expression statement. */
        struct node_t *expr_stmt;

        /* Attribute node. */
        attrib_t attrib;

        /* Type node. */
        struct {
            /* The type represented by this node. If a complex type,
             * additional information is found in the `info` union. */
            typeclass_t    tclass;
            struct node_t *elem_type;

            union {
                struct {
                    struct node_t *length;
                } array;

                struct {
                    bool mut;
                } ptr;

                struct {
                    bool mut;
                } slice;

                struct {
                    nodespan_t     params; /* Parameter types */
                    struct node_t *ret;    /* Return type */
                    nodespan_t     throws;
                } fn;
            } info;
        } type;

        /* Reference node. */
        struct {
            struct node_t *target; /* Target of the reference. */
            bool           mut;    /* If this is a mutable reference. */
        } ref;

        /* Align expression */
        struct node_t *align;

        /* Expression nodes. */
        struct {
            const char *text;     /* Original text of the number literal */
            u16         text_len; /* Length of the original text */
            imm_t       value;    /* Parsed value based on type context */
        } number;

        struct {
            const char *name;
            u16         length;
        } ident;

        struct {
            binop_t        op;
            struct node_t *left;
            struct node_t *right;
        } binop;

        struct {
            unop_t         op;
            struct node_t *expr;
        } unop;

        struct {
            struct node_t *expr; /* Expression to cast */
            struct node_t *type; /* Target type */
        } as_expr;

        struct {
            nodespan_t      stmts;
            struct scope_t *scope; /* Scope for this block */
        } block;

        struct {
            struct node_t *callee;
            nodespan_t     args;
        } call;

        struct {
            builtin_kind_t kind;
            nodespan_t     args; /* Arguments to builtin */
        } builtin;

        struct {
            struct node_t *label; /* Optional label, or `NULL`. */
            struct node_t *expr;  /* Argument expression. */
        } call_arg;

        struct {
            struct node_t *lval;
            struct node_t *rval;
        } assign;

        struct {
            struct node_t *value;
        } return_stmt;

        struct {
            struct node_t *expr;
        } throw_stmt;

        struct {
            struct node_t *message; /* Optional message expression */
        } panic_stmt;

        struct {
            struct node_t *cond;
            struct node_t *body;
            struct node_t *rbranch; /* Optional else clause */
        } while_stmt;

        struct {
            struct node_t  *var;     /* Variable to bind */
            struct node_t  *expr;    /* Optional expression to unwrap */
            struct node_t  *guard;   /* Optional guard condition */
            struct node_t  *body;    /* Loop body */
            struct node_t  *rbranch; /* Optional else clause */
            struct scope_t *scope;   /* Holds the bound variable */
        } while_let_stmt;

        struct {
            struct node_t  *var;
            struct node_t  *idx; /* Optional index variable */
            struct node_t  *iter;
            struct node_t  *body;
            struct node_t  *rbranch; /* Optional else clause */
            struct scope_t *scope;   /* Holds the temporary variable. */
        } for_stmt;

        struct {
            struct node_t *body;
        } loop_stmt;

        struct {
            struct node_t *expr;       /* Expression guarded by try */
            struct node_t *catch_expr; /* Fallback expression for catch */
            nodespan_t     handlers;
            bool           panic; /* Emit ebreak on error when true */
            bool optional;        /* Return optional instead of propagating */
        } try_expr;

        struct {
            struct node_t *cond;
            struct node_t *lbranch;
            struct node_t *rbranch;
        } if_stmt;

        struct {
            struct node_t  *var;     /* Variable to bind */
            struct node_t  *expr;    /* Optional expression to unwrap */
            struct node_t  *guard;   /* Optional guard */
            struct node_t  *lbranch; /* Then branch */
            struct node_t  *rbranch; /* Else branch (optional) */
            struct scope_t *scope;   /* Holds the bound variable */
        } if_let_stmt;

        struct {
            struct node_t *pattern; /* Pattern to match */
            struct node_t *expr;    /* Expression being tested */
            struct node_t *guard;   /* Optional guard */
            struct node_t *lbranch; /* Then branch */
            struct node_t *rbranch; /* Else branch (optional) */
        } if_case_stmt;

        struct {
            struct node_t *pattern; /* Pattern to match */
            struct node_t *expr;    /* Expression being tested */
            struct node_t *guard;   /* Optional guard */
            struct node_t *rbranch; /* Else branch */
        } guard_case_stmt;

        struct {
            struct node_t *var;     /* Bound variable */
            struct node_t *expr;    /* Optional expression to unwrap */
            struct node_t *rbranch; /* Else branch */
        } guard_let_stmt;

        struct {
            struct node_t *expr;
            nodespan_t     cases; /* Switch cases */
        } match_stmt;

        struct {
            nodespan_t     patterns;
            struct node_t *body;     /* Case body */
            struct node_t *guard;    /* Optional guard condition */
            struct node_t *variable; /* Bound variable */
        } match_case;

        struct {
            struct node_t  *binding; /* Optional bound identifier */
            struct node_t  *body;    /* Catch handler body */
            struct scope_t *scope;   /* Scope for the bound variable */
        } catch_clause;

        var_decl_t    var;
        const_decl_t  constant;
        static_decl_t static_decl;

        struct {
            struct node_t *ident;
            struct node_t *type;
        } param;

        union_decl_t    union_decl;
        union_variant_t union_variant;

        struct {
            struct node_t *type;   /* Type identifier */
            nodespan_t     fields; /* Field initializers */
            bool           etc;    /* Pattern discards remaining fields */
        } record_lit;

        /* Field initialization in a record literal. */
        record_lit_field_t record_lit_field;

        /* Record declarations. */
        record_decl_t record_decl;
        /* Anonymous record */
        record_type_t record_type;

        /* Array indexing and record field access, eg. `x.y` or `x[y]`. */
        struct {
            struct node_t *lval;
            struct node_t *rval;
        } access;

        /* Range expression, e.g. `0..5` or `..` */
        struct {
            struct node_t *start; /* Start expression (NULL if omitted) */
            struct node_t *end;   /* End expression (NULL if omitted) */
        } range;

        /* Array literal, e.g. `[1, 2, 3]` */
        struct {
            nodespan_t elems;
        } array_lit;

        /* Array repeat literal, e.g. `[0; 24]` */
        struct {
            struct node_t *value; /* Value to repeat */
            struct node_t *count; /* Number of repetitions */
        } array_repeat_lit;

        /* Use declaration, e.g. `use std::net::tcp;` */
        struct {
            struct node_t *path;     /* Path to the module being used */
            struct node_t *attribs;  /* Attributes (e.g. pub) */
            bool           wildcard; /* Whether this is a wildcard import */
        } use_decl;

        /* Module declaration, e.g. `mod util;` */
        struct {
            struct node_t *ident;   /* Name of the module */
            struct node_t *attribs; /* Attributes (e.g. pub) */
        } mod_decl;

        fn_decl_t fn_decl;
    } val;
} node_t;

/* Forward declaration for parser_t */
struct parser_t;

/* Check if the node is a binary comparison operation. */
bool node_is_comp(node_t *n);
/* Check if the node is a literal */
bool node_is_literal(node_t *n);
/* Check if the binary operator is a logical one. */
bool op_is_logical(binop_t);
/* Add a parameter to a function. */
void node_fn_add_param(struct parser_t *p, node_t *fn, node_t *param);

/* Access node pointers from a span. Requires parser context. */
struct parser_t;
node_t **nodespan_ptrs(struct parser_t *p, nodespan_t span);

/* Allocate a new span from the parser's pointer pool. */
nodespan_t nodespan_alloc(struct parser_t *p, u16 cap);

/* Append a node to a span, growing if needed. Returns false on overflow. */
bool nodespan_push(struct parser_t *p, nodespan_t *span, node_t *node);

/* Forward declaration for module_t */
struct module_t;

/* Add a statement to a block node using module's parser. */
void node_block_add_stmt(struct module_t *mod, node_t *block, node_t *stmt);

#endif