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 = (isIType) && (funct3 == 3'b000);
  wire isADD  = (isRType) && (funct3 == 3'b000) && (funct7 == 7'b0000000);
  
  // Branch instructions
  wire isBEQ  = (isBType) && (funct3 == 3'b000);
  wire isBNE  = (isBType) && (funct3 == 3'b001);
  wire isBLT  = (isBType) && (funct3 == 3'b100);
  wire isBGE  = (isBType) && (funct3 == 3'b101);
  wire isBLTU = (isBType) && (funct3 == 3'b110);
  wire isBGEU = (isBType) && (funct3 == 3'b111);

  //store load instructions
  wire isLB = (isIType) && (funct3 == 3'b000);
  wire isLH = (isIType) && (funct3 == 3'b001);
  wire isLW = (isIType) && (funct3 == 3'b010);
  wire isLBU = (isIType) && (funct3 == 3'b100);
  wire isLHU = (isIType) && (funct3 == 3'b101);

  wire isSB = (isSType) && (funct3 == 3'b000);
  wire isSH = (isSType) && (funct3 == 3'b001);
  wire isSW = (isSType) && (funct3 == 3'b010);

  //sl and sr instructions
  wire isSLT = (isRType) && (funct3 == 3'b010) && (funct7[5] == 1'b0);
  wire isSLTU = (isRType) && (funct3 == 3'b011) && (funct7[5] == 1'b0);
  wire isSLL = (isRType) && (funct3 == 3'b001) && (funct7[5] == 1'b0);
  wire isSLTI = (isIType) && (funct3 == 3'b010);
  wire isSLTIU = (isIType) && (funct3 == 3'b011);

  wire isSRA = (isRType) && (funct3 == 3'b101) && (funct7[5] == 1'b1);
  wire isSRL = (isRType) && (funct3 == 3'b101) && (funct7[5] == 1'b0);
  wire isSRAI = (isIType) && (funct3 == 3'b101) && (funct7[5] == 1'b1);
  
  wire isSLLI = (isIType) && (funct3 == 3'b001) && (funct7[5] == 1'b0);
  wire isSRLI = (isIType) && (funct3 == 3'b101) && (funct7[5] == 1'b0);

  //logic imms
  wire isANDI = (isIType) && (funct3 == 3'b111);
  wire isORI = (isIType) && (funct3 == 3'b110);
  wire isXORI = (isIType) && (funct3 == 3'b100);

  // logic
  wire isAND = (isRType) && (funct3 == 3'b111) && (funct7[5] == 1'b0);
  wire isOR = (isRType) && (funct3 == 3'b110) && (funct7[5] == 1'b0);
  wire isXOR = (isRType) && (funct3 == 3'b100) && (funct7[5] == 1'b0);
  wire isSUB = (isRType) && (funct3 == 3'b000) && (funct7[5] == 1'b1);

  //lui auipc (opcode use)
  wire isLUI = (opcode == 7'b0110111); 
  wire isAUIPC = (opcode == 7'b0010111);
  wire isJAL = (opcode == 7'b1101111);
  //jal UType

  wire isJALR = (opcode == 7'b1100111) && (funct3 == 3'b000);

  // ALU operations
  //sltu and slt
  wire [31:0] sltu_rslt = {31'b0, (rs1_val < rs2_val)};
  wire [31:0] signed_slt = (rs1_val[31] && !rs2_val[31]) ? 1'b1 :
                           (!rs1_val[31] && rs2_val[31]) ? 1'b0 :
                           (rs1_val < rs2_val);
  wire [31:0] slt_rslt = {31'b0, signed_slt};
  wire [31:0] slti_rslt = ((rs1_val[31] == Iimm[31]) ? sltu_rslt : {31'b0, rs1_val[31]});

  wire [63:0] SErs1_val = {32{rs1_val[31]}, (rs1_val < rs2_val)};

  wire [63:0] sra_rslt = {SErs1_val >> rs2_val[4:0]};
  wire [63:0] srai_rslt = {SErs1_val >> Iimm[4:0]};
  wire [31:0] sltiu_rslt = {31'b0, (rs1_val < Iimm)};
 
  wire [31:0] alu_result = (isADDI) ? (rs1_val + Iimm) :
                           (isADD) ? (rs1_val + rs2_val) : 
                           (isSLT) ? slt_rslt :
                           (isSLTU) ? sltu_rslt :
                           (isSRA) ? sra_rslt :
                           (isSRAI) ? srai_rslt :
                           (isSLTI) ? slti_rslt :
                           (isANDI) ? (rs1_val & Iimm) :
                           (isORI) ? (rs1_val | Iimm) :
                           (isXORI) ? (rs1_val ^ Iimm) :
                           (isSLLI) ? (rs1_val << Iimm[5:0]) :
                           (isSRLI) ? (rs1_val >> Iimm[5:0]) :
                           (isAND) ? (rs1_val & rs2_val) :
                           (isOR) ? (rs1_val | rs2_val) :
                           (isXOR) ? (rs1_val ^ rs2_val) :
                           (isSUB) ? (rs1_val - rs2_val) :
                           (isSLL) ? (rs1_val << rs2_val[4:0]) :
                           (isSRL) ? (rs1_val >> rs2_val[4:0]) :
                           (isSLTIU) ? (sltiu_rslt) :
                           (isLUI) ? ({Iimm[31:12], 12'b0}) :
                           (isAUIPC) ? (pc + Iimm) :
                           (isJAL) ? (pc + 32'd4) :
                           (isJALR) ? (pc + 32'd4) :
                           (isSRA) ? (sra_rslt[31:0]) :
                           (isSRAI) ? (srai_rslt[31:0]) : 
                           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