#ifndef AMALGAM_AST_H
#define AMALGAM_AST_H

#include "defs.h"
#include "nullable.h"
#include "std/defs.h"
#include "std/buffer_view.h"
#include "std/buffer.h"
#include "std/misc.h"
#include "std/arena_allocator.h"
#include "std/hash_map.h"
#include "number.h"
#include "binop_type.h"
#include "ir/ir.h"

#include <setjmp.h>

/* General error */
#define AST_ERR -1
#define AST_ERR_DEF_DUP -20

typedef struct Ast Ast;
typedef struct FunctionParameter FunctionParameter;
typedef struct FunctionCall FunctionCall;
typedef struct StructDecl StructDecl;
typedef struct StructField StructField;
typedef struct AssignmentExpr AssignmentExpr;
typedef struct Import Import;
typedef struct String String;
typedef struct Variable Variable;
typedef struct Number Number;
typedef struct AstBool AstBool;
typedef struct Binop Binop;
typedef struct IfStatement IfStatement;
typedef struct ElseIfStatement ElseIfStatement;
typedef struct WhileStatement WhileStatement;
typedef struct ReturnExpr ReturnExpr;

/* TODO: Instead of using pointers, use the data directly */
typedef union {
    void *data;
    FunctionDecl *func_decl;
    FunctionCall *func_call;
    StructDecl *struct_decl;
    StructField *struct_field;
    LhsExpr *lhs_expr;
    AssignmentExpr *assign_expr;
    Import *import;
    String *string;
    Number *number;
    AstBool *bool;
    Variable *variable;
    Binop *binop;
    IfStatement *if_stmt;
    WhileStatement *while_stmt;
    ReturnExpr *return_expr;
} AstValue;

typedef enum {
    AST_FUNCTION_DECL,
    AST_FUNCTION_CALL,
    AST_STRUCT_DECL,
    AST_STRUCT_FIELD,
    AST_LHS,
    AST_ASSIGN,
    AST_IMPORT,
    AST_STRING,
    AST_NUMBER,
    AST_BOOL,
    AST_VARIABLE,
    AST_BINOP,
    AST_IF_STATEMENT,
    AST_WHILE_STATEMENT,
    AST_RETURN
} AstType;

typedef enum {
    AST_NOT_RESOLVED,
    AST_RESOLVING,
    AST_RESOLVED,
    AST_IR_RESOLVED
} AstResolveStatus;

typedef enum {
    RESOLVED_TYPE_NONE,
    RESOLVED_TYPE_LHS_EXPR,
    RESOLVED_TYPE_FUNC_SIG
} AstResolvedTypeEnum;

typedef struct {
    union {
        void *data;
        LhsExpr *lhs_expr;
        FunctionSignature *func_sig;
    } value;
    AstResolvedTypeEnum type;
} AstResolvedType;

typedef struct {
    AstResolvedType type;
    AstResolveStatus status;
    IrRegister ir_reg;
} AstResolveData;

typedef enum {
    NAMED_OBJECT_NONE,
    NAMED_OBJECT_LHS_EXPR,
    NAMED_OBJECT_FUNC_PARAM
} ScopeNamedObjectType;

typedef struct {
    union {
        void *data;
        LhsExpr *lhs_expr;
        FunctionParameter *func_param;
    } value;
    ScopeNamedObjectType type;
    AstResolveData *resolve_data; /* Borrowed from @value.func_param or from the Ast @value.lhs_expr belongs to */
} ScopeNamedObject;

struct Ast {
    AstValue value;
    AstType type;
    AstResolveData resolve_data;
    /*
        Borrowed. This is the parser (thread) that is currently parsing the expression.
        TODO: Move this to LhsExpr
    */
    Parser *parser;
};

struct Scope {
    Buffer/*<Ast*>*/ ast_objects;
    HashMapType(BufferView, Ast*) named_objects; /* Value is always an Ast* with type LhsExpr */
    Scope *parent;
    FunctionSignature *function_signature; /* Borrowed from FunctionDecl. Only used if the scope belongs to FunctionDecl */
};

