#pragma once
/*
 This is the Lab1 skeleton cgen header. As given, it contains only basic
 functionality. You will need to add members to each of the classes
 to get them to perform their desired functions. Document your important
 design decisions below. We should be able to read your documentation and
 get a general overview of how your compiler generates code. For instance,
 how does your compiler generate structures for classes, how is inheritance
 modeled, how do you handle dynamic binding, etc.
*/

// ------------------ INSERT DESIGN DOCUMENTATION HERE --------------------- //
/*
 * Types for classes are declared when its CgenNode is installed to deal with circular dependencies
 *
 */
// ----------------------------- END DESIGN DOCS --------------------------- //

#include "cool_tree.h"
// #include "stringtab.h"
// #include "symtab.h"
#include <llvm/IR/IRBuilder.h>
#include <llvm/IR/LLVMContext.h>
#include <llvm/IR/Module.h>
#include <llvm/IR/DIBuilder.h>
#include "llvm/IR/DebugInfoMetadata.h"
#include <llvm/Support/raw_ostream.h>

#if defined _MSC_VER && !defined __clang__
#define __attribute(...)
#endif

class CgenNode;

#pragma region "CgenClassTable"
// CgenClassTable represents the top level of a Cool program, which is
// basically a list of classes. The class table is used to look up classes
// (CgenNodes) by name, and it also handles global code generation tasks.
// The CgenClassTable constructor is where you'll find the entry point for
// code generation for an entire Cool program.
class CgenClassTable : public cool::SymbolTable<CgenNode> {
public:
    // CgenClassTable constructor begins and ends the code generation process
    CgenClassTable(Classes);

private:
    // The following creates an inheritance graph from a list of classes.
    // The graph is implemented as a tree of `CgenNode', and class names
    // are placed in the base class symbol table.
    void install_basic_classes();

    void install_classes(Classes cs);

    void install_class(CgenNode *nd);

    void install_special_class(CgenNode *nd);

    void build_inheritance_tree();

    void set_relations(CgenNode *nd);

    // Create declarations for C runtime functions we need to generate code
    void setup_external_functions();

    void setup_classes(CgenNode *c, int depth);

    // Setup each class in the table and prepare for code generation phase
    void setup();

    // Code generation functions. You need to write these functions.
    void code_module();

    void code_classes(CgenNode *c);

    void code_constants();

    void code_main();

    /* Util functions */
    CgenNode *root(); // Get the root of the class Tree, i.e. Object
public:
    int get_num_classes() const { return current_tag; }

private:
    // Class lists and current class tag
    std::vector<CgenNode *> nds, special_nds;
    int current_tag;

public:
    // LLVM Util
    llvm::Function *create_llvm_function(const std::string &funcName,
                                         llvm::Type *retType,
                                         llvm::ArrayRef<llvm::Type *> argTypes,
                                         bool isVarArgs);

    llvm::Function *create_function(const std::string &name,
                                    llvm::Type *ret_type,
                                    llvm::ArrayRef<llvm::Type *> arg_types,
                                    llvm::GlobalValue::LinkageTypes linkage);

    llvm::Function *create_var_function(const std::string &name,
                                        llvm::Type *ret_type,
                                        llvm::ArrayRef<llvm::Type *> arg_types,
                                        llvm::GlobalValue::LinkageTypes
                                        linkage);

    // CgenClassTable owns the current LLVM module and everything attached.
    // One program, one class table, one module.
    llvm::LLVMContext context;
    llvm::IRBuilder<> builder;
    llvm::Module the_module;
    llvm::DIBuilder di_builder;
    llvm::DIFile* di_file;
    llvm::DICompileUnit* di_cu;

    llvm::Type *i1;
    llvm::Type *i8;
    llvm::PointerType *ptr;
    llvm::Type *i32;
    llvm::Type *void_;

    llvm::DIBasicType* di1;
    llvm::DIBasicType* di8;
    llvm::DIBasicType* di32;
    llvm::DIDerivedType* dptr;
    llvm::DIBasicType* dvoid;
};
#pragma endregion

