#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 "binop_type.h"
#include "ssa/ssa.h"

#include <setjmp.h>

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

typedef struct Ast Ast;
typedef struct FunctionSignature FunctionSignature;
typedef struct FunctionCall FunctionCall;
typedef struct StructDecl StructDecl;
typedef struct StructField StructField;
typedef struct LhsExpr LhsExpr;
typedef struct AssignmentExpr AssignmentExpr;
typedef struct Import Import;
typedef struct String String;
typedef struct Variable Variable;
typedef struct Number Number;
typedef struct Binop Binop;
typedef struct IfStatement IfStatement;
typedef struct ElseIfStatement ElseIfStatement;
typedef struct WhileStatement WhileStatement;

/* 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;
    Variable *variable;
    Binop *binop;
    IfStatement *if_stmt;
    WhileStatement *while_stmt;
} 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_VARIABLE,
    AST_BINOP,
    AST_IF_STATEMENT,
    AST_WHILE_STATEMENT
} AstType;

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

typedef struct {
    LhsExpr *type;
    AstResolveStatus status;
    Parser *parser; /* Borrowed. This is the parser that is currently parsing the expression */
} AstResolveData;

struct Ast {
    AstValue value;
    AstType type;
    AstResolveData resolve_data;
    SsaRegister ssa_reg;
};

DefineNullablePtrType(Ast);

struct Scope {
    Buffer/*<Ast*>*/ ast_objects;
    HashMap/*(key=BufferView, value=Ast<LhsExpr>*)*/ named_objects;
    Scope *parent;
    /* Is null unless the scope is a file scope, in which case this is the parser that owns the scope */
    Parser *parser;
};

struct FileScopeReference {
    Parser *parser;
    Buffer canonical_path;
};

struct Variable {
    BufferView name;
    Ast *resolved_var; /* resolved_var will always be a LhsExpr after resolved */
};

struct FunctionSignature {
    int params; /* TODO: Implement signatures */
    bool resolved;
};

struct FunctionDecl {
    FunctionSignature *signature;
    Scope body;
    SsaFuncIndex ssa_func_index;
};

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

struct StructDecl {
    Scope body;
};

struct StructField {
    BufferView name;
    Variable type;
};

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

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

/*
    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 {
    bool is_extern;
    bool is_pub;
    bool is_const;
    BufferView var_name;
    VariableType type;
    NullablePtr(Ast) rhs_expr;
};

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

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

struct String {
    BufferView str;
};

struct Number {
    union {
        i64 integer;
        f64 floating;
    } value;
    bool is_integer;
    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;
};

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),
        as such, @parser and scope_get_parser(@scope) parser may not be the same.
    */
    Scope *scope;
} AstCompilerContext;

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

void function_signature_init(FunctionSignature *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);
void structfield_init(StructField *self, BufferView name, BufferView type_name);
void lhsexpr_init(LhsExpr *self, bool is_extern, bool is_pub, bool is_const, 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, i64 value, bool is_integer, 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);

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);



CHECK_RESULT bool resolved_type_is_func_decl(Ast *self);

#endif