#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 "ast.hpp"

#include <algorithm>
#include <cstdint>
#include <cstring>
#include <sstream>
#include <string>
#include <sys/types.h>
#include <tuple>
#include <utility>
#include <vector>

// $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;

pair<string, bool>
CodeGen::getRegName(Value *v, int i) const {
    assert(i == 0 or i == 1);
    bool find;
    string name;
    bool is_float = v->get_type()->is_float_type();
    auto regmap = RA.get();
    if (regmap.find(v) == regmap.end()) {
        name = tmpregname(i, is_float);
        find = false;
    } else {
        auto regid = regmap.find(v)->second;
        name = regname(regid, is_float);
        find = true;
    }
    return {name, find};
}

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()) {
            LRA.run(&func);
            RA.LinearScan(LRA.get());

            std::cout << "register map for function: " << func.get_name()
                      << std::endl;
            RA.print([](int i) { return regname(i); });

            auto regmap = RA.get();
            for (auto [_, op] : LRA.get()) {
                if (regmap.find(op) == regmap.end())
                    std::cout << "no reg belongs to " << op->get_name()
                              << std::endl;
            }

            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();
        }
    }
    // read only data
    output.push_back(".section .rodata");
    for (auto [constant, name] : ROdata) {
        auto int_ = dynamic_cast<ConstantInt *>(constant);
        auto float_ = dynamic_cast<ConstantFP *>(constant);
        string value;
        assert((int_ || float_) && "wrong type");
        if (int_) {
            value = to_string(int_->get_value());
        } else {
            std::stringstream ss;
            float v = (float_->get_value());
            ss << std::hex << *(uint32_t *)&v;
            value = "0x" + ss.str();
        }
        output.push_back(name + ": ");
        output.push_back(".word " + value);
    }
}

void
CodeGen::ptrContent2reg(Value *ptr, string reg_name) {
    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 suff = suffix(ele_tp);

    auto [addr_reg, find] = getRegName(ptr, 1);
    if (not find)
        addr_reg = "$t1";

    if (dynamic_cast<GlobalVariable *>(ptr)) {
        output.push_back("la.local " + addr_reg + ", " + ptr->get_name());
        output.push_back(instr_ir + suff + " " + reg_name + ", " + addr_reg +
                         ", 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.at(ptr)));
    } else if (dynamic_cast<GetElementPtrInst *>(ptr)) {
        // auto GEP_instr = static_cast<GetElementPtrInst *>(ptr);
        if (not find) {
            output.push_back("ld.d " + addr_reg + ", $fp, -" +
                             to_string(off.at(ptr)));
        }
        output.push_back(instr_ir + suff + " " + reg_name + ", " + addr_reg +
                         ", 0");
    } else
        assert(false && "unknown type");
}

string
CodeGen::value2reg(Value *v, int i) {
    bool is_float = v->get_type()->is_float_type();
    string tmp_ireg = "$t" + to_string(i);
    auto [reg_name, find] = getRegName(v, i);
    if (find)
        return reg_name;
    // now is the stack allocation case
    if (dynamic_cast<Constant *>(v)) {
        if (v == ConstantInt::get(0, m))
            return "$zero";
        auto constant = static_cast<Constant *>(v);
#ifdef __RO_PART__
        if (ROdata.find(constant) == ROdata.end())
            ROdata[constant] = ".LC" + to_string(ROdata.size());
        string instr_ir, addr = ROdata[constant];
        if (dynamic_cast<ConstantInt *>(constant))
            instr_ir = "ld.w";
        else if (dynamic_cast<ConstantFP *>(constant))
            instr_ir = "fld.s";
        else
            assert(false && "wait for completion");
        // bug here: maybe
        output.push_back("la.local " + tmp_ireg + ", " + addr);
        output.push_back(instr_ir + " " + reg_name + ", " + tmp_ireg + ", 0");
#else

        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
            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");
#endif
    } 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.at(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.at(cur_func).at(id)));
    } else {
        string instr_ir = is_float ? "fld" : "ld";
        output.push_back(instr_ir + suffix(v->get_type()) + " " + reg_name +
                         ", $fp, -" + to_string(off.at(v)));
    }
    return reg_name;
}

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] = argN;
}