#pragma region "CgenNode"
// Each CgenNode corresponds to a Cool class. As such, it is responsible for
// performing code generation on the class level. This includes laying out
// the class attributes, creating the necessary Types for the class and
// generating code for each of its methods.
class CgenNode : public class__class {
public:
    enum Basicness { Basic, NotBasic };

    CgenNode(Class_ c, Basicness bstatus, CgenClassTable *class_table)
        : class__class((const class__class &) *c), parentnd(0), children(0),
          basic_status(bstatus), class_table(class_table), tag(-1), ctx(class_table->context),
          i1(class_table->i1), i8(class_table->i8), ptr(class_table->ptr),
          i32(class_table->i32), void_(class_table->void_),
          type(nullptr), vtable_type(nullptr), vtable(nullptr),
          di_builder(class_table->di_builder)
    {}

    // Relationships with other nodes in the tree
    void set_parent(CgenNode *p)
    {
        assert(this->parentnd == nullptr && p != nullptr);
        p->children.push_back(this);
        this->parentnd = p;
    }

    std::optional<CgenNode*> get_direct_parent() const
    {
        if (parentnd->basic())
            return std::nullopt;
        return parentnd;
    }

    int basic() const { return basic_status == Basic; }
    std::vector<CgenNode *> &get_children() { return children; }
    void set_max_child(int mc) { max_child = mc; }
    int get_max_child() const { return max_child; }

    // Accessors for other provided fields
    int get_tag() const { return tag; }
    CgenClassTable *get_classtable() const { return class_table; }

    llvm::StructType *type;
    llvm::StructType *vtable_type;
    llvm::GlobalVariable *vtable;

    llvm::DIType* dtype;

    std::string get_type_name() const { return name->get_string(); }

    std::string get_vtable_type_name() const
    {
        return "_" + get_type_name() + "_vtable";
    }

    std::string get_vtable_name() const
    {
        return "_" + get_type_name() + "_vtable_prototype";
    }

    std::string get_init_function_name() const { return get_type_name() + "_new"; }

    // Class setup. You need to write the body of this function.
    void setup(int tag, int depth);

    // Layout the methods and attributes for code generation
    void layout_features();

    // Class codegen. You need to write the body of this function.
    void code_class();

    // Codegen for the init function of every class
    void code_init_function(CgenEnvironment *env) const;

    void make_vtable_type(CgenEnvironment *env);

    std::vector<std::pair<method_class *, llvm::FunctionType *> > methods;
    std::vector<std::pair<attr_class *, llvm::Type *> > attributes;

    llvm::Function *create_function(const std::string &fName, llvm::FunctionType *ty) const
    {
        return llvm::Function::Create(ty, llvm::GlobalValue::ExternalLinkage,
                                      this->name->get_string() + "_" + fName,
                                      this->class_table->the_module);
    }

    static llvm::FunctionType *create_functionty(llvm::Type *retType,
                                                 llvm::ArrayRef<llvm::Type *> argTypes)
    {
        return llvm::FunctionType::get(retType, argTypes, false);
    }

    void make_vtable_prototype(CgenEnvironment *env);

    void premake_type()
    {
        assert(type == nullptr);
        this->type = llvm::StructType::create(ctx,
                                              this->name->get_string());
        assert(type != nullptr);
    }
    void premake_debug_type();

    llvm::DIBuilder &di_builder;
    std::vector<llvm::Metadata*> dfeats;
private:
    CgenNode *parentnd; // Parent of class
    std::vector<CgenNode *> children; // Children of class
    Basicness basic_status; // `Basic' or 'NotBasic'
    CgenClassTable *class_table;
    // Class tag. Should be unique for each class in the tree
    int tag, max_child;
    llvm::LLVMContext &ctx;


    llvm::Type *i1;
    llvm::Type *i8;
    llvm::PointerType *ptr;
    llvm::Type *i32;
    llvm::Type *void_;

