/*
 * Copyright (C) 2024 olang maintainers
 *
 * This program is free software: you can redistribute it and/or modify
 * it under the terms of the GNU General Public License as published by
 * the Free Software Foundation, either version 3 of the License, or
 * (at your option) any later version.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License
 * along with this program.  If not, see <https://www.gnu.org/licenses/>.
 */
#ifndef AST_H
#define AST_H

#include <stdint.h>

#include "arena.h"
#include "lexer.h"
#include "list.h"
#include "scope.h"
#include "string_view.h"
#include "type.h"

#define AST_NODE_BASE                                                          \
    {                                                                          \
        ast_node_kind_t kind;                                                  \
        token_loc_t loc;                                                       \
    }

/**
 * Every node within the *ast_node* union must include the base segment.
 *
 * This macro must be declared first within any node definition to ensure the
 * ast_node base segment's integrity.
 */

#define AST_NODE_HEAD struct AST_NODE_BASE base

typedef union ast_node ast_node_t;

typedef enum
{
    AST_NODE_TRANSLATION_UNIT,
    AST_NODE_BLOCK,
    AST_NODE_FN_DEF,
    AST_NODE_FN_CALL,
    AST_NODE_VAR_DEF,
    AST_NODE_BINARY_OP,
    AST_NODE_UNARY_OP,
    AST_NODE_RETURN_STMT,
    AST_NODE_IF_STMT,
    AST_NODE_WHILE_STMT,
    AST_NODE_LITERAL,
    AST_NODE_REF,
    AST_NODE_UNKNOWN
} ast_node_kind_t;

typedef struct ast_node_base AST_NODE_BASE ast_node_base_t;

typedef struct ast_block
{
    AST_NODE_HEAD;
    list_t *nodes;
} ast_block_t;

typedef struct ast_id
{
    string_view_t name;
    scope_t *scope;
    symbol_t *symbol;
} ast_id_t;

typedef struct ast_translation_unit
{
    AST_NODE_HEAD;
    list_t *decls;
} ast_translation_unit_t;

typedef struct ast_fn_param
{
    string_view_t id;
    type_t *type;
} ast_fn_param_t;

typedef struct ast_fn_definition
{
    AST_NODE_HEAD;
    string_view_t id;
    list_t *params;
    type_t *return_type;
    bool _extern;
    ast_node_t *block;
    scope_t *scope;
} ast_fn_definition_t;

typedef struct ast_fn_call
{
    AST_NODE_HEAD;
    ast_id_t id;
    list_t *args;
} ast_fn_call_t;

typedef struct ast_var_definition
{
    AST_NODE_HEAD;
    ast_id_t id;
    type_t *type;
    ast_node_t *value;
} ast_var_definition_t;

typedef enum
{
    AST_LITERAL_U32
} ast_literal_kind_t;

typedef struct ast_literal
{
    AST_NODE_HEAD;
    ast_literal_kind_t kind;
    union
    {
        uint32_t as_u32;
    };
} ast_literal_t;

typedef struct ast_ref
{
    AST_NODE_HEAD;
    ast_id_t id;
} ast_ref_t;

typedef enum ast_binary_op_kind
{
    AST_BINOP_ADDITION,
    AST_BINOP_SUBTRACTION,
    AST_BINOP_MULTIPLICATION,
    AST_BINOP_DIVISION,
    AST_BINOP_REMINDER,
    AST_BINOP_BITWISE_LSHIFT,
    AST_BINOP_BITWISE_RSHIFT,
    AST_BINOP_BITWISE_XOR,
    AST_BINOP_BITWISE_AND,
    AST_BINOP_BITWISE_OR,
    AST_BINOP_CMP_LT,
    AST_BINOP_CMP_GT,
    AST_BINOP_CMP_LEQ,
    AST_BINOP_CMP_GEQ,
    AST_BINOP_CMP_EQ,
    AST_BINOP_CMP_NEQ,
    AST_BINOP_LOGICAL_AND,
    AST_BINOP_LOGICAL_OR,
    AST_BINOP_ASSIGN,
} ast_binary_op_kind_t;

