`ifndef HAZARD_SV
`define HAZARD_SV

module hazard_unit (
    // (I)Decode/EXecute stage
    input  logic        idex_MemRead,
    input  logic [4:0]  idex_Rt,
    input  logic [4:0]  idex_Rs,

    // EXecute/MEMory stage
    input  logic        exmem_RegWrite,
    input  logic [4:0]  exmem_Rd,

    // MEMory/WriteBack stage
    input  logic        memwb_RegWrite,
    input  logic [4:0]  memwb_Rd,

    // (I)Fetch/(I)Decode stage
    input  logic [4:0]  ifid_Rs,
    input  logic [4:0]  ifid_Rt,

    // branch / jump control
    input  logic        branch_taken,
    input  logic        jump_taken,

    // outputs
    output logic [1:0]  forwardA,
    output logic [1:0]  forwardB,
    output logic        stallF, // Stall OC register
    output logic        stallD, // Stall Instruction Memory
    output logic        flushE, // Flush ALU inputs
    output logic        flushD  // Flush
);

    // Forwarding encoding:
    // 00 = no forwarding
    // 01 = forward from MEM/WB
    // 10 = forward from EX/MEM

    always_comb begin
        // defaults
        forwardA = 2'b00;
        forwardB = 2'b00;
        stallF   = 1'b0;
        stallD   = 1'b0;
        flushE   = 1'b0;
        flushD   = 1'b0;

        // Forwarding for ALU src A
        if (exmem_RegWrite && (exmem_Rd != 5'd0) && (exmem_Rd == idex_Rs))
            forwardA = 2'b10;
        else if (memwb_RegWrite && (memwb_Rd != 5'd0) && (memwb_Rd == idex_Rs))
            forwardA = 2'b01;

        // Forwarding for ALU src B
        if (exmem_RegWrite && (exmem_Rd != 5'd0) && (exmem_Rd == idex_Rt))
            forwardB = 2'b10;
        else if (memwb_RegWrite && (memwb_Rd != 5'd0) && (memwb_Rd == idex_Rt))
            forwardB = 2'b01;

        // -------------------------
        // Load-use hazard stall
        // if instruction in EX is a load and
        // next instruction uses its result
        // -------------------------
        if (idex_MemRead &&
            ((idex_Rt == ifid_Rs) || (idex_Rt == ifid_Rt))) begin
            stallF = 1'b1;
            stallD = 1'b1;
            flushE = 1'b1;   // insert bubble into EX
        end

        // -------------------------
        // Control hazard flush
        // branch/jump resolved in EX
        // -------------------------
        if (branch_taken || jump_taken) begin
            flushD = 1'b1;   // flush wrong instruction in decode
            flushE = 1'b1;   // flush wrong instruction in execute
        end
    end

endmodule
`endif