// the addr for opk is: fp - off[opk]
void
CodeGen::stackMemAlloc() {
    // preserved for $ra and $fp
    stackN = 16;
    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("st.d $ra, $sp," + to_string(stackN - 8));
    output.push_back("st.d $fp, $sp, " + to_string(stackN - 16));
    output.push_back("addi.d $fp, $sp, " + to_string(stackN));
}

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

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());
    string f_reg = value2reg(instr->get_operand(0));
    auto [i_reg, _] = getRegName(instr);
    output.push_back("ftintrz.w.s " + f_reg + ", " + f_reg);
    output.push_back("movfr2gr.s " + i_reg + ", " + f_reg);
    gencopy(instr, i_reg);
}
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());
    string i_reg = value2reg(instr->get_operand(0));
    auto [f_reg, _] = getRegName(instr);
    output.push_back("movgr2fr.w " + f_reg + ", " + i_reg);
    output.push_back("ffint.s.w " + f_reg + ", " + f_reg);
    gencopy(instr, f_reg);
}

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;
        }
        auto reg1 = value2reg(instr->get_operand(0), 0);
        auto reg2 = value2reg(instr->get_operand(1), 1);
        return instr_ir + " " +
               (reverse ? (reg2 + ", " + reg1) : (reg1 + ", " + reg2)) + ",";
    } else {
        assert(false && "not implemented");
        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;
        }
        auto reg1 = value2reg(instr->get_operand(0), 0);
        auto reg2 = value2reg(instr->get_operand(1), 1);
        output.push_back(instr_ir + ", " + reg1 + ", " + reg2);
        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));
    auto func_argN = func_arg_N.at(func);
    // analyze the registers that need to be stored
    int cur_i = LRA.get_instr_id().at(instr);
    auto regmap = RA.get();
    //
    int storeN = 0;
    vector<std::tuple<Value *, string, int>> store_record;
    for (auto [op, interval] : LRA.get_interval_map()) {
        if (RA::no_reg_alloca(op))
            continue;
        if (interval.i <= cur_i and cur_i <= interval.j) {
            int tplen = typeLen(op->get_type());
            storeN = ALIGN(storeN, tplen) + tplen;
            auto name = regname(regmap.at(op), op->get_type()->is_float_type());
            store_record.push_back({op, name, storeN});
        }
    }
    int totalN = STACK_ALIGN(ALIGN(storeN, 8) + func_argN);
    // cout << "debug: " << STACK_ALIGN(12) << endl;
    // stack space allocation
    output.push_back("addi.d $sp, $sp, -" + to_string(totalN));
    string instr_ir, suff, reg;
    // place the reserved regs
    for (auto [op, reg, off] : store_record) {
        instr_ir = (op->get_type()->is_float_type() ? "fst" : "st");
        suff = suffix(op->get_type());
        output.push_back(instr_ir + suff + " " + reg + ", $sp, " +
                         to_string(totalN - off));
    }
    // 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 = value2reg(arg);
        output.push_back(instr_ir + suff + " " + reg + ", $sp, " +
                         to_string(func_arg_off.at(func).at(i)));
    }
    output.push_back("bl " + func->get_name());
    // bug here: maybe
    gencopy(instr, instr->get_type()->is_float_type() ? "$fa0" : "$a0");
    // restore the reserved regs
    for (auto [op, reg, off] : store_record) {
        instr_ir = (op->get_type()->is_float_type() ? "fld" : "ld");
        suff = suffix(op->get_type());
        output.push_back(instr_ir + suff + " " + reg + ", $sp, " +
                         to_string(totalN - off));
    }
    output.push_back("addi.d $sp, $sp, " + to_string(totalN));
    // output.push_back("addi.d $fp, $sp, " + to_string(stackN));
}

void
CodeGen::IR2assem(BinaryInst *instr) {
    auto reg1 = value2reg(instr->get_operand(0), 0);
    auto reg2 = value2reg(instr->get_operand(1), 1);
    auto [reg, _] = getRegName(instr);

    string instr_ir = instr->get_instr_op_name();
    if (instr_ir == "sdiv")
        instr_ir = "div";
    string suff = suffix(instr->get_type());
    output.push_back(instr_ir + suff + " " + reg + ", " + reg1 + ", " + reg2);
    gencopy(instr, reg);
}

void
CodeGen::IR2assem(GetElementPtrInst *instr) {
    assert(instr->get_num_operand() <= 3);
    auto addr_reg = value2reg(instr->get_operand(0), 0);
    assert(addr_reg == "$t0");

    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");
        if (size != 4) {
            // getelementptr [5 x i32], [5 x i32]* @w, i32 0, i32 4
            assert(instr->get_operand(i) == ConstantInt::get(0, m) &&
                   "cminus support only 1 dimension array, so first offset is "
                   "must 0");
            continue;
        }
        auto off_reg = value2reg(instr->get_operand(i), 1);
        // value2reg(ConstantInt::get(size, m), 2);
        output.push_back("slli.d " + off_reg + ", " + off_reg + ", 2");
        output.push_back("add.d " + addr_reg + ", " + addr_reg + ", " +
                         off_reg);
    }
    gencopy(instr, addr_reg);
}

void
CodeGen::IR2assem(LoadInst *instr) {
    auto [reg, find] = getRegName(instr);
    ptrContent2reg(instr->get_lval(), reg);
    if (not find) // this if is just for logically clear
        last_reg = reg;
}

void
CodeGen::IR2assem(StoreInst *instr) {
    auto reg1 = value2reg(instr->get_rval(), 0);
    auto reg2 = 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());
    output.push_back(instr_ir + suff + " " + reg1 + ", " + reg2 + ", 0");
}

void
CodeGen::IR2assem(ReturnInst *instr) {
    if (not instr->is_void_ret()) {
        auto value = instr->get_operand(0);
        auto is_float = value->get_type()->is_float_type();
        auto reg = value2reg(value);
        if (is_float and reg != "$fa0")
            output.push_back("fmov.s $fa0, " + reg);
        else if (not is_float and reg != "$a0")
            output.push_back("or $a0, $zero " + reg);
    }

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