// 局部变量究竟保存在哪里？
// alloca会返回一个指针，这个指针指向一块可用的内存区域
// 这个该如何体现在assemb上？
// NOTE: 参数类型需要额外考虑！

#include "codegen.hpp"

#include "BasicBlock.h"
#include "Constant.h"
#include "Function.h"
#include "GlobalVariable.h"
#include "Instruction.h"
#include "Type.h"
#include "Value.h"

#include <cstdint>
#include <cstring>
#include <string>
#include <utility>

// $r0          $zero       constant 0
// $r1          $ra         return address
// $r2          $tp         thread pointer
// $r3          $sp         stack pointer
// $r4 - $r5    $a0 - $a1   argument, return value
// $r6 - $r11   $a2 - $a7   argument
// $r12 - $r20  $t0 - $t8   temporary
// $r21                     saved
// $r22         $fp         frame pointer
// $r23 - $r31  $s0 - $s8   static

using std::to_string;

class Reg {
  public:
    Reg(int index) : id(index) {}

    int id;

    string print() {
        if (id == 0)
            return "$zero";
        if (id == 1)
            return "$ra";
        if (id == 2)
            return "$tp";
        if (id == 3)
            return "$sp";
        if (4 <= id and id <= 11)
            return "$a" + to_string(id - 4);
        if (12 <= id and id <= 20)
            return "$t" + to_string(id - 12);
        if (id == 22)
            return "$fp";
        assert(false);
    }
};

void
CodeGen::getPhiMap() {
    phi_map.clear();
    for (auto &func : m->get_functions())
        for (auto &bb : func.get_basic_blocks())
            for (auto &instr : bb.get_instructions()) {
                if (not instr.is_phi())
                    continue;
                for (int i = 0; i < instr.get_num_operand() / 2; i++) {
                    auto key =
                        static_cast<BasicBlock *>(instr.get_operand(2 * i + 1));
                    auto value =
                        std::make_pair(&instr, instr.get_operand(2 * i));
                    phi_map[key].push_back(value);
                }
            }
}

void
CodeGen::run() {
    // TODO: implement
    // 以下内容生成 int main() { return 0; } 的汇编代码
    getPhiMap();
    output.push_back(".text");
    // global variables
    for (auto &globl : m->get_global_variable()) {
        auto type = globl.get_type()->get_pointer_element_type();

        output.push_back(".comm " + globl.get_name() + ", " +
                         to_string(typeLen(type)));
    }
    // arguments: stack transfer
    for (auto &func : m->get_functions())
        if (not func.is_declaration())
            compute_arg_info(&func);
    // funtions
    for (auto &func : m->get_functions()) {
        if (not func.is_declaration()) {
            cur_func = &func;
            output.push_back("");
            output.push_back(".globl " + func.get_name());
            output.push_back(".type " + func.get_name() + ", @function");
            output.push_back(func.get_name() + ":");
            // sp and fp balabala
            stackMemAlloc();

            for (auto &bb : func.get_basic_blocks()) {
                output.push_back(label_in_assem(&bb) + ":");
                for (auto &instr : bb.get_instructions()) {
                    IR2assem(instr, bb);
                }
            }
            // restore the stack
            stackMemDealloc();
        }
    }
}

void
CodeGen::ptrContent2reg(Value *ptr, int id) {
    auto ele_tp = ptr->get_type()->get_pointer_element_type();
    assert(ele_tp);
    bool is_float = ele_tp->is_float_type();
    string instr_ir = (is_float ? "fld" : "ld");
    string reg_name = (is_float ? "$fa" : "$a") + to_string(id);
    string suff = suffix(ele_tp);
    if (dynamic_cast<GlobalVariable *>(ptr)) {
        output.push_back("la.local " + reg_name + ", " + ptr->get_name());
        output.push_back(instr_ir + suff + " " + reg_name + ", " + reg_name +
                         ", 0");
    } else if (dynamic_cast<AllocaInst *>(ptr)) {
        /* auto alloc_instr = static_cast<AllocaInst *>(ptr);
         * string suff = suffix(alloc_instr->get_alloca_type()); */
        output.push_back(instr_ir + suff + " " + reg_name + ", $fp, -" +
                         to_string(off[ptr]));
    } else if (dynamic_cast<GetElementPtrInst *>(ptr)) {
        // auto GEP_instr = static_cast<GetElementPtrInst *>(ptr);
        output.push_back("ld.d " + reg_name + ", $fp, -" + to_string(off[ptr]));
        output.push_back(instr_ir + suff + " " + reg_name + ", " + reg_name +
                         ", 0");
    } else
        assert(false && "unknown type");
}

