`timescale 1ns / 1ps
module InstDecode (
    input clk,
    // From prev stage
    input [31:0] prev_fetched_instruction,
    input [31:0] prev_PC_plus_4,
    // For hazard unit
    input [1:0] IFIDSrc,
    // From WB stage
    input WB_write_enable,
    input [4:0] WB_write_address,
    input [31:0] WB_write_data,
    // To IF stage
    output [1:0] PC_jump,
    output [31:0] jump_target,
    output [31:0] jump_register_target,
    // To hazard unit
    output is_loadword,
    // To next stage
    output is_branch,
    output WB_source,
    output memory_write,
    output [4:0] ALU_function,
    output ALU_source1,
    output ALU_source2,
    output register_write_destination_source,
    output register_write,
    output [31:0] PC_plus_4,
    output [31:0] register_file_read_A,
    output [31:0] register_file_read_B,
    output [4:0] shamt,
    output [31:0] extended_immediate,
    output [4:0] rs_address,
    output [4:0] rt_address,
    output [4:0] rd_address
);

    reg  [31:0] IFID_instruction;
    reg  [31:0] IFID_PC_plus_4;

    wire [ 5:0] opcode;
    wire [ 4:0] rs;
    wire [ 4:0] rt;
    wire [ 4:0] rd;
    wire [ 5:0] funct;
    wire [25:0] j_addr;
    wire [15:0] immediate;

    // Signals derived from instruction
    assign opcode = IFID_instruction[31:26];
    assign rs = IFID_instruction[25:21];
    assign rt = IFID_instruction[20:16];
    assign rd = IFID_instruction[15:11];
    assign shamt = IFID_instruction[10:6];
    assign funct = IFID_instruction[5:0];
    assign j_addr = IFID_instruction[25:0];
    assign immediate = IFID_instruction[15:0];

    // This output is directly derived from IFID regs
    assign jump_target = {IFID_PC_plus_4[31:28], j_addr, 2'b00};
    assign rs_address = rs;
    assign rt_address = rt;
    assign rd_address = rd;

    // Signals to connect from control unit to register file and immediate extend unit
    wire write_ra;
    wire ra_addr_source;
    wire extendop;

    ControlUnit control_unit (
        .opcode(opcode),
        .funct(funct),
        .PC_jump(PC_jump),
        .is_branch(is_branch),
        .is_loadword(is_loadword),
        .write_ra(write_ra),
        .ra_addr_source(ra_addr_source),
        .WB_source(WB_source),
        .memory_write(memory_write),
        .ALU_function(ALU_function),
        .ALU_source1(ALU_source1),
        .ALU_source2(ALU_source2),
        .register_write(register_write),
        .register_write_destination_source(register_write_destination_source),
        .extendop(extendop)
    );

    // Signal for register file
    wire [31:0] RF_read_A_out;
    wire [ 4:0] write_ra_addr_after_mux;

    assign write_ra_addr_after_mux = (ra_addr_source == 1'b0) ? 5'b11111 : rd;

    RegisterFile register_file (
        .clk(clk),
        .read_addr1(rs),
        .read_addr2(rt),
        .write_enable(WB_write_enable),
        .write_addr(WB_write_address),
        .write_data(WB_write_data),
        .write_ra(write_ra),
        .write_ra_addr(write_ra_addr_after_mux),
        .write_ra_data(IFID_PC_plus_4),
        .read_output1(RF_read_A_out),
        .read_output2(register_file_read_B)
    );

    assign register_file_read_A = RF_read_A_out;
    assign jump_register_target = RF_read_A_out;

    ImmediateExtender immediate_extender(
        .immediate(immediate),
        .extendop(extendop),
        .extended_immediate(extended_immediate)
    );

    always @(posedge clk) begin
        case (IFIDSrc)
            2'b00: begin
                IFID_instruction <= prev_fetched_instruction;
                IFID_PC_plus_4   <= prev_PC_plus_4;
            end
            2'b01: begin
                IFID_instruction <= 32'h00000000;
                IFID_PC_plus_4   <= 32'h00000000;
            end
            2'b10: begin
                IFID_instruction <= IFID_instruction;
                IFID_PC_plus_4   <= IFID_PC_plus_4;
            end
            default: begin
                IFID_instruction <= 32'h00000000;
                IFID_PC_plus_4   <= 32'h00000000;
            end
        endcase
    end

endmodule