package projet

import chisel3._
import chisel3.util.switch
import chisel3.util.is
import projet.op.OpImm
import projet.op.OpR
import projet.op.OpBranch
import projet.op.OpStore
import projet.op.OpLoad
import projet.op.OpImmW

class CUInterface extends Bundle {
    val instruction = Input(UInt(32.W))
    val reg_file_we = Output(Bool())
    val alu_opcode = Output(AluOpCode())

    // TODO: Maybe bundle alu <> cu
    // Alu comparison result
    val alu_comp_result = Input(Bool())
    val alu_word_mode = Output(Bool())

    val optype = Output(OpType())

    // True if the decoded instruction is a jump
    val is_jump = Output(Bool())

    // DMem connections
    val memory_size = Output(DMemSize())
    val memory_en = Output(Bool())
    val memory_we = Output(Bool())
    val memory_sign_extend = Output(
        Bool()
    ) // Wether to signe extend the output of the memory for bytes, half loads

    // Muxers
    // The naming convention is as follows : mux_xxx_yyy.
    // If the signal is hot, xxx is selected otherwise yyy is selected

    // Sends alu or immediate to write back
    val mux_alu_imm = Output(Bool())

    // Selects regA or pc for alu a
    val mux_rega_pc = Output(Bool())

    // Selects regB or immediate for alu b
    val mux_regb_imm = Output(Bool())

    // Selects 4 or immediate for pc incrementation
    val mux_incrpc4_imm = Output(Bool())

    // Selects pc or previous pc for pc incrementation
    val mux_regpc_executepc = Output(Bool())

    // Selects pc adder or write back for pc write
    val mux_pcadder_jalrwriteback = Output(Bool())

    // Selects pc or write back for register write
    val mux_writeback_pc = Output(Bool())

    // Selects aluout or memory data out for writeback_line
    val mux_executeout_dout = Output(Bool())
}

object OpType extends ChiselEnum {
    val U, I, J, B, S = Value

    // Type IR is not defined in the RiscV specification
    // but the immediate needs to be handled differently
    // since the 31st bit indicates wether it is an
    // arithmetic or logical shift
    val IR = Value
}

object OpCode extends ChiselEnum {
    // Type U
    val LUI = "b0110111".U
    val AUIPC = "b0010111".U

    // Type I + Type IR

    // Instruction format
    //    xxx rd, rs1, imm
    // Perfomrs some operation with rs1 and immediate
    // and writes back to a register

    // addi, andi, ori, xori, slli, srli, srai,
    // slti, sltiu
    val OpImm = "b0010011".U

    // Instruction format
    //    xxx rd, rs1, rs2
    // Perfomrs some operation with rs1 and rd2
    // and writes back to a register

    // add, and, or, xor, sll, srl, sra, slt, stlu, mul, mulh, mulhsu, mulhu, div, divu, rem, remu
    val OpR = "b0110011".U;

    // Type J
    val JAL = "b1101111".U

    // Type I
    val JALR = "b1100111".U

    // Type B
    val OpBranch = "b1100011".U

    // Type S
    val OpStore = "b0100011".U

    // Type I
    val OpLoad = "b0000011".U

    // ~~ Rv64i ~~
    val OpImmW = "b0011011".U

    // TypeR
    val OpRW = "b0111011".U
}

class ControlUnit() extends Module {
    val io = IO(new CUInterface())

    io.reg_file_we := false.B
    io.alu_opcode := AluOpCode.Add
    io.alu_word_mode := false.B
    io.optype := OpType.U
    io.is_jump := false.B;
    io.memory_size := DMemSize.Word
    io.memory_en := false.B
    io.memory_we := false.B
    io.memory_sign_extend := false.B;
    io.mux_alu_imm := true.B;
    io.mux_rega_pc := true.B;
    io.mux_regb_imm := true.B;
    io.mux_incrpc4_imm := true.B;
    io.mux_regpc_executepc := true.B;
    io.mux_writeback_pc := true.B;
    io.mux_pcadder_jalrwriteback := true.B;
    io.mux_executeout_dout := true.B;

    // Decode instruction
    val opcode = io.instruction(6, 0);
    val funct3 = io.instruction(14, 12);
    val funct7 = io.instruction(31, 25);

    switch(opcode) {
        // ~~~ Single instruction OpCodes ~~

        // ~ U instructions ~
        is(OpCode.LUI) {
            io.reg_file_we := true.B // Write to regfile
            io.mux_alu_imm := false.B // Send immy directly to refile
            io.optype := OpType.U
        }
        is(OpCode.AUIPC) {
            io.alu_opcode := AluOpCode.Add
            io.reg_file_we := true.B // Write to regfile
            io.mux_rega_pc := false.B // Send PC to ALU A
            io.mux_regb_imm := false.B; // Send immu to ALU B
            io.optype := OpType.U
        }

        // ~ J instructions ~
        is(OpCode.JAL) {
            io.alu_opcode := AluOpCode.Add // ALU is unused

            io.reg_file_we := true.B // Write to regfile
            io.mux_writeback_pc := false.B;

            // io.mux_rega_pc := false.B // Send PC to ALU A
            // io.mux_regb_imm := false.B; // Send imm to ALU B

            io.is_jump := true.B
            io.mux_incrpc4_imm := false.B; // Send the immediate in PC
            io.mux_regpc_executepc := false.B; // Add the immediate to the previous PC
            io.mux_pcadder_jalrwriteback := true.B;
            io.optype := OpType.J
        }

        is(OpCode.JALR) {
            io.alu_opcode := AluOpCode.Add

            io.reg_file_we := true.B // Write to regfile
            io.mux_writeback_pc := false.B // Write next instruction PC to rd
            io.is_jump := true.B // Insert no-op to discard fetched next instruction

            io.mux_rega_pc := true.B // Send rs1 to ALU A
            io.mux_regb_imm := false.B; // Send imm to ALU B

            io.mux_alu_imm := true.B; // Send alu out to write back line
            io.mux_pcadder_jalrwriteback := false.B; // Set pc to writeback dest
            io.optype := OpType.I
        }

        is(OpCode.OpImm) {
            OpImm.opImm(io)
        }

        is(OpCode.OpR) {
            OpR.opR(io);
        }

        is(OpCode.OpBranch) {
            OpBranch.opBranch(io);
        }

        is(OpCode.OpStore) {
            OpStore.opStore(io);
        }

        is(OpCode.OpLoad) {
            OpLoad.opLoad(io);
        }

        is(OpCode.OpImmW) {
            OpImmW.opImmW(io);
        }

        is(OpCode.OpRW) {
            io.alu_word_mode := true.B
            OpR.opR(io);
        }
    }
}