#ifndef RISCV64_BACKEND_H
#define RISCV64_BACKEND_H

#include "IR.h"
#include <string>
#include <vector>
#include <map>
#include <set>
#include <memory>
#include <iostream>
#include <functional> // For std::function

extern int DEBUG;
extern int DEEPDEBUG;



namespace sysy {

// 为活跃性分析的结果定义一个结构体，以同时持有 live_in 和 live_out 集合
struct LivenessResult {
    std::map<Instruction*, std::set<std::string>> live_in;
    std::map<Instruction*, std::set<std::string>> live_out;
};

class RISCv64CodeGen {
public:
    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
    };

    // Move DAGNode and RegAllocResult to public section
    struct DAGNode {
        enum NodeKind { CONSTANT, LOAD, STORE, BINARY, CALL, RETURN, BRANCH, ALLOCA_ADDR, UNARY };
        NodeKind kind;
        Value* value = nullptr; // For IR Value
        std::string inst;       // Generated RISC-V instruction(s) for this node
        std::string result_vreg; // Virtual register assigned to this node's result
        std::vector<DAGNode*> operands;
        std::vector<DAGNode*> users; // For debugging and potentially optimizations
        DAGNode(NodeKind k) : kind(k) {}

        // Debugging / helper
        std::string getNodeKindString() const {
            switch (kind) {
                case CONSTANT: return "CONSTANT";
                case LOAD: return "LOAD";
                case STORE: return "STORE";
                case BINARY: return "BINARY";
                case CALL: return "CALL";
                case RETURN: return "RETURN";
                case BRANCH: return "BRANCH";
                case ALLOCA_ADDR: return "ALLOCA_ADDR";
                case UNARY: return "UNARY";
                default: return "UNKNOWN";
            }
        }
    };

    struct RegAllocResult {
        std::map<std::string, PhysicalReg> vreg_to_preg; // Virtual register to Physical Register mapping
        std::map<Value*, int> stack_map; // Value (AllocaInst) to stack offset
        int stack_size = 0; // Total stack frame size for locals and spills
    };

    RISCv64CodeGen(Module* mod) : module(mod) {}

    std::string code_gen();
    std::string module_gen();
    std::string function_gen(Function* func);
    // 修改 basicBlock_gen 的声明，添加 int block_idx 参数
    std::string basicBlock_gen(BasicBlock* bb, const RegAllocResult& alloc, int block_idx);

    // DAG related
    std::vector<std::unique_ptr<DAGNode>> build_dag(BasicBlock* bb);
    void select_instructions(DAGNode* node, const RegAllocResult& alloc);
    // 改变 emit_instructions 的参数，使其可以直接添加汇编指令到 main ss
    void emit_instructions(DAGNode* node, std::stringstream& ss, const RegAllocResult& alloc, std::set<DAGNode*>& emitted_nodes);

    // Register Allocation related
    LivenessResult liveness_analysis(Function* func);
    std::map<std::string, std::set<std::string>> build_interference_graph(const LivenessResult& liveness);
    void color_graph(std::map<std::string, PhysicalReg>& vreg_to_preg,
                     const std::map<std::string, std::set<std::string>>& interference_graph);
    RegAllocResult register_allocation(Function* func);
    void eliminate_phi(Function* func); // Phi elimination is typically done before DAG building

    // Utility
    std::string reg_to_string(PhysicalReg reg);
    void print_dag(const std::vector<std::unique_ptr<DAGNode>>& dag, const std::string& bb_name);

private:
    static const std::vector<PhysicalReg> allocable_regs;
    std::map<Value*, std::string> value_vreg_map; // Maps IR Value* to its virtual register name
    Module* module;
    int vreg_counter = 0; // Counter for unique virtual register names
    int alloca_offset_counter = 0; // Counter for alloca offsets

    // 新增一个成员变量来存储当前函数的所有 DAGNode，以确保其生命周期贯穿整个函数代码生成
    // 这样可以在多个 BasicBlock_gen 调用中访问到完整的 DAG 节点
    std::vector<std::unique_ptr<DAGNode>> current_function_dag_nodes;

    // 为空标签定义一个伪名称前缀，加上块索引以确保唯一性
    const std::string ENTRY_BLOCK_PSEUDO_NAME = "entry_block_";

    // !!! 修改：get_operand_node 辅助函数现在需要传入 value_to_node 和 nodes_storage 的引用
    // 因为它们是 build_dag 局部管理的
    DAGNode* get_operand_node(
        Value* val_ir,
        std::map<Value*, DAGNode*>& value_to_node,
        std::vector<std::unique_ptr<DAGNode>>& nodes_storage
    );

    // !!! 新增：create_node 辅助函数也需要传入 value_to_node 和 nodes_storage 的引用
    // 并且它应该不再是 lambda，而是一个真正的成员函数
    DAGNode* create_node(
        DAGNode::NodeKind kind,
        Value* val,
        std::map<Value*, DAGNode*>& value_to_node,
        std::vector<std::unique_ptr<DAGNode>>& nodes_storage
    );

    std::vector<std::unique_ptr<Instruction>> temp_instructions_storage; // 用于存储 build_dag 中创建的临时 BinaryInst
};

} // namespace sysy

#endif // RISCV64_BACKEND_H