typedef enum {
    VARIABLE_TYPE_FLAG_NONE      = 0,
    VARIABLE_TYPE_FLAG_OPTIONAL  = 1 << 0,
    VARIABLE_TYPE_FLAG_BORROW    = 1 << 1
} VariableTypeFlag;

typedef struct {
    enum {
        VARIABLE_TYPE_NONE,
        VARIABLE_TYPE_VARIABLE,
        VARIABLE_TYPE_SIGNATURE
    } type;

    union {
        Variable *variable;
        FunctionSignature *signature;
    } value;
    VariableTypeFlag variable_type_flags;
} VariableType;

struct FileScopeReference {
    Parser *parser; /* Borrowed, the parser that belongs to the same file as the FileScopeReference */
    Buffer canonical_path;
};

struct ParserFileScopeReference {
    FileScopeReference *file_scope_ref;
    /* Index in the parser import list. Used by bytecode to know which import an expression belongs to */
    int import_index;
};

struct Variable {
    BufferView name;
    ScopeNamedObject resolved_var;
};

struct FunctionParameter {
    BufferView name;
    VariableType type;
    AstResolveData resolve_data;
};

typedef struct {
    VariableType type;
    AstResolvedType resolved_type;
} FunctionReturnType;

struct FunctionSignature {
    Buffer/*FunctionParameter*/ parameters;
    Buffer/*FunctionReturnType*/ return_types;
    /* Borrowed from the FunctionDecl the function signature belongs to. This is used to access the scope the function signature belongs to */
    FunctionDecl *func_decl;
    bool resolved;
};

struct FunctionDecl {
    LhsExpr *lhs_expr; /* Borrowed from the LhsExpr that owns this FunctionDecl if it exists, otherwise NULL */
    FunctionSignature *signature; /* TODO: Make this a non-pointer type? */
    Scope body;
    IrFuncIndex ir_func_index;
};

struct FunctionCall {
    Variable func;
    Buffer/*<Ast*>*/ args;
};

struct StructDecl {
    /*
        TODO: Instead of having a scope body, use a body with only structdecl. This improves performance
        and simplifies the code.
    */
    Scope body;
    /* The number of fields that are pointers or a type that is different on different platforms (i.e isize, usize, etc), recursive */
    u32 fields_num_pointers;
    /* The total size of all fields which are of a type that have the same size on all platforms (i.e u8, i32, etc), recursive */
    u32 fields_fixed_size_bytes;
    bool is_signed; /* Only default arithmetic types can be signed, everything else is unsigned */
};

struct StructField {
    BufferView name;
    VariableType type;
};

typedef enum {
    DECL_FLAG_NONE        = 0,
    DECL_FLAG_EXTERN      = 1 << 0,
    DECL_FLAG_EXPORT      = 1 << 1,
    DECL_FLAG_PUB         = 1 << 2,
    DECL_FLAG_CONST       = 1 << 3
} DeclFlag;

/*
    Note: When resolving AST, multiple threads can end up resolving the same expressions at the same time.
    This is intentional. Instead of using mutex for every expression and locking/unlocking everytime
    which uses more memory and affects performance, we assume such race conditions are rare and let them happen.
    TODO: Investigate possible deadlock when multiple threads resolve the same expression and that
    expression has a recursive dependency and the threads execute @ast_resolve at the same speed,
    leading to @ast_resolve running again for the same expression.
*/
struct LhsExpr {
    u8 decl_flags;
    BufferView var_name;
    VariableType type;
    nullable Ast *rhs_expr;
    u16 extern_index; /* Index to extern func, extern variable, etc; if applicable */
};

#define LHS_EXPR_IS_EXTERN(expr)    ((expr)->decl_flags & DECL_FLAG_EXTERN)
#define LHS_EXPR_IS_EXPORT(expr)    ((expr)->decl_flags & DECL_FLAG_EXPORT)
#define LHS_EXPR_IS_PUB(expr)       ((expr)->decl_flags & DECL_FLAG_PUB)
#define LHS_EXPR_IS_CONST(expr)     ((expr)->decl_flags & DECL_FLAG_CONST)

struct AssignmentExpr {
    Ast *lhs_expr;
    Ast *rhs_expr;
};

