module RISCCore (
  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 = (pc[31:2] < 64) ? imem[pc[31:2]] : 32'h00000013;  // Word-aligned access

  //Data Mem
  reg [31:0] dmem [0:31];
  wire [31:0] ld_data;
  wire [4:0] word_addr = mem_addr[6:2];

  integer j;
  initial begin
    for(j = 0; j < 32; j = j + 1) begin
      dmem[j] = 32'b0;
    end 
  end

  assign ld_data = (word_addr < 32) ? dmem[word_addr] : 32'b0;

  // 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
   // Initialize register file (x0 always zero)
  reg [31:0] rf [0:31]; //Register file
  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  = (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 = (opcode == 7'b0000011) && (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 = (opcode == 7'b0100011) && (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  = (opcode == 7'b0110011) && (funct3 == 3'b001) && (funct7 == 7'b0000000);
  wire isSLTI = (isIType) && (funct3 == 3'b010);
  wire isSLTIU = (isIType) && (funct3 == 3'b011);

  wire isSRL  = (opcode == 7'b0110011) && (funct3 == 3'b101) && (funct7 == 7'b0000000);
  wire isSRA  = (opcode == 7'b0110011) && (funct3 == 3'b101) && (funct7 == 7'b0100000);

  wire isSLLI = (opcode == 7'b0010011) && (funct3 == 3'b001) && (funct7[6:1] == 6'b000000);
  wire isSRLI = (opcode == 7'b0010011) && (funct3 == 3'b101) && (funct7[6:1] == 6'b000000);
  wire isSRAI = (opcode == 7'b0010011) && (funct3 == 3'b101) && (funct7[6:1] == 6'b010000);

  //logic imms
  wire isANDI = (opcode == 7'b0010011) && (funct3 == 3'b111);
  wire isORI  = (opcode == 7'b0010011) && (funct3 == 3'b110);
  wire isXORI = (opcode == 7'b0010011) && (funct3 == 3'b100);
  // logic
  wire isAND  = (opcode == 7'b0110011) && (funct3 == 3'b111) && (funct7 == 7'b0000000);
  wire isOR   = (opcode == 7'b0110011) && (funct3 == 3'b110) && (funct7 == 7'b0000000);
  wire isXOR  = (opcode == 7'b0110011) && (funct3 == 3'b100) && (funct7 == 7'b0000000);
  wire isSUB = (opcode == 7'b0110011) && (funct3 == 3'b000) && (funct7 == 7'b0100000);
  //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);

  //Mem address calculation
  /*wire [31:0] mem_addr = (is_mem_op) ? alu_result : 32'b0;
  Using direct */
  wire [31:0] mem_addr = rs1_val + (isSType ? Simm : Iimm);

  //Mem address logic
  //For simplicity we only implement all loads as word loads (lw)
  //byte/halfword ignored
  wire is_load = (opcode == 7'b0000011);  // All load instructions
  wire is_store = (opcode == 7'b0100011); // All store instructions
  wire is_mem_op = is_store || is_load;


//Load operations

  reg [31:0] load_data;

  always @(*) begin 
    if (mem_addr[1:0] == 2'b00) begin 
      case (funct3)
        3'b000: load_data = {{24{ld_data[7]}}, ld_data[7:0]};
        3'b001: load_data = {{16{ld_data[15]}}, ld_data[15:0]};
        3'b010: load_data = ld_data;
        3'b100: load_data = {24'b0, ld_data[7:0]};
        3'b101: load_data = {16'b0, ld_data[15:0]};
        default: load_data = 32'b0;
      endcase
    end else begin 
      load_data = 32'b0;
      if (is_load) begin  // Only show error for actual loads
        $display("ERROR: Misaligned memory address at addr %h", mem_addr);
      end
    end 
  end

//Store operations
always @(posedge clk) begin 
  if (!rst && is_store && (word_addr < 32) && (mem_addr[1:0] == 2'b00)) begin 

    case (funct3)
      3'b000: dmem[word_addr] <= {dmem[word_addr][31:8], rs2_val[7:0]};   // SB
      3'b001: dmem[word_addr] <= {dmem[word_addr][31:16], rs2_val[15:0]}; // SH
      3'b010: dmem[word_addr] <= rs2_val;                                 // SW
    endcase
    // The $display can be simplified now as well
    $display("MEM Write: word_addr=%h, data=%h", word_addr, rs2_val);
  end
end
  // ALU operations

  //sltu and slt
  // For signed operations, cast the inputs to 'signed'.
  // Verilog will then use 2's complement arithmetic automatically.
  wire signed [31:0] signed_rs1_val = rs1_val;
  wire signed [31:0] signed_Iimm    = Iimm;

  // SLTU: Set if Less Than (Unsigned)
  // Compares rs1_val and rs2_val as unsigned integers.
  wire [31:0] sltu_rslt = (rs1_val < rs2_val) ? 32'd1 : 32'd0;

  // SLT: Set if Less Than (Signed)
  // Compares rs1_val and rs2_val as signed integers.
  wire [31:0] slt_rslt = (signed_rs1_val < $signed(rs2_val)) ? 32'd1 : 32'd0;

  // SLTIU: Set if Less Than Immediate (Unsigned)
  // Compares rs1_val (unsigned) with the sign-extended immediate (unsigned).
  wire [31:0] sltiu_rslt = (rs1_val < Iimm) ? 32'd1 : 32'd0;

  // SLTI: Set if Less Than Immediate (Signed)
  // Compares rs1_val and the immediate as signed integers.
  wire [31:0] slti_rslt = (signed_rs1_val < signed_Iimm) ? 32'd1 : 32'd0;

  // SRA: Shift Right Arithmetic
  // Shifts signed_rs1_val right by the amount in rs2_val[4:0].
  // The '>>>' operator performs an arithmetic shift, preserving the sign bit.
  wire [31:0] sra_rslt = signed_rs1_val >>> rs2_val[4:0];

  // SRAI: Shift Right Arithmetic Immediate
  // Shifts signed_rs1_val right by the immediate amount in Iimm[4:0].
  wire [31:0] srai_rslt = signed_rs1_val >>> Iimm[4:0];
 
  wire [31:0] alu_result = (isADDI) ? (rs1_val + Iimm) :
                           (isADD) ? (rs1_val + rs2_val) : 
                           (isSLT) ? slt_rslt :
                           (isSLTU) ? sltu_rslt :
                           (isSLTI) ? slti_rslt :
                           (isANDI) ? (rs1_val & Iimm) :
                           (isORI) ? (rs1_val | Iimm) :
                           (isXORI) ? (rs1_val ^ Iimm) :
                           (isSLLI) ? (rs1_val << Iimm[4:0]) :
                           (isSRLI) ? (rs1_val >> Iimm[4: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) ? Uimm :
                           (isAUIPC) ? (pc + Uimm) :
                           (isJAL || isJALR) ? (pc + 32'd4) :
                           (isSRA) ? sra_rslt :
                           (isSRAI) ? srai_rslt :
                           32'b0;


  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;
  wire [31:0] jalr_tgt_pc = (rs1_val + Iimm) & 32'hFFFFFFFE;
  wire [31:0] jal_target = pc + Jimm;
  
  assign next_pc = branch_taken ? branch_target :
                   isJAL ? jal_target :
                   isJALR ? jalr_tgt_pc :
                   next_pc_base;

  // Register write back
  wire rf_write_enable = (rd != 0) && (
                          isADDI || isADD || isSUB ||
                          isJAL || isJALR ||
                          isSLT || isSLTU || isSLTI || isSLTIU ||
                          isSLL || isSRL || isSRA || isSLLI || isSRLI || isSRAI ||
                          isLB || isLH || isLW || isLBU || isLHU ||
                          isXORI || isORI || isANDI ||
                          isXOR || isOR || isAND ||
                          is_load || isLUI || isAUIPC
                        );
  

  wire [31:0] writeback_data = is_load ? load_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;  // Use writeback_data instead of alu_result
    $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
  
  always @(posedge clk) begin
    if (!rst && (opcode == 7'b0000011)) begin
      $display("LOAD: byte_addr=%h, word_addr=%h, data=%h, funct3=%b", 
               mem_addr, word_addr, load_data, funct3);
    end
  end

endmodule