# pyright: reportInvalidTypeForm=false

from amaranth import *
from amaranth.lib import wiring

from isa import AluOp

class ALU(wiring.Component):
    # Inputs
    value1:     wiring.In(32)
    value2:     wiring.In(32)

    funct3:     wiring.In(3)
    funct7:     wiring.In(7)
    immediate:  wiring.In(1)

    # Outputs
    result:     wiring.Out(32)

    def __init__(self):
        super().__init__()

    def elaborate(self, platform):
        m = Module()

        shamt   = Signal(5)

        m.d.comb += shamt.eq(self.value2[0:5])

        with m.Switch(self.funct3):
            with m.Case(AluOp.ADD | AluOp.SUB):
                m.d.comb += self.result.eq(
                    Mux(
                        # SUBI doesn't exsit. However, since funct7[5] falls into immediate region in I type instructions
                        # we must differentiate between immediate alu op and register alu op
                        self.funct7[5] & (~self.immediate),
                        self.value1 - self.value2,
                        self.value1 + self.value2,
                    )
                )
            with m.Case(AluOp.SLL):
                m.d.comb += self.result.eq(self.value1 << shamt)
            with m.Case(AluOp.SLT):
                m.d.comb += self.result.eq(self.value1.as_signed() < self.value2.as_signed())
            with m.Case(AluOp.SLTU):
                m.d.comb += self.result.eq(self.value1 < self.value2)
            with m.Case(AluOp.XOR):
                m.d.comb += self.result.eq(self.value1 ^ self.value2)
            with m.Case(AluOp.SRA | AluOp.SRL):
                m.d.comb += self.result.eq(
                    Mux(
                        self.funct7[5],
                        self.value1.as_signed() >> shamt,
                        self.value1 >> shamt,
                    )
                )
            with m.Case(AluOp.OR):
                m.d.comb += self.result.eq(self.value1 | self.value2)
            with m.Case(AluOp.AND):
                m.d.comb += self.result.eq(self.value1 & self.value2)

        return m