    void make_type_forrealz(CgenEnvironment *env);

    void make_debug_type(CgenEnvironment* env);
};
#pragma endregion

#pragma region "CgenEnvironment"
// CgenEnvironment provides the environment for code generation of a method.
// Its main task is to provide a mapping from Cool names to LLVM Values.
// This mapping needs to be maintained as scopes are entered and exited, new
// variables are declared, and so on. CgenEnvironment is also a good place
// to put non-local information you will need during code generation. Two
// examples are the current CgenNode and the current Function.
class CgenEnvironment {
public:
    // Class CgenEnvironment should be constructed by a class prior to code
    // generation for each method. You may need to add parameters to this
    // constructor.
    CgenEnvironment(CgenNode *cur_class)
        : cur_class(cur_class),
          class_table(*cur_class->get_classtable()),
          context(class_table.context),
          builder(class_table.builder), the_module(class_table.the_module), data_layout(&the_module),
          i1(class_table.i1), i8(class_table.i8), ptr(class_table.ptr),
          i32(class_table.i32), void_(class_table.void_)
    {
        vars.enterscope();
    }

    CgenNode *get_class() const { return cur_class; }
    void set_class(CgenNode *c) { cur_class = c; }

    // Must return the CgenNode for a class given the symbol of its name
    CgenNode *type_to_class(Symbol t) const;

    std::pair<llvm::Type*, llvm::Value*> find_in_scopes(Symbol name) const;
    std::optional<llvm::Value*> find_if_exists(Symbol name) const;


    void add_binding(Symbol name, llvm::Value *val_ptr)
    {
        vars.insert(name, val_ptr);
    }

    void open_scope() { vars.enterscope(); }
    void close_scope() { vars.exitscope(); }

    // LLVM Utils:
    // Create a new llvm function in the current module with class name appended
    llvm::Function *get_function(Symbol name) const
    {
        return the_module.getFunction(name->get_string());
    }

    // Insert a new BasicBlock at the end of the current function (the function
    // that builder is in)
    llvm::BasicBlock *new_bb_at_end(const std::string &name)
    {
        return llvm::BasicBlock::Create(context, name,
                                        builder.GetInsertBlock()->getParent());
    }

    // Insert an alloca instruction in the head BasicBlock of the current
    // function, such that this alloca is available in all BasicBlocks of the
    // function.
    llvm::AllocaInst *insert_alloca_at_head(llvm::Type *ty);

    llvm::AllocaInst *insert_alloca_at_head(llvm::Type *, llvm::Twine const &);

    // Get or insert a BasicBlock with the name "abort" which calls the ::abort
    // function. This block will be inserted at the end of the given function,
    // without moving the builder.
    llvm::BasicBlock *get_or_insert_abort_block(llvm::Function *f);

    llvm::Type *as_type(Symbol) const __attribute((pure));
    llvm::Type *as_mostly_boxed_type(Symbol) const __attribute((pure));
    llvm::DIType* as_boxed_dtype(Symbol) const __attribute((pure));

    llvm::Value* default_value(llvm::Type*) const __attribute((pure));
    // Generate any code necessary to convert from given Value* to
    // dest_type, assuming it has already been checked to be compatible
    llvm::Value* conform(llvm::Value* src, llvm::Type* dest_type) __attribute((pure));
private:
    // mapping from variable names to memory locations
    cool::SymbolTable<llvm::Value> vars;
    CgenNode *cur_class;

public:
    CgenClassTable &class_table;
    // These are references to the current LLVM context and module,
    // and they point to the ones in CgenClassTable.
    llvm::LLVMContext &context;
    llvm::IRBuilder<> &builder;
    llvm::Module &the_module;
    llvm::DataLayout data_layout;

    llvm::Type *i1;
    llvm::Type *i8;
    llvm::PointerType *ptr;
    llvm::Type *i32;
    llvm::Type *void_;
};
#pragma endregion