module RISCcore2 (
  input rst,
  input clk,
  output reg [31:0] pc,
  output wire [31:0] next_pc,  // Changed to wire
  output wire [31:0] instr
);

  // IMem - reduced size for simplicity
  reg [31:0] imem [0:63];
  initial begin
    $readmemh("program.hex", imem);
  end
  assign instr = imem[pc[31:2]];  // Word-aligned access

  // Instruction decoder
  wire [6:0] opcode = instr[6:0];
  wire [4:0] rs1 = instr[19:15];
  wire [4:0] rs2 = instr[24:20];
  wire [4:0] rd = instr[11:7];
  wire [2:0] funct3 = instr[14:12];
  wire [6:0] funct7 = instr[31:25];

  // Instruction type detection
  wire isUType = (opcode == 7'b0110111) || (opcode == 7'b0010111);  // LUI, AUIPC
  wire isIType = (opcode == 7'b0000011) || (opcode == 7'b0000111) ||  // LOAD
                 (opcode == 7'b0010011) || (opcode == 7'b0011011) ||  // OP-IMM
                 (opcode == 7'b1100111);                               // JALR
  wire isRType = (opcode == 7'b0110011) || (opcode == 7'b0111011);     // OP
  wire isSType = (opcode == 7'b0100011) || (opcode == 7'b0100111);     // STORE
  wire isBType = (opcode == 7'b1100011);                               // BRANCH
  wire isJType = (opcode == 7'b1101111);                               // JAL

  // Immediate generation
  wire [31:0] Iimm = {{20{instr[31]}}, instr[31:20]};
  wire [31:0] Simm = {{20{instr[31]}}, instr[31:25], instr[11:7]};
  wire [31:0] Bimm = {{19{instr[31]}}, instr[31], instr[7], instr[30:25], instr[11:8], 1'b0};
  wire [31:0] Uimm = {instr[31:12], 12'b0};
  wire [31:0] Jimm = {{11{instr[31]}}, instr[31], instr[19:12], instr[20], instr[30:21], 1'b0};

  // Register file
  reg [31:0] rf [0:31];
  
  // Initialize register file (x0 always zero)
  integer i;
  initial begin
    for (i = 0; i < 32; i = i + 1) begin
      rf[i] = 32'b0;
    end
  end

  // Read ports
  wire [31:0] rs1_val = (rs1 != 0) ? rf[rs1] : 32'b0;
  wire [31:0] rs2_val = (rs2 != 0) ? rf[rs2] : 32'b0;

  // Instruction decoding
  wire isADDI = (opcode == 7'b0010011) && (funct3 == 3'b000);
  wire isADD  = (opcode == 7'b0110011) && (funct3 == 3'b000) && (funct7 == 7'b0000000);
  
  // Branch instructions
  wire isBEQ  = (opcode == 7'b1100011) && (funct3 == 3'b000);
  wire isBNE  = (opcode == 7'b1100011) && (funct3 == 3'b001);
  wire isBLT  = (opcode == 7'b1100011) && (funct3 == 3'b100);
  wire isBGE  = (opcode == 7'b1100011) && (funct3 == 3'b101);
  wire isBLTU = (opcode == 7'b1100011) && (funct3 == 3'b110);
  wire isBGEU = (opcode == 7'b1100011) && (funct3 == 3'b111);

  // ALU operations
  wire [31:0] alu_src2 = isADDI ? Iimm : rs2_val;
  wire [31:0] alu_result = (isADDI || isADD) ? (rs1_val + alu_src2) : 32'b0;

  // Branch condition logic
  wire signed [31:0] signed_rs1 = rs1_val;
  wire signed [31:0] signed_rs2 = rs2_val;
  
  wire branch_taken =
      isBEQ  ? (rs1_val == rs2_val) :
      isBNE  ? (rs1_val != rs2_val) :
      isBLT  ? (signed_rs1 < signed_rs2) :
      isBGE  ? (signed_rs1 >= signed_rs2) :
      isBLTU ? (rs1_val < rs2_val) :
      isBGEU ? (rs1_val >= rs2_val) :
      1'b0;

  // Next PC calculation - FIXED: using wire for continuous assignment
  wire [31:0] branch_target = pc + Bimm;
  wire [31:0] next_pc_base = pc + 32'h4;
  assign next_pc = branch_taken ? branch_target : next_pc_base;

  // Register write back
  wire rf_write_enable = (rd != 0) && (isADDI || isADD);
  wire [31:0] writeback_data = alu_result;

  // PC update
  always @(posedge clk) begin
    if (rst) begin
      pc <= 32'h0;
    end else begin
      pc <= next_pc;
    end
  end

  // Register write back
  always @(posedge clk) begin
    if (rf_write_enable && !rst) begin
      rf[rd] <= writeback_data;
      $display("RF Write: x%d = %h", rd, writeback_data);
    end
  end

  // Debug monitoring
  always @(posedge clk) begin
    if (!rst) begin
      $display("PC=%08h, Instr=%08h, rs1=x%d(%h), rs2=x%d(%h), rd=x%d, branch_taken=%b",
               pc, instr, rs1, rs1_val, rs2, rs2_val, rd, branch_taken);
      if (pc[1:0] != 2'b00) begin
        $display("WARNING: PC not word-aligned: %h", pc);
      end
    end
  end

endmodule