from .insn_lui import *
from .insn_auipc import *
from .insn_jal import *
from .insn_jalr import *
from .insn_beq import *
from .insn_bne import *
from .insn_blt import *
from .insn_bge import *
from .insn_bltu import *
from .insn_bgeu import *
from .insn_lb import *
from .insn_lh import *
from .insn_lw import *
from .insn_lbu import *
from .insn_lhu import *
from .insn_sb import *
from .insn_sh import *
from .insn_sw import *
from .insn_addi import *
from .insn_slti import *
from .insn_sltiu import *
from .insn_xori import *
from .insn_ori import *
from .insn_andi import *
from .insn_slli import *
from .insn_srli import *
from .insn_srai import *
from .insn_add import *
from .insn_sub import *
from .insn_sll import *
from .insn_slt import *
from .insn_sltu import *
from .insn_xor import *
from .insn_srl import *
from .insn_sra import *
from .insn_or import *
from .insn_and import *
from .insn_mul import *
from .insn_mulh import *
from .insn_mulhsu import *
from .insn_mulhu import *
from .insn_div import *
from .insn_divu import *
from .insn_rem import *
from .insn_remu import *

"""
RV32M Standard Extension
"""

class IsaRV32M(Elaboratable):
    def __init__(self, params):
        # Core-specific parameters
        self.params = params

        # Input ports
        self.rvfi_valid = Signal(1)
        self.rvfi_insn = Signal(self.params.ilen)
        self.rvfi_pc_rdata = Signal(self.params.xlen)
        self.rvfi_rs1_rdata = Signal(self.params.xlen)
        self.rvfi_rs2_rdata = Signal(self.params.xlen)
        self.rvfi_mem_rdata = Signal(self.params.xlen)
        if self.params.csr_misa:
            self.rvfi_csr_misa_rdata = Signal(self.params.xlen)

        # Output ports
        if self.params.csr_misa:
            self.spec_csr_misa_rmask = Signal(self.params.xlen)
        self.spec_valid = Signal(1)
        self.spec_trap = Signal(1)
        self.spec_rs1_addr = Signal(5)
        self.spec_rs2_addr = Signal(5)
        self.spec_rd_addr = Signal(5)
        self.spec_rd_wdata = Signal(self.params.xlen)
        self.spec_pc_wdata = Signal(self.params.xlen)
        self.spec_mem_addr = Signal(self.params.xlen)
        self.spec_mem_rmask = Signal(int(self.params.xlen // 8))
        self.spec_mem_wmask = Signal(int(self.params.xlen // 8))
        self.spec_mem_wdata = Signal(self.params.xlen)

    def ports(self):
        input_ports = [
            self.rvfi_valid,
            self.rvfi_insn,
            self.rvfi_pc_rdata,
            self.rvfi_rs1_rdata,
            self.rvfi_rs2_rdata,
            self.rvfi_mem_rdata
        ]
        if self.params.csr_misa:
            input_ports.append(self.rvfi_csr_misa_rdata)
        output_ports = [
            self.spec_valid,
            self.spec_trap,
            self.spec_rs1_addr,
            self.spec_rs2_addr,
            self.spec_rd_addr,
            self.spec_rd_wdata,
            self.spec_pc_wdata,
            self.spec_mem_addr,
            self.spec_mem_rmask,
            self.spec_mem_wmask,
            self.spec_mem_wdata
        ]
        if self.params.csr_misa:
            output_ports.append(self.spec_csr_misa_rmask)
        return input_ports + output_ports

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

        insn_submodules = {}

        m.submodules._lui = insn_submodules['lui'] = InsnLui(self.params)
        m.submodules._auipc = insn_submodules['auipc'] = InsnAuipc(self.params)
        m.submodules._jal = insn_submodules['jal'] = InsnJal(self.params)
        m.submodules._jalr = insn_submodules['jalr'] = InsnJalr(self.params)
        m.submodules._beq = insn_submodules['beq'] = InsnBeq(self.params)
        m.submodules._bne = insn_submodules['bne'] = InsnBne(self.params)
        m.submodules._blt = insn_submodules['blt'] = InsnBlt(self.params)
        m.submodules._bge = insn_submodules['bge'] = InsnBge(self.params)
        m.submodules._bltu = insn_submodules['bltu'] = InsnBltu(self.params)
        m.submodules._bgeu = insn_submodules['bgeu'] = InsnBgeu(self.params)
        m.submodules._lb = insn_submodules['lb'] = InsnLb(self.params)
        m.submodules._lh = insn_submodules['lh'] = InsnLh(self.params)
        m.submodules._lw = insn_submodules['lw'] = InsnLw(self.params)
        m.submodules._lbu = insn_submodules['lbu'] = InsnLbu(self.params)
        m.submodules._lhu = insn_submodules['lhu'] = InsnLhu(self.params)
        m.submodules._sb = insn_submodules['sb'] = InsnSb(self.params)
        m.submodules._sh = insn_submodules['Sh'] = InsnSh(self.params)
        m.submodules._sw = insn_submodules['sw'] = InsnSw(self.params)
        m.submodules._addi = insn_submodules['addi'] = InsnAddi(self.params)
        m.submodules._slti = insn_submodules['slti'] = InsnSlti(self.params)
        m.submodules._sltiu = insn_submodules['sltiu'] = InsnSltiu(self.params)
        m.submodules._xori = insn_submodules['xori'] = InsnXori(self.params)
        m.submodules._ori = insn_submodules['ori'] = InsnOri(self.params)
        m.submodules._andi = insn_submodules['andi'] = InsnAndi(self.params)
        m.submodules._slli = insn_submodules['slli'] = InsnSlli(self.params)
        m.submodules._srli = insn_submodules['srli'] = InsnSrli(self.params)
        m.submodules._srai = insn_submodules['srai'] = InsnSrai(self.params)
        m.submodules._add = insn_submodules['add'] = InsnAdd(self.params)
        m.submodules._sub = insn_submodules['sub'] = InsnSub(self.params)
        m.submodules._sll = insn_submodules['sll'] = InsnSll(self.params)
        m.submodules._slt = insn_submodules['slt'] = InsnSlt(self.params)
        m.submodules._sltu = insn_submodules['sltu'] = InsnSltu(self.params)
        m.submodules._xor = insn_submodules['xor'] = InsnXor(self.params)
        m.submodules._srl = insn_submodules['srl'] = InsnSrl(self.params)
        m.submodules._sra = insn_submodules['sra'] = InsnSra(self.params)
        m.submodules._or = insn_submodules['or'] = InsnOr(self.params)
        m.submodules._and = insn_submodules['and'] = InsnAnd(self.params)
        m.submodules._mul = insn_submodules['mul'] = InsnMul(self.params)
        m.submodules._mulh = insn_submodules['mulh'] = InsnMulh(self.params)
        m.submodules._mulhsu = insn_submodules['mulhsu'] = InsnMulhsu(self.params)
        m.submodules._mulhu = insn_submodules['mulhu'] = InsnMulhu(self.params)
        m.submodules._div = insn_submodules['div'] = InsnDiv(self.params)
        m.submodules._divu = insn_submodules['divu'] = InsnDivu(self.params)
        m.submodules._rem = insn_submodules['rem'] = InsnRem(self.params)
        m.submodules._remu = insn_submodules['remu'] = InsnRemu(self.params)

        for _, insn in insn_submodules.items():
            m.d.comb += insn.rvfi_valid.eq(self.rvfi_valid)
            m.d.comb += insn.rvfi_insn.eq(self.rvfi_insn)
            m.d.comb += insn.rvfi_pc_rdata.eq(self.rvfi_pc_rdata)
            m.d.comb += insn.rvfi_rs1_rdata.eq(self.rvfi_rs1_rdata)
            m.d.comb += insn.rvfi_rs2_rdata.eq(self.rvfi_rs2_rdata)
            m.d.comb += insn.rvfi_mem_rdata.eq(self.rvfi_mem_rdata)
            if self.params.csr_misa:
                m.d.comb += insn.rvfi_csr_misa_rdata.eq(self.rvfi_csr_misa_rdata)

        spec_valid = 0
        for _, insn in insn_submodules.items():
            spec_valid = Mux(insn.spec_valid, insn.spec_valid, spec_valid)
        m.d.comb += self.spec_valid.eq(spec_valid)

        spec_trap = 0
        for _, insn in insn_submodules.items():
            spec_trap = Mux(insn.spec_valid, insn.spec_trap, spec_trap)
        m.d.comb += self.spec_trap.eq(spec_trap)

        spec_rs1_addr = 0
        for _, insn in insn_submodules.items():
            spec_rs1_addr = Mux(insn.spec_valid, insn.spec_rs1_addr, spec_rs1_addr)
        m.d.comb += self.spec_rs1_addr.eq(spec_rs1_addr)

        spec_rs2_addr = 0
        for _, insn in insn_submodules.items():
            spec_rs2_addr = Mux(insn.spec_valid, insn.spec_rs2_addr, spec_rs2_addr)
        m.d.comb += self.spec_rs2_addr.eq(spec_rs2_addr)

        spec_rd_addr = 0
        for _, insn in insn_submodules.items():
            spec_rd_addr = Mux(insn.spec_valid, insn.spec_rd_addr, spec_rd_addr)
        m.d.comb += self.spec_rd_addr.eq(spec_rd_addr)

        spec_rd_wdata = 0
        for _, insn in insn_submodules.items():
            spec_rd_wdata = Mux(insn.spec_valid, insn.spec_rd_wdata, spec_rd_wdata)
        m.d.comb += self.spec_rd_wdata.eq(spec_rd_wdata)

        spec_pc_wdata = 0
        for _, insn in insn_submodules.items():
            spec_pc_wdata = Mux(insn.spec_valid, insn.spec_pc_wdata, spec_pc_wdata)
        m.d.comb += self.spec_pc_wdata.eq(spec_pc_wdata)

        spec_mem_addr = 0
        for _, insn in insn_submodules.items():
            spec_mem_addr = Mux(insn.spec_valid, insn.spec_mem_addr, spec_mem_addr)
        m.d.comb += self.spec_mem_addr.eq(spec_mem_addr)

        spec_mem_rmask = 0
        for _, insn in insn_submodules.items():
            spec_mem_rmask = Mux(insn.spec_valid, insn.spec_mem_rmask, spec_mem_rmask)
        m.d.comb += self.spec_mem_rmask.eq(spec_mem_rmask)

        spec_mem_wmask = 0
        for _, insn in insn_submodules.items():
            spec_mem_wmask = Mux(insn.spec_valid, insn.spec_mem_wmask, spec_mem_wmask)
        m.d.comb += self.spec_mem_wmask.eq(spec_mem_wmask)

        spec_mem_wdata = 0
        for _, insn in insn_submodules.items():
            spec_mem_wdata = Mux(insn.spec_valid, insn.spec_mem_wdata, spec_mem_wdata)
        m.d.comb += self.spec_mem_wdata.eq(spec_mem_wdata)

        if self.params.csr_misa:
            spec_csr_misa_rmask = 0
            for _, insn in insn_submodules.items():
                spec_csr_misa_rmask = Mux(insn.spec_valid, insn.spec_csr_misa_rmask, spec_csr_misa_rmask)
            m.d.comb += self.spec_csr_misa_rmask.eq(spec_csr_misa_rmask)

        return m