package projet

import chisel3._
import chisel3.util.switch
import chisel3.util.is
import chisel3.util.Cat
import chisel3.util.Fill

object DMemSize extends ChiselEnum {
    val Byte, Half, Word, Long = Value
}

class DMem extends Module {
    val io = IO(new Bundle {
        val dbus = Flipped(new BusInterface)
        // The size of the input data
        val size = Input(DMemSize())
        val en = Input(Bool())
        val we = Input(Bool())
        val addr = Input(UInt(64.W))
        val data_in = Input(UInt(64.W))

        val data_out = Output(UInt(64.W))
        val sign_extend = Input(Bool())
    })
    io.dbus.addr := Cat(io.addr(63, 3), 0.U(3.W));
    io.dbus.en := io.en;
    io.data_out := 0.U

    val bytes_enabled = VecInit(
        false.B,
        false.B,
        false.B,
        false.B,
        false.B,
        false.B,
        false.B,
        false.B
    )
    val data = WireInit(0.U(64.W))

    when(io.we && io.en) {
        switch(io.size) {
            is(DMemSize.Byte) {
                bytes_enabled(io.addr(2, 0)) := true.B
                data := io.data_in <<
                    (io.addr(2) * 32.U + io.addr(1) * 16.U + io.addr(0) * 8.U)
            }
            is(DMemSize.Half) {
                bytes_enabled(io.addr(2, 0)) := true.B
                bytes_enabled(io.addr(2, 0) + 1.U) := true.B
                data := io.data_in << (io.addr(2, 0) * 8.U)
            }
            is(DMemSize.Word) {
                bytes_enabled(io.addr(2, 0)) := true.B
                bytes_enabled(io.addr(2, 0) + 1.U) := true.B
                bytes_enabled(io.addr(2, 0) + 2.U) := true.B
                bytes_enabled(io.addr(2, 0) + 3.U) := true.B
                data := io.data_in << (io.addr(2) * 32.U)
            }
            is(DMemSize.Long) {
                bytes_enabled :=
                    VecInit(
                        true.B,
                        true.B,
                        true.B,
                        true.B,
                        true.B,
                        true.B,
                        true.B,
                        true.B
                    )
                data := io.data_in
            }
        }
    }

    // Data out is available one cycle after the read request
    val read_addr = Delay.Delay(io.addr, 1, 0.U);
    val read_sign_extend = Delay.Delay(io.sign_extend, 1, false.B);
    when(Delay.Delay(io.en && !io.we, 1, false.B)) {
        switch(Delay.Delay(io.size, 1, DMemSize.Byte)) {
            is(DMemSize.Byte) {
                val dout = (io.dbus.rdata >>
                    (read_addr(2) * 32.U + read_addr(1) * 16.U + read_addr(
                        0
                    ) * 8.U)) &
                    0xff.U;
                io.data_out := dout | Mux(
                    read_sign_extend,
                    Fill(56, dout(7)) << 8.U,
                    0.U
                );
            }

            is(DMemSize.Half) {
                val dout = (io.dbus.rdata >>
                    (read_addr(2, 0) * 8.U)) & 0xffff.U;
                io.data_out := dout | Mux(
                    read_sign_extend,
                    Fill(48, dout(15)) << 16.U,
                    0.U
                );
            }

            is(DMemSize.Word) {
                val dout = (io.dbus.rdata >>
                    (read_addr(2) * 32.U)) & 0xffffffffL.U;
                io.data_out := dout | Mux(
                    read_sign_extend,
                    Fill(32, dout(31)) << 32.U,
                    0.U
                );
            }
            is(DMemSize.Long) {
                io.data_out := io.dbus.rdata
            }
        }
    }
    io.dbus.be := bytes_enabled

    io.dbus.wdata := data

}
