#include "RISCv32Backend.h"
#include <sstream>
#include <algorithm>

namespace sysy {

const std::vector<RISCv32CodeGen::PhysicalReg> RISCv32CodeGen::allocable_regs = {
    PhysicalReg::T0, PhysicalReg::T1, PhysicalReg::T2, PhysicalReg::T3,
    PhysicalReg::T4, PhysicalReg::T5, PhysicalReg::T6,
    PhysicalReg::A0, PhysicalReg::A1, PhysicalReg::A2, PhysicalReg::A3,
    PhysicalReg::A4, PhysicalReg::A5, PhysicalReg::A6, PhysicalReg::A7
};

std::string RISCv32CodeGen::reg_to_string(PhysicalReg reg) {
    switch (reg) {
        case PhysicalReg::T0: return "t0"; case PhysicalReg::T1: return "t1";
        case PhysicalReg::T2: return "t2"; case PhysicalReg::T3: return "t3";
        case PhysicalReg::T4: return "t4"; case PhysicalReg::T5: return "t5";
        case PhysicalReg::T6: return "t6"; case PhysicalReg::A0: return "a0";
        case PhysicalReg::A1: return "a1"; case PhysicalReg::A2: return "a2";
        case PhysicalReg::A3: return "a3"; case PhysicalReg::A4: return "a4";
        case PhysicalReg::A5: return "a5"; case PhysicalReg::A6: return "a6";
        case PhysicalReg::A7: return "a7";
        default: return "";
    }
}

std::string RISCv32CodeGen::code_gen() {
    std::stringstream ss;
    ss << ".text\n";
    ss << module_gen();
    return ss.str();
}

std::string RISCv32CodeGen::module_gen() {
    std::stringstream ss;
    // 生成全局变量（数据段）
    for (const auto& global : module->getGlobals()) {
        ss << ".data\n";
        ss << ".globl " << global->getName() << "\n";
        ss << global->getName() << ":\n";
        ss << "  .word 0\n"; // 假设初始化为0
    }
    // 生成函数（文本段）
    ss << ".text\n";
    for (const auto& func : module->getFunctions()) {
        ss << function_gen(func.second.get());
    }
    return ss.str();
}

std::string RISCv32CodeGen::function_gen(Function* func) {
    std::stringstream ss;
    // 函数标签
    ss << ".globl " << func->getName() << "\n";
    ss << func->getName() << ":\n";
    // 序言：保存 ra，分配堆栈
    bool is_leaf = true; // 简化假设
    ss << "  addi sp, sp, -16\n";
    ss << "  sw ra, 12(sp)\n";
    // 寄存器分配
    auto alloc = register_allocation(func);
    // 生成基本块代码
    for (const auto& bb : func->getBasicBlocks()) {
        ss << basicBlock_gen(bb.get(), alloc);
    }
    // 结尾：恢复 ra，释放堆栈
    ss << "  lw ra, 12(sp)\n";
    ss << "  addi sp, sp, 16\n";
    ss << "  ret\n";
    return ss.str();
}

std::string RISCv32CodeGen::basicBlock_gen(BasicBlock* bb, const RegAllocResult& alloc) {
    std::stringstream ss;
    ss << bb->getName() << ":\n";
    for (const auto& inst : bb->getInstructions()) {
        auto riscv_insts = instruction_gen(inst.get());
        for (const auto& riscv_inst : riscv_insts) {
            ss << "  " << riscv_inst.opcode;
            for (size_t i = 0; i < riscv_inst.operands.size(); ++i) {
                if (i > 0) ss << ", ";
                if (riscv_inst.operands[i].kind == Operand::Kind::Reg) {
                    auto it = alloc.reg_map.find(riscv_inst.operands[i].value);
                    if (it != alloc.reg_map.end()) {
                        ss << reg_to_string(it->second);
                    } else {
                        auto stack_it = alloc.stack_map.find(riscv_inst.operands[i].value);
                        if (stack_it != alloc.stack_map.end()) {
                            ss << stack_it->second << "(sp)";
                        } else {
                            ss << "%" << riscv_inst.operands[i].value->getName();
                        }
                    }
                } else if (riscv_inst.operands[i].kind == Operand::Kind::Imm) {
                    ss << riscv_inst.operands[i].label;
                } else {
                    ss << riscv_inst.operands[i].label;
                }
            }
            ss << "\n";
        }
    }
    return ss.str();
}

std::vector<RISCv32CodeGen::RISCv32Inst> RISCv32CodeGen::instruction_gen(Instruction* inst) {
    std::vector<RISCv32Inst> insts;
    if (auto bin = dynamic_cast<BinaryInst*>(inst)) {
        std::string opcode;
        if (bin->getKind() == BinaryInst::kAdd) opcode = "add";
        else if (bin->getKind() == BinaryInst::kSub) opcode = "sub";
        else if (bin->getKind() == BinaryInst::kMul) opcode = "mul";
        else return insts; // 其他操作未实现
        insts.emplace_back(opcode, std::vector<Operand>{
            {Operand::Kind::Reg, bin},
            {Operand::Kind::Reg, bin->getLhs()},
            {Operand::Kind::Reg, bin->getRhs()}
        });
    } else if (auto load = dynamic_cast<LoadInst*>(inst)) {
        insts.emplace_back("lw", std::vector<Operand>{
            {Operand::Kind::Reg, load},
            {Operand::Kind::Label, load->getPointer()->getName()}
        });
    } else if (auto store = dynamic_cast<StoreInst*>(inst)) {
        insts.emplace_back("sw", std::vector<Operand>{
            {Operand::Kind::Reg, store->getValue()},
            {Operand::Kind::Label, store->getPointer()->getName()}
        });
    }
    return insts;
}

void RISCv32CodeGen::eliminate_phi(Function* func) {
    // TODO: 实现 phi 指令消除
}

std::map<Instruction*, std::set<Value*>> RISCv32CodeGen::liveness_analysis(Function* func) {
    std::map<Instruction*, std::set<Value*>> live_sets;
    // TODO: 实现活跃性分析
    return live_sets;
}

std::map<Value*, std::set<Value*>> RISCv32CodeGen::build_interference_graph(
    const std::map<Instruction*, std::set<Value*>>& live_sets) {
    std::map<Value*, std::set<Value*>> graph;
    // TODO: 实现干扰图构建
    return graph;
}

RISCv32CodeGen::RegAllocResult RISCv32CodeGen::register_allocation(Function* func) {
    RegAllocResult result;
    // TODO: 实现寄存器分配
    return result;
}

} // namespace sysy