package projet

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

object AluOpCode extends ChiselEnum {
    val Add, Sub, And, Or, Xor, ShiftLeft, ShiftRight, ShiftArithmeticRight,
        LessThanSigned, LessThanUnsigned, GreaterEqualSigned,
        GreaterEqualUnsigned, Equal = Value
}

class Alu extends Module {
    import AluOpCode._

    val io = IO(new Bundle {
        val a = Input(UInt(64.W))
        val b = Input(UInt(64.W))
        val out = Output(UInt(64.W))
        val opcode = Input(AluOpCode())
        val comp_result = Output(Bool())

        // ~~ Rv64i ~~
        // If word_mode is enabled, operations are performed
        // on 32-bit-truncated inputs and then sign extended to 64bits
        val word_mode = Input(Bool())
    })

    io.out := 0.U;
    io.comp_result := false.B;
    val out = WireInit(0.U(64.W))
    val op_a = WireInit(0.U(64.W))
    val op_b = WireInit(0.U(64.W))

    // Defaults
    op_a := io.a
    op_b := io.b
    io.out := out

    // Truncate operands, sext result
    when(io.word_mode) {
        op_a := io.a(31, 0)
        op_b := io.b(31, 0)

        io.out := SignExtend.Sext(out(31, 0), 64)
    }

    switch(io.opcode) {
        is(Add) {
            out := op_a + op_b;
        }

        is(Sub) {
            out := op_a - op_b;
        }

        is(And) {
            out := op_a & op_b;
        }

        is(Or) {
            out := op_a | op_b;
        }

        is(Xor) {
            out := op_a ^ op_b;
        }

        is(ShiftLeft) {
            out := op_a << op_b(5, 0);
        }

        is(ShiftRight) {
            out := op_a >> op_b(5, 0);
        }

        is(ShiftArithmeticRight) {
            // >> means arithmetic shift for SInts
            out := (op_a.asSInt >> op_b(5, 0)).asUInt;
            // Signed shift on 32bit requires careful handling
            when(io.word_mode) {

                out := (op_a(31, 0).asSInt >> op_b(5, 0)).asUInt
            }
        }

        is(LessThanUnsigned) {
            val less_than = op_a < op_b;
            out := less_than;
            io.comp_result := less_than;
        }

        is(LessThanSigned) {
            val less_than = op_a.asSInt < op_b.asSInt
            out := less_than
            io.comp_result := less_than;
        }

        is(GreaterEqualUnsigned) {
            val geq_than = op_a >= op_b;
            out := geq_than
            io.comp_result := geq_than;
        }

        is(GreaterEqualSigned) {
            val geq_than = op_a.asSInt >= op_b.asSInt
            out := geq_than
            io.comp_result := geq_than;
        }

        is(Equal) {
            io.comp_result := op_a === op_b;
        }
    }

}