module core(
    input clk,
    input reset,
    output dummy_out,

    output reg [31:0] mem_inst_addr,
    input [31:0] mem_inst_data,

    output reg [31:0] mem_data_addr,
    output reg [31:0] mem_data_wdata,
    input [31:0] mem_data_rdata,
    output reg mem_data_en,
    output reg mem_data_we,
    input mem_data_valid,
    input mem_data_done
);

// Register File
reg [31:0] regfile [0:31];
initial regfile[0] = 32'h00000000;

// Registers
reg [31:0] r_if_pc = 0, r_id_pc, r_ex_pc, r_mem_pc, r_wb_pc;
reg r_id_stall, r_ex_stall, r_mem_stall, r_wb_stall;
reg [31:0] r_id_inst, r_ex_inst, r_mem_inst, r_wb_inst;
reg [4:0] r_ex_rd, r_mem_rd, r_wb_rd;
reg r_ex_alu_seed;
reg [2:0] r_ex_aluop;
reg [31:0] r_ex_s1, r_ex_s2, r_mem_s1, r_mem_s2;
reg [31:0] r_mem_alu_out, r_wb_alu_out;
reg r_mem_alu_zero;
reg r_ex_jump;
reg r_ex_store, r_mem_store;
reg r_ex_load, r_mem_load;

// IF
reg s_if_halt;
reg [31:0] s_if_next_pc;
reg [31:0] s_if_inst;
reg s_if_stall;

always @(*) begin
    s_if_halt = 0;

    if (r_ex_jump) begin
        s_if_next_pc    = s_ex_alu_out;
        s_if_stall      = 1'b1;
    end else begin
        s_if_next_pc    = r_if_pc + 4;
        s_if_stall      = 1'b0;
    end

    mem_inst_addr   = r_if_pc;
    s_if_inst       = mem_inst_data;
end

// ID
reg s_id_halt;
reg [6:0] s_id_opcode;
reg [2:0] s_id_funct3;
reg [6:0] s_id_funct7;
reg [4:0] s_id_rd, s_id_rs1, s_id_rs2;
reg [31:0] s_id_immed_itype, s_id_immed_stype, s_id_immed_utype, s_id_immed_btype, s_id_immed_jtype;
reg [31:0] s_id_s1, s_id_s2;
reg [2:0] s_id_aluop;
reg s_id_alu_seed;
reg s_id_invalid;
reg s_id_jump, s_id_branch;
reg s_id_store, s_id_load;

// RV32I / RV64I / RV32M
localparam  OP_LUI      = 7'b0110111,
            OP_AUIPC    = 7'b0010111,
            OP_JAL      = 7'b1101111,
            OP_JALR     = 7'b1100111,
            OP_BRANCH   = 7'b1100011,
            OP_LOAD     = 7'b0000011,
            OP_STORE    = 7'b0100011,
            OP_IMM      = 7'b0010011,
            OP_ALU      = 7'b0110011,
            OP_FENCE    = 7'b0001111,
            OP_SYSTEM   = 7'b1110011;
// RV64M
// localparam  OP_???????? = 7'b0111011;
// RV32A / RV64A
// localparam  OP_ATOMIC   = 7'b0101111;
// TODO: add opcodes for other extensions

// ALU OPCODES
localparam  ALUOP_ADD   = 0,
            ALUOP_XOR   = 1,
            ALUOP_OR    = 2,
            ALUOP_AND   = 3,
            ALUOP_SL    = 4,
            ALUOP_SR    = 5,
            ALUOP_SLT   = 6,
            ALUOP_SLTU  = 7;