struct Import {
    BufferView path;
    ParserFileScopeReference *file_scope;
};

struct String {
    BufferView str;
};

struct Number {
    AmalNumber value;
    BufferView code_ref;
};

struct AstBool {
    bool value;
    BufferView code_ref;
};

struct Binop {
    Ast *lhs;
    Ast *rhs;
    BinopType type;
    /* Is the binop already ordered - no need to reorder it */
    bool grouped;
};

struct IfStatement {
    Ast *condition;
    Scope body;
    ElseIfStatement *else_if_stmt;
};

/* ElseIfStatement where condition == NULL is an Else statement */
struct ElseIfStatement {
    Ast *condition;
    Scope body;
    ElseIfStatement *next_else_if_stmt;
};

struct WhileStatement {
    Ast *condition;
    Scope body;
};

struct ReturnExpr {
    Ast *rhs_expr;
};

typedef struct {
    u32 num_pointers;
    u32 fixed_size;
} TypeSize;

typedef struct {
    jmp_buf env;
    amal_compiler *compiler; /* Borrowed */
    Parser *parser; /* Borrowed. This is the parser that belongs to the thread */
    /* Borrowed. This is the current scope. Note that this scope can belong to another parser (and thread) */
    Scope *scope;
} AstCompilerContext;

CHECK_RESULT int ast_create(ArenaAllocator *allocator, void *value, AstType type, Ast **result);
BufferView ast_get_name(Ast *self);

CHECK_RESULT int function_signature_init(FunctionSignature *self, ArenaAllocator *allocator);
/* Adds a copy of @param to the function signature parameter list */
CHECK_RESULT int function_signature_add_parameter(FunctionSignature *self, const FunctionParameter *param);
/* Adds a copy of @return_type to the function signature return type list */
CHECK_RESULT int function_signature_add_return_type(FunctionSignature *self, const VariableType *return_type);
CHECK_RESULT bool function_signature_equals(FunctionSignature *self, FunctionSignature *other);

void function_parameter_init(FunctionParameter *self);
CHECK_RESULT int funcdecl_init(FunctionDecl *self, FunctionSignature *signature, Scope *parent, ArenaAllocator *allocator);
CHECK_RESULT int funccall_init(FunctionCall *self, BufferView name, ArenaAllocator *allocator);
CHECK_RESULT int structdecl_init(StructDecl *self, Scope *parent, ArenaAllocator *allocator);
CHECK_RESULT int structdecl_add_field(StructDecl *self, StructField *field, ArenaAllocator *allocator);
LhsExpr* structdecl_get_field_by_name(StructDecl *self, BufferView field_name);
/* Copies @type */
void structfield_init(StructField *self, BufferView name, VariableType *type);
void lhsexpr_init(LhsExpr *self, DeclFlag decl_flag, BufferView var_name);
void assignmentexpr_init(AssignmentExpr *self, Ast *lhs_expr, Ast *rhs_expr);
void import_init(Import *self, BufferView path);
CHECK_RESULT int string_init(String *self, BufferView str);
void number_init(Number *self, AmalNumber *value, BufferView code_ref);
void ast_bool_init(AstBool *self, bool value, BufferView code_ref);
void variable_init(Variable *self, BufferView name);
void binop_init(Binop *self);
CHECK_RESULT int if_statement_init(IfStatement *self, Scope *parent, ArenaAllocator *allocator);
CHECK_RESULT int else_if_statement_init(ElseIfStatement *self, Scope *parent, ArenaAllocator *allocator);
CHECK_RESULT int while_statement_init(WhileStatement *self, Scope *parent, ArenaAllocator *allocator);
void return_expr_init(ReturnExpr *self, Ast *rhs_expr);

TypeSize resolved_type_get_byte_size(AstResolvedType *self);

CHECK_RESULT int scope_init(Scope *self, Scope *parent, ArenaAllocator *allocator);
CHECK_RESULT int file_scope_reference_init(FileScopeReference *self, BufferView canonical_path, ArenaAllocator *allocator);
CHECK_RESULT int scope_add_child(Scope *self, Ast *child);

/* longjump to compiler env on failure */
void scope_resolve(Scope *self, AstCompilerContext *context);

#endif