package projet

import chisel3._

class Rv64i(sim: Boolean = true) extends Module {
    val io = IO(new Bundle {
        val ibus = Flipped(new BusInterface)
        val dbus = Flipped(new BusInterface)
        val x31 = Output(UInt(64.W))
        val valid_x31 = if (sim) Some(Output(Bool())) else None
    })

    val reg_pc = RegInit("x10000".U(64.W));
    val reg_file = Module(new RegFile(sim));

    reg_file.io.rd_data := 0.U;
    io.x31 := reg_file.io.x31
    if (sim) {
        io.valid_x31.get := reg_file.io.valid_x31.get
    }

    val alu = Module(new Alu());
    val immediate_decoder = Module(new ImmediateDecoder());
    val control_unit = Module(new ControlUnit());
    val dmem = Module(new DMem());
    dmem.io.dbus <> io.dbus
    dmem.io.en := control_unit.io.memory_en
    dmem.io.we := control_unit.io.memory_we
    dmem.io.size := control_unit.io.memory_size
    dmem.io.sign_extend := control_unit.io.memory_sign_extend

    control_unit.io.alu_comp_result := alu.io.comp_result
    alu.io.word_mode := control_unit.io.alu_word_mode

    // PC delayed to execute stage for auipc op
    val execute_pc = Delay.Delay(reg_pc, 1, "x10000".U(64.W));
    // True if the instruction in the execute stage is a jump
    val is_jump = Delay.Delay(control_unit.io.is_jump, 1, false.B);

    if (sim) {
        val log = SimLog.file("core_pc.log");
        when(!is_jump) {
            log.printf(cf"0x${execute_pc}%x\n")
        }
    }

    // Pipelining registers for mem stage
    val reg_rd_index = RegInit(0.U);
    val reg_execute_out = RegInit(0.U);
    val reg_mux_executeout_dout = RegInit(true.B);
    val reg_regile_we = RegInit(false.B);

    // Insert no-op if jump
    val instruction =
        Mux(
            is_jump,
            "b00000000000000000000000000010011".U,
            io.ibus.rdata >> (execute_pc(2) * 32.U)
        );

    immediate_decoder.io.instruction := instruction
    immediate_decoder.io.op_type := control_unit.io.optype
    alu.io.opcode := control_unit.io.alu_opcode;

    // Reg file write-enable
    reg_regile_we := control_unit.io.reg_file_we;
    reg_file.io.we := reg_regile_we;

    // Increment PC each clock
    val pc_adder_out = Mux(
        control_unit.io.mux_regpc_executepc,
        reg_pc,
        execute_pc
    ) + Mux(
        control_unit.io.mux_incrpc4_imm,
        4.U,
        immediate_decoder.io.immediate
    );

    // Fetch instruction
    io.ibus.en := true.B;
    io.ibus.addr := reg_pc;

    control_unit.io.instruction := instruction;

    // Decode rs1 index
    val rs1_index = instruction(19, 15)
    reg_file.io.rs1_addr := rs1_index
    // Decode rs2 index
    val rs2_index = instruction(24, 20)
    reg_file.io.rs2_addr := rs2_index
    // Decode rd index
    val rd_index = instruction(11, 7)
    reg_rd_index := rd_index;
    reg_file.io.rd_addr := reg_rd_index;

    val writeback_line = Mux(
        reg_mux_executeout_dout,
        reg_execute_out,
        dmem.io.data_out
    );
    val rs1_data =
        Mux(
            rs1_index === reg_rd_index && rs1_index =/= 0.U
                && Delay.Delay(control_unit.io.reg_file_we, 1, false.B),
            writeback_line,
            reg_file.io.rs1_data
        )
    val rs2_data =
        Mux(
            rs2_index === reg_rd_index && rs2_index =/= 0.U
                && Delay.Delay(control_unit.io.reg_file_we, 1, false.B),
            writeback_line,
            reg_file.io.rs2_data
        )

    // EXECUTE

    val imm = immediate_decoder.io.immediate

    alu.io.a := Mux(
        control_unit.io.mux_rega_pc,
        rs1_data,
        execute_pc
    );
    alu.io.b := Mux(control_unit.io.mux_regb_imm, rs2_data, imm);

    // Select what to send on writeback line
    val execute_out = Mux(control_unit.io.mux_alu_imm, alu.io.out, imm);
    reg_execute_out := execute_out;
    // Jalr line is directly the execute output
    val jalr_line = execute_out

    reg_mux_executeout_dout := control_unit.io.mux_executeout_dout;

    reg_file.io.rd_data :=
        Mux(
            Delay.Delay(control_unit.io.mux_writeback_pc, 1, true.B),
            writeback_line,
            Delay.Delay(reg_pc, 1, 0.U)
        );

    // Writeback pipelining registers
    reg_pc := Mux(
        control_unit.io.mux_pcadder_jalrwriteback,
        pc_adder_out,
        jalr_line
    );

    dmem.io.data_in := rs2_data
    dmem.io.addr := execute_out

    io.ibus.wdata := 0.U;
    io.ibus.be := VecInit(
        false.B,
        false.B,
        false.B,
        false.B,
        false.B,
        false.B,
        false.B,
        false.B
    )

}