#ifndef CODEGEN_HPP
#define CODEGEN_HPP

#include "BasicBlock.h"
#include "Constant.h"
#include "Function.h"
#include "IRprinter.h"
#include "Instruction.h"
#include "Module.h"
#include "Value.h"
#include "liverange.hpp"
#include "regalloc.hpp"

#include <functional>
#include <string>

#define __PRINT_ORI__
// #define __RO_PART__
#define __PRINT_COMMENT__
// #a = 8, #t = 9, reserve $t0, $t1 for temporary
#define R_USABLE (17 - 2)
// #fa = 8, #ft=16, reserve $ft0, $ft1 for temporary
#define FR_USABLE (24 - 2)
#define ARG_R 8

#include <map>
#include <ostream>
#include <vector>

#define ALIGN(x, align) ((((x) / align) + (((x) % align) ? 1 : 0)) * align)
// #define STACK_ALIGN(x) (((x / 16) + (x % 16 ? 1 : 0)) * 16)
#define STACK_ALIGN(x) ALIGN(x, 16)

#define CONST_0 ConstantInt::get(0, m)
#define FP "$fp"
#define SP "$sp"
#define RA_reg "$ra"

using std::map;
using std::string;
using std::vector;

class CodeGen {
  public:
    CodeGen(Module *m_)
        : cmp_zext_cnt(0)
        , m(m_)
        , LRA(m_, phi_map)
        , RA_int(R_USABLE, false)
        , RA_float(FR_USABLE, true) {}

    string print() {
        string result;
        for (auto line : output) {
#ifndef __PRINT_COMMENT__
            if (line.find("#") != string::npos)
                continue;
#endif
            if (line.find(":") == string::npos and line != "")
                result += "\t"; // 添加缩进
            result += line + "\n";
        }
        return result;
    }

    void run();

  private:
    //
    // data member
    //
    std::map<Value *, unsigned int> off; // stack offset to $fp
    std::map<Function *, std::map<int, int>> func_arg_off; // to $sp
    // total space for args, NOTE: not aligned
    std::map<Function *, int> func_arg_N;
    std::map<BasicBlock *, std::vector<std::pair<Value *, Value *>>> phi_map;
    std::map<Constant *, std::string> ROdata;
    unsigned int stackN; // function local vars and so on
    uint cmp_zext_cnt;

    Function *cur_func;

    Module *m;
    vector<string> output;
    // register allocation
    LRA::LiveRangeAnalyzer LRA;
    LRA::LVITS LVITS_int, LVITS_float;
    RA::RegAllocator RA_int, RA_float;
    // some instruction has lvalue, but is stack-allocated,
    // we need this variable to track the reg name which has rvalue.
    // this variable is maintain by gencopy() and LoadInst.
    string last_reg;

    //
    // function member
    //
    void stackMemAlloc();
    void stackMemDealloc();
    // In the case of register allocation, this function will return the
    // allocated register for that value, if the value possesses no register,
    // choose from from $t0 or $t1 based on id, or the input string(not "")
    __attribute__((warn_unused_result)) string value2reg(Value *,
                                                         int i = 0,
                                                         string = "");
    // load the content in ptr to specified register.
    void ptrContent2reg(Value *, string);

    void pass_arguments(CallInst *);
    void compute_arg_info(Function *);
    string bool2branch(Instruction *);

    void getPhiMap();
    void copystmt(BasicBlock *bb) {
        // all the phi instruction is transformed to copy-stmt
        for (auto &copy : phi_map[bb]) {
            output.push_back("# " + print_as_op(copy.first, false) + " = " +
                             print_as_op(copy.second, false));
            auto lvalue = static_cast<Instruction *>(copy.first);
            auto src_reg = value2reg(copy.second);
            if (gencopy(lvalue, src_reg) == false)
                back2stack(lvalue);
        }
    }

    // this is a decorated version of push_back, it checks intermediat number in
    // the instruction, make sure imm will not overflow
    void makeSureInRange(string instr_ir,
                         string reg1,
                         string reg2,
                         int imm,
                         string tinstr,
                         int bits = 12,
                         string treg = "$t0",
                         bool u = false) {
        auto [l, h] = immRange(bits, u);
        if (l <= imm and imm <= h)
            output.push_back(instr_ir + " " + reg1 + ", " + reg2 + ", " +
                             to_string(imm));
        else {
            assert(value2reg(ConstantInt::get(imm, m), 0, treg) == treg);
            output.push_back(tinstr + " " + reg1 + ", " + reg2 + ", " + treg);
        }
    }

