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
    })

    // ### Pipeline stages
    // Fetch       : The current instruction is currently being loaded from the memory

    // Decode      : The current instruction is being read by the control unit, the register file, the immediate decoder
    // Execute     : The current instruction immediate, or register operands are currently being operated (eg. in the alu)
    // (Currently Decode and execute are the same stage)

    // Memory Read : The memory in being read/operands stored

    // ### Processing unit components

    // Generates the control signals for the following components and defines data_flow in the processor
    val control_unit = Module(new ControlUnit());

    // Arithmetic logic unit : computes operations between operands
    val alu = Module(new Alu());

    // Decodes the immediates from the instruction based on the instruction type
    val immediate_decoder = Module(new ImmediateDecoder());

    // Wraps the memory behind a interface (manages full, word, half, byte load/store operations)
    val dmem = Module(new DMem());

    // Contains the 32 GP registers, two output ports, one input port
    val reg_file = Module(new RegFile(sim));

    // Register containing the address of the instruction in the Fetch stage
    val reg_pc = RegInit(Constants.RV64I_RESET_VECTOR);

    // Register containing the address of the instruction in teh Decode stage TODO: Rename (why "execute")
    val execute_pc = Delay.Delay(reg_pc, 1, Constants.RV64I_RESET_VECTOR);

    // Wire : contains the result (i.e output of the Execution Stage) of the instruction
    val instruction_result = WireInit(0.U(Constants.RV64I_XLEN.W));

    // Wire : contains the data to be written back by the instruction in Memory Read stage (e.g. memory data)
    val writeback_line = WireInit("xdeadbeef".U(Constants.RV64I_XLEN.W));

    // Wether the executing instruction is stalling the processor (e.g. waiting for a result : division multiplication)
    val is_stalled = RegInit(false.B);

    // Holds the instruction in the decode stage if it stalls
    val stalling_insruction = RegInit(0.U(32.W));

    // Control signal that tells the core to resume execution
    val should_stop_stall =
        control_unit.io.should_stop_stall || alu.io.should_stop_stall

    // Connect memory wrapper to bus interconnect
    dmem.io.dbus <> io.dbus

    // #############################
    // ######## FETCH STAGE ########
    // #############################

    // PC ADDER
    // Sums either :
    // - PC (normal program progression)
    // - Decode PC (jump, branch)
    // with
    // - 4 constant (normal program progresion)
    // - Immediate (jump, branch)
    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
    );

    // Program counter
    // Either :
    // - Update PC according to adder (normal progression, jump, branch)
    // - Update PC according to result of instruction in Execute Stage (JALR)
    reg_pc := Mux(
        control_unit.io.mux_pcadder_jalrwriteback,
        pc_adder_out,
        instruction_result
    );

    // Handle stall state
    // We should :
    // - Stop counting as soon as we require wating
    // - Keep counting when we are stalled
    // - Start counting as soon as we requiring starting again
    when(!should_stop_stall && (is_stalled || control_unit.io.should_stall)) {
        reg_pc := reg_pc;
    }

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

    // Default data for Ibus
    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
    );

    // After this stage: the instruction will be read from memory and become available on the decode stage on next cycle

    // ######################################
    // ######## DECODE/EXECUTE STAGE ########
    // ######################################
    // Instruction is available on ibus

    // Instruction being decoded
    // Either :
    // - Stalling instructino if the cpu is waiting
    // - Data from ibus if normal operation
    // - If previous instruction being executed is a jump => Insert bubble to flush pipeline
    val decoded_instruction =
        Mux(
            is_stalled,
            stalling_insruction,
            Mux(
                Delay.Delay(control_unit.io.is_jump, 1, false.B),
                Constants.RISCV_NOP,
                io.ibus.rdata >> (execute_pc(2) * 32.U)
            )
        );
    decoded_instruction.suggestName("instruction");

    // Decide control signals for instruction
    // (Control unit does not hold any state, output signals are about "decoded_instruction")
    control_unit.io.instruction := decoded_instruction;

    // Decode immediate in the instruction based on the optype
    immediate_decoder.io.op_type := control_unit.io.optype;
    immediate_decoder.io.instruction := decoded_instruction;

    // Decode register indices

    // Decode rs1 index
    val rs1_index = decoded_instruction(19, 15)
    reg_file.io.rs1_addr := rs1_index

    // Decode rs2 index
    val rs2_index = decoded_instruction(24, 20)
    reg_file.io.rs2_addr := rs2_index

    // Decode rd index
    val rd_index = decoded_instruction(11, 7)

    // Manage the stalled state of the processor

    // Stalling behaviour
    when(!is_stalled && control_unit.io.should_stall) {
        // Set core to stalled state when instructions requires so and store it
        is_stalled := true.B;
        stalling_insruction := decoded_instruction;
    }

    when(is_stalled && should_stop_stall) {
        is_stalled := false.B
    }

    // Stalling note :
    // When instruction requires stalling it needs to be held in the decode state as its decoded information will
    // be needed in later stages when in resumes execution

    // Data gotten from register file.
    // Either :
    // - Data from register file
    // - If PREVIOUS instruction wrote to the register, get the data currently being written in the register file
    val rs1_data =
        Mux(
            rs1_index === Delay.Delay(rd_index, 1, 0.U) && 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 === Delay.Delay(rd_index, 1, 0.U) && rs2_index =/= 0.U
                && Delay.Delay(control_unit.io.reg_file_we, 1, false.B),
            writeback_line,
            reg_file.io.rs2_data
        )

    // Feed ALU

    // First operand
    // Either :
    // - Data from registers
    // - PC of next instruction (instruction being decoded) for JALR
    alu.io.a := Mux(
        control_unit.io.mux_rega_pc,
        rs1_data,
        execute_pc
    );

    // Second operand
    // Either :
    // - Data from registers
    // - Immediate from instruction
    alu.io.b := Mux(
        control_unit.io.mux_regb_imm,
        rs2_data,
        immediate_decoder.io.immediate
    );

    // Alu Op code, result, W mode
    alu.io.opcode := control_unit.io.alu_opcode;
    control_unit.io.alu_comp_result := alu.io.comp_result
    alu.io.word_mode := control_unit.io.alu_word_mode

    // Result of execute stage
    // Either :
    // - ALU Computation result
    // - Immediate from instruction
    instruction_result := Mux(
        control_unit.io.mux_alu_imm,
        alu.io.out,
        immediate_decoder.io.immediate
    );

    // Memory setup
    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;
    dmem.io.data_in := rs2_data;
    dmem.io.addr := instruction_result;

    // Debugging/ testing
    io.x31 := reg_file.io.x31
    if (sim) {
        io.valid_x31.get := reg_file.io.valid_x31.get

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

    // #############################################
    // ######## MEMORY READ/WRITEBACK STAGE ########
    // #############################################
    // One stage after Decode/Execute => Control signals should be delayed by 1

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

    // Writeback line
    // Either :
    // - Result from the exection of the instruction (previous stage)
    // - Data gotten from memory
    writeback_line := Mux(
        Delay.Delay(control_unit.io.mux_executeout_dout, 1, true.B),

        // Data from previous stage
        Delay.Delay(
            instruction_result,
            1,
            0.U(Constants.RV64I_XLEN.W)
        ),
        dmem.io.data_out
    );

    // Register writeback
    reg_file.io.rd_addr := Delay.Delay(rd_index, 1, 0.U);

    // Register write
    // Either :
    // - Writeback line (instruction result/read)
    // - This PC following this instruction's PC (JAL, JALR)
    reg_file.io.rd_data :=
        Mux(
            Delay.Delay(control_unit.io.mux_writeback_pc, 1, true.B),
            writeback_line,
            execute_pc
        );

    // Reg file write-enable
    // Do not write back if the instruction is stalling
    reg_file.io.we := Delay.Delay(
        control_unit.io.reg_file_we,
        1,
        false.B
    ) && !is_stalled;
}
