package projet

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

object AluOpCode extends ChiselEnum {
    val Add, Sub, Mul, Mulh, Mulhsu, Mulhu, Div, Divu, Rem, Remu, 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())

        // Should stop the stall for the current computation (e.g. Division)
        val should_stop_stall = 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.should_stop_stall := false.B
    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
    val op_a_sign = WireInit(op_a(63))
    val op_b_sign = WireInit(op_b(63))
    val op_a_unsigned = WireInit(op_a)
    val op_b_unsigned = WireInit(op_b)
    io.out := out

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

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

    when(op_a_sign) {
        op_a_unsigned := Utils.two_complement(op_a, io.word_mode)
    }
    when(op_b_sign) {
        op_b_unsigned := Utils.two_complement(op_b, io.word_mode)
    }

    val in_computation = RegInit(false.B)

    // Division
    val divider = Module(new Divider(64, Constants.DIVISION_CYCLES_COUNT))
    divider.io.word_mode := io.word_mode
    divider.io.dividend := op_a_unsigned
    divider.io.divisor := op_b_unsigned
    divider.io.start := false.B
    def div_by_zero() = {
        when(io.word_mode) {
            out := (-1).S(32.W).asUInt
        }.otherwise {
            out := (-1).S(64.W).asUInt
        }
    }

    // Multiplication

    // Pipeline multiplication to trim critical path
    val multiplication_res = RegNext(op_a * op_b) // Mul non signée

    switch(io.opcode) {
        is(Mul, Mulhsu, Mulh, Mulhu) {
            when(!in_computation) {
                in_computation := true.B
            }.elsewhen(in_computation) {
                in_computation := false.B
                io.should_stop_stall := true.B
            }
        }
    }

    val multiplication_low =
        Mux(io.word_mode, multiplication_res(31, 0), multiplication_res(63, 0))
    val multiplication_high =
        Mux(
            io.word_mode,
            multiplication_res(63, 32),
            multiplication_res(127, 64)
        )

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

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

        is(Mul) {
            out := multiplication_low;
        }

        is(Mulh) {
            val correction =
                Mux(op_a_sign, op_b, 0.U) + Mux(op_b_sign, op_a, 0.U)
            val multiplication_res_with_correction =
                multiplication_high - correction
            out := multiplication_res_with_correction;
        }

        is(Mulhsu) {
            val correction =
                Mux(op_a_sign, op_b, 0.U)
            val multiplication_res_with_correction =
                multiplication_high - correction
            out := multiplication_res_with_correction;
        }

        is(Mulhu) {
            out := multiplication_high;
        }

        is(Div) {
            when(op_b === 0.U) {
                div_by_zero()
            }.elsewhen(op_a_sign ^ op_b_sign) {
                out := Utils.two_complement(divider.io.quotient, io.word_mode)
            }.otherwise {
                out := divider.io.quotient
            }
        }

        is(Divu) {
            divider.io.dividend := op_a
            divider.io.divisor := op_b
            when(op_b === 0.U) {
                div_by_zero()
            }.otherwise { out := divider.io.quotient; }
        }

        is(Rem) {
            when(op_b === 0.U) {
                out := op_a
            }.elsewhen(
                op_a_sign && !op_b_sign && divider.io.remainder =/= 0.U
            ) {
                out := op_b - divider.io.remainder
            }.elsewhen(
                !op_a_sign && op_b_sign && divider.io.remainder =/= 0.U
            ) {
                out := Utils.two_complement(
                    Utils.two_complement(op_b) - divider.io.remainder,
                    io.word_mode
                )
            }.elsewhen(op_a_sign && op_b_sign) {
                out := Utils.two_complement(divider.io.remainder, io.word_mode)
            }.otherwise {
                out := divider.io.remainder
            }
        }

        is(Remu) {
            divider.io.dividend := op_a
            divider.io.divisor := op_b
            when(op_b === 0.U) {
                out := op_a
            }.otherwise {
                out := divider.io.remainder
            }
        }

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

    switch(io.opcode) {
        is(Div, Divu, Rem, Remu) {
            when(!in_computation) {
                divider.io.start := true.B
                in_computation := true.B
            }.elsewhen(divider.io.ready) {
                io.should_stop_stall := true.B
                in_computation := false.B
            }
        }
    }

    // Special overflow cases for Div and Rem
    when(
        io.word_mode &&
            op_b === ((BigInt(1) << 32) - 1).U(32.W) &&
            op_a === (-1).S(32.W).asUInt
    ) {
        switch(io.opcode) {
            is(Div) {
                out := (1.U << 31)
            }
            is(Rem) {
                out := 0.U
            }
        }
    }.elsewhen(
        !io.word_mode &&
            op_b === ((BigInt(1) << 64) - 1).U(64.W) &&
            op_a === (-1).S(64.W).asUInt
    ) {
        switch(io.opcode) {
            is(Div) {
                out := (1.U << 63)
            }
            is(Rem) {
                out := 0.U
            }
        }
    }
}
