#ifndef RISCV64_LLIR_H
#define RISCV64_LLIR_H

#include "IR.h" // 确保包含了您自己的IR头文件
#include <string>
#include <vector>
#include <memory>
#include <cstdint>
#include <map>

// 前向声明，避免循环引用
namespace sysy {
class Function;
class RISCv64ISel;
}

namespace sysy {

// 物理寄存器定义
enum class PhysicalReg {
    // --- 特殊功能寄存器 ---
    ZERO, RA, SP, GP, TP,

    // --- 整数寄存器 (按调用约定分组) ---
    // 临时寄存器 (调用者保存)
    T0, T1, T2, T3, T4, T5, T6,
    
    // 保存寄存器 (被调用者保存)
    S0, S1, S2, S3, S4, S5, S6, S7, S8, S9, S10, S11,
    
    // 参数/返回值寄存器 (调用者保存)
    A0, A1, A2, A3, A4, A5, A6, A7,

    // --- 浮点寄存器 ---
    // (保持您原有的 F0-F31 命名)
    F0, F1, F2, F3, F4, F5, F6, F7, F8, F9, F10, F11, 
    F12, F13, F14, F15, F16, F17, F18, F19, F20, F21, 
    F22, F23, F24, F25, F26, F27, F28, F29, F30, F31
};

// RISC-V 指令操作码枚举
enum class RVOpcodes {
    // 算术指令
    ADD, ADDI, ADDW, ADDIW, SUB, SUBW, MUL, MULW, DIV, DIVW, REM, REMW,
    // 逻辑指令
    XOR, XORI, OR, ORI, AND, ANDI,
    // 移位指令
    SLL, SLLI, SLLW, SLLIW, SRL, SRLI, SRLW, SRLIW, SRA, SRAI, SRAW, SRAIW,
    // 比较指令
    SLT, SLTI, SLTU, SLTIU,
    // 内存访问指令
    LW, LH, LB, LWU, LHU, LBU, SW, SH, SB, LD, SD,
    // 控制流指令
    J, JAL, JALR, RET,
    BEQ, BNE, BLT, BGE, BLTU, BGEU,
    // 伪指令
    LI, LA, MV, NEG, NEGW, SEQZ, SNEZ,
    // 函数调用
    CALL,
    // 特殊标记，非指令
    LABEL,
    // 新增伪指令，用于解耦栈帧处理
    FRAME_LOAD_W,  // 从栈帧加载 32位 Word (对应 lw)
    FRAME_LOAD_D,  // 从栈帧加载 64位 Doubleword (对应 ld)
    FRAME_STORE_W, // 保存 32位 Word 到栈帧 (对应 sw)
    FRAME_STORE_D, // 保存 64位 Doubleword 到栈帧 (对应 sd)
    FRAME_ADDR,  // 获取栈帧变量的地址
};

class MachineOperand;
class RegOperand;
class ImmOperand;
class LabelOperand;
class MemOperand;
class MachineInstr;
class MachineBasicBlock;
class MachineFunction;

// 操作数基类
class MachineOperand {
public:
    enum OperandKind { KIND_REG, KIND_IMM, KIND_LABEL, KIND_MEM };
    MachineOperand(OperandKind kind) : kind(kind) {}
    virtual ~MachineOperand() = default;
    OperandKind getKind() const { return kind; }
private:
    OperandKind kind;
};

// 寄存器操作数
class RegOperand : public MachineOperand {
public:
    // 构造虚拟寄存器
    RegOperand(unsigned vreg_num) 
        : MachineOperand(KIND_REG), vreg_num(vreg_num), is_virtual(true) {}

    // 构造物理寄存器
    RegOperand(PhysicalReg preg)
        : MachineOperand(KIND_REG), preg(preg), is_virtual(false) {}

    bool isVirtual() const { return is_virtual; }
    unsigned getVRegNum() const { return vreg_num; }
    PhysicalReg getPReg() const { return preg; }