void
CodeGen::value2reg(Value *v, int id) {
    bool is_float = v->get_type()->is_float_type();
    auto reg_name = (is_float ? "$fa" : "$a") + to_string(id);
    if (dynamic_cast<Constant *>(v)) {
        auto constant = static_cast<Constant *>(v);
        if (dynamic_cast<ConstantInt *>(constant)) {
            int k = static_cast<ConstantInt *>(constant)->get_value();
            if ((k & 0xfff) != k) {
                output.push_back("lu12i.w " + reg_name + ", " +
                                 to_string(k >> 12));
                output.push_back("ori " + reg_name + ", " + reg_name + ", " +
                                 to_string(k & 0xfff));
            } else
                output.push_back("ori " + reg_name + ", $r0, " + to_string(k));
        } else if (dynamic_cast<ConstantFP *>(constant)) {
            // move the binary code to int-reg, then use movgr2fr to move the
            // value to float-reg
            string tmp_ireg = "$t0";
            float k = static_cast<ConstantFP *>(constant)->get_value();
            int hex_int = *(uint32_t *)&k;
            if ((hex_int & 0xfff) != hex_int)
                output.push_back("lu12i.w " + tmp_ireg + ", " +
                                 to_string(hex_int >> 12));
            if (hex_int & 0xfff)
                output.push_back("ori " + tmp_ireg + ", " + tmp_ireg + ", " +
                                 to_string(hex_int & 0xfff));
            output.push_back("movgr2fr.w " + reg_name + ", " + tmp_ireg);
            // output.push_back("ffint.s.w " + reg_name + ", " + reg_name);

        } else
            assert(false && "wait for completion");
    } else if (dynamic_cast<GlobalVariable *>(v)) {
        output.push_back("la.local " + reg_name + ", " + v->get_name());
    } else if (dynamic_cast<AllocaInst *>(v)) {
        // auto alloc_instr = dynamic_cast<AllocaInst *>(v);
        // give the stack address
        output.push_back("addi.d " + reg_name + ", $fp, -" + to_string(off[v]));
    } else if (dynamic_cast<Argument *>(v)) {
        auto args = cur_func->get_args();
        int id = 1;
        for (auto iter = args.begin(); id <= args.size(); ++iter, ++id)
            if (*iter == v)
                break;
        string instr_ir = is_float ? "fld" : "ld";
        output.push_back(instr_ir + suffix(v->get_type()) + " " + reg_name +
                         ", $fp, " + to_string(func_arg_off[cur_func][id]));
    } else {
        string instr_ir = is_float ? "fld" : "ld";
        output.push_back(instr_ir + suffix(v->get_type()) + " " + reg_name +
                         ", $fp, -" + to_string(off[v]));
    }
}

void
CodeGen::compute_arg_info(Function *func) {
    if (func_arg_off.find(func) != func_arg_off.end())
        return;
    auto func_tp = func->get_function_type();
    auto &arg_off = func_arg_off[func];
    int argN = 0, arg_id = func->get_num_of_args();
    // reserve space
    for (auto iter = func_tp->param_end(); iter != func_tp->param_begin();) {
        --iter;
        auto tplen = typeLen(*iter);
        argN = ALIGN(argN, tplen) + tplen;
        arg_off[arg_id--] = argN;
    }
    for (arg_id = 1; arg_id <= func->get_num_of_args(); ++arg_id)
        arg_off[arg_id] = argN - arg_off[arg_id];
    func_arg_N[func] = STACK_ALIGN(argN);
}

void
CodeGen::stackMemDealloc() {
    output.push_back("# epilog");
    // 7: return value should be determined already!
    // output.push_back("addi.w $a0, $zero, 0");
    output.push_back(cur_func->get_name() + "_end:");
    output.push_back("ld.d $ra, $fp, -8");
    output.push_back("addi.d $sp, $sp, " + to_string(stackN));
    output.push_back("jr $ra");
}

// the addr for opk is: fp - off[opk]
void
CodeGen::stackMemAlloc() {
    // preserved for ra
    stackN = 8;
    off.clear();
    for (auto &bb : cur_func->get_basic_blocks())
        for (auto &instr : bb.get_instructions()) {
            if (no_stack_alloca(&instr))
                continue;
            int tplen;
            if (instr.is_alloca()) {
                auto alloc_instr = static_cast<AllocaInst *>(&instr);
                tplen = typeLen(alloc_instr->get_alloca_type());
            } else {
                auto type = instr.get_type();
                tplen = typeLen(type);
            }
            stackN = ALIGN(stackN, tplen) + tplen;
            off[&instr] = stackN;
        }
    stackN = STACK_ALIGN(stackN);
    output.push_back("# prolog");
    output.push_back("addi.d $sp, $sp, -" + to_string(stackN));
    output.push_back("addi.d $fp, $sp, " + to_string(stackN));
    output.push_back("st.d $ra, $fp, -8");
}
void
CodeGen::IR2assem(FpToSiInst *instr) {
    assert(instr->get_operand(0)->get_type() == m->get_float_type());
    assert(instr->get_dest_type() == m->get_int32_type());
    value2reg(instr->get_operand(0));
    output.push_back("ftint.w.s $fa0, $fa0");
    output.push_back("movfr2gr.s $a0, $fa0");
}
void
CodeGen::IR2assem(SiToFpInst *instr) {
    assert(instr->get_operand(0)->get_type() == m->get_int32_type());
    assert(instr->get_dest_type() == m->get_float_type());
    output.push_back("movgr2fr.w $fa0, $a0");
    output.push_back("ffint.s.w $fa0, $fa0");
}

