#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, S0, S1, A0, A1, A2, A3, A4, A5, A6, A7, S2, S3, S4, S5, S6, S7, S8, S9, S10, S11, T3, T4, T5, T6,
    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,  // 从栈帧加载 (AllocaInst)
    FRAME_STORE, // 保存到栈帧 (AllocaInst)
};

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, 栈偏移>
};

// 机器函数
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