    void setPReg(PhysicalReg new_preg) {
        preg = new_preg;
        is_virtual = false;
    }
private:
    unsigned vreg_num = 0;
    PhysicalReg preg = PhysicalReg::ZERO;
    bool is_virtual;
};

// 立即数操作数
class ImmOperand : public MachineOperand {
public:
    ImmOperand(int64_t value) : MachineOperand(KIND_IMM), value(value) {}
    int64_t getValue() const { return value; }
private:
    int64_t value;
};

// 标签操作数
class LabelOperand : public MachineOperand {
public:
    LabelOperand(const std::string& name) : MachineOperand(KIND_LABEL), name(name) {}
    const std::string& getName() const { return name; }
private:
    std::string name;
};

// 内存操作数, 表示 offset(base_reg)
class MemOperand : public MachineOperand {
public:
    MemOperand(std::unique_ptr<RegOperand> base, std::unique_ptr<ImmOperand> offset)
        : MachineOperand(KIND_MEM), base(std::move(base)), offset(std::move(offset)) {}
    RegOperand* getBase() const { return base.get(); }
    ImmOperand* getOffset() const { return offset.get(); }
private:
    std::unique_ptr<RegOperand> base;
    std::unique_ptr<ImmOperand> offset;
};

// 机器指令
class MachineInstr {
public:
    MachineInstr(RVOpcodes opcode) : opcode(opcode) {}

    RVOpcodes getOpcode() const { return opcode; }
    const std::vector<std::unique_ptr<MachineOperand>>& getOperands() const { return operands; }
    std::vector<std::unique_ptr<MachineOperand>>& getOperands() { return operands; }

    void addOperand(std::unique_ptr<MachineOperand> operand) {
        operands.push_back(std::move(operand));
    }
private:
    RVOpcodes opcode;
    std::vector<std::unique_ptr<MachineOperand>> operands;
};

// 机器基本块
class MachineBasicBlock {
public:
    MachineBasicBlock(const std::string& name, MachineFunction* parent) 
        : name(name), parent(parent) {}

    const std::string& getName() const { return name; }
    MachineFunction* getParent() const { return parent; }
    const std::vector<std::unique_ptr<MachineInstr>>& getInstructions() const { return instructions; }
    std::vector<std::unique_ptr<MachineInstr>>& getInstructions() { return instructions; }
    
    void addInstruction(std::unique_ptr<MachineInstr> instr) {
        instructions.push_back(std::move(instr));
    }
    
    std::vector<MachineBasicBlock*> successors;
    std::vector<MachineBasicBlock*> predecessors;
private:
    std::string name;
    std::vector<std::unique_ptr<MachineInstr>> instructions;
    MachineFunction* parent;
};

// 栈帧信息
struct StackFrameInfo {
    int locals_size = 0; // 仅为AllocaInst分配的大小
    int spill_size = 0; // 仅为溢出分配的大小
    int total_size = 0; // 总大小
    std::map<unsigned, int> alloca_offsets; // <AllocaInst的vreg, 栈偏移>
    std::map<unsigned, int> spill_offsets;  // <溢出vreg, 栈偏移>
    std::set<PhysicalReg> used_callee_saved_regs; // 使用的保存寄存器
};

// 机器函数
class MachineFunction {
public:
    MachineFunction(Function* func, RISCv64ISel* isel) : F(func), name(func->getName()), isel(isel) {}

    Function* getFunc() const { return F; }
    RISCv64ISel* getISel() const { return isel; }
    const std::string& getName() const { return name; }
    StackFrameInfo& getFrameInfo() { return frame_info; }
    const std::vector<std::unique_ptr<MachineBasicBlock>>& getBlocks() const { return blocks; }
    std::vector<std::unique_ptr<MachineBasicBlock>>& getBlocks() { return blocks; }

    void addBlock(std::unique_ptr<MachineBasicBlock> block) {
        blocks.push_back(std::move(block));
    }
private:
    Function* F;
    RISCv64ISel* isel; // 指向创建它的ISel，用于获取vreg映射等信息
    std::string name;
    std::vector<std::unique_ptr<MachineBasicBlock>> blocks;
    StackFrameInfo frame_info;
};

} // namespace sysy

#endif // RISCV64_LLIR_H