string
CodeGen::bool2branch(Instruction *instr) {
    assert(instr->get_type() == m->get_int1_type());
    auto icmp_instr = dynamic_cast<CmpInst *>(instr);
    auto fcmp_instr = dynamic_cast<FCmpInst *>(instr);
    assert(icmp_instr or fcmp_instr);

    string instr_ir;
    bool reverse = false;
    if (icmp_instr) {
        switch (icmp_instr->get_cmp_op()) {
            case CmpInst::EQ:
                instr_ir = "beq";
                break;
            case CmpInst::NE: {
                instr_ir = "bne";
                if (instr->get_operand(1) == ConstantInt::get(0, m) and
                    dynamic_cast<Instruction *>(instr->get_operand(0))
                        ->is_zext()) {
                    // something like:
                    // %op0 = icmp slt i32 1, 2 # deepest
                    // %op1 = zext i1 %op0 to i32
                    // %op2 = icmp ne i32 %op1, 0
                    // br i1 %op2, label %label3, label %label5
                    auto deepest = static_cast<Instruction *>(
                        static_cast<Instruction *>(instr->get_operand(0))
                            ->get_operand(0));
                    return bool2branch(deepest);
                }
            } break;
            case CmpInst::GT:
                instr_ir = "blt";
                reverse = true;
                break;
            case CmpInst::GE:
                instr_ir = "bge";
                break;
            case CmpInst::LT:
                instr_ir = "blt";
                break;
            case CmpInst::LE:
                instr_ir = "bge";
                reverse = true;
                break;
        }
        value2reg(instr->get_operand(0), 0);
        value2reg(instr->get_operand(1), 1);
        return instr_ir + (reverse ? " $a1, $a0," : " $a0, $a1,");
    } else {
        switch (fcmp_instr->get_cmp_op()) {
            case FCmpInst::EQ:
                instr_ir = "fcmp.ceq.s $fcc0";
                break;
            case FCmpInst::NE:
                instr_ir = "fcmp.cun.s $fcc0";
                break;
            case FCmpInst::GT:
                instr_ir = "fcmp.cle.s $fcc0";
                reverse = true;
                break;
            case FCmpInst::GE:
                instr_ir = "fcmp.clt.s $fcc0";
                reverse = true;
                break;
            case FCmpInst::LT:
                instr_ir = "fcmp.clt.s $fcc0";
                break;
            case FCmpInst::LE:
                instr_ir = "fcmp.cle.s $fcc0";
                break;
        }
        value2reg(instr->get_operand(0), 0);
        value2reg(instr->get_operand(1), 1);
        output.push_back(instr_ir + ", " + "$fa0, $fa1");
        return (reverse ? "bceqz $fcc0," : "bcnez $fcc0,");
    }
}

void
CodeGen::IR2assem(BranchInst *instr) {
    if (instr->is_cond_br()) {
        auto TBB = static_cast<BasicBlock *>(instr->get_operand(1));
        auto FBB = static_cast<BasicBlock *>(instr->get_operand(2));
        // value2reg(instr->get_operand(0));
        string instr_ir =
            bool2branch(static_cast<Instruction *>(instr->get_operand(0)));
        output.push_back(instr_ir + " " + label_in_assem(TBB));
        output.push_back("b " + label_in_assem(FBB));
    } else {
        auto bb = static_cast<BasicBlock *>(instr->get_operand(0));
        output.push_back("b " + label_in_assem(bb));
    }
}

void
CodeGen::IR2assem(CallInst *instr) {
    auto func = static_cast<Function *>(instr->get_operand(0));
    // stack space allocation
    output.push_back("addi.d $sp, $sp, -" + to_string(func_arg_N[func]));
    string instr_ir, suff, reg;
    // place the arguments
    for (int i = 1; i < instr->get_num_operand(); i++) {
        auto arg = instr->get_operand(i);
        // auto tplen = typeLen(arg->get_type());
        instr_ir = (arg->get_type()->is_float_type() ? "fst" : "st");
        suff = suffix(arg->get_type());
        reg = (arg->get_type()->is_float_type() ? "$fa0" : "$a0");
        value2reg(arg);
        output.push_back(instr_ir + suff + " " + reg + ", $sp, " +
                         to_string(func_arg_off[func][i]));
    }
    output.push_back("bl " + func->get_name());
    output.push_back("addi.d $sp, $sp, " + to_string(func_arg_N[func]));
    output.push_back("addi.d $fp, $sp, " + to_string(stackN));
}