    // if reg-allocation, store to the specific register
    // or is stack-allocation, set last_reg for back2stack()
    bool gencopy(Value *lhs, string rhs_reg) {
        auto [lhs_reg, find] = getRegName(lhs);
        if (not find) {
            // wait for back2stack() to store back
            last_reg = rhs_reg;
            return false;
        }
        auto is_float = lhs->get_type()->is_float_type();
        gencopy(lhs_reg, rhs_reg, is_float);
        return true;
    }
    void gencopy(string lhs_reg, string rhs_reg, bool is_float) {
        if (rhs_reg != lhs_reg) {
            if (is_float)
                output.push_back("fmov.s " + lhs_reg + ", " + rhs_reg);
            else
                output.push_back("or " + lhs_reg + ", $zero, " + rhs_reg);
        }
    }

    string label_in_assem(BasicBlock *bb) const {
        return cur_func->get_name() + bb->get_name().substr(5);
    }
    static int typeLen(Type *type) {
        if (type->is_float_type())
            return 4;
        else if (type->is_integer_type()) {
            if (static_cast<IntegerType *>(type)->get_num_bits() == 32)
                return 4;
            else
                return 1;
        } else if (type->is_pointer_type())
            return 8;
        else if (type->is_array_type()) {
            auto arr_tp = static_cast<ArrayType *>(type);
            int n = arr_tp->get_num_of_elements();
            return n * typeLen(arr_tp->get_element_type());
        } else {
            assert(false && "unexpected case while computing type-length");
        }
    }

    // assert the value needed to store back is in `last_reg`, according to the
    // value type
    void back2stack(Instruction *instr) {
        // std::cerr <<  instr->print() << std::endl;
        auto type = instr->get_type();
        string instr_ir = type->is_float_type() ? "fst" : "st";
        string suff = suffix(type);
        // string reg = type->is_float_type() ? "$fa0" : "$a0";
        makeSureInRange(instr_ir + suff,
                        last_reg,
                        FP,
                        -off.at(instr),
                        instr_ir + "x" + suff);
        // string addr = "$fp, -" + std::to_string(off.at(instr));
        // output.push_back(instr_ir + suff + " " + last_reg + ", " + addr);
    }

    static string suffix(Type *type) {
        int len = typeLen(type);
        switch (len) {
            case 1:
                return ".b";
            case 2:
                return ".h";
            case 4:
                return type->is_float_type() ? ".s" : ".w";
            case 8:
                return ".d";
        }
        assert(false && "no such suffix");
    }

    bool no_stack_alloca(Instruction *instr) const {
        if (instr->is_void())
            return true;
        if (instr->is_fcmp() or instr->is_cmp() or instr->is_zext())
            return true;
        auto regmap = (instr->get_type()->is_float_type() ? RA_float.get()
                                                          : RA_int.get());
        if (regmap.find(instr) != regmap.end())
            return true;

        return false;
    }
    string tmpregname(int i, bool is_float) const {
        assert(i == 0 or i == 1);
        return (is_float ? "$ft" : "$t") + to_string(i);
    }

    static string regname(uint i, bool is_float) {
        string name;
        if (is_float) {
            // assert(false && "not implemented!");
            if (1 <= i and i <= 8)
                name = "$fa" + to_string(i - 1);
            else if (9 <= i and i <= FR_USABLE)
                name = "$ft" + to_string(i - 9 + 2);
            else
                name = "WRONG_REG_" + to_string(i);
        } else {
            if (1 <= i and i <= 8)
                name = "$a" + to_string(i - 1);
            else if (9 <= i and i <= R_USABLE)
                name = "$t" + to_string(i - 9 + 2);
            else
                name = "WRONG_REG_" + to_string(i);
        }
        return name;
    }

    static pair<int, int> immRange(int bit, bool u) {
        pair<int, int> res;
        if (u) {
            res.first = 0;
            res.second = (1 << bit) - 1;
        } else {
            bit--;
            res.first = -(1 << bit);
            res.second = (1 << bit) - 1;
        }

        return res;
    };

    pair<string, bool> getRegName(Value *, int = 0) const;

    void IR2assem(Instruction &, BasicBlock &);
    void IR2assem(LoadInst *);
    void IR2assem(StoreInst *);
    void IR2assem(ReturnInst *);
    void IR2assem(GetElementPtrInst *);
    void IR2assem(CallInst *);
    void IR2assem(BranchInst *);
    void IR2assem(BinaryInst *);
    void IR2assem(FpToSiInst *);
    void IR2assem(SiToFpInst *);
    void IR2assem(PhiInst *) {}
    void IR2assem(ZextInst *);
    // The Instructions below will do nothing
    void IR2assem(AllocaInst *) {}
    // integration with BranchInst and ZextInst
    void IR2assem(CmpInst *) {}
    void IR2assem(FCmpInst *) {}
};

#endif