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
    val multiplication_res = op_a * op_b // Mul non signée
    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
            }
        }
    }
}