void
CodeGen::IR2assem(BinaryInst *instr) {
    auto is_float = instr->get_type()->is_float_type();
    value2reg(instr->get_operand(0), 0);
    value2reg(instr->get_operand(1), 1);

    string instr_ir = instr->get_instr_op_name();
    string suff = suffix(instr->get_type());
    output.push_back(instr_ir + suff +
                     (is_float ? " $fa0, $fa0, $fa1" : " $a0, $a0, $a1"));
}

void
CodeGen::IR2assem(GetElementPtrInst *instr) {
    value2reg(instr->get_operand(0), 0);
    Type *type = instr->get_operand(0)->get_type();
    for (int i = 1; i < instr->get_num_operand(); i++) {
        int size;
        if (type->is_array_type()) {
            size = type->get_array_element_type()->get_size();
            type = type->get_array_element_type();
        } else if (type->is_pointer_type()) {
            size = type->get_size();
            type = type->get_pointer_element_type();
        } else
            assert(false && "GEP translation error");
        value2reg(instr->get_operand(i), 1);
        value2reg(ConstantInt::get(size, m), 2);
        output.push_back("mul.w $a1, $a1, $a2");
        output.push_back("add.d $a0, $a0, $a1");
    }
}

void
CodeGen::IR2assem(LoadInst *instr) {
    // move the address to a0
    ptrContent2reg(instr->get_lval());

    assert(instr->get_type() == instr->get_load_type());
}

void
CodeGen::IR2assem(StoreInst *instr) {
    value2reg(instr->get_rval(), 0);
    value2reg(instr->get_lval(), 1);
    bool is_float = instr->get_rval()->get_type()->is_float_type();
    string instr_ir = (is_float ? "fst" : "st");
    string suff = suffix(instr->get_rval()->get_type());
    string reg = (is_float ? "$fa0" : "$a0");
    output.push_back(instr_ir + suff + " " + reg + ", $a1, 0");
}

void
CodeGen::IR2assem(ReturnInst *instr) {
    if (not instr->is_void_ret()) {
        auto value = instr->get_operand(0);
        value2reg(value);
    }
    output.push_back("b " + cur_func->get_name() + "_end");
}

void
CodeGen::IR2assem(Instruction &instr, BasicBlock &bb) {
    if (instr.is_br() or instr.is_ret())
        copystmt(&bb);
    output.push_back("# " + instr.print());
    switch (instr.get_instr_type()) {
        case Instruction::ret:
            IR2assem(static_cast<ReturnInst *>(&instr));
            break;
        case Instruction::br:
            IR2assem(static_cast<BranchInst *>(&instr));
            break;
            // Standard binary operators
        case Instruction::add:
        case Instruction::sub:
        case Instruction::mul:
        case Instruction::sdiv:
            // float binary operators
        case Instruction::fadd:
        case Instruction::fsub:
        case Instruction::fmul:
        case Instruction::fdiv:
            IR2assem(static_cast<BinaryInst *>(&instr));
            break;
            // Memory operators
        case Instruction::alloca:
            IR2assem(static_cast<AllocaInst *>(&instr));
            break;
        case Instruction::load:
            IR2assem(static_cast<LoadInst *>(&instr));
            break;
        case Instruction::store:
            IR2assem(static_cast<StoreInst *>(&instr));
            break;
            // Other operators
        case Instruction::cmp:
            IR2assem(static_cast<CmpInst *>(&instr));
            break;
        case Instruction::fcmp:
            IR2assem(static_cast<FCmpInst *>(&instr));
            break;
        case Instruction::phi:
            IR2assem(static_cast<PhiInst *>(&instr));
            break;
        case Instruction::call:
            IR2assem(static_cast<CallInst *>(&instr));
            break;
        case Instruction::getelementptr:
            IR2assem(static_cast<GetElementPtrInst *>(&instr));
            break;
        case Instruction::zext:
            IR2assem(static_cast<ZextInst *>(&instr));
            break;
        case Instruction::fptosi:
            IR2assem(static_cast<FpToSiInst *>(&instr));
            break;
        case Instruction::sitofp:
            IR2assem(static_cast<SiToFpInst *>(&instr));
            break;
    }
    // assert(false && "This ")
    if (not instr.is_void() and not no_stack_alloca(&instr) and
        not instr.is_alloca() and not instr.is_phi())
        back2stack(&instr);
}