#ifndef RISCV32_BACKEND_H
#define RISCV32_BACKEND_H

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

namespace sysy {

class RISCv32CodeGen {
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
    };

    // Move DAGNode and RegAllocResult to public section
    struct DAGNode {
        enum NodeKind { CONSTANT, LOAD, STORE, BINARY, CALL, RETURN, BRANCH, ALLOCA_ADDR }; // Added ALLOCA_ADDR
        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";
                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
    };

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

    std::string code_gen();
    std::string module_gen();
    std::string function_gen(Function* func);
    std::string basicBlock_gen(BasicBlock* bb, const RegAllocResult& alloc);

    // DAG related
    std::vector<std::unique_ptr<DAGNode>> build_dag(BasicBlock* bb);
    void select_instructions(DAGNode* node, const RegAllocResult& alloc);
    void emit_instructions(DAGNode* node, std::vector<std::string>& insts, const RegAllocResult& alloc, std::set<DAGNode*>& emitted_nodes); // Add emitted_nodes set

    // Register Allocation related
    std::map<Instruction*, std::set<std::string>> liveness_analysis(Function* func);
    std::map<std::string, std::set<std::string>> build_interference_graph(
        const std::map<Instruction*, std::set<std::string>>& live_sets);
    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
};

} // namespace sysy

#endif // RISCV32_BACKEND_H