/*
 * 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"

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_VAR_ASSIGN_STMT,
    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_meta {
    ast_node_kind_t kind;
    token_loc_t loc;
} ast_node_meta_t;

typedef struct ast_block {
    ast_node_meta_t meta;
    list_t *nodes;
} ast_block_t;

typedef struct ast_translation_unit {
    ast_node_meta_t meta;
    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_meta_t meta;
    string_view_t id;
    list_t *params;
    type_t *return_type;
    ast_node_t *block;
    scope_t *scope;
} ast_fn_definition_t;

typedef struct ast_fn_call {
    ast_node_meta_t meta;
    string_view_t id;
    list_t *args;
    scope_t *scope;
} ast_fn_call_t;

typedef struct ast_var_definition {
    ast_node_meta_t meta;
    string_view_t id;
    type_t *type;
    ast_node_t *value;
    scope_t *scope;
} ast_var_definition_t;

typedef enum {
    AST_LITERAL_U32
} ast_literal_kind_t;

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

typedef struct ast_ref {
    ast_node_meta_t meta;
    string_view_t id;
    scope_t *scope;
} 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_binary_op_kind_t;

typedef struct ast_binary_op {
    ast_node_meta_t meta;
    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_meta_t meta;
    ast_unary_op_kind_t kind;
    ast_node_t *expr;
} ast_unary_op_t;

typedef struct ast_var_assign_stmt {
    ast_node_meta_t meta;
    ast_node_t *ref;
    ast_node_t *expr;
} ast_var_assign_stmt_t;

typedef struct ast_return_stmt {
    ast_node_meta_t meta;
    ast_node_t *expr;
} ast_return_stmt_t;

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

typedef struct ast_while_stmt {
    ast_node_meta_t meta;
    ast_node_t *cond;
    ast_node_t *then;
} ast_while_stmt_t;

typedef union ast_node {
    // inlined ast_node_meta_t struct.
    struct {
        ast_node_kind_t kind;
        token_loc_t loc;
    };
    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,
    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 *expr);

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 */