always @(*) begin
    s_id_halt = 0;
    s_id_invalid = 0;
    s_id_store = 0;
    s_id_load = 0;

    s_id_opcode = r_id_inst[6:0];
    s_id_rd     = r_id_inst[11:7];
    s_id_rs1    = r_id_inst[19:15];
    s_id_rs2    = r_id_inst[24:20];
    s_id_funct3 = r_id_inst[14:12];
    s_id_funct7 = r_id_inst[31:25];

    s_id_immed_itype = {{20{r_id_inst[31]}}, r_id_inst[31:20]};
    s_id_immed_stype = {{20{r_id_inst[31]}}, r_id_inst[31:25], r_id_inst[11:7]};
    s_id_immed_utype = {r_id_inst[31:12], 12'b0};
    s_id_immed_btype = {{19{r_id_inst[31]}}, r_id_inst[31], r_id_inst[7], r_id_inst[30:25], r_id_inst[11:8], 1'b0};
    s_id_immed_jtype = {{11{r_id_inst[31]}}, r_id_inst[31], r_id_inst[19:12], r_id_inst[20], r_id_inst[30:21], 1'b0};

    case (s_id_opcode)
        OP_LUI: begin
            s_id_s1 = 32'h00000000;
            s_id_s2 = s_id_immed_utype;
            s_id_aluop = ALUOP_ADD;
            s_id_alu_seed = 0;
            s_id_jump = 0;
            s_id_branch = 0;
        end
        OP_AUIPC: begin
            s_id_s1 = r_id_pc;
            s_id_s2 = s_id_immed_utype;
            s_id_aluop = ALUOP_ADD;
            s_id_alu_seed = 0;
            s_id_jump = 0;
            s_id_branch = 0;
        end
        OP_JAL: begin
            s_id_s1 = r_id_pc;
            s_id_s2 = s_id_immed_jtype;
            s_id_aluop = ALUOP_ADD;
            s_id_alu_seed = 0;
            s_id_jump = 1;
            s_id_branch = 0;
        end
        OP_JALR: begin
            s_id_s1 = regfile[s_id_rs1];
            s_id_s2 = s_id_immed_itype;
            s_id_aluop = ALUOP_ADD;
            s_id_alu_seed = 0;
            s_id_jump = 1;
            s_id_branch = 0;
        end
        // OP_BRANCH: begin
            
        // end
        // OP_LOAD: begin
            
        // end
        // OP_STORE: begin
            
        // end
        OP_IMM: begin
            s_id_s1     = regfile[s_id_rs1];
            s_id_s2     = s_id_immed_itype;
            s_id_jump = 0;
            s_id_branch = 0;
            casex ({s_id_funct3, s_id_funct7})
                10'b000xxxxxxx: begin   s_id_aluop = ALUOP_ADD;     s_id_alu_seed = 1'b0;   end // ADDI
                10'b010xxxxxxx: begin   s_id_aluop = ALUOP_SLT;     s_id_alu_seed = 1'bx;   end // SLTI
                10'b011xxxxxxx: begin   s_id_aluop = ALUOP_SLTU;    s_id_alu_seed = 1'bx;   end // SLTUI
                10'b100xxxxxxx: begin   s_id_aluop = ALUOP_XOR;     s_id_alu_seed = 1'bx;   end // XORI
                10'b110xxxxxxx: begin   s_id_aluop = ALUOP_OR;      s_id_alu_seed = 1'bx;   end // ORI
                10'b111xxxxxxx: begin   s_id_aluop = ALUOP_AND;     s_id_alu_seed = 1'bx;   end // ANDI
                10'b0010000000: begin   s_id_aluop = ALUOP_SL;      s_id_alu_seed = 1'bx;   end // SLLI
                10'b1010000000: begin   s_id_aluop = ALUOP_SR;      s_id_alu_seed = 1'b0;   end // SRLI
                10'b1010100000: begin   s_id_aluop = ALUOP_SR;      s_id_alu_seed = 1'b1;   end // SRAI
                default: begin
                    s_id_s1 = 32'hxxxxxxxx;
                    s_id_s2 = 32'hxxxxxxxx;
                    s_id_invalid = 1;
                end
            endcase
        end
        OP_ALU: begin
            s_id_s1     = regfile[s_id_rs1];
            s_id_s2     = regfile[s_id_rs2];
            s_id_jump = 0;
            s_id_branch = 0;
            case ({s_id_funct3, s_id_funct7})
                10'b0000000000: begin   s_id_aluop = ALUOP_ADD;     s_id_alu_seed = 1'b0;   end // ADD
                10'b0000100000: begin   s_id_aluop = ALUOP_ADD;     s_id_alu_seed = 1'b1;   end // SUB
                10'b0010000000: begin   s_id_aluop = ALUOP_SL;      s_id_alu_seed = 1'bx;   end // SLL
                10'b0100000000: begin   s_id_aluop = ALUOP_SLT;     s_id_alu_seed = 1'bx;   end // SLT
                10'b0110000000: begin   s_id_aluop = ALUOP_SLTU;    s_id_alu_seed = 1'bx;   end // SLTU
                10'b1000000000: begin   s_id_aluop = ALUOP_XOR;     s_id_alu_seed = 1'bx;   end // XOR
                10'b1100000000: begin   s_id_aluop = ALUOP_OR;      s_id_alu_seed = 1'bx;   end // OR
                10'b1110000000: begin   s_id_aluop = ALUOP_AND;     s_id_alu_seed = 1'bx;   end // AND
                10'b1010000000: begin   s_id_aluop = ALUOP_SR;      s_id_alu_seed = 1'b0;   end // SRL
                10'b1010100000: begin   s_id_aluop = ALUOP_SR;      s_id_alu_seed = 1'b1;   end // SRA
                default: begin
                    s_id_s1 = 32'hxxxxxxxx;
                    s_id_s2 = 32'hxxxxxxxx;
                    s_id_invalid = 1;
                end
            endcase
        end
        // OP_FENCE: begin
            
        // end
        // OP_SYSTEM: begin
            
        // end
        default: begin
            s_id_jump = 0;
            s_id_branch = 0;
            s_id_s1 = 32'hxxxxxxxx;
            s_id_s2 = 32'hxxxxxxxx;
            s_id_invalid = 1;
        end
    endcase

    if (s_id_invalid) begin
        $display("Invalid instruction at PC=0x%h", r_id_pc);
        s_id_halt       = 1'b1;
        s_id_aluop      = 3'hx;
        s_id_alu_seed   = 1'bx;
    end
end

// EX
reg s_ex_halt;
reg [31:0] s_ex_data1, s_ex_data2;
reg [31:0] s_ex_alu_out;
reg s_ex_alu_zero;
reg [31:0] s_ex_ra;

always @(*) begin
    s_ex_halt = 0;

    s_ex_data1 = r_ex_s1;
    s_ex_data2 = r_ex_s2;

    case (r_ex_aluop)
        ALUOP_ADD: begin    // seed=1: subtract
            s_ex_alu_out = s_ex_data1 + (s_ex_data2 ^ {32{r_ex_alu_seed}}) + r_ex_alu_seed;
        end
        ALUOP_XOR: begin
            s_ex_alu_out = (|s_ex_data1) ^ (|s_ex_data2);
        end
        ALUOP_OR: begin
            s_ex_alu_out = (|s_ex_data1) | (|s_ex_data2);
        end
        ALUOP_AND: begin
            s_ex_alu_out = (|s_ex_data1) & (|s_ex_data2);
        end
        ALUOP_SL: begin
            s_ex_alu_out = s_ex_data1 << s_ex_data2;
        end
        ALUOP_SR: begin     // seed=1: arithmetic
            s_ex_alu_out = s_ex_data1 >> s_ex_data2;
            if (r_ex_alu_seed) begin
                s_ex_alu_out = s_ex_data1 >>> s_ex_data2;
            end
        end
        ALUOP_SLT: begin
            s_ex_alu_out = $signed(s_ex_data1) < $signed(s_ex_data2);
        end
        ALUOP_SLTU: begin
            s_ex_alu_out = s_ex_data1 < s_ex_data2;
        end
        default: begin
            s_ex_halt       = 1;
            s_ex_alu_out    = 32'hxxxxxxxx;
        end
    endcase
    s_ex_alu_zero = (s_ex_alu_out == 0);

    s_ex_ra = r_ex_pc + 4;
end

// MEM
reg s_mem_halt;
reg s_mem_bp;

always @(*) begin
    s_mem_halt = 0;
    s_mem_bp = 0;

    if (r_mem_store) begin
        mem_data_en = 1;
        mem_data_we = 1;
        s_mem_bp = !mem_data_done;
    end else if (r_mem_load) begin
        mem_data_en = 1;
        mem_data_we = 0;
        s_mem_bp = !mem_data_valid;
    end else begin
        mem_data_en = 0;
        mem_data_we = 0;
        s_mem_bp = 0;
    end
end

// WB
reg s_wb_halt;
reg [31:0] s_wb_data;
reg s_wb_write;

always @(*) begin
    s_wb_halt = 0;

    // load instructions do not use output of alu in wb
    s_wb_data = r_wb_alu_out;

    // FIXME: always writes!!!
    s_wb_write = !r_wb_stall;
end

// SYS
reg s_sys_halt;

always @(*) begin
    s_sys_halt = s_if_halt || s_id_halt || s_ex_halt || s_mem_halt || s_wb_halt;
end

// Register update
always @(posedge clk) begin
    if (reset) begin
        r_if_pc     <= 32'h00000000;
        // rather than resetting all flip-flops just stall the pipeline so values are ignored.
        r_id_stall  <= 1;
        r_ex_stall  <= 1;
        r_mem_stall <= 1;
        r_wb_stall  <= 1;
        
    end else begin
        // NOTE: halt disabled because startup causes hault
        // if (s_sys_halt && 0) begin
        //     // stay halted forever
        // end else begin
            // IF
            // if (!s_mem_bp) begin
                r_if_pc     <= s_if_next_pc;
            // end

            // ID
            // if (!s_mem_bp) begin
                r_id_stall  <= s_if_stall;
                r_id_pc     <= r_if_pc;
                r_id_inst   <= s_if_inst;
            // end

            // EX
            // if (!s_mem_bp) begin
                // TODO: also stall EX if taking branch
                r_ex_stall  <= r_id_stall;
                r_ex_pc     <= r_id_pc;
                r_ex_inst   <= r_id_inst;
                r_ex_rd     <= s_id_rd;
                r_ex_s1     <= s_id_s1;
                r_ex_s2     <= s_id_s2;
                r_ex_aluop  <= s_id_aluop;
                r_ex_alu_seed <= s_id_alu_seed;
                r_ex_jump   <= s_id_jump;
                r_ex_store  <= s_id_store;
                r_ex_load   <= s_id_load;
            // end
            

            // MEM
            // if (!s_mem_bp) begin
                r_mem_stall <= r_ex_stall;
                r_mem_pc    <= r_ex_pc;
                r_mem_inst  <= r_ex_inst;
                r_mem_rd    <= r_ex_rd;
                r_mem_s1    <= r_ex_s1;
                r_mem_s2    <= r_ex_s2;
                r_mem_alu_out <= s_ex_alu_out;
                r_mem_alu_zero <= s_ex_alu_zero;
                r_mem_store <= r_ex_store;
                r_mem_load  <= r_ex_load;
            // end

            // WB
            // if (!s_mem_bp) begin
                r_wb_stall  <= r_mem_stall;
                r_wb_pc     <= r_mem_pc;
                r_wb_rd     <= r_mem_rd;
                r_wb_alu_out <= r_mem_alu_out;
            // end

            // Register File
            if (r_wb_rd != 0 && s_wb_write) begin
                regfile[r_wb_rd] <= s_wb_data;
            end
        // end
    end
end

assign dummy_out = s_wb_data[0];

endmodule