typedef struct ast_binary_op
{
    AST_NODE_HEAD;
    ast_binary_op_kind_t kind;
    ast_node_t *lhs;
    ast_node_t *rhs;
} ast_binary_op_t;

typedef enum ast_unary_op_kind
{
    AST_UNARY_BITWISE_NOT,
    AST_UNARY_LOGICAL_NOT,
    AST_UNARY_NEGATIVE,
    AST_UNARY_POSITIVE,
    AST_UNARY_DEREFERENCE,
    AST_UNARY_ADDRESSOF,
} ast_unary_op_kind_t;

typedef struct ast_unary_op
{
    AST_NODE_HEAD;
    ast_unary_op_kind_t kind;
    ast_node_t *operand;
} ast_unary_op_t;

typedef struct ast_var_assign_stmt
{
    AST_NODE_HEAD;
    ast_node_t *ref;
    ast_node_t *value;
} ast_var_assign_stmt_t;

typedef struct ast_return_stmt
{
    AST_NODE_HEAD;
    ast_node_t *value;
} ast_return_stmt_t;

typedef struct ast_if_stmt
{
    AST_NODE_HEAD;
    ast_node_t *cond;
    ast_node_t *then;
    ast_node_t *_else;
} ast_if_stmt_t;

typedef struct ast_while_stmt
{
    AST_NODE_HEAD;
    ast_node_t *cond;
    ast_node_t *then;
} ast_while_stmt_t;

typedef union ast_node
{
    struct AST_NODE_BASE;
    ast_translation_unit_t as_translation_unit;
    ast_fn_definition_t as_fn_def;
    ast_fn_call_t as_fn_call;
    ast_var_definition_t as_var_def;
    ast_binary_op_t as_bin_op;
    ast_unary_op_t as_unary_op;
    ast_literal_t as_literal;
    ast_ref_t as_ref;
    ast_block_t as_block;
    ast_var_assign_stmt_t as_var_assign_stmt;
    ast_return_stmt_t as_return_stmt;
    ast_if_stmt_t as_if_stmt;
    ast_while_stmt_t as_while_stmt;
} ast_node_t;

ast_node_t *
ast_new_translation_unit(arena_t *arena);

ast_node_t *
ast_new_node_fn_def(arena_t *arena,
                    token_loc_t loc,
                    string_view_t id,
                    list_t *params,
                    type_t *return_type,
                    bool _extern,
                    ast_node_t *block);

ast_node_t *
ast_new_node_fn_call(arena_t *arena,
                     token_loc_t loc,
                     string_view_t id,
                     list_t *args);

ast_node_t *
ast_new_node_var_def(arena_t *arena,
                     token_loc_t loc,
                     string_view_t id,
                     type_t *type,
                     ast_node_t *value);

ast_node_t *
ast_new_node_bin_op(arena_t *arena,
                    token_loc_t loc,
                    ast_binary_op_kind_t kind,
                    ast_node_t *lhs,
                    ast_node_t *rhs);

ast_node_t *
ast_new_node_unary_op(arena_t *arena,
                      token_loc_t loc,
                      ast_unary_op_kind_t kind,
                      ast_node_t *operand);

ast_node_t *
ast_new_node_literal_u32(arena_t *arena, token_loc_t loc, uint32_t value);

ast_node_t *
ast_new_node_ref(arena_t *arena, token_loc_t loc, string_view_t id);

ast_node_t *
ast_new_node_var_assign_stmt(arena_t *arena,
                             token_loc_t loc,
                             ast_node_t *ref,
                             ast_node_t *expr);

ast_node_t *
ast_new_node_return_stmt(arena_t *arena, token_loc_t loc, ast_node_t *expr);

ast_node_t *
ast_new_node_if_stmt(arena_t *arena,
                     token_loc_t loc,
                     ast_node_t *cond,
                     ast_node_t *then,
                     ast_node_t *_else);

ast_node_t *
ast_new_node_while_stmt(arena_t *arena,
                        token_loc_t loc,
                        ast_node_t *cond,
                        ast_node_t *then);

ast_node_t *
ast_new_node_block(arena_t *arena);

ast_fn_param_t *
ast_new_fn_param(arena_t *arena, string_view_t id, type_t *type);

#endif /* AST_H */