package projet

import chisel3._
import chisel3.util.switch
import chisel3.util.is

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_writeback = 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, IR, J, B, S = 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;

    // mulw, divw, remw
    val OpMw = "b0111011".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
}

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_writeback := 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) {

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

        // OpImm instructions
        is(OpCode.OpImm) {
            io.reg_file_we := true.B // Write to regfile
            opImm(funct3, funct7, io);
        }

        // OpR instructions
        is(OpCode.OpR) {
            io.reg_file_we := true.B // Write to regfile
            opR(funct3, funct7, io);
        }

        // OpMw instructions
        is(OpCode.OpMw) {
            io.reg_file_we := true.B // Write to regfile
            io.alu_word_mode := true.B
            opM(funct3, io);
        }

        // 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_writeback := 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_writeback := false.B; // Set pc to writeback dest
            io.optype := OpType.I
        }

        is(OpCode.OpBranch) {
            io.optype := OpType.B;
            opBranch(funct3, io);
        }

        is(OpCode.OpStore) {
            io.optype := OpType.S
            opStore(funct3, io);
        }

        is(OpCode.OpLoad) {
            io.optype := OpType.I
            opLoad(funct3, io);
        }

        is(OpCode.OpImmW) {
            io.reg_file_we := true.B // Write to regfile
            io.optype := OpType.I
            opImmW(funct3, funct7, io);
        }
    }

    // Implements functions for all instruction of the opImm kind
    def opImm(funct3: UInt, funct7: UInt, io: CUInterface) = {
        // Meaning of Funct3 in the context of an OpImm instruction
        object Funct3 extends ChiselEnum {
            // Type I
            val ADDI = "b000".U
            val SLTI = "b010".U
            val SLTIU = "b011".U
            val XORI = "b100".U
            val ORI = "b110".U
            val ANDI = "b111".U

            // Type IR
            val SLLI = "b001".U
            val SRLI_SRAI = "b101".U // Right shifts
        }

        // Meaning of Funct7 in the context of SHIFTR IR instructions
        object IRFunct7 extends ChiselEnum {
            val SRLI = "b0000000".U
            val SRAI = "b0100000".U
        }

        io.mux_regb_imm := false.B; // OpImm are operations with imm, so send imm to ALU
        io.mux_alu_imm := true.B;
        io.mux_rega_pc := true.B;
        io.alu_word_mode := false.B;
        switch(funct3) {
            // Type I
            is(Funct3.ADDI) {
                io.alu_opcode := AluOpCode.Add
                io.optype := OpType.I
            }

            is(Funct3.ANDI) {
                io.alu_opcode := AluOpCode.And
                io.optype := OpType.I
            }

            is(Funct3.ORI) {
                io.alu_opcode := AluOpCode.Or
                io.optype := OpType.I
            }

            is(Funct3.XORI) {
                io.alu_opcode := AluOpCode.Xor
                io.optype := OpType.I
            }

            is(Funct3.SLLI) {
                io.alu_opcode := AluOpCode.ShiftLeft
                io.optype := OpType.I
            }

            is(Funct3.SLTI) {
                io.alu_opcode := AluOpCode.LessThanSigned
                io.optype := OpType.I
            }

            is(Funct3.SLTIU) {
                io.alu_opcode := AluOpCode.LessThanUnsigned
                io.optype := OpType.I
            }

            // Type IR
            is(Funct3.SRLI_SRAI) {
                io.optype := OpType.IR
                switch(funct7) {
                    is(IRFunct7.SRLI) {
                        io.alu_opcode := AluOpCode.ShiftRight
                    }

                    is(IRFunct7.SRAI) {
                        io.alu_opcode := AluOpCode.ShiftArithmeticRight
                    }
                }
            }
        }
    }

    // Implements functions for all instruction of the opR kind
    def opR(funct3: UInt, funct7: UInt, io: CUInterface) = {
        // Meaning of Funct3 in the context of an OpR instruction
        object Funct3 extends ChiselEnum {
            // Type I
            val ADD_SUB = "b000".U
            val AND = "b111".U
            val OR = "b110".U
            val XOR = "b100".U
            val SLL = "b001".U
            val SLT = "b010".U
            val SLTU = "b011".U

            // Arithmetic op (discriminate by Funct7)
            val SRL_SRA = "b101".U
        }

        // Meaning of Funct7 in the context of SHIFTR IR instructions
        object IRFunct7 extends ChiselEnum {
            val SRL = "b0000000".U
            val SRA = "b0100000".U
        }

        object ADDFunct7 extends ChiselEnum {
            val ADD = "b0000000".U
            val SUB = "b0100000".U
        }

        io.mux_regb_imm := true.B;
        io.mux_rega_pc := true.B;
        io.mux_alu_imm := true.B;
        switch(funct3) {
            // Type I
            is(Funct3.ADD_SUB) {
                switch(funct7) {
                    is(ADDFunct7.ADD) {
                        io.alu_opcode := AluOpCode.Add
                    }
                    is(ADDFunct7.SUB) {
                        io.alu_opcode := AluOpCode.Sub
                    }
                }
            }

            is(Funct3.AND) {
                io.alu_opcode := AluOpCode.And
            }

            is(Funct3.OR) {
                io.alu_opcode := AluOpCode.Or
            }

            is(Funct3.XOR) {
                io.alu_opcode := AluOpCode.Xor
            }

            is(Funct3.SLL) {
                io.alu_opcode := AluOpCode.ShiftLeft
            }

            is(Funct3.SLT) {
                io.alu_opcode := AluOpCode.LessThanSigned
            }

            is(Funct3.SLTU) {
                io.alu_opcode := AluOpCode.LessThanUnsigned
            }

            // Type IR
            is(Funct3.SRL_SRA) {
                switch(funct7) {
                    is(IRFunct7.SRL) {
                        io.alu_opcode := AluOpCode.ShiftRight
                    }

                    is(IRFunct7.SRA) {
                        io.alu_opcode := AluOpCode.ShiftArithmeticRight
                    }
                }
            }
        }

        // Switch to mul extension
        val RV32M_FUNCT7 = "b0000001".U
        when(funct7 === RV32M_FUNCT7) {
            opM(funct3, io);
        }
    }

    // Implements functions for all instruction of the branch kind
    def opBranch(funct3: UInt, io: CUInterface) = {
        // Meaning of Funct3 in the context of an OpBranch instruction
        object Funct3 extends ChiselEnum {
            val BEQ = "b000".U
            val BNE = "b001".U
            val BLT = "b100".U
            val BGE = "b101".U
            val BLTU = "b110".U
            val BGEU = "b111".U
        }

        val branch = WireInit(false.B);
        switch(funct3) {
            is(Funct3.BEQ) {
                // Send operands from registers to alu
                io.alu_opcode := AluOpCode.Equal
                io.mux_rega_pc := true.B;
                io.mux_regb_imm := true.B;
                branch := io.alu_comp_result
            }

            is(Funct3.BNE) {
                // Send operands from registers to alu
                io.alu_opcode := AluOpCode.Equal
                io.mux_rega_pc := true.B;
                io.mux_regb_imm := true.B;
                branch := ~io.alu_comp_result
            }

            // Signed comparison
            is(Funct3.BLT) {
                // Send operands from registers to alu
                io.alu_opcode := AluOpCode.LessThanSigned
                io.mux_rega_pc := true.B;
                io.mux_regb_imm := true.B;
                branch := io.alu_comp_result
            }

            is(Funct3.BGE) {
                // Send operands from registers to alu
                io.alu_opcode := AluOpCode.LessThanSigned
                io.mux_rega_pc := true.B;
                io.mux_regb_imm := true.B;
                branch := ~io.alu_comp_result
            }

            // Signed comparison
            is(Funct3.BLTU) {
                // Send operands from registers to alu
                io.alu_opcode := AluOpCode.LessThanUnsigned
                io.mux_rega_pc := true.B;
                io.mux_regb_imm := true.B;
                branch := io.alu_comp_result
            }

            is(Funct3.BGEU) {
                // Send operands from registers to alu
                io.alu_opcode := AluOpCode.LessThanUnsigned
                io.mux_rega_pc := true.B;
                io.mux_regb_imm := true.B;
                branch := ~io.alu_comp_result
            }
        }

        // Branch if comparison is successful
        when(branch) {
            io.mux_pcadder_writeback := true.B;
            io.mux_incrpc4_imm := false.B;
            io.mux_regpc_executepc := false.B;
            io.is_jump := true.B;
        }
    }
    // Implements functions for all instruction of the OpStore kind
    def opStore(funct3: UInt, io: CUInterface) = {
        // Meaning of Funct3 in the context of an OpStore instruction
        object Funct3 extends ChiselEnum {
            val SB = "b000".U
            val SH = "b001".U
            val SW = "b010".U
        }

        io.alu_opcode := AluOpCode.Add
        io.mux_alu_imm := true.B
        io.mux_rega_pc := true.B
        io.mux_regb_imm := false.B
        io.memory_we := true.B
        io.memory_en := true.B
        switch(funct3) {
            is(Funct3.SB) {
                io.memory_size := DMemSize.Byte
            }
            is(Funct3.SH) {
                io.memory_size := DMemSize.Half
            }
            is(Funct3.SW) {
                io.memory_size := DMemSize.Word
            }
        }
    }
    // Implements functions for all instruction of the OpLoad kind
    def opLoad(funct3: UInt, io: CUInterface) = {
        // Meaning of Funct3 in the context of an OpLoad instruction
        object Funct3 extends ChiselEnum {
            val LB = "b000".U
            val LH = "b001".U
            val LW = "b010".U
            val LBU = "b100".U
            val LHU = "b101".U
        }

        io.alu_opcode := AluOpCode.Add
        io.mux_alu_imm := true.B
        io.mux_rega_pc := true.B
        io.mux_regb_imm := false.B
        io.mux_writeback_pc := true.B;
        io.mux_executeout_dout := false.B;
        io.reg_file_we := true.B
        io.memory_en := true.B
        io.memory_we := false.B
        switch(funct3) {
            // Byte load
            is(Funct3.LB) {
                io.memory_size := DMemSize.Byte
                io.memory_sign_extend := true.B;
            }
            is(Funct3.LBU) {
                io.memory_size := DMemSize.Byte
            }

            // Half load
            is(Funct3.LH) {
                io.memory_size := DMemSize.Half
                io.memory_sign_extend := true.B;
            }
            is(Funct3.LHU) {
                io.memory_size := DMemSize.Half
            }

            // Word load
            is(Funct3.LW) {
                io.memory_size := DMemSize.Word
            }
        }
    }

    // ~~~ EXTENSIONS ~~~

    // ~~ Rv64i ~~

    // Implements functions for all instruction of the opImm, w  kind
    def opImmW(funct3: UInt, funct7: UInt, io: CUInterface) = {
        // Meaning of Funct3 in the context of an OpImm instruction
        val _ = funct7;
        object Funct3 extends ChiselEnum {
            // Type I
            val ADDIW = "b000".U
            val SLLIW = "b001".U

            // Unimplemented
            val SRLIW_SRAIW = "b101".U // Right shifts
        }

        // Meaning of Funct7 in the context of SHIFTR IR instructions
        // object IRFunct7 extends ChiselEnum {
        //     val SRLI = "b0000000".U
        //     val SRAI = "b0100000".U
        // }

        io.mux_regb_imm := false.B; // OpImm are operations with imm, so send imm to ALU
        io.mux_alu_imm := true.B;
        io.mux_rega_pc := true.B;
        switch(funct3) {
            // Type I
            is(Funct3.ADDIW) {
                io.alu_word_mode := true.B
                io.alu_opcode := AluOpCode.Add
                io.optype := OpType.I
            }
        }
    }

    // ~~ Rv32M ~~

    // Implements functions for instructions of the M extension of type R
    def opM(funct3: UInt, io: CUInterface) = {
        // Meaning of Funct3 in the context of an OpM instruction
        object Funct3 extends ChiselEnum {
            val MUL = "b000".U
            val MULH = "b001".U
            val MULHSU = "b010".U
            val MULHU = "b011".U
            val DIV = "b100".U
            val DIVU = "b101".U
            val REM = "b110".U
            val REMU = "b111".U
        }

        switch(funct3) {
            is(Funct3.MUL) {
                io.alu_opcode := AluOpCode.Mul
            }
            is(Funct3.MULH) {
                io.alu_opcode := AluOpCode.Mulh
            }
            is(Funct3.MULHSU) {
                io.alu_opcode := AluOpCode.Mulhsu
            }
            is(Funct3.MULHU) {
                io.alu_opcode := AluOpCode.Mulhu
            }
            is(Funct3.DIV) {
                io.alu_opcode := AluOpCode.Div
            }
            is(Funct3.DIVU) {
                io.alu_opcode := AluOpCode.Divu
            }
            is(Funct3.REM) {
                io.alu_opcode := AluOpCode.Rem
            }
            is(Funct3.REMU) {
                io.alu_opcode := AluOpCode.Remu
            }
        }
    }
}